fix(vllm): publish KV gauges from prefill workers

components/src/dynamo/vllm/publisher.py

@@ -89,7 +89,7 @@ def record(
     def init_publish(self) -> None:
         self.inner.publish(self.dp_rank, kv_used_blocks=0)
         dp_rank_str = str(self.dp_rank)
-        self.component_gauges.set_total_blocks(dp_rank_str, 0)
+        self.component_gauges.set_total_blocks(dp_rank_str, self.num_gpu_block)
         self.component_gauges.set_gpu_cache_usage(dp_rank_str, 0.0)
 
     def log_engine_initialized(self) -> None:

components/src/dynamo/vllm/tests/test_vllm_publisher.py

@@ -12,11 +12,14 @@
 """
 
 from types import SimpleNamespace
-from unittest.mock import AsyncMock, Mock
+from unittest.mock import AsyncMock, Mock, call
 
 import pytest
+from prometheus_client import CollectorRegistry
 
 import dynamo.vllm.publisher as publisher_mod
+from dynamo import prometheus_names
+from dynamo.common.utils.prometheus import LLMBackendMetrics
 from dynamo.vllm.publisher import (
     DynamoStatLoggerPublisher,
     NoopStatLogger,
@@ -135,18 +138,29 @@ def _fake_publisher(*args, **kwargs):
     factory = StatLoggerFactory(
         endpoint=SimpleNamespace(), component_gauges=SimpleNamespace()
     )
-    for dp_rank in range(3):
+    dp_ranks = (2, 4, 7)
+    for dp_rank in dp_ranks:
         factory.create_stat_logger(dp_rank=dp_rank)
 
     factory.set_num_gpu_blocks_all(4096)
     factory.init_publish()
 
-    assert factory.created_loggers == dict(enumerate(loggers))
+    assert factory.created_loggers == dict(zip(dp_ranks, loggers, strict=True))
     for logger in loggers:
         logger.set_num_gpu_block.assert_called_once_with(4096)
         logger.init_publish.assert_called_once_with()
 
 
+def test_factory_initialization_without_loggers_is_a_noop():
+    """Disabled vLLM stat logging leaves the factory without loggers."""
+    factory = StatLoggerFactory(endpoint=SimpleNamespace())
+
+    factory.set_num_gpu_blocks_all(4096)
+    factory.init_publish()
+
+    assert factory.created_loggers == {}
+
+
 def test_factory_binds_deferred_endpoint_to_every_dp_rank_logger(monkeypatch):
     loggers = []
 
@@ -177,18 +191,58 @@ async def test_deferred_logger_starts_with_fresh_metrics_state(monkeypatch):
         publisher_mod, "WorkerMetricsPublisher", Mock(side_effect=publishers)
     )
 
+    registry = CollectorRegistry()
+    component_gauges = LLMBackendMetrics(
+        registry=registry,
+        model_name="test-model",
+        component_name="prefill",
+    )
+
     logger = DynamoStatLoggerPublisher(
         endpoint=None,
-        component_gauges=SimpleNamespace(),
+        dp_rank=5,
+        component_gauges=component_gauges,
     )
-    logger.inner.publish(dp_rank=0, kv_used_blocks=7)
+    logger.inner.publish(dp_rank=5, kv_used_blocks=7)
 
     endpoint = SimpleNamespace()
     logger.bind_endpoint(endpoint)
     assert logger.inner is publishers[1]
 
     assert logger._endpoint_task is not None
     await logger._endpoint_task
-    publishers[0].publish.assert_called_once_with(dp_rank=0, kv_used_blocks=7)
+    publishers[0].publish.assert_called_once_with(dp_rank=5, kv_used_blocks=7)
     publishers[0].create_endpoint.assert_not_called()
     publishers[1].create_endpoint.assert_awaited_once_with(endpoint)
+
+    metric_labels = {
+        prometheus_names.labels.MODEL: "test-model",
+        prometheus_names.labels.COMPONENT: "prefill",
+        prometheus_names.labels.DP_RANK: "5",
+    }
+    total_blocks_name = (
+        f"{prometheus_names.name_prefix.COMPONENT}_"
+        f"{prometheus_names.kvstats.TOTAL_BLOCKS}"
+    )
+    cache_usage_name = (
+        f"{prometheus_names.name_prefix.COMPONENT}_"
+        f"{prometheus_names.kvstats.GPU_CACHE_USAGE_PERCENT}"
+    )
+
+    logger.set_num_gpu_block(400)
+    logger.init_publish()
+    assert registry.get_sample_value(total_blocks_name, metric_labels) == 400
+    assert registry.get_sample_value(cache_usage_name, metric_labels) == 0.0
+
+    logger.record(SimpleNamespace(kv_cache_usage=0.25), None)
+    assert registry.get_sample_value(total_blocks_name, metric_labels) == 400
+    assert registry.get_sample_value(cache_usage_name, metric_labels) == 0.25
+
+    logger.record(SimpleNamespace(kv_cache_usage=0.0), None)
+    assert registry.get_sample_value(total_blocks_name, metric_labels) == 400
+    assert registry.get_sample_value(cache_usage_name, metric_labels) == 0.0
+    assert publishers[1].publish.call_args_list == [
+        call(5, kv_used_blocks=0),
+        call(5, kv_used_blocks=100),
+        call(5, kv_used_blocks=0),
+    ]

components/src/dynamo/vllm/tests/test_vllm_worker_factory.py

@@ -1114,21 +1114,32 @@ async def _noop(*_args, **_kwargs) -> None:
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("lora_enabled", [True, False])
-async def test_prefill_serves_lora_lifecycle_endpoints_when_enabled(
+@pytest.mark.parametrize("snapshot_mode", [True, False])
+async def test_prefill_initializes_metrics_and_serves_lora_lifecycle(
     monkeypatch: pytest.MonkeyPatch,
+    tmp_path,
     lora_enabled: bool,
+    snapshot_mode: bool,
 ) -> None:
     engine_client = Mock()
-    vllm_config = Mock(additional_config={})
+    vllm_config = Mock(
+        additional_config={},
+        cache_config=SimpleNamespace(num_gpu_blocks=96),
+    )
     engine_tuple: EngineSetupResult = (
         engine_client,
         vllm_config,
         Mock(),
-        "/tmp/prom",
+        str(tmp_path / "prometheus"),
         Mock(),
     )
+    stat_logger = Mock()
+    snapshot_engine: SnapshotEngineSetupResult | None = (
+        (engine_tuple, stat_logger) if snapshot_mode else None
+    )
+    setup_vllm_engine = Mock(return_value=engine_tuple)
     factory = WorkerFactory(
-        setup_vllm_engine_fn=Mock(return_value=engine_tuple),
+        setup_vllm_engine_fn=setup_vllm_engine,
         setup_kv_event_publisher_fn=Mock(return_value=None),
         register_vllm_model_fn=AsyncMock(),
         setup_fpm_relay_fn=Mock(return_value=None),
@@ -1155,6 +1166,13 @@ async def _noop(*_args, **_kwargs) -> None:
     monkeypatch.setattr(
         "dynamo.vllm.worker_factory.configure_kv_event_block_size", _noop
     )
+    monkeypatch.setattr(
+        "dynamo.vllm.worker_factory.get_dp_range_for_worker", lambda _config: (3, 2)
+    )
+    stat_logger_factory = Mock(return_value=stat_logger)
+    monkeypatch.setattr(
+        "dynamo.vllm.worker_factory.StatLoggerFactory", stat_logger_factory
+    )
 
     endpoints: dict[str, Mock] = {}
 
@@ -1187,8 +1205,26 @@ def _endpoint(name: str) -> Mock:
         config,
         asyncio.Event(),
         shutdown_endpoints,
+        snapshot_engine=snapshot_engine,
     )
 
+    if snapshot_mode:
+        stat_logger_factory.assert_not_called()
+        setup_vllm_engine.assert_not_called()
+        stat_logger.bind_endpoint.assert_called_once_with(
+            endpoints["dyn.prefill.generate"]
+        )
+    else:
+        stat_logger_factory.assert_called_once_with(
+            endpoint=endpoints["dyn.prefill.generate"]
+        )
+        setup_vllm_engine.assert_called_once_with(
+            config, stat_logger, fpm_worker_id="cid"
+        )
+        stat_logger.bind_endpoint.assert_not_called()
+    stat_logger.set_num_gpu_blocks_all.assert_called_once_with(48)
+    stat_logger.init_publish.assert_called_once_with()
+
     lifecycle_names = {
         "dyn.prefill.load_lora",
         "dyn.prefill.unload_lora",

components/src/dynamo/vllm/worker_factory.py

@@ -1557,40 +1557,39 @@ async def _run_prefill_worker(
 
         # Use pre-created engine if provided (checkpoint mode), otherwise create new
         fpm_worker_id = str(generate_endpoint.connection_id())
-        snapshot_factory: Optional[StatLoggerFactory] = None
         if snapshot_engine is not None:
-            engine_setup, snapshot_factory = snapshot_engine
+            engine_setup, factory = snapshot_engine
             (
                 engine_client,
                 vllm_config,
                 default_sampling_params,
                 prometheus_temp_dir,
                 _component_gauges,
             ) = engine_setup
-            snapshot_factory.bind_endpoint(generate_endpoint)
+            factory.bind_endpoint(generate_endpoint)
             # TODO: The scheduler in the child process still has worker_id=""
             # because the engine was forked before the runtime existed.
             # Propagating the new ID to the child requires shared memory or
             # a restart of the EngineCore process.
             os.environ[ENV_FPM_WORKER_ID] = fpm_worker_id
         else:
+            factory = StatLoggerFactory(endpoint=generate_endpoint)
             (
                 engine_client,
                 vllm_config,
                 default_sampling_params,
                 prometheus_temp_dir,
                 _component_gauges,
-            ) = self.setup_vllm_engine(config, fpm_worker_id=fpm_worker_id)
+            ) = self.setup_vllm_engine(config, factory, fpm_worker_id=fpm_worker_id)
         await configure_kv_event_block_size(engine_client, vllm_config)
 
-        if snapshot_factory is not None:
-            _, dp_size = get_dp_range_for_worker(vllm_config)
-            per_rank_num_gpu_blocks = per_rank_kv_blocks(
-                vllm_config.cache_config.num_gpu_blocks,
-                dp_size,
-            )
-            snapshot_factory.set_num_gpu_blocks_all(per_rank_num_gpu_blocks or 0)
-            snapshot_factory.init_publish()
+        _, dp_size = get_dp_range_for_worker(vllm_config)
+        per_rank_num_gpu_blocks = per_rank_kv_blocks(
+            vllm_config.cache_config.num_gpu_blocks,
+            dp_size,
+        )
+        factory.set_num_gpu_blocks_all(per_rank_num_gpu_blocks or 0)
+        factory.init_publish()
 
         encode_worker_client = await self._maybe_get_encode_worker_client(
             runtime, config

tests/serve/test_vllm.py

@@ -365,11 +365,17 @@ class VLLMConfig(EngineConfig):
             pytest.mark.core,
             pytest.mark.gpu_2,
             pytest.mark.pre_merge,
-        ],  # TODO: profile to get max_vram and timeout
+            pytest.mark.timeout(720),  # 3x ~240s local runtime
+        ],  # TODO: profile to get max_vram
         model="Qwen/Qwen3-0.6B",
         request_payloads=[
             chat_payload_default(),
             completion_payload_default(),
+            metric_payload_default(
+                min_num_requests=1,
+                backend="vllm",
+                port=DefaultPort.SYSTEM2.value,
+            ),
         ],
     ),
     "disaggregated_same_gpu": VLLMConfig(