fix: inject prometheus multiproc_dir into env var when worker starts (#…

…4820) * fix: inject prometheus multiproc_dir into env var when worker starts Signed-off-by: Frost Ming <[email protected]> * fix: remove unneeded tests Signed-off-by: Frost Ming <[email protected]> * fix: tests Signed-off-by: Frost Ming <[email protected]> * fix: replace bentoml.metrics with prometheus_client Signed-off-by: Frost Ming <[email protected]> --------- Signed-off-by: Frost Ming <[email protected]>
bentoml · Jun 23, 2024 · 65e7176 · 65e7176
1 parent da22538
commit 65e7176
Show file tree

Hide file tree

Showing 12 changed files with 51 additions and 476 deletions.
diff --git a/docs/source/guides/snippets/metrics/metric_defs.py b/docs/source/guides/snippets/metrics/metric_defs.py
@@ -1,8 +1,8 @@
 from __future__ import annotations
 
-import bentoml
+import prometheus_client
 
-inference_duration = bentoml.metrics.Histogram(
+inference_duration = prometheus_client.Histogram(
     name="inference_duration",
     documentation="Duration of inference",
     labelnames=["nltk_version", "sentiment_cls"],
@@ -25,7 +25,7 @@
     ),
 )
 
-polarity_counter = bentoml.metrics.Counter(
+polarity_counter = prometheus_client.Counter(
     name="polarity_total",
     documentation="Count total number of analysis by polarity scores",
     labelnames=["polarity"],

diff --git a/src/_bentoml_impl/server/app.py b/src/_bentoml_impl/server/app.py
@@ -29,14 +29,14 @@
 
 if t.TYPE_CHECKING:
     from opentelemetry.sdk.trace import Span
+    from prometheus_client import Histogram
     from starlette.applications import Starlette
     from starlette.requests import Request
     from starlette.routing import BaseRoute
 
     from bentoml._internal import external_typing as ext
     from bentoml._internal.context import ServiceContext
     from bentoml._internal.types import LifecycleHook
-    from bentoml.metrics import Histogram
 
 R = t.TypeVar("R")
 

diff --git a/src/_bentoml_impl/worker/service.py b/src/_bentoml_impl/worker/service.py
@@ -174,6 +174,9 @@ def main(
 
     if prometheus_dir is not None:
         BentoMLContainer.prometheus_multiproc_dir.set(prometheus_dir)
+    os.environ["PROMETHEUS_MULTIPROC_DIR"] = (
+        BentoMLContainer.prometheus_multiproc_dir.get()
+    )
     server_context.service_name = service.name
     if service.bento is None:
         server_context.bento_name = service.name

diff --git a/src/bentoml/_internal/server/metrics/prometheus.py b/src/bentoml/_internal/server/metrics/prometheus.py
@@ -2,12 +2,13 @@
 
 import logging
 import os
-import sys
 import typing as t
 from functools import partial
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+    from prometheus_client import Metric
+
     from ... import external_typing as ext
 
 logger = logging.getLogger(__name__)
@@ -45,20 +46,6 @@ def __init__(
 
     @property
     def prometheus_client(self):
-        if self.multiproc and not self._imported:
-            # step 1: check environment
-            assert (
-                "prometheus_client" not in sys.modules
-            ), "prometheus_client is already imported, multiprocessing will not work properly"
-
-            assert (
-                self.multiproc_dir
-            ), f"Invalid prometheus multiproc directory: {self.multiproc_dir}"
-            assert os.path.isdir(self.multiproc_dir)
-
-            os.environ["PROMETHEUS_MULTIPROC_DIR"] = self.multiproc_dir
-
-        # step 2:
         import prometheus_client
         import prometheus_client.exposition
         import prometheus_client.metrics
@@ -173,6 +160,10 @@ def Metric(self):
         """
         A Metric family and its samples.
 
-        This is a base class to be used by instrumentation client. Custom collectors should use ``bentoml.metrics.metrics_core.GaugeMetricFamily``, ``bentoml.metrics.metrics_core.CounterMetricFamily``, ``bentoml.metrics.metrics_core.SummaryMetricFamily`` instead.
+        This is a base class to be used by instrumentation client.
+        Custom collectors should use
+        ``prometheus_client.metrics_core.GaugeMetricFamily``,
+        ``prometheus_client.metrics_core.CounterMetricFamily``,
+        ``prometheus_client.metrics_core.SummaryMetricFamily`` instead.
         """
         return partial(self.prometheus_client.Metric, registry=self.registry)