Tiny support sglang_routing_keys_active in engine (#16570)

python/sglang/srt/utils/common.py

@@ -1501,6 +1501,26 @@ def add_prometheus_middleware(app):
     app.routes.append(metrics_route)
 
 
+class RefCountedGauge:
+    def __init__(self, gauge):
+        self._gauge = gauge
+        self._refcount: Dict[str, int] = {}
+
+    def inc(self, key: str):
+        if key in self._refcount:
+            self._refcount[key] += 1
+        else:
+            self._refcount[key] = 1
+            self._gauge.inc()
+
+    def dec(self, key: str):
+        if key in self._refcount:
+            self._refcount[key] -= 1
+            if self._refcount[key] == 0:
+                del self._refcount[key]
+                self._gauge.dec()
+
+
 def add_prometheus_track_response_middleware(app):
     from prometheus_client import Counter, Gauge
 
@@ -1523,15 +1543,26 @@ def add_prometheus_track_response_middleware(app):
         multiprocess_mode="livesum",
     )
 
+    routing_keys_active = RefCountedGauge(
+        Gauge(
+            name="sglang:routing_keys_active",
+            documentation="Number of unique routing keys with active requests",
+            multiprocess_mode="livesum",
+        )
+    )
+
     @app.middleware("http")
     async def track_http_status_code(request, call_next):
         # With recording all requests, we have the risk of high cardinality if requests have arbitrary unhandled paths.
         # But given that SGLang engines with metrics enabled are usually behind routers this looks safe.
         path, is_handled_path = _get_fastapi_request_path(request)
         method = request.method
+        routing_key = request.headers.get("x-smg-routing-key")
 
         http_request_counter.labels(endpoint=path, method=method).inc()
         http_requests_active.labels(endpoint=path, method=method).inc()
+        if routing_key:
+            routing_keys_active.inc(routing_key)
 
         try:
             response = await call_next(request)
@@ -1545,6 +1576,8 @@ def add_prometheus_track_response_middleware(app):
             return response
         finally:
             http_requests_active.labels(endpoint=path, method=method).dec()
+            if routing_key:
+                routing_keys_active.dec(routing_key)
 
 
 # https://github.com/blueswen/fastapi-observability/blob/132a3c576f8b09e5311c68bd553215013bc75685/fastapi_app/utils.py#L98

test/registered/metrics/test_metrics.py

@@ -101,6 +101,16 @@ class TestEnableMetrics(CustomTestCase):
             for _ in response.iter_lines(decode_unicode=False):
                 pass
 
+            response = requests.post(
+                f"{DEFAULT_URL_FOR_TEST}/generate",
+                json={
+                    "text": "Hello",
+                    "sampling_params": {"temperature": 0, "max_new_tokens": 5},
+                },
+                headers={"x-smg-routing-key": "test-key"},
+            )
+            self.assertEqual(response.status_code, 200)
+
             # Get metrics
             metrics_response = requests.get(f"{DEFAULT_URL_FOR_TEST}/metrics")
             self.assertEqual(metrics_response.status_code, 200)
@@ -133,6 +143,7 @@ class TestEnableMetrics(CustomTestCase):
             "sglang:inter_token_latency_seconds",
             "sglang:e2e_request_latency_seconds",
             "sglang:http_requests_active",
+            "sglang:routing_keys_active",
         ]
         for metric in essential_metrics:
             self.assertIn(metric, metrics_text, f"Missing metric: {metric}")