Feature/prometheus metrics

pyproject.toml

@@ -159,6 +159,10 @@ convention = "google"
 # Unlike Flake8, default to a complexity level of 10.
 max-complexity = 10
 
+[project.optional-dependencies]
+prometheus = ["prometheus_client"]
+all = ["prometheus_client"]
+
 [dependency-groups]
 dev = [
     "asgi-lifespan>=2.1.0",

src/litserve/__init__.py

@@ -35,3 +35,10 @@
     "set_trace_if_debug",
     "test_examples",
 ]
+
+try:
+    from litserve.metrics import PrometheusLogger  # noqa: F401
+
+    __all__.append("PrometheusLogger")
+except ImportError:
+    pass

src/litserve/loggers.py

@@ -160,11 +160,11 @@ def run(self, lit_server: "LitServer"):
 
         module_logger.debug(f"Starting logger process with {len(logger_proxies)} loggers")
         ctx = mp.get_context("spawn")
-        process = ctx.Process(
+        self._process = ctx.Process(
             target=_LoggerConnector._process_logger_queue,
             args=(
                 logger_proxies,
                 queue,
             ),
         )
-        process.start()
+        self._process.start()

src/litserve/metrics.py

@@ -0,0 +1,99 @@
+# Copyright The Lightning AI team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import logging
+import time
+
+from starlette.middleware.base import BaseHTTPMiddleware
+
+from litserve.loggers import Logger
+
+logger = logging.getLogger(__name__)
+
+try:
+    import prometheus_client  # noqa: F401
+    from prometheus_client import CollectorRegistry, Counter, Histogram, make_asgi_app, multiprocess
+
+    _PROMETHEUS_AVAILABLE = True
+except ImportError:  # pragma: no cover
+    _PROMETHEUS_AVAILABLE = False
+
+
+def _check_prometheus_available():
+    if not _PROMETHEUS_AVAILABLE:  # pragma: no cover
+        raise ImportError(
+            "prometheus_client is not installed. Please install it with `pip install litserve[prometheus]` "
+            "to use PrometheusLogger."
+        )
+
+
+class PrometheusMiddleware(BaseHTTPMiddleware):
+    """Middleware to track HTTP request metrics for Prometheus."""
+
+    def __init__(self, app):
+        super().__init__(app)
+        _check_prometheus_available()
+
+        # Use singleton pattern to avoid hitting private registry internals or re-creating metrics
+        if not hasattr(PrometheusMiddleware, "_metrics_initialized"):
+            PrometheusMiddleware.http_requests_total = Counter(
+                "litserve_http_requests_total", "Total requests", ["method", "endpoint", "status"]
+            )
+            PrometheusMiddleware.http_request_duration_seconds = Histogram(
+                "litserve_http_request_duration_seconds", "Request latency", ["endpoint"]
+            )
+            PrometheusMiddleware._metrics_initialized = True
+
+    async def dispatch(self, request, call_next):
+        if request.url.path in ["/metrics", "/health", "/info"]:
+            return await call_next(request)
+
+        start_time = time.time()
+        response = await call_next(request)
+        duration = time.time() - start_time
+
+        PrometheusMiddleware.http_request_duration_seconds.labels(endpoint=request.url.path).observe(duration)
+        PrometheusMiddleware.http_requests_total.labels(
+            method=request.method, endpoint=request.url.path, status=response.status_code
+        ).inc()
+        return response
+
+
+class PrometheusLogger(Logger):
+    """A built-in logger that exposes Prometheus metrics in a multiprocess environment."""
+
+    def __init__(self):
+        super().__init__()
+        _check_prometheus_available()
+
+        async def lazy_asgi_app(scope, receive, send):
+            if not hasattr(self, "_asgi_app"):
+                registry = CollectorRegistry()
+                multiprocess.MultiProcessCollector(registry)
+                self._asgi_app = make_asgi_app(registry=registry)
+            await self._asgi_app(scope, receive, send)
+
+        self.mount("/metrics", lazy_asgi_app)
+
+        # We don't initialize the Histogram here because PROMETHEUS_MULTIPROC_DIR
+        # is not guaranteed to be set yet. It will be initialized lazily.
+        self._batch_size_histogram = None
+
+    def process(self, key, value):
+        """Handle self.log() calls from LitAPI workers."""
+        if self._batch_size_histogram is None:
+            # By the time process() is called, the server has started and set PROMETHEUS_MULTIPROC_DIR
+            self._batch_size_histogram = Histogram("litserve_batch_size", "Inference batch size")
+
+        if key == "batch_size":
+            self._batch_size_histogram.observe(value)

src/litserve/server.py

@@ -753,6 +753,27 @@ def __init__(
             )
             raise ValueError(_msg)
 
+        self._metrics_dir = None
+        try:
+            from litserve.metrics import PrometheusLogger
+
+            if loggers is not None:
+                _loggers_list = loggers if isinstance(loggers, list) else [loggers]
+                if any(isinstance(logger, PrometheusLogger) for logger in _loggers_list):
+                    import os
+                    import tempfile
+
+                    from litserve.metrics import PrometheusMiddleware
+
+                    self._metrics_dir = tempfile.mkdtemp(prefix="litserve_prom_")
+                    os.environ["PROMETHEUS_MULTIPROC_DIR"] = self._metrics_dir
+                    import prometheus_client.values
+
+                    prometheus_client.values.ValueClass = prometheus_client.values.MultiProcessValue()
+                    middlewares.append(PrometheusMiddleware)
+        except ImportError:
+            pass
+
         # Handle 0.3.0 migration
         if api_path is not None:
             _migration_warning("api_path")
@@ -1190,11 +1211,12 @@ def _perform_graceful_shutdown(
         logger.info("Shutting down LitServe...")
 
         # Handle transport closure based on shutdown reason
-        if shutdown_reason == "keyboard_interrupt":
-            logger.debug("KeyboardInterrupt detected - skipping transport cleanup to avoid hanging")
-            self._transport.close(send_sentinel=False)
-        else:
-            self._transport.close(send_sentinel=True)
+        if hasattr(self, "_transport"):
+            if shutdown_reason == "keyboard_interrupt":  # pragma: no cover
+                logger.debug("KeyboardInterrupt detected - skipping transport cleanup to avoid hanging")
+                self._transport.close(send_sentinel=False)
+            else:
+                self._transport.close(send_sentinel=True)
 
         # terminate Uvicorn server workers tracked by LitServe (the master processes/threads)
         if len(uvicorn_workers) > 0:
@@ -1234,6 +1256,25 @@ def _perform_graceful_shutdown(
             except Exception as e:
                 logger.error(f"Error while terminating worker {worker_name} (PID: {worker_pid}): {e}")
 
+        # terminate logger process
+        if hasattr(self, "_logger_connector") and hasattr(self._logger_connector, "_process"):
+            lp = self._logger_connector._process
+            if lp and lp.is_alive():  # pragma: no cover
+                logger.debug(f"Terminating logger process (PID: {lp.pid})...")
+                try:
+                    lp.terminate()
+                    lp.join(timeout=5)
+                    if lp.is_alive():
+                        logger.warning(f"Logger process (PID: {lp.pid}) did not terminate gracefully. Killing.")
+                        lp.kill()
+                except Exception as e:
+                    logger.error(f"Error during termination of logger process: {e}")
+
+        if getattr(self, "_metrics_dir", None) and os.path.exists(self._metrics_dir):
+            import shutil
+
+            shutil.rmtree(self._metrics_dir, ignore_errors=True)
+
         manager.shutdown()
 
     def _resolve_workers_per_device_config(self, workers_per_device):

src/litserve/utils.py

@@ -120,10 +120,27 @@ def wrap_litserve_start(server: "LitServer", worker_monitor: bool = False):
     finally:
         server._shutdown_event.set()
         # First close the transport to signal to the response_queue_to_buffer task that it should stop
-        server._transport.close()
+        if hasattr(server, "_transport"):
+            server._transport.close()
         for p in server.inference_workers:
             p.terminate()
             p.join()
+
+        if hasattr(server, "_logger_connector") and hasattr(server._logger_connector, "_process"):
+            lp = server._logger_connector._process
+            if lp and lp.is_alive():  # pragma: no cover
+                lp.terminate()
+                lp.join(timeout=1)
+                if lp.is_alive():
+                    lp.kill()
+
+        if getattr(server, "_metrics_dir", None) and os.path.exists(server._metrics_dir):
+            import contextlib
+            import shutil
+
+            with contextlib.suppress(Exception):
+                shutil.rmtree(server._metrics_dir)
+
         server.manager.shutdown()
 
 

tests/unit/test_prometheus.py

@@ -0,0 +1,149 @@
+import os
+
+import pytest
+from fastapi.testclient import TestClient
+
+import litserve as ls
+import litserve.metrics
+from litserve.metrics import _PROMETHEUS_AVAILABLE
+
+# Only run tests if prometheus is available
+pytestmark = pytest.mark.skipif(not _PROMETHEUS_AVAILABLE, reason="prometheus_client is not installed")
+
+
+class SimpleAPI(ls.LitAPI):
+    def setup(self, device):
+        pass
+
+    def predict(self, x):
+        self.log("batch_size", len(x))
+        return x
+
+
+def test_prometheus_logger_initialization():
+    api = SimpleAPI()
+    logger = ls.metrics.PrometheusLogger()
+    server = ls.LitServer(api, loggers=[logger])
+
+    assert server._metrics_dir is not None
+    assert os.environ.get("PROMETHEUS_MULTIPROC_DIR") == server._metrics_dir
+    assert os.path.exists(server._metrics_dir)
+    assert os.path.basename(server._metrics_dir).startswith("litserve_prom_")
+
+
+def test_prometheus_metrics_endpoint():
+    api = SimpleAPI()
+    logger = ls.metrics.PrometheusLogger()
+    server = ls.LitServer(api, loggers=[logger])
+
+    # Use wrap_litserve_start to initialize internal states
+    with ls.utils.wrap_litserve_start(server):
+        client = TestClient(server.app)
+
+        # Hit healthcheck multiple times
+        client.get("/health")
+        client.get("/health")
+        client.get("/health")
+
+        # Hit an invalid route to test error statuses
+        client.get("/invalid_route")
+
+        # Hit metrics
+        response = client.get("/metrics")
+        assert response.status_code == 200
+        text = response.text
+
+        # Parse metrics instead of brittle string matching
+        from prometheus_client.parser import text_string_to_metric_families
+
+        metrics = list(text_string_to_metric_families(text))
+
+        http_reqs = next((m for m in metrics if m.name == "litserve_http_requests"), None)
+        assert http_reqs is not None
+
+        # Validate healthcheck route is ignored (since it's skipped in middleware)
+        health_sample = next((s for s in http_reqs.samples if s.labels.get("endpoint") == "/health"), None)
+        assert health_sample is None
+
+        # Validate invalid route
+        invalid_sample = next((s for s in http_reqs.samples if s.labels.get("endpoint") == "/invalid_route"), None)
+        assert invalid_sample is not None
+        assert invalid_sample.value == 1.0
+
+
+def test_prometheus_inference_metrics():
+    api = SimpleAPI()
+    logger = ls.metrics.PrometheusLogger()
+    server = ls.LitServer(api, loggers=[logger])
+
+    # Use wrap_litserve_start to initialize internal states
+    with ls.utils.wrap_litserve_start(server):
+        # We must explicitly start the logger connector process since wrap_litserve_start skips it
+        server._logger_connector.run(server)
+
+        with TestClient(server.app) as client:
+            # Post to predict to trigger full pipeline: API -> Queue -> Worker -> process() -> /metrics
+            client.post("/predict", json=[1, 2, 3, 4])
+            client.post("/predict", json=[1, 2])
+
+            import time
+
+            # Retry loop to wait deterministically for background processes instead of a fixed sleep
+            for _ in range(20):
+                response = client.get("/metrics")
+                text = response.text
+
+                from prometheus_client.parser import text_string_to_metric_families
+
+                metrics = list(text_string_to_metric_families(text))
+                batch_size = next((m for m in metrics if m.name == "litserve_batch_size"), None)
+
+                if batch_size is not None:
+                    batch_count = next((s for s in batch_size.samples if s.name == "litserve_batch_size_count"), None)
+                    if batch_count is not None and batch_count.value == 2.0:
+                        break
+                time.sleep(0.1)
+            else:
+                pytest.fail("Timeout waiting for metrics to propagate from background processes")
+
+            batch_sum = next((s for s in batch_size.samples if s.name == "litserve_batch_size_sum"), None)
+            assert batch_sum is not None
+            assert batch_sum.value == 6.0
+
+
+def test_prometheus_cleanup_on_shutdown():
+    api = SimpleAPI()
+    logger = ls.metrics.PrometheusLogger()
+    server = ls.LitServer(api, loggers=[logger])
+    metrics_dir = server._metrics_dir
+
+    # Fake the logger connector setup
+    server._logger_connector = ls.loggers._LoggerConnector(server, [logger])
+    server.inference_workers = []
+
+    # Assert dir exists
+    assert os.path.exists(metrics_dir)
+
+    # Call shutdown logic
+    class FakeManager:
+        def shutdown(self):
+            pass
+
+    server._perform_graceful_shutdown(FakeManager(), {})
+
+    # Assert dir is deleted
+    assert not os.path.exists(metrics_dir)
+
+
+def test_prometheus_process_method():
+    api = SimpleAPI()
+    logger = ls.metrics.PrometheusLogger()
+    _ = ls.LitServer(api, loggers=[logger])
+
+    # Explicitly call process to ensure coverage on main thread
+    logger.process("batch_size", 4)
+    logger.process("batch_size", 8)
+
+    # We can inspect the internal histogram to verify it tracked correctly
+    assert logger._batch_size_histogram is not None
+    # No direct assert on metric values needed, just ensuring it doesn't crash