46 lines
1.2 KiB
Python
46 lines
1.2 KiB
Python
"""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)
|
|
assert 'MINIO_BUCKET_NAME' in os.environ
|
|
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
|