Repository: bentoml/BentoML
Pull request #5110 (2024-12-04 03:15:51)

fix: ensure all models to exists on deployment
Signed-off-by: Frost Ming <me@frostming.com>

## What does this PR address?

<!--
Thanks for sending a pull request!

Congrats for making it this far! Here's a 🍱 for you. There are still a few steps ahead.

Please make sure to read the contribution guidelines, then fill out the blanks below before requesting a code review.

Name your Pull Request with one of the following prefixes, e.g. "feat: add support for PyTorch", to indicate the type of changes proposed. This is based on the [Conventional Commits specification](https://www.conventionalcommits.org/en/v1.0.0/#summary).
  - feat: (new feature for the user, not a new feature for build script)
  - fix: (bug fix for the user, not a fix to a build script)
  - docs: (changes to the documentation)
  - style: (formatting, missing semicolons, etc; no production code change)
  - refactor: (refactoring production code, eg. renaming a variable)
  - perf: (code changes that improve performance)
  - test: (adding missing tests, refactoring tests; no production code change)
  - chore: (updating grunt tasks etc; no production code change)
  - build: (changes that affect the build system or external dependencies)
  - ci: (changes to configuration files and scripts)
  - revert: (reverts a previous commit)

Describe your changes in detail. Attach screenshots here if appropriate.

Once you're done with this, someone from BentoML team or community member will help review your PR (see "Who can help review?" section for potential reviewers.). If no one has reviewed your PR after a week have passed, don't hesitate to post a new comment and ping @-the same person. Notifications sometimes get lost 🥲.
-->

<!-- Remove if not applicable -->

Fixes #(issue)

## Before submitting:

<!--- Go over all the following points, and put an `x` in all the boxes that apply. -->
<!--- If you're unsure about any of these, don't hesitate to ask. We're here to help! -->
<!--- If you plan to update documentation or tests in follow-up, please note -->

- [ ] Does the Pull Request follow [Conventional Commits specification](https://www.conventionalcommits.org/en/v1.0.0/#summary) naming? Here are [GitHub's
      guide](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-a-pull-request) on how to create a pull request.
- [ ] Does the code follow BentoML's code style, `pre-commit run -a` script has passed ([instructions](https://github.com/bentoml/BentoML/blob/main/DEVELOPMENT.md#style-check-auto-formatting-type-checking))?
- [ ] Did you read through [contribution guidelines](https://github.com/bentoml/BentoML/blob/main/CONTRIBUTING.md#ways-to-contribute) and follow [development guidelines](https://github.com/bentoml/BentoML/blob/main/DEVELOPMENT.md#start-developing)?
- [ ] Did your changes require updates to the documentation? Have you updated
      those accordingly? Here are [documentation guidelines](https://github.com/bentoml/BentoML/tree/main/docs) and [tips on writting docs](https://github.com/bentoml/BentoML/tree/main/docs#writing-documentation).
- [ ] Did you write tests to cover your changes?



Patch:
diff --git a/src/_bentoml_impl/server/serving.py b/src/_bentoml_impl/server/serving.py
index 69a2b655..ace69ab9 100644
--- a/src/_bentoml_impl/server/serving.py
+++ b/src/_bentoml_impl/server/serving.py
@@ -140,8 +140,12 @@ def server_on_deployment(
         if callable(member) and getattr(member, "__bentoml_deployment_hook__", False):
             member()
     # Resolve models before server starts.
-    for model in svc.models:
-        model.resolve()
+    if bento := svc.bento:
+        for model in bento.info.all_models:
+            model.to_model().resolve()
+    else:
+        for model in svc.models:
+            model.resolve()
     if os.path.exists(result_file):
         os.remove(result_file)
 
diff --git a/src/_bentoml_sdk/__init__.py b/src/_bentoml_sdk/__init__.py
index 52d3e686..f965d58c 100644
--- a/src/_bentoml_sdk/__init__.py
+++ b/src/_bentoml_sdk/__init__.py
@@ -1,4 +1,10 @@
+from __future__ import annotations
 from bentoml._internal.utils.pkg import pkg_version_info
+from typing_extensions import deprecated
+from typing import TypeVar, Callable, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from bentoml._internal.external_typing import ASGIApp
 
 if (ver := pkg_version_info("pydantic")) < (2,):
     raise ImportError(
@@ -8,7 +14,7 @@ if (ver := pkg_version_info("pydantic")) < (2,):
 
 # ruff: noqa
 
-from .decorators import api, on_shutdown, mount_asgi_app, on_deployment, task
+from .decorators import api, on_shutdown, asgi_app, on_deployment, task
 from .service import get_current_service
 from .service import depends
 from .service import Service, ServiceConfig
@@ -16,11 +22,22 @@ from .service import service
 from .service import runner_service
 from .io_models import IODescriptor
 
+T = TypeVar("T")
+
+
+@deprecated("Deprecated in favor of `bentoml.asgi_app`")
+def mount_asgi_app(
+    app: ASGIApp, *, path: str = "/", name: str | None = None
+) -> Callable[[T], T]:
+    return asgi_app(app, path=path, name=name)
+
+
 __all__ = [
     "api",
     "task",
     "on_shutdown",
     "on_deployment",
+    "asgi_app",
     "mount_asgi_app",
     "depends",
     "Service",
diff --git a/src/_bentoml_sdk/decorators.py b/src/_bentoml_sdk/decorators.py
index c10ae391..1cedbdfc 100644
--- a/src/_bentoml_sdk/decorators.py
+++ b/src/_bentoml_sdk/decorators.py
@@ -103,7 +103,7 @@ def api(
     return wrapper
 
 
-def mount_asgi_app(
+def asgi_app(
     app: ASGIApp, *, path: str = "/", name: str | None = None
 ) -> t.Callable[[R], R]:
     """Mount an ASGI app to the service.
diff --git a/src/_bentoml_sdk/gradio.py b/src/_bentoml_sdk/gradio.py
index ef7bae7d..a1c915a8 100644
--- a/src/_bentoml_sdk/gradio.py
+++ b/src/_bentoml_sdk/gradio.py
@@ -5,7 +5,7 @@ from pathlib import Path
 
 from bentoml.exceptions import MissingDependencyException
 
-from .decorators import mount_asgi_app
+from .decorators import asgi_app
 
 R = t.TypeVar("R")
 
@@ -55,7 +55,7 @@ def mount_gradio_app(blocks: Blocks, path: str, name: str = "gradio_ui"):
         blocks.root_path = path
         blocks.favicon_path = favicon_path
         gradio_app = gr.routes.App.create_app(blocks, app_kwargs={"root_path": path})
-        mount_asgi_app(gradio_app, path=path, name=name)(obj)
+        asgi_app(gradio_app, path=path, name=name)(obj)
 
         # @bentoml.service() decorator returns a wrapper instead of the original class
         # Check if the object is an instance of Service
diff --git a/src/bentoml/__init__.py b/src/bentoml/__init__.py
index 6fbcdf35..54618fe0 100644
--- a/src/bentoml/__init__.py
+++ b/src/bentoml/__init__.py
@@ -104,6 +104,7 @@ if TYPE_CHECKING:
     from _bentoml_impl.loader import importing
     from _bentoml_sdk import IODescriptor
     from _bentoml_sdk import api
+    from _bentoml_sdk import asgi_app
     from _bentoml_sdk import depends
     from _bentoml_sdk import get_current_service
     from _bentoml_sdk import images
@@ -252,6 +253,7 @@ else:
         "depends",
         "on_shutdown",
         "on_deployment",
+        "asgi_app",
         "mount_asgi_app",
         "get_current_service",
         "IODescriptor",
@@ -363,6 +365,7 @@ __all__ = [
     "validators",
     "Field",
     "get_current_service",
+    "asgi_app",
     "mount_asgi_app",
     # new implementation
     "SyncHTTPClient",
diff --git a/src/bentoml/_internal/bento/bento.py b/src/bentoml/_internal/bento/bento.py
index 4adef598..6b301cfb 100644
--- a/src/bentoml/_internal/bento/bento.py
+++ b/src/bentoml/_internal/bento/bento.py
@@ -51,6 +51,7 @@ if TYPE_CHECKING:
     from fs.base import FS
 
     from _bentoml_sdk import Service as NewService
+    from _bentoml_sdk.models import Model
     from _bentoml_sdk.service import ServiceConfig
     from _bentoml_sdk.service.dependency import Dependency
 
@@ -663,6 +664,17 @@ class BentoModelInfo:
             alias=alias,
         )
 
+    def to_model(self) -> Model[t.Any]:
+        from bentoml.models import BentoModel
+        from bentoml.models import HuggingFaceModel
+
+        if self.registry == "bentoml":
+            return BentoModel.from_info(self)
+        elif self.registry == "huggingface":
+            return HuggingFaceModel.from_info(self)
+        else:
+            raise ValueError(f"Unknown model registry: {self.registry}")
+
 
 @attr.frozen
 class BentoDependencyInfo:
diff --git a/src/bentoml/_internal/cloud/bento.py b/src/bentoml/_internal/cloud/bento.py
index 6112ee5b..9c16dc19 100644
--- a/src/bentoml/_internal/cloud/bento.py
+++ b/src/bentoml/_internal/cloud/bento.py
@@ -102,7 +102,7 @@ class BentoAPI:
             if model.registry == "huggingface":
                 models_to_push.append(HuggingFaceModel.from_info(model))
             else:
-                model = BentoModel(model.tag)
+                model = BentoModel.from_info(model)
                 if model.stored is not None:
                     models_to_push.append(model)
         model_api = ModelAPI(self._client, self.spinner)
diff --git a/src/bentoml/_internal/container/__init__.py b/src/bentoml/_internal/container/__init__.py
index 1746216d..832cc2a6 100644
--- a/src/bentoml/_internal/container/__init__.py
+++ b/src/bentoml/_internal/container/__init__.py
@@ -168,7 +168,7 @@ def construct_containerfile(
                 model_ref = HuggingFaceModel.from_info(model)
                 model_ref.resolve(hf_model_dir)
             else:
-                model_ref = BentoModel(model.tag)
+                model_ref = BentoModel.from_info(model)
                 model_ref.resolve(bento_model_dir)
 
         # NOTE: dockerfile_template is already included in the
diff --git a/src/bentoml/_internal/container/frontend/dockerfile/entrypoint.sh b/src/bentoml/_internal/container/frontend/dockerfile/entrypoint.sh
index 3ec4e7c7..2e26e72f 100755
--- a/src/bentoml/_internal/container/frontend/dockerfile/entrypoint.sh
+++ b/src/bentoml/_internal/container/frontend/dockerfile/entrypoint.sh
@@ -10,7 +10,11 @@ _is_sourced() {
 }
 
 _main() {
-	[ -d .venv ] && source .venv/bin/activate
+	if [ -d "/app/.venv" ]; then
+		source /app/.venv/bin/activate
+	elif [ -d ".venv" ]; then
+		source .venv/bin/activate
+	fi
 	# For backwards compatibility with the yatai<1.0.0, adapting the old "yatai" command to the new "start" command.
 	if [ "${#}" -gt 0 ] && [ "${1}" = 'python' ] && [ "${2}" = '-m' ] && { [ "${3}" = 'bentoml._internal.server.cli.runner' ] || [ "${3}" = "bentoml._internal.server.cli.api_server" ]; }; then # SC2235, use { } to avoid subshell overhead
 		if [ "${3}" = 'bentoml._internal.server.cli.runner' ]; then
diff --git a/src/bentoml/_internal/server/http/access.py b/src/bentoml/_internal/server/http/access.py
index 9a7e1222..ce240e01 100644
--- a/src/bentoml/_internal/server/http/access.py
+++ b/src/bentoml/_internal/server/http/access.py
@@ -65,19 +65,24 @@ class AccessLogMiddleware:
         receive: ext.ASGIReceive,
         send: ext.ASGISend,
     ) -> None:
-        if not scope["type"].startswith("http"):
+        if not scope["type"].startswith(("http", "websocket")):
             await self.app(scope, receive, send)
             return
 
         start = default_timer()
         client = scope.get("client")
         scheme = scope["scheme"]
-        method = scope["method"]
+        method = scope.get("method", "")
         path = scope["path"]
         if path.startswith(tuple(self.skip_paths)):
             await self.app(scope, receive, send)
             return
 
+        if client:
+            address = f"{client[0]}:{client[1]}"
+        else:
+            address = "_"
+
         if self.has_request_content_length or self.has_request_content_type:
             for key, value in scope["headers"]:
                 if key == CONTENT_LENGTH:
@@ -85,7 +90,23 @@ class AccessLogMiddleware:
                 elif key == CONTENT_TYPE:
                     request_content_type.set(value)
 
+        async def wrapped_receive() -> "ext.ASGIMessage":
+            message = await receive()
+            if message["type"] == "websocket.connect":
+                self.logger.info(
+                    "%s (scheme=%s,path=%s) - Client connected", address, scheme, path
+                )
+            elif message["type"] == "websocket.disconnect":
+                self.logger.info(
+                    "%s (scheme=%s,path=%s) - Client disconnected",
+                    address,
+                    scheme,
+                    path,
+                )
+            return message
+
         async def wrapped_send(message: "ext.ASGIMessage") -> None:
+            latency = max(default_timer() - start, 0) * 1000
             if message["type"] == "http.response.start":
                 status.set(message["status"])
                 if self.has_response_content_length or self.has_response_content_type:
@@ -94,17 +115,23 @@ class AccessLogMiddleware:
                             response_content_length.set(value)
                         elif key == CONTENT_TYPE:
                             response_content_type.set(value)
+            elif message["type"] == "websocket.close":
+                self.logger.info(
+                    "%s (scheme=%s,path=%s) - Connection closed", address, scheme, path
+                )
+            elif message["type"] == "websocket.accept":
+                self.logger.info(
+                    "%s (scheme=%s,path=%s) - Connection established",
+                    address,
+                    scheme,
+                    path,
+                )
 
             elif message["type"] == "http.response.body":
                 if "more_body" in message and message["more_body"]:
                     await send(message)
                     return
 
-                if client:
-                    address = f"{client[0]}:{client[1]}"
-                else:
-                    address = "_"
-
                 request = [f"scheme={scheme}", f"method={method}", f"path={path}"]
                 if self.has_request_content_type:
                     request.append(f"type={request_content_type.get().decode()}")
@@ -117,7 +144,6 @@ class AccessLogMiddleware:
                 if self.has_response_content_length:
                     response.append(f"length={response_content_length.get().decode()}")
 
-                latency = max(default_timer() - start, 0) * 1000
                 await send(message)
                 self.logger.info(
                     "%s (%s) (%s) %.3fms",
@@ -130,4 +156,4 @@ class AccessLogMiddleware:
 
             await send(message)
 
-        await self.app(scope, receive, wrapped_send)
+        await self.app(scope, wrapped_receive, wrapped_send)
diff --git a/src/bentoml/_internal/server/http/instruments.py b/src/bentoml/_internal/server/http/instruments.py
index e5b717dd..15bf7e6d 100644
--- a/src/bentoml/_internal/server/http/instruments.py
+++ b/src/bentoml/_internal/server/http/instruments.py
@@ -188,6 +188,25 @@ class RunnerTrafficMetricsMiddleware:
             labelnames=["endpoint", "service_name", "service_version", "runner_name"],
             multiprocess_mode="livesum",
         )
+        self.metrics_websocket_connections = metrics_client.Gauge(
+            namespace=self.namespace,
+            name="websocket_connections",
+            documentation="Total number of websocket connections",
+            labelnames=["endpoint", "service_name", "service_version", "runner_name"],
+            multiprocess_mode="livesum",
+        )
+        self.metrics_websocket_data_received = metrics_client.Summary(
+            namespace=self.namespace,
+            name="websocket_data_received",
+            documentation="Total number of bytes received from websocket",
+            labelnames=["endpoint", "service_name", "service_version", "runner_name"],
+        )
+        self.metrics_websocket_data_sent = metrics_client.Summary(
+            namespace=self.namespace,
+            name="websocket_data_sent",
+            documentation="Total number of bytes sent to websocket",
+            labelnames=["endpoint", "service_name", "service_version", "runner_name"],
+        )
         self._is_setup = True
 
     async def __call__(
@@ -198,11 +217,11 @@ class RunnerTrafficMetricsMiddleware:
     ) -> None:
         if not self._is_setup:
             self._setup()
-        if not scope["type"].startswith("http"):
+        if not scope["type"].startswith(("http", "websocket")):
             await self.app(scope, receive, send)
             return
 
-        if scope["path"] == "/metrics":
+        if scope["type"].startswith("http") and scope["path"] == "/metrics":
             from starlette.responses import Response
 
             response = Response(
@@ -214,37 +233,81 @@ class RunnerTrafficMetricsMiddleware:
             return
 
         endpoint = scope["path"]
-        START_TIME_VAR.set(default_timer())
+        start_time = default_timer()
+        status_code = 0
+
+        async def wrapped_receive() -> "ext.ASGIMessage":
+            message = await receive()
+            if message["type"] == "websocket.disconnect":
+                self.metrics_websocket_connections.labels(
+                    endpoint=endpoint,
+                    service_name=server_context.bento_name,
+                    service_version=server_context.bento_version,
+                    runner_name=server_context.service_name,
+                ).dec()
+            elif message["type"] == "websocket.receive":
+                if message.get("bytes") is not None:
+                    data_len = len(message["bytes"])
+                else:
+                    data_len = len(message["text"])
+                self.metrics_websocket_data_received.labels(
+                    endpoint=endpoint,
+                    service_name=server_context.bento_name,
+                    service_version=server_context.bento_version,
+                    runner_name=server_context.service_name,
+                ).observe(data_len)
+            return message
 
         async def wrapped_send(message: "ext.ASGIMessage") -> None:
+            nonlocal status_code
             if message["type"] == "http.response.start":
-                STATUS_VAR.set(message["status"])
+                status_code = message["status"]
             elif message["type"] == "http.response.body":
                 if ("more_body" not in message) or not message["more_body"]:
-                    assert START_TIME_VAR.get() != 0
-                    assert STATUS_VAR.get() != 0
-
                     # instrument request total count
                     self.metrics_request_total.labels(
                         endpoint=endpoint,
                         service_name=server_context.bento_name,
                         service_version=server_context.bento_version,
-                        http_response_code=STATUS_VAR.get(),
+                        http_response_code=status_code,
                         runner_name=server_context.service_name,
                     ).inc()
 
                     # instrument request duration
-                    total_time = max(default_timer() - START_TIME_VAR.get(), 0)
+                    total_time = max(default_timer() - start_time, 0)
                     self.metrics_request_duration.labels(  # type: ignore
                         endpoint=endpoint,
                         service_name=server_context.bento_name,
                         service_version=server_context.bento_version,
-                        http_response_code=STATUS_VAR.get(),
+                        http_response_code=status_code,
                         runner_name=server_context.service_name,
                     ).observe(total_time)
+            elif message["type"] == "websocket.send":
+                if message.get("bytes") is not None:
+                    data_len = len(message["bytes"])
+                else:
+                    data_len = len(message["text"])
+                self.metrics_websocket_data_sent.labels(
+                    endpoint=endpoint,
+                    service_name=server_context.bento_name,
+                    service_version=server_context.bento_version,
+                    runner_name=server_context.service_name,
+                ).observe(data_len)
+            elif message["type"] == "websocket.accept":
+                self.metrics_websocket_connections.labels(
+                    endpoint=endpoint,
+                    service_name=server_context.bento_name,
+                    service_version=server_context.bento_version,
+                    runner_name=server_context.service_name,
+                ).inc()
+            elif message["type"] == "websocket.close":
+                self.metrics_websocket_connections.labels(
+                    endpoint=endpoint,
+                    service_name=server_context.bento_name,
+                    service_version=server_context.bento_version,
+                    runner_name=server_context.service_name,
+                ).dec()
 
-                    START_TIME_VAR.set(0)
-                    STATUS_VAR.set(0)
             await send(message)
 
         with self.metrics_request_in_progress.labels(
@@ -253,5 +316,5 @@ class RunnerTrafficMetricsMiddleware:
             service_version=server_context.bento_version,
             runner_name=server_context.service_name,
         ).track_inprogress():
-            await self.app(scope, receive, wrapped_send)
+            await self.app(scope, wrapped_receive, wrapped_send)
             return


Test patch:
diff --git a/tests/unit/bentoml_io/test_decorators.py b/tests/unit/bentoml_io/test_decorators.py
index 9f8a9b3a..f3c3b44f 100644
--- a/tests/unit/bentoml_io/test_decorators.py
+++ b/tests/unit/bentoml_io/test_decorators.py
@@ -18,7 +18,7 @@ def test_mount_asgi_app():
 
     app = FastAPI()
 
-    @bentoml.mount_asgi_app(app, path="/test")
+    @bentoml.asgi_app(app, path="/test")
     @bentoml.service(metrics={"enabled": False})
     class TestService:
         @app.get("/hello")
@@ -37,7 +37,7 @@ def test_mount_asgi_app_later():
     app = FastAPI()
 
     @bentoml.service(metrics={"enabled": False})
-    @bentoml.mount_asgi_app(app, path="/test")
+    @bentoml.asgi_app(app, path="/test")
     class TestService:
         @app.get("/hello")
         def hello(self):