moved files into src folder

shared/datastore/src/__init__.py

@@ -0,0 +1,8 @@
+from .connect_minio import connect_minio
+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/datastore/src/connect_minio.py

@@ -0,0 +1,42 @@
+"""Definition of connect function."""
+
+import logging
+import os
+
+from minio import Minio
+
+
+def connect_minio() -> Minio:
+    """Connect to MinIO server."""
+    # ensure necessary env vars available
+    env_var_list = [
+        'MINIO_ENDPOINT',
+        'MINIO_ACCESS_KEY',
+        'MINIO_SECRET_KEY',
+        'MINIO_BUCKET_NAME',
+    ]
+    for env_var in env_var_list:
+        # ensure env var set
+        assert (
+            env_var in os.environ
+        ), f"environment variable not set: {
+            env_var
+        }"
+    # prepare arguments
+    minio_endpoint = os.getenv('MINIO_ENDPOINT', default='')
+    minio_access_key = os.getenv('MINIO_ACCESS_KEY', default='')
+    minio_secret_key = os.getenv('MINIO_SECRET_KEY', default='')
+    minio_bucket_name = os.getenv('MINIO_BUCKET_NAME', default='')
+    # connect client
+    client = Minio(
+        endpoint=minio_endpoint,
+        access_key=minio_access_key,
+        secret_key=minio_secret_key,
+        secure=False,
+    )
+    # ensure bucket exists
+    if not client.bucket_exists(bucket_name=minio_bucket_name):
+        logging.info('creating bucket: %s', minio_bucket_name)
+        client.make_bucket(bucket_name=minio_bucket_name)
+    logging.debug('finished')
+    return client

shared/datastore/src/delete.py

@@ -0,0 +1,30 @@
+"""Definition of delete function."""
+
+from __future__ import annotations
+
+import logging
+import os
+
+from minio import Minio
+
+
+def delete(
+    client: Minio,
+    object_name: str,
+) -> None:
+    """Delete object from MinIO."""
+    assert isinstance(client, Minio)
+    assert isinstance(object_name, str)
+    bucket_name = os.getenv('MINIO_BUCKET_NAME', default='')
+    assert len(bucket_name) > 0
+    # remove object
+    try:
+        client.remove_object(
+            bucket_name=bucket_name,
+            object_name=object_name,
+        )
+    except Exception as exc:
+        logging.debug(exc)
+        logging.error('failed deleting %s', object_name)
+    else:
+        logging.debug('deleted %s', object_name)

shared/datastore/src/get.py

@@ -0,0 +1,39 @@
+"""Definition of get function."""
+
+import logging
+import os
+from io import BytesIO
+from traceback import print_exc
+
+from minio import Minio
+
+
+def get(
+    client: Minio,
+    object_name: str,
+) -> BytesIO:
+    """Get buffer from bucket in MinIO."""
+    assert isinstance(client, Minio)
+    assert isinstance(object_name, str)
+    bucket_name = os.getenv('MINIO_BUCKET_NAME', default='')
+    # get buffer
+    try:
+        response = client.get_object(
+            bucket_name=bucket_name,
+            object_name=object_name,
+        )
+        assert response.status == 200
+        buffer = BytesIO()
+        chunk_size = 2**14
+        while chunk := response.read(chunk_size):
+            buffer.write(chunk)
+        buffer.seek(0)
+        logging.debug('got %s', object_name)
+        return buffer
+    except Exception as exc:
+        logging.error('failed getting data from MinIO')
+        print_exc()
+        raise exc
+    finally:
+        response.close()
+        response.release_conn()

shared/datastore/src/get_image.py

@@ -0,0 +1,41 @@
+"""Definition of get_image function."""
+
+import logging
+import os
+from io import BytesIO
+
+from minio import Minio
+from PIL import Image
+
+
+def get_image(
+    client: Minio,
+    object_name: str,
+) -> Image.Image:
+    """Get image from image subfolder in bucket in Minio."""
+    assert isinstance(client, Minio)
+    assert isinstance(object_name, str)
+    assert len(object_name) > 0
+    # prepare arguments
+    assert 'MINIO_BUCKET_NAME' in os.environ
+    bucket_name = os.getenv('MINIO_BUCKET_NAME', default='')
+    subfolder = 'images'
+    # get object from bucket
+    object_name = f'{subfolder}/{object_name}'
+    try:
+        response = client.get_object(
+            bucket_name=bucket_name,
+            object_name=object_name,
+        )
+        buffer = BytesIO(response.data)
+    except Exception as exc:
+        logging.error('failed getting data from MinIO')
+        raise exc
+    finally:
+        response.close()
+        response.release_conn()
+    buffer.seek(0)
+    # convert data to image
+    image = Image.open(buffer)
+    logging.debug('got data from %s', object_name)
+    return image

shared/datastore/src/get_model.py

@@ -0,0 +1,30 @@
+"""Definition of get_model function."""
+
+import logging
+from collections import OrderedDict
+
+import torch
+from minio import Minio
+
+from .get import get
+
+
+def get_model(
+    client: Minio,
+    object_name: str,
+) -> OrderedDict:
+    """Get model from model subfolder in bucket in Minio."""
+    assert isinstance(client, Minio)
+    assert isinstance(object_name, str)
+    subfolder = 'models'
+    object_name = f'{subfolder}/{object_name}'
+    # get buffer
+    buffer = get(
+        client=client,
+        object_name=object_name,
+    )
+    # convert data to model checkpoint
+    buffer.seek(0)
+    model_content = torch.load(buffer)
+    logging.debug('finished')
+    return model_content

shared/datastore/src/put.py

@@ -0,0 +1,39 @@
+"""Definition of put function."""
+
+from __future__ import annotations
+
+import logging
+import os
+from hashlib import md5
+from io import BytesIO
+
+from minio import Minio
+
+
+def put(
+    client: Minio,
+    buffer: BytesIO,
+) -> str:
+    """Put buffer in bucket in MinIO and return MD5 checksum as object name."""
+    assert isinstance(client, Minio)
+    assert isinstance(buffer, BytesIO)
+    bucket_name = os.getenv('MINIO_BUCKET_NAME', default='')
+    assert len(bucket_name) > 0
+    # get md5 of image
+    checksum = md5(buffer.getbuffer()).hexdigest()
+    # prepare for saving
+    num_bytes = buffer.tell()
+    buffer.seek(0)
+    # send data to bucket
+    try:
+        client.put_object(
+            bucket_name=bucket_name,
+            object_name=checksum,
+            length=num_bytes,
+            data=buffer,
+        )
+    except Exception as exc:
+        logging.error('failed saving data to MinIO')
+        raise exc
+    logging.debug('saved data to %s', checksum)
+    return checksum

shared/datastore/src/put_image.py

@@ -0,0 +1,44 @@
+"""Definition of put_image function."""
+
+import logging
+import os
+from hashlib import md5
+from io import BytesIO
+
+from minio import Minio
+from PIL import Image
+
+
+def put_image(
+    client: Minio,
+    image: Image.Image,
+) -> str:
+    """Put image in image subfolder in bucket in Minio and return MD5 checksum
+    used as object name."""
+    assert isinstance(client, Minio)
+    assert isinstance(image, Image.Image)
+    bucket_name = os.getenv('MINIO_BUCKET_NAME', default='')
+    assert len(bucket_name) > 0
+    subfolder = 'images'
+    # save image to buffer
+    buffer = BytesIO()
+    image.save(buffer, 'png')
+    # 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

shared/datastore/src/put_model.py

@@ -0,0 +1,45 @@
+"""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', default='')
+    assert len(bucket_name) > 0
+    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