[Observability] Add startup, memory, and hybrid SWA diagnostics (#33375)

test/registered/core/test_srt_endpoint.py

@@ -41,6 +41,7 @@ class TestSRTEndpoint(CustomTestCase):
     # Extra server-launch env; subclasses override to run the same suite
     # against a different server flavor (e.g. SGLANG_RUST_SERVER=1).
     env = {}
+    expect_startup_observability = True
 
     @classmethod
     def setUpClass(cls):
@@ -562,6 +563,45 @@ class TestSRTEndpoint(CustomTestCase):
         version = response_json["version"]
         self.assertIsInstance(version, str)
 
+        if not self.expect_startup_observability:
+            return
+
+        startup_time = response_json["startup_time"]
+        for phase in (
+            "load_weight",
+            "kv_cache_allocation",
+            "scheduler_e2e",
+            "tokenizer_e2e",
+        ):
+            self.assertIsInstance(startup_time[phase], float)
+            self.assertGreater(startup_time[phase], 0)
+
+        graph_phases = {
+            "prefill",
+            "decode",
+            "target_verify",
+            "draft_prefill",
+            "draft_decode",
+            "draft_extend",
+        }
+        self.assertTrue(graph_phases.issubset(startup_time["cuda_graph"]))
+        for phase in graph_phases:
+            self.assertIsInstance(startup_time["cuda_graph"][phase], float)
+            self.assertGreaterEqual(startup_time["cuda_graph"][phase], 0)
+        self.assertGreater(startup_time["cuda_graph"]["decode"], 0)
+
+        memory_usage = response_json["internal_states"][0]["memory_usage"]
+        self.assertIsInstance(memory_usage["weight"], float)
+        self.assertIsInstance(memory_usage["kvcache"], float)
+        self.assertEqual(memory_usage["token_capacity"], max_total_num_tokens)
+        self.assertIsNone(memory_usage["token_capacity_swa"])
+        self.assertIsInstance(memory_usage["startup_available"], float)
+        self.assertGreater(memory_usage["startup_available"], 0)
+        self.assertTrue(graph_phases.issubset(memory_usage["graph"]))
+        for phase in graph_phases:
+            self.assertIsInstance(memory_usage["graph"][phase], float)
+            self.assertGreaterEqual(memory_usage["graph"][phase], 0)
+
     def test_logit_bias(self):
         """Test that a very high logit bias forces sampling of a specific token."""
         # Choose a token ID to bias (using 5 as an example)
@@ -864,6 +904,7 @@ class TestTokenizeDetokenize(CustomTestCase):
 )
 class TestRustServerEndpoint(TestSRTEndpoint):
     env = {"SGLANG_RUST_SERVER": "1"}
+    expect_startup_observability = False
 
     _RUST_TODO = "not implemented by the embedded Rust server yet"
 

test/registered/observability/test_metrics.py

@@ -29,6 +29,14 @@ register_cuda_ci(est_time=74, stage="base-b", runner_config="1-gpu-small")
 register_amd_ci(est_time=32, suite="stage-b-test-1-gpu-small-amd")
 
 _MODEL_NAME = "Qwen/Qwen3-0.6B"
+_GRAPH_PHASES = {
+    "prefill",
+    "decode",
+    "target_verify",
+    "draft_prefill",
+    "draft_decode",
+    "draft_extend",
+}
 
 
 class TestEnableMetrics(CustomTestCase):
@@ -66,14 +74,6 @@ class TestEnableMetrics(CustomTestCase):
                     "sglang:dp_cooperation_realtime_tokens_total",
                     {"mode": "decode"},
                 ),
-                (
-                    "sglang:dp_cooperation_forward_execution_seconds_total",
-                    {"category": "extend"},
-                ),
-                (
-                    "sglang:dp_cooperation_forward_execution_seconds_total",
-                    {"category": "decode"},
-                ),
             ]
             _check_metrics_positive(self, metrics, metrics_to_check)
 
@@ -139,8 +139,8 @@ class TestEnableMetrics(CustomTestCase):
             for _ in response.iter_lines(decode_unicode=False):
                 pass
 
