fix: load model aliases before loading new SDK service (#4727)

* fix: load model aliases before loading new SDK service Signed-off-by: Frost Ming <[email protected]> * fix: skip pre-commit for certain jobs Signed-off-by: Frost Ming <[email protected]>
bentoml · May 14, 2024 · 3c33b4d · 3c33b4d
1 parent 30b1992
commit 3c33b4d
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -2,6 +2,9 @@ ci:
   autoupdate_schedule: monthly
   autofix_commit_msg: "ci: auto fixes from pre-commit.ci\n\nFor more information, see https://pre-commit.ci"
   autoupdate_commit_msg: 'ci: pre-commit autoupdate [skip ci]'
+  skip:  # exceeds tier max size
+    - buf-format
+    - buf-lint
 exclude: '(.*\.(css|js|svg))|(.*/(snippets|grpc|proto)/.*)$'
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit

diff --git a/src/_bentoml_impl/loader.py b/src/_bentoml_impl/loader.py
@@ -11,6 +11,7 @@
     from _bentoml_sdk import Service
 
 BENTO_YAML_FILENAME = "bento.yaml"
+BENTO_BUILD_CONFIG_FILENAME = "bentofile.yaml"
 
 
 def normalize_identifier(
@@ -104,6 +105,9 @@ def import_service(
     `normalize_identifier` function.
     """
     from _bentoml_sdk import Service
+    from bentoml._internal.bento.bento import BentoInfo
+    from bentoml._internal.bento.build_config import BentoBuildConfig
+    from bentoml._internal.configuration.containers import BentoMLContainer
 
     if bento_path is None:
         bento_path = pathlib.Path(".")
@@ -123,7 +127,6 @@ def import_service(
         bento_path.parent.joinpath(BENTO_YAML_FILENAME).exists()
         and bento_path.parent.joinpath("models").exists()
     ):
-        from bentoml._internal.configuration.containers import BentoMLContainer
         from bentoml._internal.models import ModelStore
 
         original_model_store = BentoMLContainer.model_store.get()
@@ -134,6 +137,19 @@ def import_service(
     else:
         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}
+    elif (bentofile := bento_path.joinpath(BENTO_BUILD_CONFIG_FILENAME)).exists():
+        with open(bentofile, encoding="utf-8") as f:
+            build_config = BentoBuildConfig.from_yaml(f)
+        model_aliases = build_config.model_aliases
+    else:
+        model_aliases = {}
+    BentoMLContainer.model_aliases.set(model_aliases)
+
     try:
         module_name, _, attrs_str = service_identifier.partition(":")