Coverage for src/python/dandeliion/client/simulator.py: 96%

53 statements  

« prev     ^ index     » next       coverage.py v7.6.12, created at 2025-02-22 10:38 +0100

1""" 

2@file python/dandeliion/client/simulator.py 

3 

4module containing Dandeliion Simulator class 

5""" 

6 

7# 

8# Copyright (C) 2024-2025 Dandeliion Team 

9# 

10# This library is free software; you can redistribute it and/or modify it under 

11# the terms of the GNU Lesser General Public License as published by the Free 

12# Software Foundation; either version 3.0 of the License, or (at your option) 

13# any later version. 

14# 

15# This library is distributed in the hope that it will be useful, but WITHOUT 

16# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS 

17# FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more 

18# details. 

19# 

20# You should have received a copy of the GNU Lesser General Public License 

21# along with this library; if not, write to the Free Software Foundation, Inc., 

22# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA 

23# 

24 

25# built-in modules 

26import logging 

27import requests 

28import threading 

29from dataclasses import dataclass 

30from typing import Optional 

31 

32# custom modules 

33from .tools.misc import update_dict 

34from .websocket import SimulatorWebSocketClient 

35from .exceptions import DandeliionAPIException 

36from .solution import Solution 

37 

38logger = logging.getLogger(__name__) 

39 

40 

41@dataclass 

42class Simulator: 

43 

44 """ 

45 Simulator class that stores authentication details and deals with job submission and result acquisition 

46 """ 

47 

48 api_url: str 

49 api_key: str 

50 

51 def submit(self, parameters: dict, is_blocking: bool = True): 

52 """ 

53 Submit parameters to Simulator instance 

54 """ 

55 

56 # submit simulation to rest api 

57 headers = {'Authorization': f'Token {self.api_key}'} # TODO adapt to server 

58 response = requests.post(url=self.api_url, json=parameters, headers=headers) 

59 if response.status_code >= 400: 

60 raise DandeliionAPIException(f"Your request has failed: {response.reason}") 

61 response_json = response.json() 

62 

63 run_id = response_json['Run']['id'] 

64 data = update_dict(parameters, response_json, inline=False) 

65 

66 if is_blocking: 

67 cond = threading.Condition() 

68 

69 def task_update_signal_hook(updates): 

70 with cond: 

71 data['Run']['status'] = updates['status'] 

72 logger.info(updates['log_message']) 

73 cond.notify_all() 

74 

75 client = SimulatorWebSocketClient( 

76 url=data['Run']['ws_status_url'], 

77 api_key=self.api_key, 

78 on_update=task_update_signal_hook, 

79 ) 

80 client.subscribe(run_id) 

81 while data['Run']['status'] in ['queued', 'running']: 

82 # block until task update signalled 

83 with cond: 

84 cond.wait() 

85 # closing connection again 

86 client.close() 

87 

88 return Solution(self, data) 

89 

90 def update_results(self, prefetched_data: dict, keys: list = None, inline: bool = False) -> Optional[dict]: 

91 """ 

92 Function to (pre)fetch result columns from server and update/append prefetched_data 

93 """ 

94 params = [('key', key) for key in keys] if keys is not None else [] 

95 params.append(('id', prefetched_data['Run']['id'])) 

96 

97 headers = {'Authorization': f'Token {self.api_key}'} # TODO adapt to server 

98 response = requests.get(url=self.api_url, params=params, headers=headers) 

99 if response.status_code >= 400: 

100 raise DandeliionAPIException(f"Your request has failed: {response.reason}. Try again?") 

101 

102 response_json = response.json() 

103 # sanity check if id for sim returned is same as the one requested 

104 if response_json['Run']['id'] != prefetched_data['Run']['id']: 

105 raise DandeliionAPIException( 

106 "Something went wrong." 

107 f" Reported run id is {response_json['Run']['id']}" 

108 f" (requested: {prefetched_data['Run']['id']})" 

109 ) 

110 

111 if inline: 

112 update_dict(prefetched_data, response.json()) 

113 else: 

114 return response.json() 

115 

116 def get_status(self, prefetched_data: dict) -> str: 

117 """ 

118 Returns current status of a simulation (as either stored in prefetched_data 

119 if finished/failed or retrieved from server if potentially still queued/running) 

120 """ 

121 if prefetched_data['Run']['status'] in ['queued', 'running']: 

122 self.update_results(prefetched_data, inline=True) 

123 return prefetched_data['Run']['status']