From 692a28690847bbc2f2394d5b7f8c8dd02965296e Mon Sep 17 00:00:00 2001 From: Frost Ming Date: Fri, 14 Jun 2024 10:52:09 +0800 Subject: [PATCH] fix: populate bento name and version of server context (#4787) Signed-off-by: Frost Ming --- src/_bentoml_impl/loader.py | 18 ++++++++++-------- src/_bentoml_impl/worker/service.py | 6 ++++++ 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/src/_bentoml_impl/loader.py b/src/_bentoml_impl/loader.py index 230710843fe..84c1d623e92 100644 --- a/src/_bentoml_impl/loader.py +++ b/src/_bentoml_impl/loader.py @@ -106,7 +106,7 @@ def import_service( `normalize_identifier` function. """ from _bentoml_sdk import Service - from bentoml._internal.bento.bento import BentoInfo + from bentoml._internal.bento.bento import Bento from bentoml._internal.bento.build_config import BentoBuildConfig from bentoml._internal.configuration.containers import BentoMLContainer @@ -144,10 +144,10 @@ def import_service( original_model_store = None # load model aliases - if (bento_yaml := bento_path.with_name(BENTO_YAML_FILENAME)).exists(): - with open(bento_yaml, encoding="utf-8") as f: - info = BentoInfo.from_yaml_file(f) - model_aliases = {m.alias: str(m.tag) for m in info.all_models if m.alias} + bento: Bento | None = None + if bento_path.with_name(BENTO_YAML_FILENAME).exists(): + bento = Bento.from_path(str(bento_path.parent)) + model_aliases = {m.alias: str(m.tag) for m in bento.info.all_models if m.alias} elif (bentofile := bento_path.joinpath(BENTO_BUILD_CONFIG_FILENAME)).exists(): with open(bentofile, encoding="utf-8") as f: build_config = BentoBuildConfig.from_yaml(f) @@ -165,16 +165,18 @@ def import_service( module = importlib.import_module(module_name) root_service_name, _, depend_path = attrs_str.partition(".") - root_service = getattr(module, root_service_name) + root_service = t.cast("Service[t.Any]", getattr(module, root_service_name)) assert isinstance( root_service, Service ), f'import target "{module_name}:{attrs_str}" is not a bentoml.Service instance' if not depend_path: - return root_service # type: ignore + svc = root_service else: - return root_service.find_dependent(depend_path) + svc = root_service.find_dependent(depend_path) + svc.bento = bento + return svc except (ImportError, AttributeError, KeyError, AssertionError) as e: sys_path = sys.path.copy() diff --git a/src/_bentoml_impl/worker/service.py b/src/_bentoml_impl/worker/service.py index 75ed939feaa..103563df5e0 100644 --- a/src/_bentoml_impl/worker/service.py +++ b/src/_bentoml_impl/worker/service.py @@ -175,6 +175,12 @@ def main( if prometheus_dir is not None: BentoMLContainer.prometheus_multiproc_dir.set(prometheus_dir) server_context.service_name = service.name + if service.bento is None: + server_context.bento_name = service.name + server_context.bento_version = "not available" + else: + server_context.bento_name = service.bento.tag.name + server_context.bento_version = service.bento.tag.version or "not available" asgi_app = service.to_asgi( is_main=server_context.service_type == "entry_service", init=False