-            for i in range(2):
-                # Send the request twice to trigger cached token metrics
+            for _ in range(3):
+                # The third request returns to the first rank under DP round-robin.
                 response = requests.post(
                     f"{DEFAULT_URL_FOR_TEST}/generate",
                     json={
@@ -187,6 +187,12 @@ class TestEnableMetrics(CustomTestCase):
             "sglang:num_unique_running_routing_keys",
             "sglang:routing_key_running_req_count",
             "sglang:routing_key_all_req_count",
+            "sglang:weight_memory_usage_gb",
+            "sglang:kv_cache_memory_usage_gb",
+            "sglang:graph_memory_usage_gb",
+            "sglang:startup_available_gpu_memory_gb",
+            "sglang:startup_time_seconds",
+            "sglang:startup_cuda_graph_time_seconds",
         ]
         mfu_metrics = [
             "sglang:estimated_flops_per_gpu_total",
@@ -224,9 +230,34 @@ class TestEnableMetrics(CustomTestCase):
             ("sglang:forward_execution_seconds_total", {"category": "extend"}),
             ("sglang:forward_execution_seconds_total", {"category": "decode"}),
             ("sglang:process_cpu_seconds_total", {"component": "tokenizer"}),
+            ("sglang:weight_memory_usage_gb", {"model_name": _MODEL_NAME}),
+            ("sglang:kv_cache_memory_usage_gb", {"model_name": _MODEL_NAME}),
+            (
+                "sglang:startup_available_gpu_memory_gb",
+                {"model_name": _MODEL_NAME},
+            ),
+            ("sglang:startup_time_seconds", {"phase": "load_weight"}),
+            ("sglang:startup_time_seconds", {"phase": "kv_cache_allocation"}),
+            ("sglang:startup_time_seconds", {"phase": "scheduler_e2e"}),
+            ("sglang:startup_time_seconds", {"phase": "tokenizer_e2e"}),
+            ("sglang:startup_cuda_graph_time_seconds", {"phase": "decode"}),
         ]
         _check_metrics_positive(self, metrics, metrics_to_check)
 
+        for metric_name in (
+            "sglang:graph_memory_usage_gb",
+            "sglang:startup_cuda_graph_time_seconds",
+        ):
+            phases = {
+                sample.labels.get("phase")
+                for sample in metrics[metric_name]
+                if sample.labels.get("model_name") == _MODEL_NAME
+            }
+            self.assertTrue(
+                _GRAPH_PHASES.issubset(phases),
+                f"{metric_name}: missing graph phases {_GRAPH_PHASES - phases}",
+            )
+
         if expect_mfu_metrics:
             # Estimated perf metrics may have multiple series (e.g., by rank). Ensure
             # that at least one series for this model has a positive accumulated value.

test/registered/unit/entrypoints/test_server_info.py

@@ -57,6 +57,7 @@ def _call_server_info_with(
     tokenizer_manager.server_args = server_args
     tokenizer_manager.model_path = server_args.model_path
     tokenizer_manager.served_model_name = server_args.served_model_name
+    tokenizer_manager.startup_time = None
     tokenizer_manager._config_updates = (
         [("test", dict(config_updates))] if config_updates else []
     )

test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py

@@ -77,12 +77,12 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
             "check_cuda_graph_backend",
             return_value=False,
         ):
-            runner = capture_prefill_graph(
+            capture = capture_prefill_graph(
                 model_runner=model_runner,
                 eager_runner=eager_runner,
             )
 
-        self.assertIs(runner, eager_runner)
+        self.assertIs(capture.runner, eager_runner)
 
     def test_prefix_chunk_capacity_is_aggregate_and_can_be_overridden(self):
         model_runner = SimpleNamespace(

test/registered/unit/observability/test_ray_wrappers.py

@@ -225,11 +225,11 @@ class TestCollectorSubclassWiring(TestRayWrapperBase):
         self.assertIs(cls._histogram_cls, self.rw.RayHistogramWrapper)
         self.assertIs(cls._summary_cls, self.rw.RaySummaryWrapper)
 
-    def test_tokenizer_overrides_counter_histogram_only(self):
+    def test_tokenizer_overrides_counter_gauge_histogram(self):
         cls = self.rw.RayTokenizerMetricsCollector
         self.assertIs(cls._counter_cls, self.rw.RayCounterWrapper)
+        self.assertIs(cls._gauge_cls, self.rw.RayGaugeWrapper)
         self.assertIs(cls._histogram_cls, self.rw.RayHistogramWrapper)
-        self.assertIsNone(cls._gauge_cls)
         self.assertIsNone(cls._summary_cls)
 
     def test_storage_overrides_counter_histogram_only(self):