removed newline char

model/docker-compose.server.yml

@@ -9,3 +9,8 @@ services:
       - ../server.env
     environment:
       - ENV=TEST
+    deploy:
+      resources:
+        limits:
+          cpus: '2.000'
+          memory: 8G

model/src/main.py

@@ -1,45 +1,42 @@
 """Main script to be run by service."""
-from __future__ import annotations
 
-from hashlib import md5
-from io import BytesIO
+import logging
+from pathlib import Path
 
 import torch
-from data_store import connect
-from data_store import put
+from dotenv import load_dotenv
 from models import VisualCommunicationModel
 from torchinfo import summary
-from utils import check_env
+from utils import get_model_name
 
+from shared.data_store import connect, get_model
 from shared.utils import setup_logging
 
-
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 if __name__ == '__main__':
-    # check that environment variables are set
-    check_env()
+    # load in env file
+    env_path = Path(__file__).parent.parent.parent / 'server.env'
+    assert env_path.exists()
+    load_dotenv(env_path)
     # setup logging
     setup_logging()
+    # connect to minio
+    minio_client = connect()
     # instantiate model
-    model = VisualCommunicationModel().to(DEVICE)
+    model = VisualCommunicationModel()
+    # get model object name
+    model_name_path = Path(__file__).parent / 'model_name.txt'
+    model_object_name = get_model_name(path=model_name_path)
+    logging.info('using model: %s', model_object_name)
+    # load model from minio
+    model_checkpoint = get_model(
+        client=minio_client,
+        object_name=model_object_name,
+    )
+    model.load_state_dict(model_checkpoint)
+    # move model to selected device
+    model = model.to(DEVICE)
     # show model weights
     summary(model)
-    # for param in model.parameters():
-    #     logging.info(param.data)
-    # save model to buffer
-    buffer = BytesIO()
-    torch.save(model.state_dict(), buffer)
-    print(f"buffer size: {len(buffer.getvalue())}")
-    # calculate buffer hash
-    hash_str = md5(buffer.getbuffer()).hexdigest()
-    print(f"hash string: {hash_str}")
-    # connect to minio
-    client = connect()
-    # put buffer in minio bucket
-    put(
-        client=client,
-        buffer=buffer,
-        object_name=hash_str,
-    )
-    print('saved to Minio')
+    print('loaded model')

model/src/model_name.txt

@@ -0,0 +1 @@
+fb75db772ba8b280bdaf1774d5f73f58

model/src/utils/__init__.py

@@ -0,0 +1 @@
+from .get_model_name import get_model_name

model/src/utils/get_model_name.py

@@ -0,0 +1,19 @@
+"""Definition of get_model_name function."""
+
+import logging
+from pathlib import Path
+
+
+def get_model_name(
+    path: Path,
+) -> str:
+    """Get model name from model_name file."""
+    assert isinstance(path, Path)
+    assert path.exists(), f'{path} does not exist'
+    # read file contents
+    with open(file=path, encoding='utf-8') as fh:
+        model_name = fh.read()
+    # remove newline character
+    model_name = model_name.strip()
+    logging.debug('finished')
+    return model_name

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/connect.py

@@ -34,14 +34,9 @@ def connect() -> Minio:
         secret_key=minio_secret_key,
         secure=False,
     )
-    # ensure buckets exists
-    bucket_name_list = [
-        'images',
-        'models',
-    ]
-    for bucket_name in bucket_name_list:
-        if not client.bucket_exists(bucket_name=bucket_name):
-            logging.info('creating bucket: %s', bucket_name)
+    # 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/data_store/get.py

@@ -1,7 +1,9 @@
 """Definition of get function."""
 
+import logging
 import os
 from io import BytesIO
+from traceback import print_exc
 
 from minio import Minio
 
@@ -16,10 +18,22 @@ def get(
     bucket_name = os.getenv('MINIO_BUCKET_NAME')
     # get buffer
     try:
-        response = client.get_object(bucket_name, object_name)
-        buffer = BytesIO(response.data)
+        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()
-    buffer.seek(0)
-    return buffer

shared/data_store/get_model.py

@@ -0,0 +1,31 @@
+"""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,
+    )
+    logging.debug('buffer size: %s', buffer.getbuffer().nbytes)
+    # convert data to model checkpoint
+    buffer.seek(0)
+    model_content = torch.load(buffer)
+    logging.debug('finished')
+    return model_content

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

tests/model_io_test.py

@@ -1,37 +1,34 @@
 """Test that a model can be saved and loaded again."""
 
-from io import BytesIO
+from pathlib import Path
 
 import torch
+from dotenv import load_dotenv
 from torchinfo import summary
 
 from model.src.models import VisualCommunicationModel
-from shared.data_store import connect, put
-from shared.utils import check_env, setup_logging
+from shared.data_store import connect, put_model
+from shared.utils import setup_logging
 
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 if __name__ == '__main__':
-    # check that environment variables are set
-    check_env()
+    # load in env file
+    env_path = Path(__file__).parent.parent / 'server.env'
+    assert env_path.exists()
+    load_dotenv(env_path)
     # setup logging
     setup_logging()
+    # connect to minio
+    client = connect()
     # instantiate model
     model = VisualCommunicationModel().to(DEVICE)
     # show model weights
     summary(model)
-    # for param in model.parameters():
-    #     logging.info(param.data)
-    # save model to buffer
-    buffer = BytesIO()
-    torch.save(model.state_dict(), buffer)
-    print(f"buffer size: {len(buffer.getvalue())}")
-    # connect to minio
-    client = connect(bucket='models')
     # put buffer in minio bucket
-    hash_str = put(
+    hash_str = put_model(
         client=client,
-        buffer=buffer,
+        model=model,
     )
     print(f"hash string: {hash_str}")
     print('saved to Minio')