object_based_datastore #65
@@ -1,10 +1 @@
|
||||
from .src import (
|
||||
connect_minio,
|
||||
delete,
|
||||
get,
|
||||
get_image,
|
||||
get_model,
|
||||
put,
|
||||
put_image,
|
||||
put_model,
|
||||
)
|
||||
from .src import Datastore
|
||||
|
||||
@@ -1,8 +1 @@
|
||||
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
|
||||
from .datastore_minio import DatastoreMinio as Datastore
|
||||
|
||||
@@ -1,38 +0,0 @@
|
||||
"""Definition of connect function."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
from minio import Minio
|
||||
|
||||
from shared.utils import check_env
|
||||
|
||||
|
||||
def connect_minio() -> Minio:
|
||||
"""Connect to MinIO server."""
|
||||
# ensure necessary env vars available
|
||||
var_list = {
|
||||
'MINIO_ENDPOINT',
|
||||
'MINIO_ACCESS_KEY',
|
||||
'MINIO_SECRET_KEY',
|
||||
'MINIO_BUCKET_NAME',
|
||||
}
|
||||
check_env(var_list)
|
||||
# prepare arguments
|
||||
minio_endpoint = str(os.getenv('MINIO_ENDPOINT'))
|
||||
minio_access_key = str(os.getenv('MINIO_ACCESS_KEY'))
|
||||
minio_secret_key = str(os.getenv('MINIO_SECRET_KEY'))
|
||||
minio_bucket_name = str(os.getenv('MINIO_BUCKET_NAME'))
|
||||
# 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
|
||||
@@ -1,31 +0,0 @@
|
||||
"""Definition of delete function."""
|
||||
|
||||
import logging
|
||||
from traceback import print_exc
|
||||
|
||||
from minio import Minio
|
||||
|
||||
|
||||
def delete(
|
||||
client: Minio,
|
||||
bucket_name: str,
|
||||
object_name: str,
|
||||
) -> None:
|
||||
"""Delete object from MinIO."""
|
||||
assert isinstance(client, Minio)
|
||||
assert isinstance(bucket_name, str)
|
||||
assert len(bucket_name) > 0
|
||||
assert isinstance(object_name, str)
|
||||
assert len(object_name) > 0
|
||||
# remove object
|
||||
try:
|
||||
client.remove_object(
|
||||
bucket_name=bucket_name,
|
||||
object_name=object_name,
|
||||
)
|
||||
except Exception as exc:
|
||||
logging.error('failed deleting %s', object_name)
|
||||
print_exc()
|
||||
raise exc
|
||||
else:
|
||||
logging.debug('deleted %s', object_name)
|
||||
@@ -1,44 +0,0 @@
|
||||
"""Definition of get function."""
|
||||
|
||||
import logging
|
||||
from io import BytesIO
|
||||
from traceback import print_exc
|
||||
|
||||
from minio import Minio
|
||||
|
||||
|
||||
def get(
|
||||
client: Minio,
|
||||
bucket_name: str,
|
||||
object_name: str,
|
||||
) -> BytesIO:
|
||||
"""Get buffer from bucket in MinIO."""
|
||||
assert isinstance(client, Minio)
|
||||
assert isinstance(bucket_name, str)
|
||||
assert len(bucket_name) > 0
|
||||
assert isinstance(object_name, str)
|
||||
assert len(object_name) > 0
|
||||
try:
|
||||
# make request
|
||||
response = client.get_object(
|
||||
bucket_name=bucket_name,
|
||||
object_name=object_name,
|
||||
)
|
||||
assert response.status == 200
|
||||
# get buffer
|
||||
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:
|
||||
# close connection if established
|
||||
if 'response' in locals():
|
||||
response.close()
|
||||
response.release_conn()
|
||||
@@ -1,35 +0,0 @@
|
||||
"""Definition of get_image function."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
from minio import Minio
|
||||
from PIL import Image
|
||||
|
||||
from shared.utils import check_env
|
||||
|
||||
from .get import get
|
||||
|
||||
|
||||
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
|
||||
check_env({'MINIO_BUCKET_NAME'})
|
||||
bucket_name = str(os.getenv('MINIO_BUCKET_NAME'))
|
||||
object_name = f'images/{object_name}'
|
||||
# get object from bucket
|
||||
buffer = get(
|
||||
client=client,
|
||||
bucket_name=bucket_name,
|
||||
object_name=object_name,
|
||||
)
|
||||
# convert data to image
|
||||
image = Image.open(buffer)
|
||||
logging.debug('got data from %s', object_name)
|
||||
return image
|
||||
@@ -1,36 +0,0 @@
|
||||
"""Definition of get_model function."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from collections import OrderedDict
|
||||
|
||||
import torch
|
||||
from minio import Minio
|
||||
|
||||
from shared.utils import check_env
|
||||
|
||||
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)
|
||||
assert len(object_name) > 0
|
||||
# prepare arguments
|
||||
check_env({'MINIO_BUCKET_NAME'})
|
||||
bucket_name = str(os.getenv('MINIO_BUCKET_NAME'))
|
||||
object_name = f'models/{object_name}'
|
||||
# get object from bucket
|
||||
buffer = get(
|
||||
client=client,
|
||||
bucket_name=bucket_name,
|
||||
object_name=object_name,
|
||||
)
|
||||
# convert data to model checkpoint
|
||||
model_content = torch.load(buffer)
|
||||
logging.debug('finished')
|
||||
return model_content
|
||||
@@ -1,36 +0,0 @@
|
||||
"""Definition of put function."""
|
||||
|
||||
import logging
|
||||
from io import BytesIO
|
||||
|
||||
from minio import Minio
|
||||
|
||||
|
||||
def put(
|
||||
client: Minio,
|
||||
buffer: BytesIO,
|
||||
bucket_name: str,
|
||||
object_name: str,
|
||||
) -> None:
|
||||
"""Put buffer in bucket in MinIO and return MD5 checksum as object name."""
|
||||
assert isinstance(client, Minio)
|
||||
assert isinstance(buffer, BytesIO)
|
||||
assert isinstance(bucket_name, str)
|
||||
assert len(bucket_name) > 0
|
||||
assert isinstance(object_name, str)
|
||||
assert len(object_name) > 0
|
||||
# prepare for saving
|
||||
num_bytes = len(buffer.getvalue())
|
||||
buffer.seek(0)
|
||||
# send data to bucket
|
||||
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)
|
||||
@@ -1,40 +0,0 @@
|
||||
"""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
|
||||
|
||||
from .put import put
|
||||
|
||||
|
||||
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)
|
||||
# get bucket name from env
|
||||
bucket_name = os.getenv('MINIO_BUCKET_NAME', default='')
|
||||
assert len(bucket_name) > 0
|
||||
# save data to buffer
|
||||
buffer = BytesIO()
|
||||
image.save(buffer, 'png')
|
||||
# get md5 of buffer
|
||||
checksum = md5(buffer.getbuffer()).hexdigest()
|
||||
# set object name
|
||||
object_name = f'images/{checksum}'
|
||||
# send data to bucket
|
||||
put(
|
||||
client=client,
|
||||
buffer=buffer,
|
||||
bucket_name=bucket_name,
|
||||
object_name=object_name,
|
||||
)
|
||||
logging.debug('finished')
|
||||
return checksum
|
||||
@@ -1,43 +0,0 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user