def _bento_orm_obj_to_pb(bento_obj, labels=None):
# Backwards compatible support loading saved bundle created before 0.8.0
if ('apis' in bento_obj.bento_service_metadata
and bento_obj.bento_service_metadata['apis']):
for api in bento_obj.bento_service_metadata['apis']:
if 'handler_type' in api:
api['input_type'] = api['handler_type']
del api['handler_type']
if 'handler_config' in api:
api['input_config'] = api['handler_config']
del api['handler_config']
if 'output_type' not in api:
api['output_type'] = 'DefaultOutput'
bento_service_metadata_pb = ParseDict(bento_obj.bento_service_metadata,
BentoServiceMetadata())
bento_uri = BentoUri(uri=bento_obj.uri,
type=BentoUri.StorageType.Value(bento_obj.uri_type))
if not bento_obj.upload_status:
upload_status = DEFAULT_UPLOAD_STATUS
else:
upload_status = UploadStatus(status=UploadStatus.Status.Value(
bento_obj.upload_status['status']))
if labels is not None:
bento_service_metadata_pb.labels.update(labels)
return BentoPB(
name=bento_obj.name,
version=bento_obj.version,
uri=bento_uri,
bento_service_metadata=bento_service_metadata_pb,
status=upload_status,
)
def _update_bento_upload_progress(
self, bento_service_metadata, status=UploadStatus.DONE, percentage=None
):
upload_status = UploadStatus(status=status, percentage=percentage)
upload_status.updated_at.GetCurrentTime()
update_bento_req = UpdateBentoRequest(
bento_name=bento_service_metadata.name,
bento_version=bento_service_metadata.version,
upload_status=upload_status,
service_metadata=bento_service_metadata,
)
self.yatai_service.UpdateBento(update_bento_req)
from google.protobuf.json_format import ParseDict
from bentoml.utils import ProtoMessageToDict
from bentoml.exceptions import YataiRepositoryException
from bentoml.yatai.db import Base, create_session
from bentoml.yatai.proto.repository_pb2 import (
UploadStatus,
BentoUri,
BentoServiceMetadata,
Bento as BentoPB,
ListBentoRequest,
)
logger = logging.getLogger(__name__)
DEFAULT_UPLOAD_STATUS = UploadStatus(status=UploadStatus.UNINITIALIZED)
class Bento(Base):
__tablename__ = 'bentos'
__table_args__ = tuple(
UniqueConstraint('name', 'version', name='_name_version_uc'))
id = Column(Integer, primary_key=True)
name = Column(String, nullable=False)
version = Column(String, nullable=False)
# Storage URI for this Bento
uri = Column(String, nullable=False)
# Name is is required for PostgreSQL and any future supported database which
def UploadBento(self, request_iterator, context=None):
if not is_file_system_repo(self.repo):
logger.error(
"UploadBento RPC only works with File System based repository, "
"for other types of repositories(s3, gcs, minio), "
"use pre-signed URL for upload")
return UploadBentoResponse(status=Status.INTERNAL(''))
try:
with self.db.create_session() as sess:
lock_obj = None
bento_pb = None
with TempDirectory() as temp_dir:
temp_tar_path = os.path.join(
temp_dir, f'{uuid.uuid4().hex[:12]}.tar')
file = open(temp_tar_path, 'wb+')
for request in request_iterator:
# Initial request is without bundle
if not request.bento_bundle:
bento_name = request.bento_name
bento_version = request.bento_version
bento_pb = self.db.metadata_store.get(
sess, bento_name, bento_version)
if not bento_pb:
result_status = Status.NOT_FOUND(
"BentoService `{}:{}` is not found".
format(bento_name, bento_version))
return UploadBentoResponse(
status=result_status)
if bento_pb.status:
if bento_pb.status.status == UploadStatus.DONE:
return UploadBentoResponse(
status=Status.CANCELLED(
f"Bento bundle `{bento_name}:"
f"{bento_version}` is uploaded"
))
if bento_pb.status.status == UploadStatus.UPLOADING:
return UploadBentoResponse(
status=Status.CANCELLED(
f"Bento bundle `{bento_name}:"
f"{bento_version}` is currently "
f"uploading"))
if lock_obj is None:
lock_obj = LockStore.acquire(
sess=sess,
lock_type=LockType.WRITE,
resource_id=
f'{bento_name}_{bento_version}',
ttl_min=DEFAULT_TTL_MIN,
)
else:
if (bento_name == request.bento_name
and bento_version
== request.bento_version):
file.write(request.bento_bundle)
else:
lock_obj.release(sess)
raise BadInput(
f"Incoming stream request doesn't match "
f"with initial request info "
f"{bento_name}:{bento_version} - "
f"{request.bento_name}:"
f"{request.bento_version}")
file.seek(0)
with tarfile.open(fileobj=file, mode='r') as tar:
tar.extractall(path=bento_pb.uri.uri)
upload_status = UploadStatus(status=UploadStatus.DONE)
upload_status.updated_at.GetCurrentTime()
self.db.metadata_store.update_upload_status(
sess, bento_name, bento_version, upload_status)
lock_obj.release(sess)
return UploadBentoResponse(status=Status.OK())
except BentoMLException as e:
logger.error("RPC ERROR UploadBento: %s", e)
return UploadBentoResponse(status=e.status_proto)
except Exception as e: # pylint: disable=broad-except
logger.error("RPC ERROR UploadBento: %s", e)
return UploadBentoResponse(status=Status.INTERNAL())
finally:
if file is not None:
file.close()