fix(vllm): seed prefill KV cache gauges

components/src/dynamo/vllm/publisher.py

@@ -67,51 +67,29 @@ def record(
         dp_rank_str = str(self.dp_rank)
         self.component_gauges.set_total_blocks(dp_rank_str, self.num_gpu_block)
 
-        # Set GPU cache usage percentage directly from scheduler_stats
-        # Note: vLLM's scheduler_stats.kv_cache_usage returns very small values
-        # (e.g., 0.0000834 for ~0.08% usage), which Prometheus outputs in scientific
-        # notation (8.34e-05). This is the correct value and will be properly parsed.
         self.component_gauges.set_gpu_cache_usage(
             dp_rank_str, scheduler_stats.kv_cache_usage
         )
 
     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:
         pass
 
 
 class NoopStatLogger(StatLoggerBase):
-    """Stat logger that drops every record.
-
-    vLLM's ``AsyncLLM`` always invokes a ``StatLoggerBase`` subclass
-    during engine init, but for some worker shapes the chat-style
-    publish path (KV cache usage, scheduler gauges) is meaningless --
-    embedding/pooling workers are the current driver, but the same
-    no-op semantics fit any future worker that wants to satisfy vLLM's
-    factory contract without registering Prometheus collectors. Reach
-    for this class instead of writing a new throwaway subclass each
-    time.
-    """
+    """Drops stats for workers without scheduler metrics."""
 
     def __init__(
         self,
         vllm_config: Optional[VllmConfig] = None,
         engine_index: int = 0,
     ) -> None:
-        # vLLM's ``StatLoggerBase`` declares ``__init__`` as abstract, so
-        # subclasses must provide one even when they hold no state.
-        # Without this, instantiation raises ``TypeError: Can't
-        # instantiate abstract class NoopStatLogger without an
-        # implementation for abstract method '__init__'``. The
-        # ``(vllm_config, engine_index)`` signature mirrors what vLLM
-        # passes to concrete ``StatLoggerBase`` subclasses, so this
-        # logger remains a drop-in if vLLM ever wires the factory call
-        # to invoke the constructor directly.
+        # StatLoggerBase requires this vLLM-compatible constructor.
         del vllm_config, engine_index
 
     def record(
@@ -141,16 +119,12 @@ def __init__(
         self.endpoint = endpoint
         self.component_gauges = component_gauges
         self.embedding_worker = embedding_worker
-        self.created_logger: Optional[DynamoStatLoggerPublisher] = None
+        self.created_loggers: list[DynamoStatLoggerPublisher] = []
 
     def create_stat_logger(self, dp_rank: int) -> StatLoggerBase:
-        # Embedding workers have no KV cache and no scheduler stats worth
-        # publishing -- short-circuit before constructing the chat-shaped
-        # WorkerMetricsPublisher and skipping the component_gauges check.
         if self.embedding_worker:
             return NoopStatLogger()
-        # component_gauges must be set by setup_vllm_engine() before vLLM
-        # calls create_stat_logger() during engine initialization.
+        # Engine setup installs gauges before vLLM creates its loggers.
         assert (
             self.component_gauges is not None
         ), "component_gauges must be set before creating stat loggers"
@@ -159,7 +133,7 @@ def create_stat_logger(self, dp_rank: int) -> StatLoggerBase:
             dp_rank=dp_rank,
             component_gauges=self.component_gauges,
         )
-        self.created_logger = logger
+        self.created_loggers.append(logger)
 
         return logger
 
@@ -168,9 +142,9 @@ def __call__(self, vllm_config: VllmConfig, dp_rank: int) -> StatLoggerBase:
 
     # TODO Remove once we publish metadata to shared storage
     def set_num_gpu_blocks_all(self, num_blocks: int) -> None:
-        if self.created_logger:
-            self.created_logger.set_num_gpu_block(num_blocks)
+        for logger in self.created_loggers:
+            logger.set_num_gpu_block(num_blocks)
 
     def init_publish(self) -> None:
-        if self.created_logger:
-            self.created_logger.init_publish()
+        for logger in self.created_loggers:
+            logger.init_publish()

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

@@ -1,15 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-"""Unit tests for the vLLM stat-logger factory.
-
-These tests focus on the embedding-worker gating path: vLLM workers run
-either a chat/decode engine or a pooling (embedding) engine, and the
-chat-shaped Prometheus collectors are only meaningful on the former.
-The factory is the single seam where vLLM calls into dynamo per dp_rank,
-so it is also the seam where the embedding worker must short-circuit
-the chat-shaped pipeline.
-"""
+"""Unit tests for the vLLM stat-logger factory."""
 
 from types import SimpleNamespace
 from unittest.mock import Mock
@@ -33,17 +25,6 @@
 
 
 def test_factory_returns_noop_logger_for_embedding_worker(monkeypatch):
-    """``create_stat_logger`` returns a ``NoopStatLogger`` on the
-    embedding path -- no ``DynamoStatLoggerPublisher`` /
-    ``WorkerMetricsPublisher`` / NATS endpoint construction.
-
-    Why this matters: ``DynamoStatLoggerPublisher.__init__`` schedules a
-    ``create_endpoint`` task on the runtime and registers chat-shaped
-    publish callbacks. On a pooling engine there is no kv_cache_usage to
-    publish (vLLM never emits ``SchedulerStats``) and the endpoint is
-    never queried -- so the factory must not construct it at all.
-    """
-
     def _explode(*_a, **_kw):
         raise AssertionError(
             "embedding-worker path must not construct DynamoStatLoggerPublisher"
@@ -59,22 +40,12 @@ def _explode(*_a, **_kw):
     logger = factory.create_stat_logger(dp_rank=0)
 
     assert isinstance(logger, NoopStatLogger)
-    # Embedding factory never tracks a created chat logger, so the
-    # downstream ``init_publish`` / ``set_num_gpu_blocks_all`` calls in
-    # the chat path are safe no-ops if anyone ever wires them on the
-    # embedding branch by mistake.
-    assert factory.created_logger is None
+    assert factory.created_loggers == []
 
 
 def test_noop_stat_logger_record_is_safe_with_none_stats():
-    """vLLM calls ``record`` every iteration even when the engine has
-    nothing useful to report. The chat-path publisher returns early on
-    ``scheduler_stats is None``; the embedding noop must accept the same
-    shape (and the variadic mm/engine_idx args vLLM passes) without
-    raising."""
     logger = NoopStatLogger()
 
-    # Mirrors the call shape from vllm/v1/metrics/loggers.py.
     logger.record(None, None)
     logger.record(None, None, None, 0)
     logger.record(
@@ -87,26 +58,17 @@ def test_noop_stat_logger_record_is_safe_with_none_stats():
 
 
 def test_factory_embedding_flag_skips_component_gauges_assert():
-    """On the chat path the factory asserts
-    ``component_gauges is not None`` because ``setup_vllm_engine`` is
-    responsible for setting it before vLLM invokes the factory. The
-    embedding path skips that step entirely (no chat-shaped gauges to
-    register), so the factory must not blow up when it stays None."""
     factory = StatLoggerFactory(
         endpoint=SimpleNamespace(),
         embedding_worker=True,
     )
     assert factory.component_gauges is None
 
-    # Would AssertionError on the chat path; must succeed here.
     logger = factory.create_stat_logger(dp_rank=0)
     assert isinstance(logger, NoopStatLogger)
 
 
 def test_factory_default_is_chat_path(monkeypatch):
-    """Sibling check: the default (``embedding_worker=False``) still
-    constructs ``DynamoStatLoggerPublisher`` so the gating doesn't
-    accidentally swallow the chat path."""
     constructed = []
 
     def _fake_publisher(*args, **kwargs):
@@ -125,3 +87,47 @@ def _fake_publisher(*args, **kwargs):
     assert constructed[0]["endpoint"] is endpoint
     assert constructed[0]["dp_rank"] == 3
     assert constructed[0]["component_gauges"] is component_gauges
+
+
+@pytest.mark.pre_merge
+@pytest.mark.gpu_0
+@pytest.mark.unit
+@pytest.mark.vllm
+@pytest.mark.core
+def test_factory_initializes_every_data_parallel_logger(monkeypatch):
+    loggers = [Mock(spec=DynamoStatLoggerPublisher) for _ in range(2)]
+    monkeypatch.setattr(
+        publisher_mod,
+        "DynamoStatLoggerPublisher",
+        Mock(side_effect=loggers),
+    )
+    factory = StatLoggerFactory(
+        endpoint=SimpleNamespace(),
+        component_gauges=SimpleNamespace(),
+    )
+
+    factory.create_stat_logger(dp_rank=0)
+    factory.create_stat_logger(dp_rank=1)
+    factory.set_num_gpu_blocks_all(9)
+    factory.init_publish()
+
+    for logger in loggers:
+        logger.set_num_gpu_block.assert_called_once_with(9)
+        logger.init_publish.assert_called_once_with()
+
+
+@pytest.mark.pre_merge
+@pytest.mark.gpu_0
+@pytest.mark.unit
+@pytest.mark.vllm
+@pytest.mark.core
+def test_initial_publish_uses_configured_gpu_block_count():
+    logger = DynamoStatLoggerPublisher.__new__(DynamoStatLoggerPublisher)
+    logger.inner = Mock()
+    logger.dp_rank = 1
+    logger.component_gauges = Mock()
+    logger.num_gpu_block = 9
+
+    logger.init_publish()
+
+    logger.component_gauges.set_total_blocks.assert_called_once_with("1", 9)

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

@@ -6,6 +6,7 @@
 import asyncio
 import json
 import logging
+from dataclasses import dataclass
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, Mock
 
@@ -1068,6 +1069,141 @@ async def _noop(*_args, **_kwargs) -> None:
         assert len(shutdown_endpoints) == 6
 
 
+@pytest.mark.asyncio
+@pytest.mark.pre_merge
+@pytest.mark.gpu_0
+@pytest.mark.unit
+@pytest.mark.vllm
+@pytest.mark.core
+async def test_prefill_worker_wires_stat_logger_and_seeds_kv_gauges(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    stop_after_register = RuntimeError("stop-after-register")
+
+    class RecordingStatLoggerFactory:
+        def __init__(self, *, endpoint, component_gauges=None):
+            captured.factory_endpoint = endpoint
+            captured.component_gauges = component_gauges
+            self.num_gpu_blocks: list[int] = []
+            self.init_publish_calls = 0
+            captured.factory = self
+
+        def set_num_gpu_blocks_all(self, num_blocks: int) -> None:
+            self.num_gpu_blocks.append(num_blocks)
+
+        def init_publish(self) -> None:
+            self.init_publish_calls += 1
+
+    @dataclass
+    class PrefillCapture:
+        factory_endpoint: object | None = None
+        component_gauges: object | None = None
+        factory: RecordingStatLoggerFactory | None = None
+        model_input: ModelInput | None = None
+        model_type: ModelType | None = None
+        worker_type: WorkerType | None = None
+        needs: list[list[WorkerType]] | None = None
+
+    captured = PrefillCapture()
+
+    async def fake_register_vllm_model(
+        model_input,
+        model_type,
+        endpoint,
+        config,
+        engine_client,
+        vllm_config,
+        worker_type,
+        needs,
+    ) -> None:
+        captured.model_input = model_input
+        captured.model_type = model_type
+        captured.worker_type = worker_type
+        captured.needs = needs
+        raise stop_after_register
+
+    engine_client = Mock()
+    vllm_config = Mock()
+    vllm_config.additional_config = {}
+    vllm_config.cache_config.num_gpu_blocks = 9
+    engine_tuple: EngineSetupResult = (
+        engine_client,
+        vllm_config,
+        Mock(),
+        "/tmp/prom",
+        Mock(),
+    )
+    setup_vllm_engine = Mock(return_value=engine_tuple)
+
+    factory = WorkerFactory(
+        setup_vllm_engine_fn=setup_vllm_engine,
+        setup_kv_event_publisher_fn=Mock(return_value=None),
+        register_vllm_model_fn=fake_register_vllm_model,
+        setup_fpm_relay_fn=Mock(return_value=None),
+        setup_metrics_collection_fn=Mock(),
+    )
+    factory._maybe_get_encode_worker_client = AsyncMock(return_value=None)  # type: ignore[assignment]
+    factory._maybe_wait_for_failover_lock = AsyncMock()  # type: ignore[assignment]
+    factory.register_engine_routes = Mock()  # type: ignore[assignment]
+
+    mock_handler = Mock(embedding_cache_manager=None)
+    monkeypatch.setattr(
+        "dynamo.vllm.worker_factory.PrefillWorkerHandler",
+        Mock(return_value=mock_handler),
+    )
+    monkeypatch.setattr(
+        "dynamo.vllm.worker_factory.StatLoggerFactory",
+        RecordingStatLoggerFactory,
+    )
+    monkeypatch.setattr(
+        "dynamo.vllm.worker_factory.get_dp_range_for_worker", lambda _config: (0, 2)
+    )
+
+    async def _noop(*_args, **_kwargs) -> None:
+        return None
+
+    monkeypatch.setattr(
+        "dynamo.vllm.worker_factory.configure_kv_event_block_size", _noop
+    )
+
+    config = _make_config(
+        disaggregation_mode=DisaggregationMode.PREFILL,
+        route_to_encoder=False,
+        use_vllm_tokenizer=False,
+        namespace="dyn",
+        component="prefill",
+        endpoint="generate",
+        served_model_name="m",
+        model="m",
+        frontend_decoding=False,
+        enable_multimodal=False,
+        enable_rl=False,
+        engine_args=SimpleNamespace(enable_lora=False),
+    )
+    endpoint = Mock(connection_id=Mock(return_value="cid"))
+    runtime = Mock()
+    runtime.endpoint.return_value = endpoint
+
+    with pytest.raises(RuntimeError, match="stop-after-register"):
+        await factory._create_prefill_worker(
+            runtime,
+            config,
+            asyncio.Event(),
+            [],
+        )
+
+    recording_factory = captured.factory
+    assert recording_factory is not None
+    setup_vllm_engine.assert_called_once_with(
+        config,
+        recording_factory,
+        fpm_worker_id="cid",
+    )
+    assert captured.factory_endpoint is endpoint
+    assert recording_factory.num_gpu_blocks == [4]
+    assert recording_factory.init_publish_calls == 1
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("lora_enabled", [True, False])
 async def test_prefill_serves_lora_lifecycle_endpoints_when_enabled(