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 = '/fake/path/to/bundle'
bento_pb.uri.type = repo_storage_type
api = BentoServiceMetadata.BentoServiceApi(name='predict')
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):
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 load_bento_service_metadata(archive_path):
config = load_bentoml_config(archive_path)
bento_service_metadata = BentoServiceMetadata()
bento_service_metadata.name = config["metadata"]["service_name"]
bento_service_metadata.version = config["metadata"]["service_version"]
bento_service_metadata.created_at.FromDatetime(config["metadata"]["created_at"])
if "env" in config:
if "setup_sh" in config["env"]:
bento_service_metadata.env.setup_sh = config["env"]["setup_sh"]
if "conda_env" in config["env"]:
bento_service_metadata.env.conda_env = dump_to_yaml_str(
config["env"]["conda_env"]
)
if "pip_dependencies" in config["env"]:
bento_service_metadata.env.pip_dependencies = "\n".join(
config["env"]["pip_dependencies"]
)
if "python_version" in config["env"]:
bento_service_metadata.env.python_version = config["env"]["python_version"]
if "apis" in config:
for api_config in config["apis"]:
api_metadata = BentoServiceMetadata.BentoServiceApi()
if "name" in api_config:
api_metadata.name = api_config["name"]
if "handler_type" in api_config:
api_metadata.handler_type = api_config["handler_type"]
if "docs" in api_config:
api_metadata.docs = api_config["docs"]
bento_service_metadata.apis.extend([api_metadata])
if "artifacts" in config:
for artifact_config in 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