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

1"""

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

3

4 module 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

26 import logging

27 import requests

28 import threading

29 from dataclasses import dataclass

30 from typing import Optional

31

32# custom modules

33 from .tools.misc import update_dict

34 from .websocket import SimulatorWebSocketClient

35 from .exceptions import DandeliionAPIException

36 from .solution import Solution

37

38 logger = logging.getLogger(__name__)

39

40

41@dataclass

42 class 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']