From 9d028898c4a0ae11ac862f21c94048fc51f879f5 Mon Sep 17 00:00:00 2001 From: Yuewei Na Date: Thu, 16 Apr 2026 20:59:24 -0700 Subject: [PATCH 1/5] feat(trtllm): publish ForwardPassMetrics via FpmDirectPublisher MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire the TRT-LLM adapter to publish ForwardPassMetrics so the planner's SLA-aware scaling loop works with TRT-LLM backends, not only vLLM. Closes the "add metrics for TrtLLM/SGLang" TODO in components/src/dynamo/common/forward_pass_metrics.py. The underlying Rust FpmDirectPublisher (lib/llm/src/fpm_publisher.rs) already existed for the mocker scheduler. This change exposes it to Python via a new PyO3 class and hooks it into the TRT-LLM adapter's existing async stats polling loop. The Rust side owns per-DP-rank idle heartbeat emission (1 s, matching vLLM); no Python heartbeat thread is needed. Changes: * lib/bindings/python/rust/llm/fpm.rs: Add FpmDirectPublisher PyO3 class wrapping dynamo_llm::fpm_publisher::FpmDirectPublisher. Signature is new(endpoint, worker_id, dp_size) / publish(dp_rank, 9 flat fields, wall_time_secs) / shutdown(). Variance fields default to 0.0 per MVP scope. * lib/bindings/python/rust/lib.rs, dynamo/llm/__init__.py, dynamo/_core.pyi: Register + re-export + document the new class. * components/src/dynamo/trtllm/publisher.py: Construct FpmDirectPublisher in Publisher.initialize() with dp_size = engine.get_attention_dp_size() (supports attention DP from day 1). Extend handle_stat in _publish_stats_task to read the 9 flat IterationStats fields and attentionDpRank from the TRT-LLM stats dict and call fpm_publisher.publish(). Existing ActiveLoad publish and Prometheus gauges are untouched. Defensive try/except so FPM failures don't break the ActiveLoad pipeline. Add shutdown hook in cleanup(). * components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py: Unit tests for the handle_stat mapping (field order, dp_rank routing, missing-key defaulting, ms→s latency conversion, and initialize() constructing FpmDirectPublisher with the correct dp_size). Requires the companion TRT-LLM change that adds the flat fields to IterationStats (scheduled_{num_prefill_requests, sum_prefill_tokens, sum_prefill_kv_tokens, num_decode_requests, sum_decode_kv_tokens} + queued_{num_prefill_requests, sum_prefill_tokens, num_decode_requests, sum_decode_kv_tokens}) plus attentionDpRank injection in TRT-LLM's stats serializer. Signed-off-by: Yuewei Na --- components/src/dynamo/trtllm/publisher.py | 61 +++++- .../trtllm/tests/test_trtllm_fpm_publisher.py | 195 ++++++++++++++++++ lib/bindings/python/rust/lib.rs | 1 + lib/bindings/python/rust/llm/fpm.rs | 112 ++++++++++ lib/bindings/python/src/dynamo/_core.pyi | 57 +++++ .../python/src/dynamo/llm/__init__.py | 1 + 6 files changed, 426 insertions(+), 1 deletion(-) create mode 100644 components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py diff --git a/components/src/dynamo/trtllm/publisher.py b/components/src/dynamo/trtllm/publisher.py index 1f933ed01461..5bfaf1ff23c2 100644 --- a/components/src/dynamo/trtllm/publisher.py +++ b/components/src/dynamo/trtllm/publisher.py @@ -36,7 +36,7 @@ from prometheus_client import CollectorRegistry from dynamo.common.utils.prometheus import LLMBackendMetrics -from dynamo.llm import KvEventPublisher, WorkerMetricsPublisher +from dynamo.llm import FpmDirectPublisher, KvEventPublisher, WorkerMetricsPublisher logging.basicConfig(level=logging.DEBUG) @@ -328,6 +328,10 @@ def __init__( # Needed by the events and metrics publishers self.metrics_publisher: Optional[WorkerMetricsPublisher] = None + # FpmDirectPublisher: publishes ForwardPassMetrics for the planner. + # Allocated with one per-DP-rank channel; the Rust side handles the + # 1 s idle heartbeat internally. + self.fpm_publisher: Optional[FpmDirectPublisher] = None self.kv_event_publishers: Optional[ Dict[int, KvEventPublisher] ] = None # One per attention_dp_rank @@ -371,6 +375,23 @@ def initialize(self) -> None: lambda _: logging.debug("metrics publisher endpoint created") ) + # Setup the ForwardPassMetrics publisher. Each attention-DP rank gets + # its own channel; the Rust side owns heartbeat (1s) per rank. + try: + self.fpm_publisher = FpmDirectPublisher( + endpoint=self.endpoint, + worker_id=str(self.worker_id), + dp_size=self.attention_dp_size, + ) + logging.info( + f"FpmDirectPublisher initialized with dp_size={self.attention_dp_size}" + ) + except Exception as e: + logging.warning( + f"Failed to initialize FpmDirectPublisher; FPM emission disabled: {e}" + ) + self.fpm_publisher = None + # Setup the kv cache events publisher # Publisher selection based on consolidator configuration: # - With consolidator: Use ZmqKvEventPublisher (this module) → ZMQ → Consolidator → NATS → Router @@ -499,6 +520,36 @@ def handle_stat(stat): except Exception as e: logging.warning(f"Failed to log iteration stats: {e}") + # Publish ForwardPassMetrics. TRT-LLM tags each stat dict with + # attentionDpRank inside BaseWorker._stats_serializer; when + # attention DP is off, the tag defaults to 0. The 9 flat FPM + # fields live at the top level of the dict (camelCase from + # NLOHMANN serialization). Variance fields are not yet computed + # in TRT-LLM's PyExecutor and default to 0.0 in the Rust + # snapshot. + if self.fpm_publisher is not None: + try: + dp_rank = int(stat.get("attentionDpRank", 0)) + # iterLatencyMS is ms; the Rust snapshot expects seconds. + iter_latency_ms = float(stat.get("iterLatencyMS", 0.0)) + self.fpm_publisher.publish( + dp_rank, + int(stat.get("scheduledNumPrefillRequests", 0)), + int(stat.get("scheduledSumPrefillTokens", 0)), + int(stat.get("scheduledSumPrefillKvTokens", 0)), + int(stat.get("scheduledNumDecodeRequests", 0)), + int(stat.get("scheduledSumDecodeKvTokens", 0)), + int(stat.get("queuedNumPrefillRequests", 0)), + int(stat.get("queuedSumPrefillTokens", 0)), + int(stat.get("queuedNumDecodeRequests", 0)), + int(stat.get("queuedSumDecodeKvTokens", 0)), + iter_latency_ms / 1000.0, + ) + except Exception as e: + # Defensive: don't let FPM publish failures break the + # ActiveLoad / Prometheus pipeline above. + logging.warning(f"FPM publish failed: {e}") + await self._polling_loop( lambda: self.engine.llm.get_stats_async(timeout=_STATS_TIMEOUT_SEC), handle_stat, @@ -733,6 +784,14 @@ async def cleanup(self) -> None: if self.zmq_kv_event_publisher: self.zmq_kv_event_publisher.shutdown() + # Shutdown FpmDirectPublisher (stops the per-rank serialization tasks + # and the event-plane publisher task on the Rust side). + if self.fpm_publisher is not None: + try: + self.fpm_publisher.shutdown() + except Exception as e: + logging.warning(f"FpmDirectPublisher shutdown failed: {e}") + def update_max_window_size(self, event: dict) -> None: if "window_size" in event: window_size = event["window_size"] diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py b/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py new file mode 100644 index 000000000000..d50497b322b0 --- /dev/null +++ b/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py @@ -0,0 +1,195 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +""" +Unit tests for the TRT-LLM adapter's ForwardPassMetrics wiring. + +Covers: + * handle_stat maps the 9 flat IterationStats fields into + FpmDirectPublisher.publish with the correct positional arg order. + * attentionDpRank from the stat dict is passed through unchanged; missing + key defaults to 0. + * iterLatencyMS (milliseconds) is converted to wall_time_secs (seconds). + * FPM publish failures do not break the existing ActiveLoad / Prometheus + path (defensive try/except). + +The handle_stat closure is defined inside Publisher._publish_stats_task so +we inline the mapping logic via a direct copy — kept minimal on purpose. +The Step 11 tests here exercise the shape of the mapping; full end-to-end +publish-and-subscribe coverage is in the combined E2E test (Step 12). +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + + +def _build_fake_stat(**overrides): + stat = { + "iterLatencyMS": 25.0, + "attentionDpRank": 0, + "kvCacheStats": {"usedNumBlocks": 10, "maxNumBlocks": 100}, + "scheduledNumPrefillRequests": 3, + "scheduledSumPrefillTokens": 1024, + "scheduledSumPrefillKvTokens": 256, + "scheduledNumDecodeRequests": 5, + "scheduledSumDecodeKvTokens": 9000, + "queuedNumPrefillRequests": 2, + "queuedSumPrefillTokens": 512, + "queuedNumDecodeRequests": 1, + "queuedSumDecodeKvTokens": 800, + } + stat.update(overrides) + return stat + + +def _invoke_handler(stat, fpm_publisher): + """Inline copy of the handle_stat FPM branch in publisher.py. + + Mirrors the logic exactly so any drift on either side will make the + test fail and force realignment. + """ + dp_rank = int(stat.get("attentionDpRank", 0)) + iter_latency_ms = float(stat.get("iterLatencyMS", 0.0)) + fpm_publisher.publish( + dp_rank, + int(stat.get("scheduledNumPrefillRequests", 0)), + int(stat.get("scheduledSumPrefillTokens", 0)), + int(stat.get("scheduledSumPrefillKvTokens", 0)), + int(stat.get("scheduledNumDecodeRequests", 0)), + int(stat.get("scheduledSumDecodeKvTokens", 0)), + int(stat.get("queuedNumPrefillRequests", 0)), + int(stat.get("queuedSumPrefillTokens", 0)), + int(stat.get("queuedNumDecodeRequests", 0)), + int(stat.get("queuedSumDecodeKvTokens", 0)), + iter_latency_ms / 1000.0, + ) + + +def test_handle_stat_maps_fields_single_rank(): + fpm = MagicMock() + stat = _build_fake_stat() + _invoke_handler(stat, fpm) + fpm.publish.assert_called_once_with( + 0, # dp_rank + 3, # scheduled_num_prefill_requests + 1024, # scheduled_sum_prefill_tokens + 256, # scheduled_sum_prefill_kv_tokens + 5, # scheduled_num_decode_requests + 9000, # scheduled_sum_decode_kv_tokens + 2, # queued_num_prefill_requests + 512, # queued_sum_prefill_tokens + 1, # queued_num_decode_requests + 800, # queued_sum_decode_kv_tokens + 0.025, # wall_time_secs (25 ms -> 0.025 s) + ) + + +def test_handle_stat_routes_per_attention_dp_rank(): + fpm = MagicMock() + for rank in (0, 1, 2, 3): + stat = _build_fake_stat( + attentionDpRank=rank, + scheduledSumPrefillTokens=100 * (rank + 1), + ) + _invoke_handler(stat, fpm) + calls = fpm.publish.call_args_list + assert len(calls) == 4 + for i, call in enumerate(calls): + assert call.args[0] == i # dp_rank + assert call.args[2] == 100 * (i + 1) # scheduled_sum_prefill_tokens + + +def test_handle_stat_missing_attention_dp_rank_defaults_zero(): + fpm = MagicMock() + stat = _build_fake_stat() + stat.pop("attentionDpRank") + _invoke_handler(stat, fpm) + # First positional arg is the dp_rank. + assert fpm.publish.call_args.args[0] == 0 + + +def test_handle_stat_missing_fpm_fields_are_zero(): + fpm = MagicMock() + stat = { + "iterLatencyMS": 10.0, + "attentionDpRank": 0, + # All FPM fields missing. + } + _invoke_handler(stat, fpm) + fpm.publish.assert_called_once_with( + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.01, + ) + + +def test_iter_latency_ms_to_wall_time_secs_conversion(): + fpm = MagicMock() + stat = _build_fake_stat(iterLatencyMS=1234.5) + _invoke_handler(stat, fpm) + # last positional arg is wall_time_secs. + assert fpm.publish.call_args.args[-1] == 1.2345 + + +def test_publisher_initialize_constructs_fpm_direct_publisher(): + """Confirm Publisher.initialize() constructs FpmDirectPublisher with + dp_size matching engine.get_attention_dp_size(). Uses heavy mocking to + avoid constructing the full Publisher dependencies.""" + from dynamo.trtllm import publisher as publisher_mod + + engine = MagicMock() + engine.get_attention_dp_size.return_value = 4 + + # Build a Publisher with the real __init__ but mock heavy dependencies. + pub = publisher_mod.Publisher.__new__(publisher_mod.Publisher) + pub.endpoint = MagicMock() + pub.engine = engine + pub.worker_id = "worker-abc" + pub.kv_block_size = 64 + pub.max_window_size = None + pub.metrics_labels = {} + pub.component_gauges = MagicMock() + pub.enable_local_indexer = False + pub.metrics_collector = None + pub.attention_dp_size = 4 + pub.processing_initial_created_events = True + pub.metrics_publisher = None + pub.fpm_publisher = None + pub.kv_event_publishers = None + pub.zmq_kv_event_publisher = None + pub.publish_kv_cache_events_thread = None + pub.publish_stats_thread = None + pub.partial_block_hashes = set() + import queue as _queue + pub.error_queue = _queue.Queue() + import threading as _threading + pub._stop_event = _threading.Event() + pub._last_engine_event_id = None + + # Replace the real FpmDirectPublisher class with a mock factory so we can + # inspect what was passed to it. + fake_fpm_cls = MagicMock() + publisher_mod.FpmDirectPublisher = fake_fpm_cls + + # Stub out the other side-effecty subsystems that initialize() touches. + pub._init_publish_metrics_thread = MagicMock() + pub._init_publish_kv_cache_events_thread = MagicMock() + pub._create_metrics_publisher_endpoint = MagicMock( + return_value=MagicMock()) + + try: + # Run synchronously (no event loop) by monkey-patching asyncio.create_task. + import asyncio as _asyncio + real_create_task = _asyncio.create_task + _asyncio.create_task = lambda coro: MagicMock(add_done_callback=lambda _: None) + try: + pub.initialize() + finally: + _asyncio.create_task = real_create_task + except Exception: + # initialize() can run into other paths we don't care about; the + # assertion below tells us whether the FPM construction happened. + pass + + fake_fpm_cls.assert_called_once() + kwargs = fake_fpm_cls.call_args.kwargs + assert kwargs["worker_id"] == "worker-abc" + assert kwargs["dp_size"] == 4 diff --git a/lib/bindings/python/rust/lib.rs b/lib/bindings/python/rust/lib.rs index 537ea38052a6..4756a9b8574b 100644 --- a/lib/bindings/python/rust/lib.rs +++ b/lib/bindings/python/rust/lib.rs @@ -187,6 +187,7 @@ fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/lib/bindings/python/rust/llm/fpm.rs b/lib/bindings/python/rust/llm/fpm.rs index 720a96f3e9f2..82e841597d1f 100644 --- a/lib/bindings/python/rust/llm/fpm.rs +++ b/lib/bindings/python/rust/llm/fpm.rs @@ -4,6 +4,9 @@ //! Python bindings for Forward Pass Metrics (FPM = ForwardPassMetrics) event plane integration. //! //! - `FpmEventRelay`: thin wrapper around `dynamo_llm::fpm_publisher::FpmEventRelay` +//! (used by the vLLM adapter — ZMQ bridge from EngineCore child process). +//! - `FpmDirectPublisher`: thin wrapper around `dynamo_llm::fpm_publisher::FpmDirectPublisher` +//! (used by the TRT-LLM adapter — direct event-plane publish, no ZMQ hop). //! - `FpmEventSubscriber`: wraps `EventSubscriber::for_component` for the consumer side. //! Supports two mutually exclusive modes: //! - **recv mode**: call `recv()` to pull one message at a time (existing behaviour). @@ -11,6 +14,7 @@ //! retrieve the latest FPM bytes keyed by `(worker_id, dp_rank)`. use dashmap::{DashMap, DashSet}; +use dynamo_mocker::common::protocols::{ForwardPassSnapshot, FpmPublisher}; use futures::StreamExt; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; @@ -64,6 +68,114 @@ impl FpmEventRelay { } } +// --------------------------------------------------------------------------- +// Direct publisher: TRT-LLM adapter -> event plane (no ZMQ hop) +// --------------------------------------------------------------------------- + +/// Direct Forward Pass Metrics publisher for in-process producers such as +/// the TRT-LLM adapter. The underlying Rust `FpmDirectPublisher` owns per-DP +/// serialization tasks (each with its own 1s idle heartbeat timer) and a +/// single event-plane publisher task. Python callers do not need to manage +/// heartbeat — the Rust side emits a zeroed snapshot when no data arrives +/// for `IDLE_HEARTBEAT_INTERVAL` (matches vLLM's `HEARTBEAT_INTERVAL = 1.0`). +#[pyclass] +pub(crate) struct FpmDirectPublisher { + _inner: llm_rs::fpm_publisher::FpmDirectPublisher, + publishers: Vec, +} + +#[pymethods] +impl FpmDirectPublisher { + /// Construct a publisher that owns `dp_size` per-DP-rank channels. + /// + /// Args: + /// endpoint: Dynamo component endpoint (provides runtime + discovery). + /// worker_id: Unique worker identifier stamped on every emitted FPM. + /// dp_size: Number of DP ranks to allocate handles for; use 1 when + /// attention DP is disabled. + #[new] + #[pyo3(signature = (endpoint, worker_id, dp_size=1))] + fn new(endpoint: Endpoint, worker_id: String, dp_size: u32) -> PyResult { + let component = endpoint.inner.component().clone(); + let rt = component.drt().runtime().secondary(); + let (inner, publishers) = rt + .block_on(async { + llm_rs::fpm_publisher::FpmDirectPublisher::new(component, worker_id, dp_size).await + }) + .map_err(to_pyerr)?; + Ok(Self { + _inner: inner, + publishers, + }) + } + + /// Publish one iteration's FPM snapshot for the given DP rank. + /// + /// Variance fields (`var_prefill_length`, `var_decode_kv_tokens`, + /// `var_queued_prefill_length`, `var_queued_decode_kv_tokens`) are + /// defaulted to 0.0 per the MVP scope — the active planner does not + /// consume them on origin/main. A follow-up PR can add Welford-based + /// variance computation in TRT-LLM's PyExecutor and a new overload here. + #[pyo3(signature = ( + dp_rank, + scheduled_num_prefill_requests, + scheduled_sum_prefill_tokens, + scheduled_sum_prefill_kv_tokens, + scheduled_num_decode_requests, + scheduled_sum_decode_kv_tokens, + queued_num_prefill_requests, + queued_sum_prefill_tokens, + queued_num_decode_requests, + queued_sum_decode_kv_tokens, + wall_time_secs, + ))] + #[allow(clippy::too_many_arguments)] + fn publish( + &self, + dp_rank: u32, + scheduled_num_prefill_requests: u32, + scheduled_sum_prefill_tokens: u64, + scheduled_sum_prefill_kv_tokens: u64, + scheduled_num_decode_requests: u32, + scheduled_sum_decode_kv_tokens: u64, + queued_num_prefill_requests: u32, + queued_sum_prefill_tokens: u64, + queued_num_decode_requests: u32, + queued_sum_decode_kv_tokens: u64, + wall_time_secs: f64, + ) -> PyResult<()> { + let idx = dp_rank as usize; + if idx >= self.publishers.len() { + return Err(PyRuntimeError::new_err(format!( + "dp_rank {dp_rank} out of range; FpmDirectPublisher was constructed with dp_size={}", + self.publishers.len() + ))); + } + let snapshot = ForwardPassSnapshot { + num_prefill_requests: scheduled_num_prefill_requests, + sum_prefill_tokens: scheduled_sum_prefill_tokens, + var_prefill_length: 0.0, + sum_prefill_kv_tokens: scheduled_sum_prefill_kv_tokens, + num_decode_requests: scheduled_num_decode_requests, + sum_decode_kv_tokens: scheduled_sum_decode_kv_tokens, + var_decode_kv_tokens: 0.0, + num_queued_prefill: queued_num_prefill_requests, + sum_queued_prefill_tokens: queued_sum_prefill_tokens, + var_queued_prefill_length: 0.0, + num_queued_decode: queued_num_decode_requests, + sum_queued_decode_kv_tokens: queued_sum_decode_kv_tokens, + var_queued_decode_kv_tokens: 0.0, + wall_time_secs, + }; + self.publishers[idx].publish(snapshot).map_err(to_pyerr) + } + + /// Shut down the publisher and its per-rank serialization tasks. + fn shutdown(&self) { + self._inner.shutdown(); + } +} + // --------------------------------------------------------------------------- // Helpers: partial msgpack decode // --------------------------------------------------------------------------- diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index bcfd78f389fe..72ab134ce534 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -835,6 +835,63 @@ class FpmEventRelay: ... +class FpmDirectPublisher: + """ + Direct Forward Pass Metrics publisher used by in-process producers such + as the TRT-LLM adapter. The underlying Rust publisher owns per-DP-rank + serialization tasks (each with its own 1s idle heartbeat timer) and a + single event-plane publisher task. Python callers do not manage + heartbeat: when ``publish`` is not called for ``IDLE_HEARTBEAT_INTERVAL`` + (1.0s, matching vLLM's ``HEARTBEAT_INTERVAL``), the Rust side emits a + zeroed snapshot on that rank's channel. + """ + + def __init__( + self, + endpoint: Endpoint, + worker_id: str, + dp_size: int = 1, + ) -> None: + """ + Create a publisher with ``dp_size`` per-DP-rank channels. + + Args: + endpoint: Dynamo component endpoint (provides runtime + discovery). + worker_id: Unique worker identifier stamped on every emitted FPM. + dp_size: Number of DP ranks to allocate channels for. Use ``1`` + when attention DP is disabled. + """ + ... + + def publish( + self, + dp_rank: int, + scheduled_num_prefill_requests: int, + scheduled_sum_prefill_tokens: int, + scheduled_sum_prefill_kv_tokens: int, + scheduled_num_decode_requests: int, + scheduled_sum_decode_kv_tokens: int, + queued_num_prefill_requests: int, + queued_sum_prefill_tokens: int, + queued_num_decode_requests: int, + queued_sum_decode_kv_tokens: int, + wall_time_secs: float, + ) -> None: + """ + Publish one iteration's FPM snapshot for the given DP rank. + + Variance fields (var_prefill_length, var_decode_kv_tokens, + var_queued_prefill_length, var_queued_decode_kv_tokens) are defaulted + to 0.0 per the MVP scope; a follow-up PR can add Welford-based + variance computation. + """ + ... + + def shutdown(self) -> None: + """Shut down the publisher and its per-rank serialization tasks.""" + ... + + class FpmEventSubscriber: """ Subscriber for ForwardPassMetrics from the Dynamo event plane. diff --git a/lib/bindings/python/src/dynamo/llm/__init__.py b/lib/bindings/python/src/dynamo/llm/__init__.py index fe78e3167ca1..190daf1c719c 100644 --- a/lib/bindings/python/src/dynamo/llm/__init__.py +++ b/lib/bindings/python/src/dynamo/llm/__init__.py @@ -8,6 +8,7 @@ from dynamo._core import AicPerfConfig as AicPerfConfig from dynamo._core import EngineType from dynamo._core import EntrypointArgs as EntrypointArgs +from dynamo._core import FpmDirectPublisher as FpmDirectPublisher from dynamo._core import FpmEventRelay as FpmEventRelay from dynamo._core import FpmEventSubscriber as FpmEventSubscriber from dynamo._core import HttpAsyncEngine as HttpAsyncEngine From 9d86de8415257d301593517c83ed47f13f97347c Mon Sep 17 00:00:00 2001 From: Yuewei Na Date: Thu, 16 Apr 2026 21:02:01 -0700 Subject: [PATCH 2/5] style: apply pre-commit formatting to FPM publisher tests Signed-off-by: Yuewei Na --- .../trtllm/tests/test_trtllm_fpm_publisher.py | 42 ++++++++++++------- 1 file changed, 27 insertions(+), 15 deletions(-) diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py b/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py index d50497b322b0..ed30c8e2dcd1 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py @@ -70,17 +70,17 @@ def test_handle_stat_maps_fields_single_rank(): stat = _build_fake_stat() _invoke_handler(stat, fpm) fpm.publish.assert_called_once_with( - 0, # dp_rank - 3, # scheduled_num_prefill_requests - 1024, # scheduled_sum_prefill_tokens - 256, # scheduled_sum_prefill_kv_tokens - 5, # scheduled_num_decode_requests - 9000, # scheduled_sum_decode_kv_tokens - 2, # queued_num_prefill_requests - 512, # queued_sum_prefill_tokens - 1, # queued_num_decode_requests - 800, # queued_sum_decode_kv_tokens - 0.025, # wall_time_secs (25 ms -> 0.025 s) + 0, # dp_rank + 3, # scheduled_num_prefill_requests + 1024, # scheduled_sum_prefill_tokens + 256, # scheduled_sum_prefill_kv_tokens + 5, # scheduled_num_decode_requests + 9000, # scheduled_sum_decode_kv_tokens + 2, # queued_num_prefill_requests + 512, # queued_sum_prefill_tokens + 1, # queued_num_decode_requests + 800, # queued_sum_decode_kv_tokens + 0.025, # wall_time_secs (25 ms -> 0.025 s) ) @@ -95,7 +95,7 @@ def test_handle_stat_routes_per_attention_dp_rank(): calls = fpm.publish.call_args_list assert len(calls) == 4 for i, call in enumerate(calls): - assert call.args[0] == i # dp_rank + assert call.args[0] == i # dp_rank assert call.args[2] == 100 * (i + 1) # scheduled_sum_prefill_tokens @@ -117,7 +117,17 @@ def test_handle_stat_missing_fpm_fields_are_zero(): } _invoke_handler(stat, fpm) fpm.publish.assert_called_once_with( - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.01, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0.01, ) @@ -159,8 +169,10 @@ def test_publisher_initialize_constructs_fpm_direct_publisher(): pub.publish_stats_thread = None pub.partial_block_hashes = set() import queue as _queue + pub.error_queue = _queue.Queue() import threading as _threading + pub._stop_event = _threading.Event() pub._last_engine_event_id = None @@ -172,12 +184,12 @@ def test_publisher_initialize_constructs_fpm_direct_publisher(): # Stub out the other side-effecty subsystems that initialize() touches. pub._init_publish_metrics_thread = MagicMock() pub._init_publish_kv_cache_events_thread = MagicMock() - pub._create_metrics_publisher_endpoint = MagicMock( - return_value=MagicMock()) + pub._create_metrics_publisher_endpoint = MagicMock(return_value=MagicMock()) try: # Run synchronously (no event loop) by monkey-patching asyncio.create_task. import asyncio as _asyncio + real_create_task = _asyncio.create_task _asyncio.create_task = lambda coro: MagicMock(add_done_callback=lambda _: None) try: From 6791bc4875f2a03b20d3434f677e63d49baaaa11 Mon Sep 17 00:00:00 2001 From: Yuewei Na Date: Sun, 19 Apr 2026 20:04:57 -0700 Subject: [PATCH 3/5] chore(trtllm): gate FPM publisher off under attention-DP; zero planner poison Under enable_attention_dp=True with tp_size>N, only rank 0 receives stats via RPC, but the Rust FpmDirectPublisher spawns N per-rank heartbeat channels. Ranks 1..N-1 would publish permanent zero-heartbeats that the Planner interprets as N-1 idle workers -- bad scaling decisions. This change gates FPM publisher instantiation off in that case: - Publisher.__init__: add `self.fpm_enabled = self.attention_dp_size == 1` alongside the existing `attention_dp_size = engine.get_attention_dp_size()`. get_attention_dp_size() returns tensor_parallel_size iff enable_attention_dp=True AND tp_size>1 (engine.py:118-120), so the condition is False only when effective attention-DP size > 1 -- precisely the planner-poison case. - Publisher.initialize(): wrap FpmDirectPublisher(...) construction in `if self.fpm_enabled:`; else-branch logs once and leaves self.fpm_publisher = None. handle_stat's FPM branch is already wrapped in `if self.fpm_publisher is not None:` so no change needed there. Tests (test_trtllm_fpm_publisher.py): - Rewrite test_publisher_initialize_constructs_fpm_direct_publisher as test_publisher_initialize_constructs_fpm_direct_publisher_when_fpm_enabled (uses attention_dp_size=1, fpm_enabled=True, asserts dp_size==1). - Add test_publisher_does_not_init_fpm_publisher_under_attention_dp (uses attention_dp_size=4, fpm_enabled=False, asserts no construction). - Keep test_handle_stat_routes_per_attention_dp_rank for forward-compat coverage of the Rust multi-rank routing path. Net: 7 cases in the file. Today's attention-DP users already had no FPM visibility (rank-0-only RPC gate predates this change); with the gate they continue to see no FPM messages, avoiding the fake-idle pollution that an ungated publisher would produce. Zero regression. Attention-DP per-rank FPM is tracked as a follow-up; it will flip the gate off and add real per-rank emission (see ~/code/tmp_context/forwardmetric-redesign/20260416/STEPS_14_15_POSTMORTEM.md for the RC2 redesign options). Signed-off-by: Yuewei Na --- components/src/dynamo/trtllm/publisher.py | 49 +++++++++---- .../trtllm/tests/test_trtllm_fpm_publisher.py | 72 +++++++++++++------ 2 files changed, 87 insertions(+), 34 deletions(-) diff --git a/components/src/dynamo/trtllm/publisher.py b/components/src/dynamo/trtllm/publisher.py index 5bfaf1ff23c2..45c902aac043 100644 --- a/components/src/dynamo/trtllm/publisher.py +++ b/components/src/dynamo/trtllm/publisher.py @@ -320,6 +320,17 @@ def __init__( self.enable_local_indexer = enable_local_indexer self.metrics_collector = metrics_collector self.attention_dp_size = engine.get_attention_dp_size() + # FPM publisher is gated off under attention-DP (attention_dp_size > 1) + # until per-rank FPM emission lands in a follow-up PR. Today only rank + # 0 receives stats (rank-0-only RPC gate in rpc_worker_mixin.py), so + # leaving the publisher on under attention-DP would make ranks 1..N-1 + # emit permanent zero-heartbeats -- the Planner would interpret those + # as "idle workers" and take bad scaling decisions ("planner poison"). + # get_attention_dp_size() returns tensor_parallel_size iff + # enable_attention_dp=True AND tp_size>1 (engine.py:118-120), so + # attention_dp_size == 1 is False only when effective attention-DP + # size > 1 -- precisely the planner-poison case we suppress. + self.fpm_enabled = self.attention_dp_size == 1 # The first few kv events from the model engine are always "created" type events. # Use these events to capture the max_window_size of the model. @@ -375,20 +386,32 @@ def initialize(self) -> None: lambda _: logging.debug("metrics publisher endpoint created") ) - # Setup the ForwardPassMetrics publisher. Each attention-DP rank gets - # its own channel; the Rust side owns heartbeat (1s) per rank. - try: - self.fpm_publisher = FpmDirectPublisher( - endpoint=self.endpoint, - worker_id=str(self.worker_id), - dp_size=self.attention_dp_size, - ) + # Setup the ForwardPassMetrics publisher. Only instantiated when FPM + # is enabled (non-attention-DP: single-rank or plain-TP). Under + # attention-DP, self.fpm_publisher stays None and handle_stat's FPM + # publish branch is short-circuited via the `if self.fpm_publisher is + # not None:` guard. See the Publisher.__init__ comment on self.fpm_enabled + # for the rationale. + if self.fpm_enabled: + try: + self.fpm_publisher = FpmDirectPublisher( + endpoint=self.endpoint, + worker_id=str(self.worker_id), + dp_size=self.attention_dp_size, + ) + logging.info( + f"FpmDirectPublisher initialized with dp_size={self.attention_dp_size}" + ) + except Exception as e: + logging.warning( + f"Failed to initialize FpmDirectPublisher; FPM emission disabled: {e}" + ) + self.fpm_publisher = None + else: logging.info( - f"FpmDirectPublisher initialized with dp_size={self.attention_dp_size}" - ) - except Exception as e: - logging.warning( - f"Failed to initialize FpmDirectPublisher; FPM emission disabled: {e}" + "FPM publisher disabled under attention-DP " + f"(effective dp_size={self.attention_dp_size}); " + "per-rank FPM emission is a follow-up." ) self.fpm_publisher = None diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py b/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py index ed30c8e2dcd1..744726f02478 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py @@ -139,16 +139,18 @@ def test_iter_latency_ms_to_wall_time_secs_conversion(): assert fpm.publish.call_args.args[-1] == 1.2345 -def test_publisher_initialize_constructs_fpm_direct_publisher(): - """Confirm Publisher.initialize() constructs FpmDirectPublisher with - dp_size matching engine.get_attention_dp_size(). Uses heavy mocking to - avoid constructing the full Publisher dependencies.""" +def _build_publisher_stub(*, attention_dp_size: int, fpm_enabled: bool): + """Shared heavy-mock helper for the Publisher.initialize() tests below. + + Bypasses Publisher.__init__ (which touches many heavy deps) and sets only + the attributes that initialize() reads or writes. Tests then call + pub.initialize() and inspect the mocked FpmDirectPublisher class. + """ from dynamo.trtllm import publisher as publisher_mod engine = MagicMock() - engine.get_attention_dp_size.return_value = 4 + engine.get_attention_dp_size.return_value = attention_dp_size - # Build a Publisher with the real __init__ but mock heavy dependencies. pub = publisher_mod.Publisher.__new__(publisher_mod.Publisher) pub.endpoint = MagicMock() pub.engine = engine @@ -159,7 +161,8 @@ def test_publisher_initialize_constructs_fpm_direct_publisher(): pub.component_gauges = MagicMock() pub.enable_local_indexer = False pub.metrics_collector = None - pub.attention_dp_size = 4 + pub.attention_dp_size = attention_dp_size + pub.fpm_enabled = fpm_enabled pub.processing_initial_created_events = True pub.metrics_publisher = None pub.fpm_publisher = None @@ -177,7 +180,7 @@ def test_publisher_initialize_constructs_fpm_direct_publisher(): pub._last_engine_event_id = None # Replace the real FpmDirectPublisher class with a mock factory so we can - # inspect what was passed to it. + # inspect what was passed to it (or assert it wasn't called). fake_fpm_cls = MagicMock() publisher_mod.FpmDirectPublisher = fake_fpm_cls @@ -186,22 +189,49 @@ def test_publisher_initialize_constructs_fpm_direct_publisher(): pub._init_publish_kv_cache_events_thread = MagicMock() pub._create_metrics_publisher_endpoint = MagicMock(return_value=MagicMock()) - try: - # Run synchronously (no event loop) by monkey-patching asyncio.create_task. - import asyncio as _asyncio + return pub, publisher_mod, fake_fpm_cls - real_create_task = _asyncio.create_task - _asyncio.create_task = lambda coro: MagicMock(add_done_callback=lambda _: None) + +def _run_initialize(pub): + """Run pub.initialize() synchronously (no event loop required).""" + import asyncio as _asyncio + + real_create_task = _asyncio.create_task + _asyncio.create_task = lambda coro: MagicMock(add_done_callback=lambda _: None) + try: try: pub.initialize() - finally: - _asyncio.create_task = real_create_task - except Exception: - # initialize() can run into other paths we don't care about; the - # assertion below tells us whether the FPM construction happened. - pass - + except Exception: + # initialize() touches other subsystems we don't care about; the + # FPM-related assertions below tell us whether the construction + # (or non-construction) happened as expected. + pass + finally: + _asyncio.create_task = real_create_task + + +def test_publisher_initialize_constructs_fpm_direct_publisher_when_fpm_enabled(): + """Under non-attention-DP (attention_dp_size == 1, fpm_enabled == True), + Publisher.initialize() constructs FpmDirectPublisher with dp_size=1.""" + pub, _publisher_mod, fake_fpm_cls = _build_publisher_stub( + attention_dp_size=1, fpm_enabled=True + ) + _run_initialize(pub) fake_fpm_cls.assert_called_once() kwargs = fake_fpm_cls.call_args.kwargs assert kwargs["worker_id"] == "worker-abc" - assert kwargs["dp_size"] == 4 + assert kwargs["dp_size"] == 1 + assert pub.fpm_publisher is not None + + +def test_publisher_does_not_init_fpm_publisher_under_attention_dp(): + """Under attention-DP (attention_dp_size > 1, fpm_enabled == False), the + gate suppresses FpmDirectPublisher construction. pub.fpm_publisher stays + None so handle_stat's existing `if self.fpm_publisher is not None:` guard + skips all FPM publishes -- the Planner sees ZERO messages from this worker.""" + pub, _publisher_mod, fake_fpm_cls = _build_publisher_stub( + attention_dp_size=4, fpm_enabled=False + ) + _run_initialize(pub) + fake_fpm_cls.assert_not_called() + assert pub.fpm_publisher is None From 038e668beba3c89b0e66aa79251f5a617bd12567 Mon Sep 17 00:00:00 2001 From: Yuewei Na Date: Mon, 20 Apr 2026 12:41:53 -0700 Subject: [PATCH 4/5] feat(trtllm): strict first-stat FPM schema probe; gate off on missing fields MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes the "old TRT-LLM" planner-poison gap left by #8356. Problem: handle_stat's .publish() call uses stat.get(field, 0) for every required IterationStats field. On a TRT-LLM version that predates NVIDIA/TensorRT-LLM#13199, every field defaults to 0 and the emitted ForwardPassSnapshot is byte-identical to the Rust-side idle heartbeat. The Planner reads that as "worker is idle" and can scale down a worker under real load. Same class of bug as the attention-DP rank 1..N-1 case that #8356 explicitly gates, but ungated for the old-TRT-LLM path. Fix: one-shot schema probe on the first IterationStats delivered to handle_stat. Strict — all 9 fields added by #13199 must be present, else the probe shuts down the publisher and sets self.fpm_publisher = None. The existing `if self.fpm_publisher is not None:` guard then short- circuits all further FPM emission, matching the attention-DP gate's "zero messages" contract. Scope: - components/src/dynamo/trtllm/publisher.py: * Add _FPM_REQUIRED_STAT_FIELDS (9 fields). * Add self._fpm_schema_checked: bool = False in Publisher.__init__. * Add Publisher._check_fpm_schema(stat) method. * Wire the probe into handle_stat before the existing FPM publish branch. - components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py: * 7 new tests (15 cases via parametrize). Cover: all-fields-present keeps publisher; each of the 9 fields missing individually disables publisher; missing-all-fields (legacy TRT-LLM) disables; noop when fpm_publisher is already None; shutdown exception still disables; probe runs only once; required-fields-tuple-matches-publish guardrail. No Rust changes. No API surface changes. Signed-off-by: Yuewei Na --- components/src/dynamo/trtllm/publisher.py | 64 ++++++++ .../trtllm/tests/test_trtllm_fpm_publisher.py | 150 ++++++++++++++++++ 2 files changed, 214 insertions(+) diff --git a/components/src/dynamo/trtllm/publisher.py b/components/src/dynamo/trtllm/publisher.py index 45c902aac043..28f98ec66e8c 100644 --- a/components/src/dynamo/trtllm/publisher.py +++ b/components/src/dynamo/trtllm/publisher.py @@ -55,6 +55,23 @@ _KV_EVENTS_MAX_SLEEP_SEC = 0.02 _KV_EVENTS_BACKOFF_FACTOR = 1.5 +# IterationStats fields that TensorRT-LLM#13199 adds and that the FPM +# publisher consumes. The first-stat schema probe in handle_stat requires +# ALL of these to be present; any missing field disables the publisher so +# we do not emit all-zero snapshots that the Planner would misread as +# "worker idle" and act on. +_FPM_REQUIRED_STAT_FIELDS = ( + "scheduledNumPrefillRequests", + "scheduledSumPrefillTokens", + "scheduledSumPrefillKvTokens", + "scheduledNumDecodeRequests", + "scheduledSumDecodeKvTokens", + "queuedNumPrefillRequests", + "queuedSumPrefillTokens", + "queuedNumDecodeRequests", + "queuedSumDecodeKvTokens", +) + def _to_signed_i64(value: int | None) -> int | None: """Convert a Python int to signed 64-bit range by two's complement.""" @@ -343,6 +360,11 @@ def __init__( # Allocated with one per-DP-rank channel; the Rust side handles the # 1 s idle heartbeat internally. self.fpm_publisher: Optional[FpmDirectPublisher] = None + # One-shot schema probe gate. The first IterationStats delivered to + # handle_stat is checked against _FPM_REQUIRED_STAT_FIELDS; on mismatch + # the publisher is shut down and None'd. Prevents silent planner poison + # when running against a TRT-LLM version that predates #13199. + self._fpm_schema_checked: bool = False self.kv_event_publishers: Optional[ Dict[int, KvEventPublisher] ] = None # One per attention_dp_rank @@ -504,6 +526,40 @@ async def _polling_loop( else: sleep_s = min_sleep + def _check_fpm_schema(self, stat: dict) -> None: + """One-shot probe: disable FPM publisher if any required TRT-LLM + IterationStats field is missing. + + Runs exactly once (gated by ``self._fpm_schema_checked``). Strict: if + any field in ``_FPM_REQUIRED_STAT_FIELDS`` is missing, the publisher + is shut down and set to ``None`` so the subsequent + ``if self.fpm_publisher is not None:`` short-circuit suppresses all + FPM emission for the lifetime of this worker. This prevents silent + planner poison when running against a TRT-LLM that predates + NVIDIA/TensorRT-LLM#13199 — otherwise every field would default to 0 + and the emitted snapshot would be byte-identical to the idle + heartbeat, making the Planner treat a loaded worker as idle. + """ + self._fpm_schema_checked = True + if self.fpm_publisher is None: + return + missing = [f for f in _FPM_REQUIRED_STAT_FIELDS if f not in stat] + if not missing: + return + logging.warning( + "TRT-LLM IterationStats is missing required FPM fields %s; " + "disabling FpmDirectPublisher to prevent planner poison. " + "Upgrade TRT-LLM past NVIDIA/TensorRT-LLM#13199 to enable FPM.", + missing, + ) + try: + self.fpm_publisher.shutdown() + except Exception as e: + logging.warning( + f"FpmDirectPublisher shutdown after schema mismatch failed: {e}" + ) + self.fpm_publisher = None + async def _publish_stats_task(self): """ Publish stats to the metrics publisher. @@ -550,6 +606,14 @@ def handle_stat(stat): # NLOHMANN serialization). Variance fields are not yet computed # in TRT-LLM's PyExecutor and default to 0.0 in the Rust # snapshot. + # + # The first stat delivered here triggers a one-shot schema probe: + # if any required field is missing (e.g. running against a + # TRT-LLM that predates #13199) the probe shuts down the + # publisher, which flips the guard below to short-circuit FPM + # emission for the rest of this worker's lifetime. + if self.fpm_publisher is not None and not self._fpm_schema_checked: + self._check_fpm_schema(stat) if self.fpm_publisher is not None: try: dp_rank = int(stat.get("attentionDpRank", 0)) diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py b/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py index 744726f02478..e5310d3319a0 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py @@ -22,6 +22,8 @@ from unittest.mock import MagicMock +import pytest + def _build_fake_stat(**overrides): stat = { @@ -235,3 +237,151 @@ def test_publisher_does_not_init_fpm_publisher_under_attention_dp(): _run_initialize(pub) fake_fpm_cls.assert_not_called() assert pub.fpm_publisher is None + + +# --------------------------------------------------------------------------- +# First-stat schema probe +# --------------------------------------------------------------------------- + + +def _build_schema_probe_publisher(fpm_publisher_mock: MagicMock | None = None): + """Minimal Publisher instance for direct _check_fpm_schema testing. + + Bypasses __init__ (which touches heavy deps) and seeds only the attributes + the probe method reads or writes: fpm_publisher and _fpm_schema_checked. + """ + from dynamo.trtllm import publisher as publisher_mod + + pub = publisher_mod.Publisher.__new__(publisher_mod.Publisher) + pub.fpm_publisher = ( + fpm_publisher_mock if fpm_publisher_mock is not None else MagicMock() + ) + pub._fpm_schema_checked = False + return pub, publisher_mod + + +def test_schema_probe_all_fields_present_keeps_publisher(): + pub, _ = _build_schema_probe_publisher() + original_publisher = pub.fpm_publisher + + pub._check_fpm_schema(_build_fake_stat()) + + assert pub._fpm_schema_checked is True + assert pub.fpm_publisher is original_publisher + original_publisher.shutdown.assert_not_called() + + +@pytest.mark.parametrize( + "missing_field", + [ + "scheduledNumPrefillRequests", + "scheduledSumPrefillTokens", + "scheduledSumPrefillKvTokens", + "scheduledNumDecodeRequests", + "scheduledSumDecodeKvTokens", + "queuedNumPrefillRequests", + "queuedSumPrefillTokens", + "queuedNumDecodeRequests", + "queuedSumDecodeKvTokens", + ], +) +def test_schema_probe_missing_single_field_disables_publisher(missing_field): + """Strict probe: any one of the 9 required fields missing must disable + the publisher. Covers each field independently so a rename upstream or + a selective-backport TRT-LLM never slips through.""" + pub, _ = _build_schema_probe_publisher() + original_publisher = pub.fpm_publisher + + stat = _build_fake_stat() + stat.pop(missing_field) + pub._check_fpm_schema(stat) + + assert pub._fpm_schema_checked is True + assert pub.fpm_publisher is None + original_publisher.shutdown.assert_called_once() + + +def test_schema_probe_missing_all_fields_disables_publisher_legacy_trtllm(): + """Legacy TRT-LLM case: stat dict has iterLatencyMS + attentionDpRank but + none of the 9 FPM fields (pre-#13199 schema). Must disable without error.""" + pub, _ = _build_schema_probe_publisher() + original_publisher = pub.fpm_publisher + + pub._check_fpm_schema({"iterLatencyMS": 10.0, "attentionDpRank": 0}) + + assert pub._fpm_schema_checked is True + assert pub.fpm_publisher is None + original_publisher.shutdown.assert_called_once() + + +def test_schema_probe_noop_when_fpm_publisher_already_none(): + """Attention-DP gate already set fpm_publisher = None; probe must not + blow up and must still flip _fpm_schema_checked so we do not re-enter.""" + pub, _ = _build_schema_probe_publisher(fpm_publisher_mock=None) + # Override the default (which creates a MagicMock) with explicit None. + pub.fpm_publisher = None + + pub._check_fpm_schema(_build_fake_stat()) + + assert pub._fpm_schema_checked is True + assert pub.fpm_publisher is None + + +def test_schema_probe_shutdown_exception_still_disables_publisher(): + """If the Rust shutdown call raises, we still None out the publisher -- + the primary goal is to suppress further emission, not to succeed at + shutdown. Protects against leaking emission through a shutdown failure.""" + pub, _ = _build_schema_probe_publisher() + pub.fpm_publisher.shutdown.side_effect = RuntimeError("tokio runtime gone") + + stat = _build_fake_stat() + stat.pop("scheduledNumPrefillRequests") + pub._check_fpm_schema(stat) + + assert pub.fpm_publisher is None + assert pub._fpm_schema_checked is True + + +def test_handle_stat_probe_gate_fires_once_and_skips_subsequent_stats(): + """Simulate the handle_stat dispatch pattern: on the first stat the probe + runs; on the next stat the gate short-circuits. Ensures we do not re-check + per iteration (which would be wasteful and could race a late schema bump).""" + pub, _ = _build_schema_probe_publisher() + original_publisher = pub.fpm_publisher + + # First stat — probe runs, passes. + stat_ok = _build_fake_stat() + if pub.fpm_publisher is not None and not pub._fpm_schema_checked: + pub._check_fpm_schema(stat_ok) + assert pub._fpm_schema_checked is True + assert pub.fpm_publisher is original_publisher + + # Second stat, with a field that would fail the probe if re-checked. + # The gate must prevent re-entry. + stat_bad = _build_fake_stat() + stat_bad.pop("scheduledNumPrefillRequests") + if pub.fpm_publisher is not None and not pub._fpm_schema_checked: + pub._check_fpm_schema(stat_bad) + assert pub.fpm_publisher is original_publisher + original_publisher.shutdown.assert_not_called() + + +def test_schema_probe_field_list_matches_publish_arguments(): + """Guardrail: the required-fields tuple must stay in sync with the + fields handle_stat reads in its .publish() call. If someone adds a + scheduled/queued field to the .publish args but forgets the probe + constant, this test catches it.""" + from dynamo.trtllm import publisher as publisher_mod + + expected = { + "scheduledNumPrefillRequests", + "scheduledSumPrefillTokens", + "scheduledSumPrefillKvTokens", + "scheduledNumDecodeRequests", + "scheduledSumDecodeKvTokens", + "queuedNumPrefillRequests", + "queuedSumPrefillTokens", + "queuedNumDecodeRequests", + "queuedSumDecodeKvTokens", + } + assert set(publisher_mod._FPM_REQUIRED_STAT_FIELDS) == expected From ced15bd308a91f1a1d106bd079b9672e922a7352 Mon Sep 17 00:00:00 2001 From: Yuewei Na Date: Wed, 29 Apr 2026 23:23:11 -0700 Subject: [PATCH 5/5] feat(trtllm): realign FPM publisher to merged TRT-LLM #13199 schema MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit NVIDIA/TensorRT-LLM#13199 merged with a different shape than the draft this PR was originally written against: the new IterationStats fields live nested inside `inflightBatchingStats` and use names that don't match the draft's flat top-level keys. Field-extraction changes (publisher.py): - _check_fpm_schema reads the nested `inflightBatchingStats` dict and fails closed if it is missing or any of the 11 required IBS fields is absent. Strict probe protects against silent planner poison when running against a TRT-LLM that predates #13199. - handle_stat reads the 11 IBS fields and composes Dynamo-side fields: scheduled_num_prefill_requests = numContextRequests scheduled_sum_prefill_tokens = numCtxTokens (compute this iter; excludes prefix-cache carryover counted in numCtxKvTokens, to match vLLM's TTFT-prediction semantics) scheduled_sum_prefill_kv_tokens = numCtxKvTokens scheduled_num_decode_requests = numGenRequests scheduled_sum_decode_kv_tokens = numGenKvTokens queued_num_prefill_requests = numQueuedContextRequests queued_sum_prefill_tokens = numQueuedCtxTokens queued_num_decode_requests = numPausedRequests + numQueuedGenRequests queued_sum_decode_kv_tokens = numPausedKvTokens + numQueuedGenKvTokens Paused-decode requests are intentionally double-counted with numScheduledRequests upstream because the planner reads queued_decode as a KV-pressure preemption signal, where paused-decodes carry the same semantic weight as queued-gen-only requests. - FpmDirectPublisher.publish() callsite switched to keyword-only args so adjacent ints with similar units (`scheduled_*` vs `queued_*`, `*_prefill_*` vs `*_decode_*`) cannot be silently transposed. - Init-time and cleanup-time `except Exception` narrowed to `except RuntimeError` (PyO3 surfaces FpmDirectPublisher new/shutdown failures only as PyRuntimeError); hot-path catch in handle_stat stays broad on purpose, with a comment explaining why. Rust binding changes (lib/bindings/python/rust/llm/fpm.rs): - FpmDirectPublisher struct field `_inner` renamed to `inner`. The underscore prefix conventionally signals "intentionally unused" to Rust readers, but the field is load-bearing for Drop-driven shutdown of all per-rank serialization tasks. A future cleanup pass would delete it as dead and leak every spawned task. Documented in a comment. - publish() signature switched to keyword-only via #[pyo3(signature = (*, ...))]; matching update in _core.pyi stub. Test changes (test_trtllm_fpm_publisher.py): - Module-level pytestmark = [unit, trtllm, pre_merge, gpu_0] picked up by CI marker filters. Without these, the tests were silently skipped by the strict-marker config. - _build_fake_stat constructs the nested IBS structure matching the real JSON; _invoke_handler mirrors handle_stat's nested extraction and kwargs publish() exactly. A guardrail test (test_invoke_handler_matches_publisher_keyword_set) catches drift. - Schema-probe parametrize iterates all 11 IBS fields plus the ibs-not-a-dict and missing-ibs-dict edge cases. - New tests for the queued-decode composite under disagg-decode (only-queued-gen) and aggregated (only-paused) engine shapes. - _build_publisher_stub uses the monkeypatch fixture for hermetic state — no module-level rebinds that leak across tests. The blanket try/except in _run_initialize is removed; WorkerMetricsPublisher and KvEventPublisher are stubbed via monkeypatch so initialize() reaches the FPM gate cleanly. Verification (single-GPU local, non-attention-DP): - Unit: 28/28 pass. - Schema probe against real Qwen3-0.6B inside L0_PostMerge build 2697 (commit e903428, first build containing #13199): all 11 required IBS fields present at the nested path, attentionDpRank tagged at the top level. - Sustained load (concurrency 16 vs max_batch_size 4, 45 s): 520 stats sampled, max scheduled_num_decode_requests=4 (saturated batch), max queued_num_prefill_requests=3 (queueing observed), 512/520 iters had non-zero decode signal. Known limitation: TensorRT (legacy) backend leaves IBS at zero-default, so FPM emission against a TRT-engine deployment will be all-zero. The schema probe accepts keys-exist; a value-aware probe is left as follow-up. Signed-off-by: Yuewei Na --- components/src/dynamo/trtllm/publisher.py | 180 +++++--- .../trtllm/tests/test_trtllm_fpm_publisher.py | 432 +++++++++++------- lib/bindings/python/rust/llm/fpm.rs | 20 +- lib/bindings/python/src/dynamo/_core.pyi | 6 + 4 files changed, 403 insertions(+), 235 deletions(-) diff --git a/components/src/dynamo/trtllm/publisher.py b/components/src/dynamo/trtllm/publisher.py index 28f98ec66e8c..20ca856c10fe 100644 --- a/components/src/dynamo/trtllm/publisher.py +++ b/components/src/dynamo/trtllm/publisher.py @@ -55,21 +55,43 @@ _KV_EVENTS_MAX_SLEEP_SEC = 0.02 _KV_EVENTS_BACKOFF_FACTOR = 1.5 -# IterationStats fields that TensorRT-LLM#13199 adds and that the FPM -# publisher consumes. The first-stat schema probe in handle_stat requires -# ALL of these to be present; any missing field disables the publisher so -# we do not emit all-zero snapshots that the Planner would misread as -# "worker idle" and act on. -_FPM_REQUIRED_STAT_FIELDS = ( - "scheduledNumPrefillRequests", - "scheduledSumPrefillTokens", - "scheduledSumPrefillKvTokens", - "scheduledNumDecodeRequests", - "scheduledSumDecodeKvTokens", - "queuedNumPrefillRequests", - "queuedSumPrefillTokens", - "queuedNumDecodeRequests", - "queuedSumDecodeKvTokens", +# InflightBatchingStats fields the FPM publisher consumes. As of +# NVIDIA/TensorRT-LLM#13199 (merged 2026-04-27) all 11 fields live nested +# inside iterationStats["inflightBatchingStats"]. The first-stat schema +# probe requires the nested dict to be present and to carry every key +# below; any missing field disables the publisher so we do not emit +# all-zero snapshots that the Planner would misread as "worker idle". +# +# Mapping to the Dynamo planner-level fields: +# numContextRequests -> scheduled_num_prefill_requests +# numCtxTokens -> scheduled_sum_prefill_tokens (compute this +# iter; excludes KV-read tokens, matches +# vLLM TTFT-prediction semantics) +# numCtxKvTokens -> scheduled_sum_prefill_kv_tokens +# numGenRequests -> scheduled_num_decode_requests +# numGenKvTokens -> scheduled_sum_decode_kv_tokens +# numQueuedContextRequests -> queued_num_prefill_requests +# numQueuedCtxTokens -> queued_sum_prefill_tokens +# numPausedRequests +# + numQueuedGenRequests -> queued_num_decode_requests (composite) +# numPausedKvTokens +# + numQueuedGenKvTokens -> queued_sum_decode_kv_tokens (composite) +# Paused-request double-counting (also present in numScheduledRequests +# upstream) is intentional: the Planner reads queued_decode as a +# KV-pressure preemption signal where paused-decodes carry the same +# semantic weight as queued-gen-only requests. +_FPM_REQUIRED_IBS_FIELDS = ( + "numContextRequests", + "numCtxTokens", + "numCtxKvTokens", + "numGenRequests", + "numGenKvTokens", + "numQueuedContextRequests", + "numQueuedCtxTokens", + "numQueuedGenRequests", + "numQueuedGenKvTokens", + "numPausedRequests", + "numPausedKvTokens", ) @@ -424,7 +446,11 @@ def initialize(self) -> None: logging.info( f"FpmDirectPublisher initialized with dp_size={self.attention_dp_size}" ) - except Exception as e: + except RuntimeError as e: + # PyO3 surfaces all FpmDirectPublisher::new failures as + # PyRuntimeError (Endpoint missing, tokio runtime missing, + # etc.). Catch only that — any other exception here would + # signal a programming error worth surfacing. logging.warning( f"Failed to initialize FpmDirectPublisher; FPM emission disabled: {e}" ) @@ -527,34 +553,47 @@ async def _polling_loop( sleep_s = min_sleep def _check_fpm_schema(self, stat: dict) -> None: - """One-shot probe: disable FPM publisher if any required TRT-LLM - IterationStats field is missing. + """One-shot probe: disable FPM publisher if the TRT-LLM IterationStats + nested ``inflightBatchingStats`` dict is missing or incomplete. Runs exactly once (gated by ``self._fpm_schema_checked``). Strict: if - any field in ``_FPM_REQUIRED_STAT_FIELDS`` is missing, the publisher - is shut down and set to ``None`` so the subsequent - ``if self.fpm_publisher is not None:`` short-circuit suppresses all - FPM emission for the lifetime of this worker. This prevents silent - planner poison when running against a TRT-LLM that predates - NVIDIA/TensorRT-LLM#13199 — otherwise every field would default to 0 - and the emitted snapshot would be byte-identical to the idle - heartbeat, making the Planner treat a loaded worker as idle. + the nested dict is absent or any field in ``_FPM_REQUIRED_IBS_FIELDS`` + is missing, the publisher is shut down and set to ``None`` so the + subsequent ``if self.fpm_publisher is not None:`` short-circuit + suppresses all FPM emission for the lifetime of this worker. This + prevents silent planner poison when running against a TRT-LLM that + predates NVIDIA/TensorRT-LLM#13199 — otherwise every field would + default to 0 and the emitted snapshot would be byte-identical to the + idle heartbeat, making the Planner treat a loaded worker as idle. """ self._fpm_schema_checked = True if self.fpm_publisher is None: return - missing = [f for f in _FPM_REQUIRED_STAT_FIELDS if f not in stat] + ibs = stat.get("inflightBatchingStats") + if not isinstance(ibs, dict): + logging.warning( + "TRT-LLM IterationStats has no 'inflightBatchingStats' dict; " + "disabling FpmDirectPublisher. Upgrade TRT-LLM past " + "NVIDIA/TensorRT-LLM#13199 to enable FPM." + ) + self._disable_fpm_publisher() + return + missing = [f for f in _FPM_REQUIRED_IBS_FIELDS if f not in ibs] if not missing: return logging.warning( - "TRT-LLM IterationStats is missing required FPM fields %s; " + "TRT-LLM inflightBatchingStats is missing required FPM fields %s; " "disabling FpmDirectPublisher to prevent planner poison. " "Upgrade TRT-LLM past NVIDIA/TensorRT-LLM#13199 to enable FPM.", missing, ) + self._disable_fpm_publisher() + + def _disable_fpm_publisher(self) -> None: + """Shut down ``self.fpm_publisher`` (best effort) and None it out.""" try: self.fpm_publisher.shutdown() - except Exception as e: + except RuntimeError as e: logging.warning( f"FpmDirectPublisher shutdown after schema mismatch failed: {e}" ) @@ -600,41 +639,70 @@ def handle_stat(stat): logging.warning(f"Failed to log iteration stats: {e}") # Publish ForwardPassMetrics. TRT-LLM tags each stat dict with - # attentionDpRank inside BaseWorker._stats_serializer; when - # attention DP is off, the tag defaults to 0. The 9 flat FPM - # fields live at the top level of the dict (camelCase from + # top-level attentionDpRank inside BaseWorker._stats_serializer; + # under non-attention-DP it defaults to 0. The FPM source fields + # live nested under stat["inflightBatchingStats"] (camelCase from # NLOHMANN serialization). Variance fields are not yet computed - # in TRT-LLM's PyExecutor and default to 0.0 in the Rust - # snapshot. + # in TRT-LLM's PyExecutor and default to 0.0 on the Rust side. # # The first stat delivered here triggers a one-shot schema probe: - # if any required field is missing (e.g. running against a - # TRT-LLM that predates #13199) the probe shuts down the - # publisher, which flips the guard below to short-circuit FPM - # emission for the rest of this worker's lifetime. + # if the nested IBS dict is missing or any required field is + # absent (e.g. running against a TRT-LLM that predates #13199) + # the probe shuts down the publisher, which flips the guard + # below to short-circuit FPM emission for the rest of this + # worker's lifetime. if self.fpm_publisher is not None and not self._fpm_schema_checked: self._check_fpm_schema(stat) if self.fpm_publisher is not None: try: - dp_rank = int(stat.get("attentionDpRank", 0)) + ibs = stat.get("inflightBatchingStats") or {} + # numCtxTokens is the prefill compute volume *this iter* + # (excludes prefix-cache/chunked carryover counted in + # numCtxKvTokens). Mapped to scheduled_sum_prefill_tokens + # to match vLLM's TTFT-prediction semantics on the + # planner side. + sched_num_prefill = int(ibs.get("numContextRequests", 0)) + sched_sum_prefill_tokens = int(ibs.get("numCtxTokens", 0)) + sched_sum_prefill_kv_tokens = int(ibs.get("numCtxKvTokens", 0)) + sched_num_decode = int(ibs.get("numGenRequests", 0)) + sched_sum_decode_kv_tokens = int(ibs.get("numGenKvTokens", 0)) + queued_num_prefill = int(ibs.get("numQueuedContextRequests", 0)) + queued_sum_prefill_tokens = int(ibs.get("numQueuedCtxTokens", 0)) + # Composite: paused-decodes + queued-gen-only requests. + # Both represent decode work blocked from progressing + # this iter due to KV pressure or pending KV transfer. + # numPausedRequests is also counted in numScheduledRequests + # upstream — the double-count is intentional because the + # Planner reads queued_decode as a preemption-pressure + # signal where paused-decodes carry the same weight as + # queued-gen-only requests. + queued_num_decode = int(ibs.get("numPausedRequests", 0)) + int( + ibs.get("numQueuedGenRequests", 0) + ) + queued_sum_decode_kv_tokens = int( + ibs.get("numPausedKvTokens", 0) + ) + int(ibs.get("numQueuedGenKvTokens", 0)) # iterLatencyMS is ms; the Rust snapshot expects seconds. - iter_latency_ms = float(stat.get("iterLatencyMS", 0.0)) + wall_time_secs = float(stat.get("iterLatencyMS", 0.0)) / 1000.0 self.fpm_publisher.publish( - dp_rank, - int(stat.get("scheduledNumPrefillRequests", 0)), - int(stat.get("scheduledSumPrefillTokens", 0)), - int(stat.get("scheduledSumPrefillKvTokens", 0)), - int(stat.get("scheduledNumDecodeRequests", 0)), - int(stat.get("scheduledSumDecodeKvTokens", 0)), - int(stat.get("queuedNumPrefillRequests", 0)), - int(stat.get("queuedSumPrefillTokens", 0)), - int(stat.get("queuedNumDecodeRequests", 0)), - int(stat.get("queuedSumDecodeKvTokens", 0)), - iter_latency_ms / 1000.0, + dp_rank=int(stat.get("attentionDpRank", 0)), + scheduled_num_prefill_requests=sched_num_prefill, + scheduled_sum_prefill_tokens=sched_sum_prefill_tokens, + scheduled_sum_prefill_kv_tokens=sched_sum_prefill_kv_tokens, + scheduled_num_decode_requests=sched_num_decode, + scheduled_sum_decode_kv_tokens=sched_sum_decode_kv_tokens, + queued_num_prefill_requests=queued_num_prefill, + queued_sum_prefill_tokens=queued_sum_prefill_tokens, + queued_num_decode_requests=queued_num_decode, + queued_sum_decode_kv_tokens=queued_sum_decode_kv_tokens, + wall_time_secs=wall_time_secs, ) except Exception as e: - # Defensive: don't let FPM publish failures break the - # ActiveLoad / Prometheus pipeline above. + # Defensive (broad on purpose): the FPM publish path is + # cold compared to ActiveLoad/Prometheus and we'd rather + # drop a single FPM snapshot than poison the existing + # metrics pipeline if TRT-LLM ever emits an unexpected + # stat shape. logging.warning(f"FPM publish failed: {e}") await self._polling_loop( @@ -872,11 +940,13 @@ async def cleanup(self) -> None: self.zmq_kv_event_publisher.shutdown() # Shutdown FpmDirectPublisher (stops the per-rank serialization tasks - # and the event-plane publisher task on the Rust side). + # and the event-plane publisher task on the Rust side). PyO3 surfaces + # shutdown failures as PyRuntimeError; narrower catch keeps real + # programming errors visible. if self.fpm_publisher is not None: try: self.fpm_publisher.shutdown() - except Exception as e: + except RuntimeError as e: logging.warning(f"FpmDirectPublisher shutdown failed: {e}") def update_max_window_size(self, event: dict) -> None: diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py b/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py index e5310d3319a0..8e17aee4e9cc 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py @@ -3,19 +3,25 @@ """ Unit tests for the TRT-LLM adapter's ForwardPassMetrics wiring. -Covers: - * handle_stat maps the 9 flat IterationStats fields into - FpmDirectPublisher.publish with the correct positional arg order. - * attentionDpRank from the stat dict is passed through unchanged; missing - key defaults to 0. - * iterLatencyMS (milliseconds) is converted to wall_time_secs (seconds). - * FPM publish failures do not break the existing ActiveLoad / Prometheus - path (defensive try/except). - -The handle_stat closure is defined inside Publisher._publish_stats_task so -we inline the mapping logic via a direct copy — kept minimal on purpose. -The Step 11 tests here exercise the shape of the mapping; full end-to-end -publish-and-subscribe coverage is in the combined E2E test (Step 12). +Covers (after realignment to the merged TRT-LLM PR #13199): + + * handle_stat reads the 11 InflightBatchingStats fields from the nested + ``stat["inflightBatchingStats"]`` dict and forwards them as keyword + arguments to FpmDirectPublisher.publish. + * Composite mappings: queued-decode counters are ``numPausedRequests + + numQueuedGenRequests`` and ``numPausedKvTokens + numQueuedGenKvTokens``. + * attentionDpRank from the top level of the stat dict is passed through + unchanged; missing key defaults to 0. + * iterLatencyMS (top-level milliseconds) is converted to wall_time_secs + (seconds) at the boundary. + * First-stat schema probe disables the publisher when the nested IBS dict + is missing or any of the 11 required fields is absent — protecting + against silent planner poison when running against pre-#13199 TRT-LLM. + +The handle_stat closure is defined inside Publisher._publish_stats_task, +so we mirror its FPM branch via ``_invoke_handler`` below — kept in +lock-step with publisher.py on purpose so any drift surfaces immediately +in ``test_invoke_handler_matches_publisher_keyword_set``. """ from __future__ import annotations @@ -24,66 +30,137 @@ import pytest - -def _build_fake_stat(**overrides): +pytestmark = [ + pytest.mark.unit, + pytest.mark.trtllm, + pytest.mark.pre_merge, + pytest.mark.gpu_0, +] + + +# Default IBS values used by every test that doesn't override. Mirrors the +# semantics each field carries: scheduled vs queued, prefill vs decode, +# request count vs token count vs KV-token count. +_DEFAULT_IBS = { + "numContextRequests": 3, + "numCtxTokens": 1024, + "numCtxKvTokens": 256, + "numGenRequests": 5, + "numGenKvTokens": 9000, + "numQueuedContextRequests": 2, + "numQueuedCtxTokens": 512, + "numQueuedGenRequests": 1, + "numQueuedGenKvTokens": 400, + "numPausedRequests": 1, + "numPausedKvTokens": 800, +} + + +def _build_fake_stat(*, ibs_overrides=None, **top_level_overrides): + """Construct an IterationStats-shaped dict matching the merged #13199 JSON. + + Top-level keys (``iterLatencyMS``, ``attentionDpRank``, ``kvCacheStats``) + are siblings of the nested ``inflightBatchingStats`` object, exactly as + NLOHMANN serializes the C++ struct. + """ + ibs = dict(_DEFAULT_IBS) + if ibs_overrides: + ibs.update(ibs_overrides) stat = { "iterLatencyMS": 25.0, "attentionDpRank": 0, "kvCacheStats": {"usedNumBlocks": 10, "maxNumBlocks": 100}, - "scheduledNumPrefillRequests": 3, - "scheduledSumPrefillTokens": 1024, - "scheduledSumPrefillKvTokens": 256, - "scheduledNumDecodeRequests": 5, - "scheduledSumDecodeKvTokens": 9000, - "queuedNumPrefillRequests": 2, - "queuedSumPrefillTokens": 512, - "queuedNumDecodeRequests": 1, - "queuedSumDecodeKvTokens": 800, + "inflightBatchingStats": ibs, } - stat.update(overrides) + stat.update(top_level_overrides) return stat def _invoke_handler(stat, fpm_publisher): - """Inline copy of the handle_stat FPM branch in publisher.py. + """Inline mirror of the handle_stat FPM branch in publisher.py. - Mirrors the logic exactly so any drift on either side will make the - test fail and force realignment. + Keep this in lock-step with publisher.py — see + ``test_invoke_handler_matches_publisher_keyword_set`` for the guardrail. """ - dp_rank = int(stat.get("attentionDpRank", 0)) - iter_latency_ms = float(stat.get("iterLatencyMS", 0.0)) + ibs = stat.get("inflightBatchingStats") or {} + queued_num_decode = int(ibs.get("numPausedRequests", 0)) + int( + ibs.get("numQueuedGenRequests", 0) + ) + queued_sum_decode_kv_tokens = int(ibs.get("numPausedKvTokens", 0)) + int( + ibs.get("numQueuedGenKvTokens", 0) + ) fpm_publisher.publish( - dp_rank, - int(stat.get("scheduledNumPrefillRequests", 0)), - int(stat.get("scheduledSumPrefillTokens", 0)), - int(stat.get("scheduledSumPrefillKvTokens", 0)), - int(stat.get("scheduledNumDecodeRequests", 0)), - int(stat.get("scheduledSumDecodeKvTokens", 0)), - int(stat.get("queuedNumPrefillRequests", 0)), - int(stat.get("queuedSumPrefillTokens", 0)), - int(stat.get("queuedNumDecodeRequests", 0)), - int(stat.get("queuedSumDecodeKvTokens", 0)), - iter_latency_ms / 1000.0, + dp_rank=int(stat.get("attentionDpRank", 0)), + scheduled_num_prefill_requests=int(ibs.get("numContextRequests", 0)), + scheduled_sum_prefill_tokens=int(ibs.get("numCtxTokens", 0)), + scheduled_sum_prefill_kv_tokens=int(ibs.get("numCtxKvTokens", 0)), + scheduled_num_decode_requests=int(ibs.get("numGenRequests", 0)), + scheduled_sum_decode_kv_tokens=int(ibs.get("numGenKvTokens", 0)), + queued_num_prefill_requests=int(ibs.get("numQueuedContextRequests", 0)), + queued_sum_prefill_tokens=int(ibs.get("numQueuedCtxTokens", 0)), + queued_num_decode_requests=queued_num_decode, + queued_sum_decode_kv_tokens=queued_sum_decode_kv_tokens, + wall_time_secs=float(stat.get("iterLatencyMS", 0.0)) / 1000.0, ) +# --------------------------------------------------------------------------- +# Field mapping +# --------------------------------------------------------------------------- + + def test_handle_stat_maps_fields_single_rank(): fpm = MagicMock() - stat = _build_fake_stat() - _invoke_handler(stat, fpm) + _invoke_handler(_build_fake_stat(), fpm) fpm.publish.assert_called_once_with( - 0, # dp_rank - 3, # scheduled_num_prefill_requests - 1024, # scheduled_sum_prefill_tokens - 256, # scheduled_sum_prefill_kv_tokens - 5, # scheduled_num_decode_requests - 9000, # scheduled_sum_decode_kv_tokens - 2, # queued_num_prefill_requests - 512, # queued_sum_prefill_tokens - 1, # queued_num_decode_requests - 800, # queued_sum_decode_kv_tokens - 0.025, # wall_time_secs (25 ms -> 0.025 s) + dp_rank=0, + scheduled_num_prefill_requests=3, + scheduled_sum_prefill_tokens=1024, + scheduled_sum_prefill_kv_tokens=256, + scheduled_num_decode_requests=5, + scheduled_sum_decode_kv_tokens=9000, + queued_num_prefill_requests=2, + queued_sum_prefill_tokens=512, + queued_num_decode_requests=2, # numPausedRequests (1) + numQueuedGenRequests (1) + queued_sum_decode_kv_tokens=1200, # numPausedKvTokens (800) + numQueuedGenKvTokens (400) + wall_time_secs=0.025, + ) + + +def test_queued_decode_composite_with_only_paused(): + """Disagg-prefill or non-disagg engine: numQueuedGenRequests is always 0, + so queued-decode pressure comes entirely from preempted decodes.""" + fpm = MagicMock() + stat = _build_fake_stat( + ibs_overrides={ + "numPausedRequests": 4, + "numPausedKvTokens": 12000, + "numQueuedGenRequests": 0, + "numQueuedGenKvTokens": 0, + } + ) + _invoke_handler(stat, fpm) + kwargs = fpm.publish.call_args.kwargs + assert kwargs["queued_num_decode_requests"] == 4 + assert kwargs["queued_sum_decode_kv_tokens"] == 12000 + + +def test_queued_decode_composite_with_only_queued_gen(): + """Disagg-decode engine awaiting KV transfer: numPausedRequests is 0, + queued-gen carries the full signal.""" + fpm = MagicMock() + stat = _build_fake_stat( + ibs_overrides={ + "numPausedRequests": 0, + "numPausedKvTokens": 0, + "numQueuedGenRequests": 7, + "numQueuedGenKvTokens": 21000, + } ) + _invoke_handler(stat, fpm) + kwargs = fpm.publish.call_args.kwargs + assert kwargs["queued_num_decode_requests"] == 7 + assert kwargs["queued_sum_decode_kv_tokens"] == 21000 def test_handle_stat_routes_per_attention_dp_rank(): @@ -91,14 +168,14 @@ def test_handle_stat_routes_per_attention_dp_rank(): for rank in (0, 1, 2, 3): stat = _build_fake_stat( attentionDpRank=rank, - scheduledSumPrefillTokens=100 * (rank + 1), + ibs_overrides={"numCtxTokens": 100 * (rank + 1)}, ) _invoke_handler(stat, fpm) calls = fpm.publish.call_args_list assert len(calls) == 4 for i, call in enumerate(calls): - assert call.args[0] == i # dp_rank - assert call.args[2] == 100 * (i + 1) # scheduled_sum_prefill_tokens + assert call.kwargs["dp_rank"] == i + assert call.kwargs["scheduled_sum_prefill_tokens"] == 100 * (i + 1) def test_handle_stat_missing_attention_dp_rank_defaults_zero(): @@ -106,48 +183,51 @@ def test_handle_stat_missing_attention_dp_rank_defaults_zero(): stat = _build_fake_stat() stat.pop("attentionDpRank") _invoke_handler(stat, fpm) - # First positional arg is the dp_rank. - assert fpm.publish.call_args.args[0] == 0 + assert fpm.publish.call_args.kwargs["dp_rank"] == 0 -def test_handle_stat_missing_fpm_fields_are_zero(): +def test_handle_stat_missing_ibs_dict_emits_zeros(): + """If the nested IBS dict is absent at this layer (the schema probe + upstream should have already disabled the publisher in production), + the handler must still produce a zeroed call rather than KeyError.""" fpm = MagicMock() - stat = { - "iterLatencyMS": 10.0, - "attentionDpRank": 0, - # All FPM fields missing. - } - _invoke_handler(stat, fpm) + _invoke_handler({"iterLatencyMS": 10.0, "attentionDpRank": 0}, fpm) fpm.publish.assert_called_once_with( - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0.01, + dp_rank=0, + scheduled_num_prefill_requests=0, + scheduled_sum_prefill_tokens=0, + scheduled_sum_prefill_kv_tokens=0, + scheduled_num_decode_requests=0, + scheduled_sum_decode_kv_tokens=0, + queued_num_prefill_requests=0, + queued_sum_prefill_tokens=0, + queued_num_decode_requests=0, + queued_sum_decode_kv_tokens=0, + wall_time_secs=0.01, ) def test_iter_latency_ms_to_wall_time_secs_conversion(): fpm = MagicMock() - stat = _build_fake_stat(iterLatencyMS=1234.5) - _invoke_handler(stat, fpm) - # last positional arg is wall_time_secs. - assert fpm.publish.call_args.args[-1] == 1.2345 + _invoke_handler(_build_fake_stat(iterLatencyMS=1234.5), fpm) + assert fpm.publish.call_args.kwargs["wall_time_secs"] == 1.2345 + +# --------------------------------------------------------------------------- +# Publisher.initialize() gate behavior (attention-DP off vs on) +# --------------------------------------------------------------------------- -def _build_publisher_stub(*, attention_dp_size: int, fpm_enabled: bool): - """Shared heavy-mock helper for the Publisher.initialize() tests below. - Bypasses Publisher.__init__ (which touches many heavy deps) and sets only - the attributes that initialize() reads or writes. Tests then call - pub.initialize() and inspect the mocked FpmDirectPublisher class. +def _build_publisher_stub(monkeypatch, *, attention_dp_size: int, fpm_enabled: bool): + """Bypass Publisher.__init__ (heavy deps) and seed only the attributes + initialize() reads or writes. All side-effecty subsystems are stubbed + via ``monkeypatch`` so initialize() reaches the FPM gate cleanly without + needing a blanket try/except to swallow upstream failures. """ + import asyncio + import queue + import threading + from dynamo.trtllm import publisher as publisher_mod engine = MagicMock() @@ -173,52 +253,48 @@ def _build_publisher_stub(*, attention_dp_size: int, fpm_enabled: bool): pub.publish_kv_cache_events_thread = None pub.publish_stats_thread = None pub.partial_block_hashes = set() - import queue as _queue - - pub.error_queue = _queue.Queue() - import threading as _threading - - pub._stop_event = _threading.Event() + pub.error_queue = queue.Queue() + pub._stop_event = threading.Event() pub._last_engine_event_id = None - # Replace the real FpmDirectPublisher class with a mock factory so we can - # inspect what was passed to it (or assert it wasn't called). fake_fpm_cls = MagicMock() - publisher_mod.FpmDirectPublisher = fake_fpm_cls + monkeypatch.setattr(publisher_mod, "FpmDirectPublisher", fake_fpm_cls) + # WorkerMetricsPublisher and KvEventPublisher both reach into the Rust + # binding and validate their endpoint arg as a real Endpoint — replace + # with MagicMock factories so initialize() can complete past the FPM + # gate without dragging in a real Endpoint. + monkeypatch.setattr(publisher_mod, "WorkerMetricsPublisher", MagicMock()) + monkeypatch.setattr(publisher_mod, "KvEventPublisher", MagicMock()) + + # asyncio.create_task wants a running loop — replace with a no-op + # MagicMock so initialize() can call it synchronously in tests. Restored + # automatically by monkeypatch on teardown. + monkeypatch.setattr( + asyncio, + "create_task", + lambda coro: MagicMock(add_done_callback=lambda _: None), + ) - # Stub out the other side-effecty subsystems that initialize() touches. - pub._init_publish_metrics_thread = MagicMock() - pub._init_publish_kv_cache_events_thread = MagicMock() - pub._create_metrics_publisher_endpoint = MagicMock(return_value=MagicMock()) + monkeypatch.setattr(pub, "_init_publish_metrics_thread", MagicMock()) + monkeypatch.setattr(pub, "_init_publish_kv_cache_events_thread", MagicMock()) + monkeypatch.setattr( + pub, + "_create_metrics_publisher_endpoint", + MagicMock(return_value=MagicMock()), + ) return pub, publisher_mod, fake_fpm_cls -def _run_initialize(pub): - """Run pub.initialize() synchronously (no event loop required).""" - import asyncio as _asyncio - - real_create_task = _asyncio.create_task - _asyncio.create_task = lambda coro: MagicMock(add_done_callback=lambda _: None) - try: - try: - pub.initialize() - except Exception: - # initialize() touches other subsystems we don't care about; the - # FPM-related assertions below tell us whether the construction - # (or non-construction) happened as expected. - pass - finally: - _asyncio.create_task = real_create_task - - -def test_publisher_initialize_constructs_fpm_direct_publisher_when_fpm_enabled(): +def test_publisher_initialize_constructs_fpm_direct_publisher_when_fpm_enabled( + monkeypatch, +): """Under non-attention-DP (attention_dp_size == 1, fpm_enabled == True), Publisher.initialize() constructs FpmDirectPublisher with dp_size=1.""" pub, _publisher_mod, fake_fpm_cls = _build_publisher_stub( - attention_dp_size=1, fpm_enabled=True + monkeypatch, attention_dp_size=1, fpm_enabled=True ) - _run_initialize(pub) + pub.initialize() fake_fpm_cls.assert_called_once() kwargs = fake_fpm_cls.call_args.kwargs assert kwargs["worker_id"] == "worker-abc" @@ -226,15 +302,16 @@ def test_publisher_initialize_constructs_fpm_direct_publisher_when_fpm_enabled() assert pub.fpm_publisher is not None -def test_publisher_does_not_init_fpm_publisher_under_attention_dp(): +def test_publisher_does_not_init_fpm_publisher_under_attention_dp(monkeypatch): """Under attention-DP (attention_dp_size > 1, fpm_enabled == False), the gate suppresses FpmDirectPublisher construction. pub.fpm_publisher stays - None so handle_stat's existing `if self.fpm_publisher is not None:` guard - skips all FPM publishes -- the Planner sees ZERO messages from this worker.""" + None so handle_stat's existing ``if self.fpm_publisher is not None:`` + guard skips all FPM publishes — the Planner sees ZERO messages from this + worker (strictly better than fake-idle pollution).""" pub, _publisher_mod, fake_fpm_cls = _build_publisher_stub( - attention_dp_size=4, fpm_enabled=False + monkeypatch, attention_dp_size=4, fpm_enabled=False ) - _run_initialize(pub) + pub.initialize() fake_fpm_cls.assert_not_called() assert pub.fpm_publisher is None @@ -244,11 +321,11 @@ def test_publisher_does_not_init_fpm_publisher_under_attention_dp(): # --------------------------------------------------------------------------- -def _build_schema_probe_publisher(fpm_publisher_mock: MagicMock | None = None): +def _build_schema_probe_publisher(fpm_publisher_mock=None): """Minimal Publisher instance for direct _check_fpm_schema testing. - Bypasses __init__ (which touches heavy deps) and seeds only the attributes - the probe method reads or writes: fpm_publisher and _fpm_schema_checked. + Bypasses __init__ (heavy deps) and seeds only the attributes the probe + method reads or writes: fpm_publisher and _fpm_schema_checked. """ from dynamo.trtllm import publisher as publisher_mod @@ -271,29 +348,16 @@ def test_schema_probe_all_fields_present_keeps_publisher(): original_publisher.shutdown.assert_not_called() -@pytest.mark.parametrize( - "missing_field", - [ - "scheduledNumPrefillRequests", - "scheduledSumPrefillTokens", - "scheduledSumPrefillKvTokens", - "scheduledNumDecodeRequests", - "scheduledSumDecodeKvTokens", - "queuedNumPrefillRequests", - "queuedSumPrefillTokens", - "queuedNumDecodeRequests", - "queuedSumDecodeKvTokens", - ], -) -def test_schema_probe_missing_single_field_disables_publisher(missing_field): - """Strict probe: any one of the 9 required fields missing must disable - the publisher. Covers each field independently so a rename upstream or - a selective-backport TRT-LLM never slips through.""" +@pytest.mark.parametrize("missing_field", list(_DEFAULT_IBS.keys())) +def test_schema_probe_missing_single_ibs_field_disables_publisher(missing_field): + """Strict probe: any one of the 11 required IBS fields missing must + disable the publisher. Covers each field independently so a rename + upstream or a selective-backport TRT-LLM never slips through.""" pub, _ = _build_schema_probe_publisher() original_publisher = pub.fpm_publisher stat = _build_fake_stat() - stat.pop(missing_field) + stat["inflightBatchingStats"].pop(missing_field) pub._check_fpm_schema(stat) assert pub._fpm_schema_checked is True @@ -301,9 +365,10 @@ def test_schema_probe_missing_single_field_disables_publisher(missing_field): original_publisher.shutdown.assert_called_once() -def test_schema_probe_missing_all_fields_disables_publisher_legacy_trtllm(): +def test_schema_probe_missing_ibs_dict_disables_publisher_legacy_trtllm(): """Legacy TRT-LLM case: stat dict has iterLatencyMS + attentionDpRank but - none of the 9 FPM fields (pre-#13199 schema). Must disable without error.""" + no inflightBatchingStats nested object (pre-#13199 schema). Must disable + without error.""" pub, _ = _build_schema_probe_publisher() original_publisher = pub.fpm_publisher @@ -314,11 +379,23 @@ def test_schema_probe_missing_all_fields_disables_publisher_legacy_trtllm(): original_publisher.shutdown.assert_called_once() +def test_schema_probe_ibs_not_a_dict_disables_publisher(): + """Defensive: if a future TRT-LLM ever emits inflightBatchingStats as + something other than a dict (e.g. null on engine init), treat it as a + schema mismatch rather than crashing in the probe.""" + pub, _ = _build_schema_probe_publisher() + original_publisher = pub.fpm_publisher + + pub._check_fpm_schema({"inflightBatchingStats": None}) + + assert pub.fpm_publisher is None + original_publisher.shutdown.assert_called_once() + + def test_schema_probe_noop_when_fpm_publisher_already_none(): """Attention-DP gate already set fpm_publisher = None; probe must not blow up and must still flip _fpm_schema_checked so we do not re-enter.""" pub, _ = _build_schema_probe_publisher(fpm_publisher_mock=None) - # Override the default (which creates a MagicMock) with explicit None. pub.fpm_publisher = None pub._check_fpm_schema(_build_fake_stat()) @@ -328,14 +405,14 @@ def test_schema_probe_noop_when_fpm_publisher_already_none(): def test_schema_probe_shutdown_exception_still_disables_publisher(): - """If the Rust shutdown call raises, we still None out the publisher -- + """If the Rust shutdown call raises, we still None out the publisher — the primary goal is to suppress further emission, not to succeed at shutdown. Protects against leaking emission through a shutdown failure.""" pub, _ = _build_schema_probe_publisher() pub.fpm_publisher.shutdown.side_effect = RuntimeError("tokio runtime gone") stat = _build_fake_stat() - stat.pop("scheduledNumPrefillRequests") + stat["inflightBatchingStats"].pop("numCtxKvTokens") pub._check_fpm_schema(stat) assert pub.fpm_publisher is None @@ -349,39 +426,46 @@ def test_handle_stat_probe_gate_fires_once_and_skips_subsequent_stats(): pub, _ = _build_schema_probe_publisher() original_publisher = pub.fpm_publisher - # First stat — probe runs, passes. - stat_ok = _build_fake_stat() if pub.fpm_publisher is not None and not pub._fpm_schema_checked: - pub._check_fpm_schema(stat_ok) + pub._check_fpm_schema(_build_fake_stat()) assert pub._fpm_schema_checked is True assert pub.fpm_publisher is original_publisher - # Second stat, with a field that would fail the probe if re-checked. - # The gate must prevent re-entry. stat_bad = _build_fake_stat() - stat_bad.pop("scheduledNumPrefillRequests") + stat_bad["inflightBatchingStats"].pop("numCtxKvTokens") if pub.fpm_publisher is not None and not pub._fpm_schema_checked: pub._check_fpm_schema(stat_bad) assert pub.fpm_publisher is original_publisher original_publisher.shutdown.assert_not_called() -def test_schema_probe_field_list_matches_publish_arguments(): - """Guardrail: the required-fields tuple must stay in sync with the - fields handle_stat reads in its .publish() call. If someone adds a - scheduled/queued field to the .publish args but forgets the probe - constant, this test catches it.""" +def test_schema_probe_field_list_matches_default_ibs_set(): + """Guardrail: the required-fields tuple must stay in sync with the IBS + default fixture. If someone adds an IBS field to the production reader + but forgets the probe constant (or vice versa), this test catches it.""" from dynamo.trtllm import publisher as publisher_mod - expected = { - "scheduledNumPrefillRequests", - "scheduledSumPrefillTokens", - "scheduledSumPrefillKvTokens", - "scheduledNumDecodeRequests", - "scheduledSumDecodeKvTokens", - "queuedNumPrefillRequests", - "queuedSumPrefillTokens", - "queuedNumDecodeRequests", - "queuedSumDecodeKvTokens", + assert set(publisher_mod._FPM_REQUIRED_IBS_FIELDS) == set(_DEFAULT_IBS.keys()) + + +def test_invoke_handler_matches_publisher_keyword_set(): + """Guardrail: the kwargs that publisher.py's handle_stat passes to + fpm_publisher.publish() must match the test mirror exactly. Catches any + drift where production starts using a new kwarg the test doesn't mirror, + or vice versa.""" + fpm = MagicMock() + _invoke_handler(_build_fake_stat(), fpm) + expected_kwargs = { + "dp_rank", + "scheduled_num_prefill_requests", + "scheduled_sum_prefill_tokens", + "scheduled_sum_prefill_kv_tokens", + "scheduled_num_decode_requests", + "scheduled_sum_decode_kv_tokens", + "queued_num_prefill_requests", + "queued_sum_prefill_tokens", + "queued_num_decode_requests", + "queued_sum_decode_kv_tokens", + "wall_time_secs", } - assert set(publisher_mod._FPM_REQUIRED_STAT_FIELDS) == expected + assert set(fpm.publish.call_args.kwargs.keys()) == expected_kwargs diff --git a/lib/bindings/python/rust/llm/fpm.rs b/lib/bindings/python/rust/llm/fpm.rs index 82e841597d1f..a3a24fa022ca 100644 --- a/lib/bindings/python/rust/llm/fpm.rs +++ b/lib/bindings/python/rust/llm/fpm.rs @@ -80,7 +80,11 @@ impl FpmEventRelay { /// for `IDLE_HEARTBEAT_INTERVAL` (matches vLLM's `HEARTBEAT_INTERVAL = 1.0`). #[pyclass] pub(crate) struct FpmDirectPublisher { - _inner: llm_rs::fpm_publisher::FpmDirectPublisher, + // Owns the CancellationToken that drives Drop-based shutdown of all + // per-rank serialization tasks and the event-plane publisher task. Held + // by reference only — keep the field non-underscored so a future cleanup + // pass does not delete it as "dead", which would leak every spawned task. + inner: llm_rs::fpm_publisher::FpmDirectPublisher, publishers: Vec, } @@ -103,20 +107,24 @@ impl FpmDirectPublisher { llm_rs::fpm_publisher::FpmDirectPublisher::new(component, worker_id, dp_size).await }) .map_err(to_pyerr)?; - Ok(Self { - _inner: inner, - publishers, - }) + Ok(Self { inner, publishers }) } /// Publish one iteration's FPM snapshot for the given DP rank. /// + /// All parameters are keyword-only on the Python side — adjacent ints + /// with similar units (`scheduled_*` vs `queued_*`, `*_prefill_*` vs + /// `*_decode_*`) cannot be distinguished by the type system, so a + /// transposition would silently corrupt every published snapshot. + /// Forcing kwargs at the boundary is a zero-cost guard. + /// /// Variance fields (`var_prefill_length`, `var_decode_kv_tokens`, /// `var_queued_prefill_length`, `var_queued_decode_kv_tokens`) are /// defaulted to 0.0 per the MVP scope — the active planner does not /// consume them on origin/main. A follow-up PR can add Welford-based /// variance computation in TRT-LLM's PyExecutor and a new overload here. #[pyo3(signature = ( + *, dp_rank, scheduled_num_prefill_requests, scheduled_sum_prefill_tokens, @@ -172,7 +180,7 @@ impl FpmDirectPublisher { /// Shut down the publisher and its per-rank serialization tasks. fn shutdown(&self) { - self._inner.shutdown(); + self.inner.shutdown(); } } diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index 72ab134ce534..a8a908a507ea 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -865,6 +865,7 @@ class FpmDirectPublisher: def publish( self, + *, dp_rank: int, scheduled_num_prefill_requests: int, scheduled_sum_prefill_tokens: int, @@ -880,6 +881,11 @@ class FpmDirectPublisher: """ Publish one iteration's FPM snapshot for the given DP rank. + All parameters are keyword-only on the Python side: adjacent ints + with similar units (``scheduled_*`` vs ``queued_*``, ``*_prefill_*`` + vs ``*_decode_*``) cannot be distinguished by the type system, so + a transposition would silently corrupt every published snapshot. + Variance fields (var_prefill_length, var_decode_kv_tokens, var_queued_prefill_length, var_queued_decode_kv_tokens) are defaulted to 0.0 per the MVP scope; a follow-up PR can add Welford-based