def create_yatai_service_mock(repo_storage_type=BentoUri.LOCAL):
bento_pb = Bento(name='bento_test_name', version='version1.1.1')
if repo_storage_type == BentoUri.LOCAL:
bento_pb.uri.uri = '/tmp/path/to/bundle'
bento_pb.uri.type = repo_storage_type
api = BentoServiceMetadata.BentoServiceApi(name='predict')
bento_pb.bento_service_metadata.apis.extend([api])
bento_pb.bento_service_metadata.env.python_version = '3.7.0'
get_bento_response = GetBentoResponse(bento=bento_pb)
yatai_service_mock = MagicMock()
yatai_service_mock.GetBento.return_value = get_bento_response
return yatai_service_mock
def get_bento_service_metadata_pb(self):
bento_service_metadata = BentoServiceMetadata()
bento_service_metadata.name = self.config["metadata"]["service_name"]
bento_service_metadata.version = self.config["metadata"][
"service_version"]
bento_service_metadata.created_at.FromDatetime(
self.config["metadata"]["created_at"])
if "env" in self.config:
if "setup_sh" in self.config["env"]:
bento_service_metadata.env.setup_sh = self.config["env"][
"setup_sh"]
if "conda_env" in self.config["env"]:
bento_service_metadata.env.conda_env = dump_to_yaml_str(
self.config["env"]["conda_env"])
if "pip_dependencies" in self.config["env"]:
bento_service_metadata.env.pip_dependencies = "\n".join(
self.config["env"]["pip_dependencies"])
if "python_version" in self.config["env"]:
bento_service_metadata.env.python_version = self.config["env"][
"python_version"]
if "docker_base_image" in self.config["env"]:
bento_service_metadata.env.docker_base_image = self.config[
"env"]["docker_base_image"]
if "apis" in self.config:
for api_config in self.config["apis"]:
api_metadata = BentoServiceMetadata.BentoServiceApi(
name=api_config["name"],
docs=api_config["docs"],
handler_type=api_config.get("handler_type", "unknown"),
)
if "handler_config" in api_config:
for k, v in api_config["handler_config"].items():
api_metadata.handler_config[k] = v
bento_service_metadata.apis.extend([api_metadata])
if "artifacts" in self.config:
for artifact_config in self.config["artifacts"]:
artifact_metadata = BentoServiceMetadata.BentoArtifact()
if "name" in artifact_config:
artifact_metadata.name = artifact_config["name"]
if "artifact_type" in artifact_config:
artifact_metadata.artifact_type = artifact_config[
"artifact_type"]
bento_service_metadata.artifacts.extend([artifact_metadata])
return bento_service_metadata
def create_yatai_service_mock(repo_storage_type=BentoUri.LOCAL):
bento_pb = Bento(name=TEST_DEPLOYMENT_BENTO_NAME,
version=TEST_DEPLOYMENT_BENTO_VERSION)
if repo_storage_type == BentoUri.LOCAL:
bento_pb.uri.uri = TEST_DEPLOYMENT_BENTO_LOCAL_URI
bento_pb.uri.type = repo_storage_type
api = BentoServiceMetadata.BentoServiceApi(name=TEST_BENTO_API_NAME)
bento_pb.bento_service_metadata.apis.extend([api])
get_bento_response = GetBentoResponse(bento=bento_pb)
yatai_service_mock = MagicMock()
yatai_service_mock.GetBento.return_value = get_bento_response
return yatai_service_mock
def get_bento_service_metadata_pb(self):
from bentoml.yatai.proto.repository_pb2 import BentoServiceMetadata
bento_service_metadata = BentoServiceMetadata()
bento_service_metadata.name = self.config["metadata"]["service_name"]
bento_service_metadata.version = self.config["metadata"][
"service_version"]
bento_service_metadata.created_at.FromDatetime(
self.config["metadata"]["created_at"])
if "env" in self.config:
if "setup_sh" in self.config["env"]:
bento_service_metadata.env.setup_sh = self.config["env"][
"setup_sh"]
if "conda_env" in self.config["env"]:
bento_service_metadata.env.conda_env = dump_to_yaml_str(
self.config["env"]["conda_env"])
if "pip_packages" in self.config["env"]:
for pip_package in self.config["env"]["pip_packages"]:
bento_service_metadata.env.pip_packages.append(pip_package)
if "python_version" in self.config["env"]:
bento_service_metadata.env.python_version = self.config["env"][
"python_version"]
if "docker_base_image" in self.config["env"]:
bento_service_metadata.env.docker_base_image = self.config[
"env"]["docker_base_image"]
if "apis" in self.config:
for api_config in self.config["apis"]:
if 'handler_type' in api_config:
# Convert handler type to input type for saved bundle created
# before version 0.8.0
input_type = api_config.get('handler_type')
elif 'input_type' in api_config:
input_type = api_config.get('input_type')
else:
input_type = "unknown"
if 'output_type' in api_config:
output_type = api_config.get('output_type')
else:
output_type = "DefaultOutput"
api_metadata = BentoServiceMetadata.BentoServiceApi(
name=api_config["name"],
docs=api_config["docs"],
input_type=input_type,
output_type=output_type,
)
if "handler_config" in api_config:
# Supports viewing API input config info for saved bundle created
# before version 0.8.0
for k, v in api_config["handler_config"].items():
if k in {'mb_max_latency', 'mb_max_batch_size'}:
setattr(api_metadata, k, v)
else:
api_metadata.input_config[k] = v
else:
if 'mb_max_latency' in api_config:
api_metadata.mb_max_latency = api_config[
"mb_max_latency"]
else:
api_metadata.mb_max_latency = DEFAULT_MAX_LATENCY
if 'mb_max_batch_size' in api_config:
api_metadata.mb_max_batch_size = api_config[
"mb_max_batch_size"]
else:
api_metadata.mb_max_batch_size = DEFAULT_MAX_BATCH_SIZE
if "input_config" in api_config:
for k, v in api_config["input_config"].items():
api_metadata.input_config[k] = v
if "output_config" in api_config:
for k, v in api_config["output_config"].items():
api_metadata.output_config[k] = v
api_metadata.batch = api_config.get("batch", False)
bento_service_metadata.apis.extend([api_metadata])
if "artifacts" in self.config:
for artifact_config in self.config["artifacts"]:
artifact_metadata = BentoServiceMetadata.BentoArtifact()
if "name" in artifact_config:
artifact_metadata.name = artifact_config["name"]
if "artifact_type" in artifact_config:
artifact_metadata.artifact_type = artifact_config[
"artifact_type"]
if "metadata" in artifact_config:
if isinstance(artifact_config["metadata"], dict):
s = Struct()
s.update(artifact_config["metadata"])
artifact_metadata.metadata.CopyFrom(s)
else:
logger.warning(
"Tried to get non-dictionary metadata for artifact "
f"{artifact_metadata.name}. Ignoring metadata...")
bento_service_metadata.artifacts.extend([artifact_metadata])
return bento_service_metadata
def get_bento_service_metadata_pb(self):
from bentoml.yatai.proto.repository_pb2 import BentoServiceMetadata
bento_service_metadata = BentoServiceMetadata()
bento_service_metadata.name = self.config["metadata"]["service_name"]
bento_service_metadata.version = self.config["metadata"][
"service_version"]
bento_service_metadata.created_at.FromDatetime(
self.config["metadata"]["created_at"])
if "env" in self.config:
if "setup_sh" in self.config["env"]:
bento_service_metadata.env.setup_sh = self.config["env"][
"setup_sh"]
if "conda_env" in self.config["env"]:
bento_service_metadata.env.conda_env = dump_to_yaml_str(
self.config["env"]["conda_env"])
if "pip_dependencies" in self.config["env"]:
bento_service_metadata.env.pip_dependencies = "\n".join(
self.config["env"]["pip_dependencies"])
if "python_version" in self.config["env"]:
bento_service_metadata.env.python_version = self.config["env"][
"python_version"]
if "docker_base_image" in self.config["env"]:
bento_service_metadata.env.docker_base_image = self.config[
"env"]["docker_base_image"]
if "apis" in self.config:
for api_config in self.config["apis"]:
if 'handler_type' in api_config:
# Convert handler type to input type for saved bundle created
# before version 0.8.0
input_type = api_config.get('handler_type')
elif 'input_type' in api_config:
input_type = api_config.get('input_type')
else:
input_type = "unknown"
if 'output_type' in api_config:
output_type = api_config.get('output_type')
else:
output_type = "DefaultOutput"
api_metadata = BentoServiceMetadata.BentoServiceApi(
name=api_config["name"],
docs=api_config["docs"],
input_type=input_type,
output_type=output_type,
)
if "handler_config" in api_config:
# Supports viewing API input config info for saved bundle created
# before version 0.8.0
for k, v in api_config["handler_config"].items():
if k in {'mb_max_latency', 'mb_max_batch_size'}:
setattr(api_metadata, k, v)
else:
api_metadata.input_config[k] = v
else:
if 'mb_max_latency' in api_config:
api_metadata.mb_max_latency = api_config[
"mb_max_latency"]
else:
api_metadata.mb_max_latency = DEFAULT_MAX_LATENCY
if 'mb_max_batch_size' in api_config:
api_metadata.mb_max_batch_size = api_config[
"mb_max_batch_size"]
else:
api_metadata.mb_max_batch_size = DEFAULT_MAX_BATCH_SIZE
if "input_config" in api_config:
for k, v in api_config["input_config"].items():
api_metadata.input_config[k] = v
if "output_config" in api_config:
for k, v in api_config["output_config"].items():
api_metadata.output_config[k] = v
bento_service_metadata.apis.extend([api_metadata])
if "artifacts" in self.config:
for artifact_config in self.config["artifacts"]:
artifact_metadata = BentoServiceMetadata.BentoArtifact()
if "name" in artifact_config:
artifact_metadata.name = artifact_config["name"]
if "artifact_type" in artifact_config:
artifact_metadata.artifact_type = artifact_config[
"artifact_type"]
bento_service_metadata.artifacts.extend([artifact_metadata])
return bento_service_metadata