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)
def upload_from_dir(self,
saved_bento_path: str,
labels: Dict = None) -> "BentoUri":
from bentoml.yatai.db.stores.label import _validate_labels
bento_service_metadata = load_bento_service_metadata(saved_bento_path)
if labels:
_validate_labels(labels)
bento_service_metadata.labels.update(labels)
get_bento_response = self.yatai_service.GetBento(
GetBentoRequest(
bento_name=bento_service_metadata.name,
bento_version=bento_service_metadata.version,
))
if get_bento_response.status.status_code == status_pb2.Status.OK:
raise BentoMLException(
"BentoService bundle {}:{} already registered in repository. Reset "
"BentoService version with BentoService#set_version or bypass BentoML's"
" model registry feature with BentoService#save_to_dir".format(
bento_service_metadata.name,
bento_service_metadata.version))
elif get_bento_response.status.status_code != status_pb2.Status.NOT_FOUND:
raise BentoMLException(
'Failed accessing YataiService. {error_code}:'
'{error_message}'.format(
error_code=Status.Name(
get_bento_response.status.status_code),
error_message=get_bento_response.status.error_message,
))
request = AddBentoRequest(
bento_name=bento_service_metadata.name,
bento_version=bento_service_metadata.version,
)
response = self.yatai_service.AddBento(request)
if response.status.status_code != status_pb2.Status.OK:
raise BentoMLException(
"Error adding BentoService bundle to repository: {}:{}".format(
Status.Name(response.status.status_code),
response.status.error_message,
))
if response.uri.type == BentoUri.LOCAL:
# When using Yatai backed by File System repository,
# if Yatai is a local instance, copy the files directly.
# Otherwise, use UploadBento RPC to stream files to remote Yatai server
if is_remote_yatai(self.yatai_service):
self._upload_bento(
bento_service_metadata.name,
bento_service_metadata.version,
saved_bento_path,
)
update_bento_service = UpdateBentoRequest(
bento_name=bento_service_metadata.name,
bento_version=bento_service_metadata.version,
service_metadata=bento_service_metadata,
)
self.yatai_service.UpdateBento(update_bento_service)
else:
if os.path.exists(response.uri.uri):
raise BentoMLException(
f'Bento bundle directory {response.uri.uri} already exist'
)
shutil.copytree(saved_bento_path, response.uri.uri)
upload_status = UploadStatus.DONE
self._update_bento_upload_progress(bento_service_metadata,
status=upload_status)
logger.info(
"BentoService bundle '%s:%s' saved to: %s",
bento_service_metadata.name,
bento_service_metadata.version,
response.uri.uri,
)
# Return URI to saved bento in repository storage
return response.uri.uri
elif response.uri.type == BentoUri.S3 or response.uri.type == BentoUri.GCS:
uri_type = 'S3' if response.uri.type == BentoUri.S3 else 'GCS'
self._update_bento_upload_progress(bento_service_metadata,
UploadStatus.UPLOADING, 0)
fileobj = io.BytesIO()
with tarfile.open(mode="w:gz", fileobj=fileobj) as tar:
tar.add(saved_bento_path, arcname=bento_service_metadata.name)
fileobj.seek(0, 0)
if response.uri.type == BentoUri.S3:
http_response = requests.put(response.uri.s3_presigned_url,
data=fileobj)
elif response.uri.type == BentoUri.GCS:
http_response = requests.put(response.uri.gcs_presigned_url,
data=fileobj)
if http_response.status_code != 200:
self._update_bento_upload_progress(bento_service_metadata,
UploadStatus.ERROR)
raise BentoMLException(
f"Error saving BentoService bundle to {uri_type}."
f"{http_response.status_code}: {http_response.text}")
self._update_bento_upload_progress(bento_service_metadata)
logger.info(
"Successfully saved BentoService bundle '%s:%s' to {uri_type}: %s",
bento_service_metadata.name,
bento_service_metadata.version,
response.uri.uri,
)
return response.uri.uri
else:
raise BentoMLException(
f"Error saving Bento to target repository, URI type {response.uri.type}"
f" at {response.uri.uri} not supported")