fix: tests for metrics

Signed-off-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>

src/bentoml/_internal/runner/runner_handle/remote.py

@@ -138,6 +138,8 @@ async def async_run_method(
         *args: P.args,
         **kwargs: P.kwargs,
     ) -> R | tuple[R, ...]:
+        import aiohttp
+
         from ...runner.container import AutoContainer
 
         inp_batch_dim = __bentoml_method.config.batch_dim[0]

src/bentoml/metrics.py

@@ -52,10 +52,14 @@ def __dir__() -> list[str]:
 
 
 def __getattr__(item: t.Any):
+    if item in _NOT_SUPPORTED:
+        raise NotImplementedError(
+            f"{item} is not supported when using '{__name__}'. See https://docs.bentoml.org/en/latest/reference/metrics.html."
+        )
     # This is the entrypoint for all bentoml.metrics.*
     if item in _PROPERTY:
         logger.warning(
-            "'%s' is a '@property', which means there is no lazy loading. See https://docs.bentoml.org/en/latest/guides/metrics.html.",
+            "'%s' is a '@property', which means there is no lazy loading. See https://docs.bentoml.org/en/latest/reference/metrics.html.",
             item,
         )
         return getattr(_LazyMetric(item), item)
@@ -66,8 +70,6 @@ class _LazyMetric:
     __slots__ = ("_attr", "_proxy", "_initialized", "_args", "_kwargs")
 
     def __init__(self, attr: str):
-        if attr in _NOT_SUPPORTED:
-            raise NotImplementedError(f"{attr} is not yet supported.")
         self._attr = attr
         self._proxy = None
         self._initialized = False

src/bentoml/serve.py

@@ -182,6 +182,8 @@ def serve_http_development(
     ssl_ciphers: str | None = Provide[BentoMLContainer.ssl.ciphers],
     reload: bool = False,
 ) -> None:
+    prometheus_dir = ensure_prometheus_dir()
+
     from circus.sockets import CircusSocket
 
     from bentoml import load
