44 lines
1.0 KiB
Python
44 lines
1.0 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
|
|
|
|
from shared.utils import check_env
|
|
|
|
from .put import put
|
|
|
|
|
|
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)
|
|
# get bucket name from env
|
|
check_env({'MINIO_BUCKET_NAME'})
|
|
bucket_name = str(os.getenv('MINIO_BUCKET_NAME'))
|
|
# save data to buffer
|
|
buffer = BytesIO()
|
|
torch.save(model.state_dict(), buffer)
|
|
# get md5 of image
|
|
checksum = md5(buffer.getbuffer()).hexdigest()
|
|
# set object name
|
|
object_name = f'models/{checksum}'
|
|
# send data to bucket
|
|
put(
|
|
client=client,
|
|
buffer=buffer,
|
|
bucket_name=bucket_name,
|
|
object_name=object_name,
|
|
)
|
|
logging.debug('finished')
|
|
return checksum
|