diff --git a/components/src/dynamo/trtllm/publisher.py b/components/src/dynamo/trtllm/publisher.py index 1f933ed01461..6f7b40cf4d3d 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) @@ -55,6 +55,45 @@ _KV_EVENTS_MAX_SLEEP_SEC = 0.02 _KV_EVENTS_BACKOFF_FACTOR = 1.5 +# 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", +) + def _to_signed_i64(value: int | None) -> int | None: """Convert a Python int to signed 64-bit range by two's complement.""" @@ -320,6 +359,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. @@ -328,6 +378,15 @@ 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 + # 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 @@ -371,6 +430,39 @@ def initialize(self) -> None: lambda _: logging.debug("metrics publisher endpoint created") ) + # 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 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}" + ) + self.fpm_publisher = None + else: + logging.info( + "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 + # Setup the kv cache events publisher # Publisher selection based on consolidator configuration: # - With consolidator: Use ZmqKvEventPublisher (this module) → ZMQ → Consolidator → NATS → Router @@ -460,6 +552,56 @@ async def _polling_loop( else: sleep_s = min_sleep + def _check_fpm_schema(self, stat: dict) -> None: + """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 + 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 + 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 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.""" + publisher = self.fpm_publisher + if publisher is None: + return + try: + publisher.shutdown() + except RuntimeError 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. @@ -499,6 +641,73 @@ 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 + # 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 on the Rust side. + # + # The first stat delivered here triggers a one-shot schema probe: + # 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: + 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. + wall_time_secs = float(stat.get("iterLatencyMS", 0.0)) / 1000.0 + self.fpm_publisher.publish( + 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 (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( lambda: self.engine.llm.get_stats_async(timeout=_STATS_TIMEOUT_SEC), handle_stat, @@ -733,6 +942,16 @@ 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). 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 RuntimeError 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..8e17aee4e9cc --- /dev/null +++ b/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py @@ -0,0 +1,471 @@ +# 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 (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 + +from unittest.mock import MagicMock + +import pytest + +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}, + "inflightBatchingStats": ibs, + } + stat.update(top_level_overrides) + return stat + + +def _invoke_handler(stat, fpm_publisher): + """Inline mirror of the handle_stat FPM branch in publisher.py. + + Keep this in lock-step with publisher.py — see + ``test_invoke_handler_matches_publisher_keyword_set`` for the guardrail. + """ + 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("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() + _invoke_handler(_build_fake_stat(), fpm) + fpm.publish.assert_called_once_with( + 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(): + fpm = MagicMock() + for rank in (0, 1, 2, 3): + stat = _build_fake_stat( + attentionDpRank=rank, + 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.kwargs["dp_rank"] == i + assert call.kwargs["scheduled_sum_prefill_tokens"] == 100 * (i + 1) + + +def test_handle_stat_missing_attention_dp_rank_defaults_zero(): + fpm = MagicMock() + stat = _build_fake_stat() + stat.pop("attentionDpRank") + _invoke_handler(stat, fpm) + assert fpm.publish.call_args.kwargs["dp_rank"] == 0 + + +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() + _invoke_handler({"iterLatencyMS": 10.0, "attentionDpRank": 0}, fpm) + fpm.publish.assert_called_once_with( + 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() + _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(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() + engine.get_attention_dp_size.return_value = attention_dp_size + + 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 = attention_dp_size + pub.fpm_enabled = fpm_enabled + 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() + pub.error_queue = queue.Queue() + pub._stop_event = threading.Event() + pub._last_engine_event_id = None + + fake_fpm_cls = MagicMock() + 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), + ) + + 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 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( + monkeypatch, attention_dp_size=1, fpm_enabled=True + ) + pub.initialize() + fake_fpm_cls.assert_called_once() + kwargs = fake_fpm_cls.call_args.kwargs + assert kwargs["worker_id"] == "worker-abc" + assert kwargs["dp_size"] == 1 + assert pub.fpm_publisher is not None + + +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 (strictly better than fake-idle pollution).""" + pub, _publisher_mod, fake_fpm_cls = _build_publisher_stub( + monkeypatch, attention_dp_size=4, fpm_enabled=False + ) + pub.initialize() + fake_fpm_cls.assert_not_called() + assert pub.fpm_publisher is None + + +# --------------------------------------------------------------------------- +# First-stat schema probe +# --------------------------------------------------------------------------- + + +def _build_schema_probe_publisher(fpm_publisher_mock=None): + """Minimal Publisher instance for direct _check_fpm_schema testing. + + 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 + + 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", 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["inflightBatchingStats"].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_ibs_dict_disables_publisher_legacy_trtllm(): + """Legacy TRT-LLM case: stat dict has iterLatencyMS + attentionDpRank but + no inflightBatchingStats nested object (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_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) + 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["inflightBatchingStats"].pop("numCtxKvTokens") + 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 + + if pub.fpm_publisher is not None and not pub._fpm_schema_checked: + pub._check_fpm_schema(_build_fake_stat()) + assert pub._fpm_schema_checked is True + assert pub.fpm_publisher is original_publisher + + stat_bad = _build_fake_stat() + 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_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 + + 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(fpm.publish.call_args.kwargs.keys()) == expected_kwargs diff --git a/docs/README.md b/docs/README.md index e7cb883b0fc3..d6f7fd78bcb0 100644 --- a/docs/README.md +++ b/docs/README.md @@ -54,7 +54,7 @@ the navigation in `docs/index.yml`, and running `fern check` to validate. | Skill | Description | |-------|-------------| -| [dynamo-docs](https://github.com/ai-dynamo/dynamo/blob/main/.claude/skills/dynamo-docs/SKILL.md) | Add, update, move, or remove a docs page | +| [dynamo-docs](https://github.com/ai-dynamo/dynamo/blob/main/.agents/skills/dynamo-docs/SKILL.md) | Add, update, move, or remove a docs page | --- diff --git a/docs/blogs/agentic-inference/agentic-inference.md b/docs/blogs/agentic-inference/agentic-inference.md index 0fe0dfc58194..236ce08608f1 100644 --- a/docs/blogs/agentic-inference/agentic-inference.md +++ b/docs/blogs/agentic-inference/agentic-inference.md @@ -51,7 +51,7 @@ Agent harnesses are increasingly adopting `v1/responses` and `v1/messages` over -We have also invested in day-0 tool call and reasoning parsing support for various open-source models. If you find that a model is not supported, please [open an issue](https://github.com/ai-dynamo/dynamo/issues) or use the [tool-call-parser-generator](https://github.com/ai-dynamo/dynamo/blob/main/.claude/skills/tool-parser-generator/SKILL.md) skill to generate it with your harness of choice. +We have also invested in day-0 tool call and reasoning parsing support for various open-source models. If you find that a model is not supported, please [open an issue](https://github.com/ai-dynamo/dynamo/issues) or use the [tool-call-parser-generator](https://github.com/ai-dynamo/dynamo/blob/main/.agents/skills/tool-parser-generator/SKILL.md) skill to generate it with your harness of choice. ### Agent Hints: The Harness-Orchestrator Interface diff --git a/lib/bindings/python/rust/lib.rs b/lib/bindings/python/rust/lib.rs index 14c0a7e51241..d358794dba13 100644 --- a/lib/bindings/python/rust/lib.rs +++ b/lib/bindings/python/rust/lib.rs @@ -188,6 +188,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 f8878a5fa447..d198450ddddf 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,122 @@ 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 { + // 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, +} + +#[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, 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, + 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 e891668d4d5a..1c0a5c820605 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -835,6 +835,69 @@ 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. + + 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 + 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 97ae41e5f912..757e2a6ff012 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