@@ -193,7 +195,6 @@ def serve_http_development(
     working_dir = os.path.realpath(os.path.expanduser(working_dir))
     svc = load(bento_identifier, working_dir=working_dir)
 
-    prometheus_dir = ensure_prometheus_dir()
     watchers: list[Watcher] = []
     circus_sockets: list[CircusSocket] = [
         CircusSocket(name=API_SERVER_NAME, host=host, port=port, backlog=backlog)
@@ -301,6 +302,8 @@ def serve_http_production(
     ssl_ca_certs: str | None = Provide[BentoMLContainer.ssl.ca_certs],
     ssl_ciphers: str | None = Provide[BentoMLContainer.ssl.ciphers],
 ) -> None:
+    prometheus_dir = ensure_prometheus_dir()
+
     from circus.sockets import CircusSocket
 
     from bentoml import load
@@ -316,7 +319,6 @@ def serve_http_production(
     circus_socket_map: t.Dict[str, CircusSocket] = {}
     runner_bind_map: t.Dict[str, str] = {}
     uds_path = None
-    prometheus_dir = ensure_prometheus_dir()
 
     if psutil.POSIX:
         # use AF_UNIX sockets for Circus
@@ -496,6 +498,8 @@ def serve_grpc_development(
     channelz: bool = Provide[BentoMLContainer.grpc.channelz.enabled],
     reflection: bool = Provide[BentoMLContainer.grpc.reflection.enabled],
 ) -> None:
+    prometheus_dir = ensure_prometheus_dir()
+
     from circus.sockets import CircusSocket
 
     from bentoml import load
@@ -508,7 +512,6 @@ def serve_grpc_development(
     working_dir = os.path.realpath(os.path.expanduser(working_dir))
     svc = load(bento_identifier, working_dir=working_dir)
 
-    prometheus_dir = ensure_prometheus_dir()
     watchers: list[Watcher] = []
     circus_sockets: list[CircusSocket] = []
 
@@ -673,6 +676,8 @@ def serve_grpc_production(
     channelz: bool = Provide[BentoMLContainer.grpc.channelz.enabled],
     reflection: bool = Provide[BentoMLContainer.grpc.reflection.enabled],
 ) -> None:
+    prometheus_dir = ensure_prometheus_dir()
+
     from bentoml import load
     from bentoml.exceptions import UnprocessableEntity
 
@@ -691,8 +696,6 @@ def serve_grpc_production(
     runner_bind_map: dict[str, str] = {}
     uds_path = None
 
-    prometheus_dir = ensure_prometheus_dir()
-
     # Check whether users are running --grpc on windows
     # also raising warning if users running on MacOS or FreeBSD
     if psutil.WINDOWS:

src/bentoml/start.py

@@ -35,9 +35,12 @@ def start_runner_server(
     """
     Experimental API for serving a BentoML runner.
     """
+    from .serve import ensure_prometheus_dir
+
+    prometheus_dir = ensure_prometheus_dir()
+
     from bentoml import load
 
-    from .serve import ensure_prometheus_dir
     from ._internal.utils import reserve_free_port
     from ._internal.utils.circus import create_standalone_arbiter
     from ._internal.utils.analytics import track_serve
@@ -51,8 +54,6 @@ def start_runner_server(
     watchers: t.List[Watcher] = []
     circus_socket_map: t.Dict[str, CircusSocket] = {}
 
-    ensure_prometheus_dir()
-
     with contextlib.ExitStack() as port_stack:
         for runner in svc.runners:
             if runner.name == runner_name:
@@ -84,6 +85,8 @@ def start_runner_server(
                             "--no-access-log",
                             "--worker-id",
                             "$(circus.wid)",
+                            "--prometheus-dir",
+                            prometheus_dir,
                         ],
                         copy_env=True,
                         stop_children=True,
@@ -130,6 +133,10 @@ def start_http_server(
     ssl_ca_certs: str | None = Provide[BentoMLContainer.ssl.ca_certs],
     ssl_ciphers: str | None = Provide[BentoMLContainer.ssl.ciphers],
 ) -> None:
+    from .serve import ensure_prometheus_dir
+
+    prometheus_dir = ensure_prometheus_dir()
+
     from circus.sockets import CircusSocket
     from circus.watcher import Watcher
 
@@ -139,7 +146,6 @@ def start_http_server(
     from .serve import API_SERVER_NAME
     from .serve import construct_ssl_args
     from .serve import PROMETHEUS_MESSAGE
-    from .serve import ensure_prometheus_dir
     from ._internal.utils.circus import create_standalone_arbiter
     from ._internal.utils.analytics import track_serve
 
@@ -152,7 +158,6 @@ def start_http_server(
         )
     watchers: t.List[Watcher] = []
     circus_socket_map: t.Dict[str, CircusSocket] = {}
-    prometheus_dir = ensure_prometheus_dir()
     logger.debug("Runner map: %s", runner_map)
     circus_socket_map[API_SERVER_NAME] = CircusSocket(
         name=API_SERVER_NAME,
@@ -237,6 +242,10 @@ def start_grpc_server(
     ssl_keyfile: str | None = Provide[BentoMLContainer.ssl.keyfile],
     ssl_ca_certs: str | None = Provide[BentoMLContainer.ssl.ca_certs],
 ) -> None:
+    from .serve import ensure_prometheus_dir
+
+    prometheus_dir = ensure_prometheus_dir()
+
     from circus.sockets import CircusSocket
     from circus.watcher import Watcher
 
@@ -245,7 +254,6 @@ def start_grpc_server(
     from .serve import create_watcher
     from .serve import construct_ssl_args
     from .serve import PROMETHEUS_MESSAGE
-    from .serve import ensure_prometheus_dir
     from .serve import PROMETHEUS_SERVER_NAME
     from ._internal.utils import reserve_free_port
     from ._internal.utils.circus import create_standalone_arbiter
@@ -260,7 +268,6 @@ def start_grpc_server(
         )
     watchers: list[Watcher] = []
     circus_socket_map: dict[str, CircusSocket] = {}
-    prometheus_dir = ensure_prometheus_dir()
     logger.debug("Runner map: %s", runner_map)
     ssl_args = construct_ssl_args(
         ssl_certfile=ssl_certfile,

tests/e2e/bento_server_grpc/service.py

@@ -206,3 +206,13 @@ async def predict_multi_images(original: Image, compared: Image):
     )
     img = PIL.Image.fromarray(output_array)
     return {"meta": "success", "result": img}
+
+
+@svc.api(input=bentoml.io.Text(), output=bentoml.io.Text())
+def ensure_metrics_are_registered(data: str) -> str:  # pylint: disable=unused-argument
+    histograms = [
+        m.name
+        for m in bentoml.metrics.text_string_to_metric_families()
+        if m.type == "histogram"
+    ]
+    assert "inference_latency" in histograms

tests/e2e/bento_server_grpc/tests/test_descriptors.py

@@ -8,7 +8,6 @@
 
 import pytest
 
-from bentoml import metrics
 from bentoml.grpc.utils import import_grpc
 from bentoml.grpc.utils import import_generated_stubs
 from bentoml.testing.grpc import create_channel
@@ -407,11 +406,3 @@ async def test_multipart(host: str, img_file: str):
                 assert_multi_images, method="pred_multi_images", im_file=img_file
             ),
         )
-
-    # Test for metrics
-    histograms = [
-        m.name
-        for m in metrics.text_string_to_metric_families()
-        if m.type == "histogram"
-    ]
-    assert "inference_latency" in histograms

tests/e2e/bento_server_grpc/tests/test_metrics.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+
+from bentoml.grpc.utils import import_generated_stubs
+from bentoml.testing.grpc import create_channel
+from bentoml.testing.grpc import async_client_call
+
+if TYPE_CHECKING:
+    from bentoml.grpc.v1alpha1 import service_pb2 as pb
+else:
+    pb, _ = import_generated_stubs()
+
+
+@pytest.mark.asyncio
+async def test_multipart(host: str, img_file: str):
+    with open(str(img_file), "rb") as f:
+        fb = f.read()
+
+    async with create_channel(host) as channel:
+        await async_client_call(
+            "predict_multi_images",
+            channel=channel,
+            data={
+                "multipart": {
+                    "fields": {
+                        "original": pb.Part(
+                            file=pb.File(kind=pb.File.FILE_TYPE_BMP, content=fb)
+                        ),
+                        "compared": pb.Part(
+                            file=pb.File(kind=pb.File.FILE_TYPE_BMP, content=fb)
+                        ),
+                    }
+                }
+            },
+        )
+        await async_client_call(
+            "ensure_metrics_are_registered",
+            channel=channel,
+            data={"text": "input_string"},
+        )

tests/e2e/bento_server_http/service.py

@@ -46,6 +46,16 @@ def echo_data_metric(data: str) -> str:
     return data
 
 
+@svc.api(input=bentoml.io.Text(), output=bentoml.io.Text())
+def ensure_metrics_are_registered(data: str) -> str:  # pylint: disable=unused-argument
+    counters = [
+        m.name
+        for m in bentoml.metrics.text_string_to_metric_families()
+        if m.type == "counter"
+    ]
+    assert "test_metrics" in counters
+
+
 @svc.api(input=JSON(), output=JSON())
 async def echo_json(json_obj: JSONSerializable) -> JSONSerializable:
     batch_ret = await py_model.echo_json.async_run([json_obj])

tests/e2e/bento_server_http/tests/test_meta.py

@@ -4,7 +4,6 @@
 
 import time
 import asyncio
-from typing import TYPE_CHECKING
 from pathlib import Path
 
 import pytest
@@ -117,11 +116,27 @@ def test_dunder_string():
 
 
 @pytest.mark.asyncio
-async def test_metrics_type(host: str):
+async def test_metrics_type(host: str, deployment_mode: str):
     await async_request(
         "POST",
         f"http://{host}/echo_data_metric",
         headers={"Content-Type": "application/json"},
-        data="[[1,2],[3,4]]",
+        data="input_string",
     )
-    assert "test_metrics_total" in bentoml.metrics.generate_latest().decode("utf-8")
+    if deployment_mode == "docker":
+        # The reason we have to do this is that there is no way
+        # to access the metrics inside a running container.
+        await async_request(
+            "POST",
+            f"http://{host}/ensure_metrics_are_registered",
+            headers={"Content-Type": "application/json"},
+            data="input_string",
+            assert_status=200,
+        )
+    else:
+        counters = [
+            m.name
+            for m in bentoml.metrics.text_string_to_metric_families()
+            if m.type == "counter"
+        ]
+        assert "test_metrics" in counters