Skip to content

Commit

Permalink
fix: inject prometheus multiproc_dir into env var when worker starts (#…
Browse files Browse the repository at this point in the history
…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]>
  • Loading branch information
frostming authored Jun 23, 2024
1 parent da22538 commit 65e7176
Show file tree
Hide file tree
Showing 12 changed files with 51 additions and 476 deletions.
6 changes: 3 additions & 3 deletions docs/source/guides/snippets/metrics/metric_defs.py
Original file line number Diff line number Diff line change
@@ -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"],
Expand All @@ -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"],
Expand Down
2 changes: 1 addition & 1 deletion src/_bentoml_impl/server/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand Down
3 changes: 3 additions & 0 deletions src/_bentoml_impl/worker/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
23 changes: 7 additions & 16 deletions src/bentoml/_internal/server/metrics/prometheus.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Loading

0 comments on commit 65e7176

Please sign in to comment.