def _download_model(self, location: dict):
if not location:
return None
model = ModelWrapper.unwrap(location, self.context.model_serializer())
if model.blob:
#Embedded model
model = model.blob
else:
xsum = model.wrapping['xsum']
if not xsum:
raise fflabc.MalformedResponseException(f"Malformed wrapping (xsum): {model.wrapping}")
#Download from bin store
url = model.wrapping.get('url', None)
if not url:
raise fflabc.MalformedResponseException(f"Malformed wrapping: {model.wrapping}")
#Download from bin store
if self.context.download_models():
with rabbitmq.RabbitHeartbeat([self.subscriber, self.publisher]):
buff = requests.get(url)
#Now compare checksums
xsum2 = model.xsum(buff.content)
if xsum != xsum2:
raise fflabc.MalformedResponseException(f"Checksum mismatch!")
model = self.context.model_serializer().deserialize(buff.content)
else:
#Let user decide what to do
model = model.wrapping
return model
def task_notification(self, timeout: int = 0, flavours: list = None) -> dict:
"""
Wait for a message to arrive or until timeout.
If message is received, check whether its notification type matches
element in given list of notification flavours.
Throws: An exception on failure
:param timeout: timeout in seconds
:type timeout: `int`
:param flavours: expected notification types
:type flavours: `list`
:return: received message
:rtype: `dict`
"""
msg = self.receive(timeout)
if 'notification' not in msg:
raise fflabc.BadNotificationException(f"Malformed object: {msg}")
if 'type' not in msg['notification']:
raise fflabc.BadNotificationException(f"Malformed object: {msg['notification']}")
try:
if fflabc.Notification(msg['notification']['type']) not in flavours:
raise ValueError
except:
raise fflabc.BadNotificationException(f"Unexpected notification " \
f"{msg['notification']['type']}, expecting {flavours}")
if 'params' not in msg:
raise fflabc.BadNotificationException(f"Malformed payload: {msg}")
model = None
if msg['params']:
model = ModelWrapper.unwrap(msg['params'], self.context.model_serializer())
if model.blob:
#Embedded model
model = model.blob
else:
#Download from bin store
url = model.wrapping.get('url', None)
if not url:
raise fflabc.MalformedResponseException(f"Malformed wrapping: {model.wrapping}")
#Download from bin store
if self.context.download_models():
self.model_files.append(utils.FileDownloader(url))
with open(self.model_files[-1].name(), 'rb') as model_file:
buff = model_file.read()
model = self.context.model_serializer().deserialize(buff)
else:
#Let user decide what to do
model = model.wrapping
return fflabc.Response(msg['notification'], model)
def _invoke_service(self, message: dict, timeout: int = 0) -> dict:
"""
Send a message and wait for a reply or until timeout.
Throws: An exception on failure
:param message: message to be sent
:type message: `dict`
:param timeout: timeout in seconds
:type timeout: `int`
:return: received message
:rtype: `dict`
"""
if not timeout:
timeout = self.timeout
try:
#Need a reply, so add this to the request message
message = self.catalog.msg_assign_reply(message,
self.command_queue.name)
message = self.context.serializer().serialize(message)
result = super().invoke_service(message,
timeout,
queue=self.command_queue)
except rabbitmq.RabbitTimedOutException as exc:
raise TimedOutException(exc) from exc
except rabbitmq.RabbitConsumerException as exc:
raise ConsumerException(exc) from exc
if not result:
raise fflabc.MalformedResponseException("Malformed object: None")
result = self.context.serializer().deserialize(result)
if 'error' in result:
raise fflabc.ServerException(
f"Server Error ({result['activation']}): {result['error']}")
if 'calls' not in result:
raise fflabc.MalformedResponseException(
f"Malformed object: {result}")
results = result['calls'][0]['count'] # calls[0] will always succeed
return result['calls'][0]['data'] if results else []
def _dispatch_model(self, task_name: str = None, model: dict = None) -> dict:
"""
Dispatch a model and determine its download location.
Throws: An exception on failure
:param model: model to be sent
:type model: `dict`
:return: download location information
:rtype: `dict`
"""
wrapper = ModelWrapper.wrap(model, self.context.model_serializer())
if not model:
return wrapper.wrapping
# First, obtain the upload location/keys
if task_name:
message = self.catalog.msg_bin_upload_object(task_name)
else:
if len(wrapper.blob) > self.context.dispatch_threshold():
message = self.catalog.msg_bin_uploader()
else:
#Small model - embed it
return wrapper.wrapping
upload_info = self._invoke_service(message)
if 'key' not in upload_info['fields']:
raise fflabc.MalformedResponseException('Update Error: Malformed URL')
key = upload_info['fields']['key']
try:
with rabbitmq.RabbitHeartbeat(self.subscriber):
# And then perform the upload
response = requests.post(upload_info['url'],
files={'file': wrapper.blob},
data=upload_info['fields'],
headers=None)
response.raise_for_status()
except requests.exceptions.RequestException as err:
raise fflabc.DispatchException(err) from err
except:
raise fflabc.DispatchException(f'General Update Error')
# Now obtain the download location/keys
if task_name:
message = self.catalog.msg_bin_download_object(key)
else:
message = self.catalog.msg_bin_downloader(key)
download_info = self._invoke_service(message)
wrapper = ModelWrapper.wrap({'url': download_info, 'key': key})
return wrapper.wrapping
def _dispatch_model(self, task_name: str = None, model: dict = None) -> dict:
"""
Dispatch a model and determine its download location.
Throws: An exception on failure
:param model: model to be sent
:type model: `dict`
:return: download location information
:rtype: `dict`
"""
wrapper = ModelWrapper.wrap(model, self.context.model_serializer())
if not model:
return wrapper.wrapping
# First, obtain the upload location/keys
if task_name:
message = self.catalog.msg_bin_upload_object(task_name)
else:
if len(wrapper.blob) > self.context.dispatch_threshold():
message = self.catalog.msg_bin_uploader()
else:
#Small model - embed it
return wrapper.wrapping
upload_info = self._invoke_service(message)
if 'key' not in upload_info['fields']:
raise fflabc.MalformedResponseException('Update Error: Malformed URL')
key = upload_info['fields']['key']
try:
with rabbitmq.RabbitHeartbeat([self.subscriber, self.publisher]):
# And then perform the upload
if len(wrapper.blob) > self.context.model_threshold():
self._post_big_model(wrapper.blob, upload_info)
else:
self._post_model(wrapper.blob, upload_info)
except (requests.exceptions.RequestException, Exception) as err:
raise fflabc.DispatchException(err) from err
#Now add the checksum to the inline message
if isinstance(model, bytes):
xsum = model.xsum(model)
else:
xsum = wrapper.xsum()
wrapper = ModelWrapper.wrap({'key': key, 'xsum': xsum})
return wrapper.wrapping
