diff --git a/src/bentoml/_internal/server/metrics/prometheus.py b/src/bentoml/_internal/server/metrics/prometheus.py index 9583e2d87fa..b9295b0a441 100644 --- a/src/bentoml/_internal/server/metrics/prometheus.py +++ b/src/bentoml/_internal/server/metrics/prometheus.py @@ -96,9 +96,25 @@ def start_http_server(self, port: int, addr: str = "") -> None: registry=self.registry, ) + start_wsgi_server = start_http_server + + def write_to_textfile(self, path: str) -> None: + """ + Write metrics to given path. This is intended to be used with + the Node expoerter textfile collector. + + Args: + path: path to write the metrics to. This file must end + with '.prom' for the textfile collector to process it. + """ + self.prometheus_client.write_to_textfile(path, registry=self.registry) + def make_wsgi_app(self) -> ext.WSGIApp: return self.prometheus_client.make_wsgi_app(registry=self.registry) # type: ignore (unfinished prometheus types) + def make_asgi_app(self) -> ext.ASGIApp: + return self.prometheus_client.make_asgi_app(registry=self.registry) # type: ignore (unfinished prometheus types) + def generate_latest(self): if self.multiproc: registry = self.prometheus_client.CollectorRegistry() diff --git a/src/bentoml/metrics.py b/src/bentoml/metrics.py index b8df86691ca..cbed5b723d9 100644 --- a/src/bentoml/metrics.py +++ b/src/bentoml/metrics.py @@ -20,18 +20,18 @@ # This sets of functions are implemented in the PrometheusClient class _INTERNAL_IMPL = [ "start_http_server", + "start_wsgi_server", "make_wsgi_app", + "make_asgi_app", "generate_latest", "text_string_to_metric_families", + "write_to_textfile", ] _NOT_IMPLEMENTED = [ "delete_from_gateway", "instance_ip_grouping_key", - "make_asgi_app", "push_to_gateway", "pushadd_to_gateway", - "start_wsgi_server", - "write_to_textfile", ] _NOT_SUPPORTED = [ "GC_COLLECTOR", diff --git a/src/bentoml_cli/worker/grpc_prometheus_server.py b/src/bentoml_cli/worker/grpc_prometheus_server.py index 47a1d8a4e62..ab0286272e2 100644 --- a/src/bentoml_cli/worker/grpc_prometheus_server.py +++ b/src/bentoml_cli/worker/grpc_prometheus_server.py @@ -68,10 +68,11 @@ def main(fd: int, backlog: int, prometheus_dir: str | None): configure_server_logging() BentoMLContainer.development_mode.set(False) - metrics_client = BentoMLContainer.metrics_client.get() if prometheus_dir is not None: BentoMLContainer.prometheus_multiproc_dir.set(prometheus_dir) + metrics_client = BentoMLContainer.metrics_client.get() + # create a ASGI app that wraps around the default HTTP prometheus server. prom_app = Starlette( debug=get_debug_mode(), middleware=[Middleware(GenerateLatestMiddleware)]