diff --git a/shared/data_store/__init__.py b/shared/data_store/__init__.py index 58cba31..5f37533 100644 --- a/shared/data_store/__init__.py +++ b/shared/data_store/__init__.py @@ -2,5 +2,7 @@ from .connect import connect from .delete import delete from .get import get from .get_image import get_image +from .get_model import get_model from .put import put from .put_image import put_image +from .put_model import put_model diff --git a/shared/data_store/get_model.py b/shared/data_store/get_model.py new file mode 100644 index 0000000..3610430 --- /dev/null +++ b/shared/data_store/get_model.py @@ -0,0 +1,14 @@ +"""Definition of get_model function.""" + +from minio import Minio +from torch.nn import Module + + +def get_model( + client: Minio, + object_name: str, +) -> Module: + """Get model from model subfolder in bucket in Minio.""" + assert isinstance(client, Minio) + assert isinstance(object_name, str) + raise NotImplementedError() diff --git a/shared/data_store/put_model.py b/shared/data_store/put_model.py new file mode 100644 index 0000000..06be4fc --- /dev/null +++ b/shared/data_store/put_model.py @@ -0,0 +1,44 @@ +"""Definition of put_model function.""" + +import logging +import os +from hashlib import md5 +from io import BytesIO + +import torch +from minio import Minio +from torch.nn import Module + + +def put_model( + client: Minio, + model: Module, +) -> str: + """Put model in model subfolder in bucket in Minio and return MD5 checksum + used as object name.""" + assert isinstance(client, Minio) + assert isinstance(model, Module) + bucket_name = os.getenv('MINIO_BUCKET_NAME') + subfolder = 'models' + # save image to buffer + buffer = BytesIO() + torch.save(model.state_dict(), buffer) + # get md5 of image + checksum = md5(buffer.getbuffer()).hexdigest() + # prepare for saving + num_bytes = buffer.tell() + buffer.seek(0) + # send data to bucket + object_name = f'{subfolder}/{checksum}' + try: + client.put_object( + bucket_name=bucket_name, + object_name=object_name, + length=num_bytes, + data=buffer, + ) + except Exception as exc: + logging.error('failed saving data to MinIO') + raise exc + logging.debug('saved data to %s', object_name) + return checksum