chore: smooth some edge cases [skip ci]

Signed-off-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>
bentoml · Oct 27, 2022 · ec3839d · ec3839d
1 parent f2bd797
commit ec3839d
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 4 deletions.
diff --git 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
@@ -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
@@ -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)]