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

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()

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