From f674dd8573fe389d87ed1f597593a48a3a49cf6c Mon Sep 17 00:00:00 2001 From: Indrajit Bhosale Date: Thu, 30 Apr 2026 11:52:15 -0700 Subject: [PATCH 1/4] Atn-dp fpm followup draft Signed-off-by: Indrajit Bhosale --- components/src/dynamo/trtllm/publisher.py | 77 ++++------- .../trtllm/tests/test_trtllm_fpm_publisher.py | 120 ++++++++++-------- docs/backends/trtllm/README.md | 1 + 3 files changed, 99 insertions(+), 99 deletions(-) diff --git a/components/src/dynamo/trtllm/publisher.py b/components/src/dynamo/trtllm/publisher.py index e341928c0ce9..ed064c5fbbe0 100644 --- a/components/src/dynamo/trtllm/publisher.py +++ b/components/src/dynamo/trtllm/publisher.py @@ -392,17 +392,6 @@ 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. @@ -411,9 +400,9 @@ 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. + # FPM is emitted as one logical channel per TRT-LLM engine. Under + # attention-DP, TRT-LLM rank 0 emits an engine-wide aggregate; ranks + # 1..N-1 have no independent planner-visible worker identity. 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 @@ -463,36 +452,24 @@ 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." + # Setup the ForwardPassMetrics publisher. Always use a single FPM + # channel: non-attention-DP has one logical rank, and attention-DP + # reports one rank-0 engine aggregate. If rank 0 stalls, the whole + # TRT-LLM engine stalls, so FPM emission is intentionally all-or-nothing. + try: + self.fpm_publisher = FpmDirectPublisher( + endpoint=self.endpoint, + worker_id=str(self.worker_id), + dp_size=1, + ) + logging.info("FpmDirectPublisher initialized with dp_size=1") + 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 @@ -602,6 +579,7 @@ def _check_fpm_schema(self, stat: dict) -> None: self._fpm_schema_checked = True if self.fpm_publisher is None: return + ibs = stat.get("inflightBatchingStats") if not isinstance(ibs, dict): logging.warning( @@ -675,11 +653,12 @@ def handle_stat(stat): 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. + # top-level attentionDpRank inside BaseWorker._stats_serializer. + # Under attention-DP, the companion TRT-LLM change emits one + # rank-0 engine aggregate. 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 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 8e17aee4e9cc..9a40e4e3e58b 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py @@ -26,6 +26,10 @@ from __future__ import annotations +import importlib.util +import sys +import types +from pathlib import Path from unittest.mock import MagicMock import pytest @@ -38,6 +42,37 @@ ] +@pytest.fixture() +def publisher_mod(monkeypatch): + prometheus_mod = types.ModuleType("dynamo.common.utils.prometheus") + prometheus_mod.LLMBackendMetrics = object + + llm_mod = types.ModuleType("dynamo.llm") + llm_mod.FpmDirectPublisher = MagicMock(name="FpmDirectPublisher") + llm_mod.KvEventPublisher = MagicMock(name="KvEventPublisher") + llm_mod.WorkerMetricsPublisher = MagicMock(name="WorkerMetricsPublisher") + + monkeypatch.setitem(sys.modules, "dynamo", types.ModuleType("dynamo")) + monkeypatch.setitem(sys.modules, "dynamo.common", types.ModuleType("dynamo.common")) + monkeypatch.setitem( + sys.modules, "dynamo.common.utils", types.ModuleType("dynamo.common.utils") + ) + monkeypatch.setitem(sys.modules, "dynamo.common.utils.prometheus", prometheus_mod) + monkeypatch.setitem(sys.modules, "dynamo.llm", llm_mod) + + repo_root = Path(__file__).resolve().parents[5] + publisher_path = repo_root / "components/src/dynamo/trtllm/publisher.py" + spec = importlib.util.spec_from_file_location( + "_trtllm_publisher_under_test", publisher_path + ) + assert spec is not None + assert spec.loader is not None + + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + # 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. @@ -214,22 +249,20 @@ def test_iter_latency_ms_to_wall_time_secs_conversion(): # --------------------------------------------------------------------------- -# Publisher.initialize() gate behavior (attention-DP off vs on) +# Publisher.initialize() FPM channel behavior # --------------------------------------------------------------------------- -def _build_publisher_stub(monkeypatch, *, attention_dp_size: int, fpm_enabled: bool): +def _build_publisher_stub(monkeypatch, publisher_mod, *, attention_dp_size: int): """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. + via ``monkeypatch`` so initialize() reaches FPM initialization 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 @@ -244,7 +277,6 @@ def _build_publisher_stub(monkeypatch, *, attention_dp_size: int, fpm_enabled: b 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 @@ -261,8 +293,8 @@ def _build_publisher_stub(monkeypatch, *, attention_dp_size: int, fpm_enabled: b 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. + # with MagicMock factories so initialize() can complete FPM initialization + # without dragging in a real Endpoint. monkeypatch.setattr(publisher_mod, "WorkerMetricsPublisher", MagicMock()) monkeypatch.setattr(publisher_mod, "KvEventPublisher", MagicMock()) @@ -286,13 +318,16 @@ def _build_publisher_stub(monkeypatch, *, attention_dp_size: int, fpm_enabled: b return pub, publisher_mod, fake_fpm_cls -def test_publisher_initialize_constructs_fpm_direct_publisher_when_fpm_enabled( - monkeypatch, +@pytest.mark.parametrize("attention_dp_size", [1, 4]) +def test_publisher_initialize_constructs_single_fpm_channel( + monkeypatch, publisher_mod, attention_dp_size ): - """Under non-attention-DP (attention_dp_size == 1, fpm_enabled == True), - Publisher.initialize() constructs FpmDirectPublisher with dp_size=1.""" + """TRT-LLM publishes one logical FPM stream per engine. Under + attention-DP, rank 0 emits an engine-wide aggregate, so the Dynamo side + still creates exactly one FPM channel and does not allocate heartbeat-only + channels for ranks 1..N-1.""" pub, _publisher_mod, fake_fpm_cls = _build_publisher_stub( - monkeypatch, attention_dp_size=1, fpm_enabled=True + monkeypatch, publisher_mod, attention_dp_size=attention_dp_size ) pub.initialize() fake_fpm_cls.assert_called_once() @@ -302,33 +337,17 @@ 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(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): +def _build_schema_probe_publisher(publisher_mod, 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() @@ -337,8 +356,8 @@ def _build_schema_probe_publisher(fpm_publisher_mock=None): return pub, publisher_mod -def test_schema_probe_all_fields_present_keeps_publisher(): - pub, _ = _build_schema_probe_publisher() +def test_schema_probe_all_fields_present_keeps_publisher(publisher_mod): + pub, _ = _build_schema_probe_publisher(publisher_mod) original_publisher = pub.fpm_publisher pub._check_fpm_schema(_build_fake_stat()) @@ -349,11 +368,13 @@ def test_schema_probe_all_fields_present_keeps_publisher(): @pytest.mark.parametrize("missing_field", list(_DEFAULT_IBS.keys())) -def test_schema_probe_missing_single_ibs_field_disables_publisher(missing_field): +def test_schema_probe_missing_single_ibs_field_disables_publisher( + publisher_mod, 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() + pub, _ = _build_schema_probe_publisher(publisher_mod) original_publisher = pub.fpm_publisher stat = _build_fake_stat() @@ -365,11 +386,11 @@ def test_schema_probe_missing_single_ibs_field_disables_publisher(missing_field) original_publisher.shutdown.assert_called_once() -def test_schema_probe_missing_ibs_dict_disables_publisher_legacy_trtllm(): +def test_schema_probe_missing_ibs_dict_disables_publisher_legacy_trtllm(publisher_mod): """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() + pub, _ = _build_schema_probe_publisher(publisher_mod) original_publisher = pub.fpm_publisher pub._check_fpm_schema({"iterLatencyMS": 10.0, "attentionDpRank": 0}) @@ -379,11 +400,11 @@ def test_schema_probe_missing_ibs_dict_disables_publisher_legacy_trtllm(): original_publisher.shutdown.assert_called_once() -def test_schema_probe_ibs_not_a_dict_disables_publisher(): +def test_schema_probe_ibs_not_a_dict_disables_publisher(publisher_mod): """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() + pub, _ = _build_schema_probe_publisher(publisher_mod) original_publisher = pub.fpm_publisher pub._check_fpm_schema({"inflightBatchingStats": None}) @@ -392,10 +413,11 @@ def test_schema_probe_ibs_not_a_dict_disables_publisher(): 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) +def test_schema_probe_noop_when_fpm_publisher_already_none(publisher_mod): + """If initialization failed and fpm_publisher is already None, the probe + must not blow up and must still flip _fpm_schema_checked so we do not + re-enter.""" + pub, _ = _build_schema_probe_publisher(publisher_mod, fpm_publisher_mock=None) pub.fpm_publisher = None pub._check_fpm_schema(_build_fake_stat()) @@ -404,11 +426,11 @@ def test_schema_probe_noop_when_fpm_publisher_already_none(): assert pub.fpm_publisher is None -def test_schema_probe_shutdown_exception_still_disables_publisher(): +def test_schema_probe_shutdown_exception_still_disables_publisher(publisher_mod): """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, _ = _build_schema_probe_publisher(publisher_mod) pub.fpm_publisher.shutdown.side_effect = RuntimeError("tokio runtime gone") stat = _build_fake_stat() @@ -419,11 +441,11 @@ def test_schema_probe_shutdown_exception_still_disables_publisher(): assert pub._fpm_schema_checked is True -def test_handle_stat_probe_gate_fires_once_and_skips_subsequent_stats(): +def test_handle_stat_probe_gate_fires_once_and_skips_subsequent_stats(publisher_mod): """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() + pub, _ = _build_schema_probe_publisher(publisher_mod) original_publisher = pub.fpm_publisher if pub.fpm_publisher is not None and not pub._fpm_schema_checked: @@ -439,12 +461,10 @@ def test_handle_stat_probe_gate_fires_once_and_skips_subsequent_stats(): original_publisher.shutdown.assert_not_called() -def test_schema_probe_field_list_matches_default_ibs_set(): +def test_schema_probe_field_list_matches_default_ibs_set(publisher_mod): """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()) diff --git a/docs/backends/trtllm/README.md b/docs/backends/trtllm/README.md index 6b36aae2951e..2e439f4c6cdf 100644 --- a/docs/backends/trtllm/README.md +++ b/docs/backends/trtllm/README.md @@ -86,5 +86,6 @@ You can deploy TensorRT-LLM with Dynamo on Kubernetes using a `DynamoGraphDeploy - **[Examples](trtllm-examples.md)**: All deployment patterns with launch scripts - **[KV Cache Transfer](trtllm-kv-cache-transfer.md)**: KV cache transfer methods for disaggregated serving - **[Observability](trtllm-observability.md)**: Metrics and monitoring +- **[Attention-DP FPM Metrics Repro](trtllm-attention-dp-fpm-repro.md)**: Local build, launch, and metrics capture for attention-DP FPM validation - **[Multinode Examples](multinode/trtllm-multinode-examples.md)**: Multi-node deployment with SLURM - **[Deploying TensorRT-LLM with Dynamo on Kubernetes](https://github.com/ai-dynamo/dynamo/tree/main/examples/backends/trtllm/deploy/README.md)**: Kubernetes deployment guide From 0004aa589400ad625697e781bab80481a0bb4bb6 Mon Sep 17 00:00:00 2001 From: Indrajit Bhosale Date: Mon, 4 May 2026 10:49:18 -0700 Subject: [PATCH 2/4] Update tests Signed-off-by: Indrajit Bhosale --- .../trtllm/tests/test_trtllm_fpm_publisher.py | 120 ++++++++---------- docs/backends/trtllm/README.md | 1 - 2 files changed, 50 insertions(+), 71 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 9a40e4e3e58b..47d856cb89f6 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py @@ -26,10 +26,6 @@ from __future__ import annotations -import importlib.util -import sys -import types -from pathlib import Path from unittest.mock import MagicMock import pytest @@ -42,37 +38,6 @@ ] -@pytest.fixture() -def publisher_mod(monkeypatch): - prometheus_mod = types.ModuleType("dynamo.common.utils.prometheus") - prometheus_mod.LLMBackendMetrics = object - - llm_mod = types.ModuleType("dynamo.llm") - llm_mod.FpmDirectPublisher = MagicMock(name="FpmDirectPublisher") - llm_mod.KvEventPublisher = MagicMock(name="KvEventPublisher") - llm_mod.WorkerMetricsPublisher = MagicMock(name="WorkerMetricsPublisher") - - monkeypatch.setitem(sys.modules, "dynamo", types.ModuleType("dynamo")) - monkeypatch.setitem(sys.modules, "dynamo.common", types.ModuleType("dynamo.common")) - monkeypatch.setitem( - sys.modules, "dynamo.common.utils", types.ModuleType("dynamo.common.utils") - ) - monkeypatch.setitem(sys.modules, "dynamo.common.utils.prometheus", prometheus_mod) - monkeypatch.setitem(sys.modules, "dynamo.llm", llm_mod) - - repo_root = Path(__file__).resolve().parents[5] - publisher_path = repo_root / "components/src/dynamo/trtllm/publisher.py" - spec = importlib.util.spec_from_file_location( - "_trtllm_publisher_under_test", publisher_path - ) - assert spec is not None - assert spec.loader is not None - - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - # 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. @@ -249,20 +214,22 @@ def test_iter_latency_ms_to_wall_time_secs_conversion(): # --------------------------------------------------------------------------- -# Publisher.initialize() FPM channel behavior +# Publisher.initialize() gate behavior (attention-DP off vs on) # --------------------------------------------------------------------------- -def _build_publisher_stub(monkeypatch, publisher_mod, *, attention_dp_size: int): +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 FPM initialization cleanly - without needing a blanket try/except to swallow upstream failures. + 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 @@ -277,6 +244,7 @@ def _build_publisher_stub(monkeypatch, publisher_mod, *, attention_dp_size: int) 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 @@ -293,8 +261,8 @@ def _build_publisher_stub(monkeypatch, publisher_mod, *, attention_dp_size: int) 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 FPM initialization - # without dragging in a real Endpoint. + # 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()) @@ -318,16 +286,27 @@ def _build_publisher_stub(monkeypatch, publisher_mod, *, attention_dp_size: int) return pub, publisher_mod, fake_fpm_cls -@pytest.mark.parametrize("attention_dp_size", [1, 4]) -def test_publisher_initialize_constructs_single_fpm_channel( - monkeypatch, publisher_mod, attention_dp_size +def test_publisher_initialize_constructs_fpm_direct_publisher_when_fpm_enabled( + monkeypatch, ): - """TRT-LLM publishes one logical FPM stream per engine. Under - attention-DP, rank 0 emits an engine-wide aggregate, so the Dynamo side - still creates exactly one FPM channel and does not allocate heartbeat-only - channels for ranks 1..N-1.""" + """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_initializes_fpm_publisher_under_attention_dp(monkeypatch): + """Under attention-DP (attention_dp_size > 1), Publisher.initialize() + still constructs the single engine-wide FpmDirectPublisher channel.""" pub, _publisher_mod, fake_fpm_cls = _build_publisher_stub( - monkeypatch, publisher_mod, attention_dp_size=attention_dp_size + monkeypatch, attention_dp_size=4, fpm_enabled=False ) pub.initialize() fake_fpm_cls.assert_called_once() @@ -342,12 +321,14 @@ def test_publisher_initialize_constructs_single_fpm_channel( # --------------------------------------------------------------------------- -def _build_schema_probe_publisher(publisher_mod, fpm_publisher_mock=None): +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() @@ -356,8 +337,8 @@ def _build_schema_probe_publisher(publisher_mod, fpm_publisher_mock=None): return pub, publisher_mod -def test_schema_probe_all_fields_present_keeps_publisher(publisher_mod): - pub, _ = _build_schema_probe_publisher(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()) @@ -368,13 +349,11 @@ def test_schema_probe_all_fields_present_keeps_publisher(publisher_mod): @pytest.mark.parametrize("missing_field", list(_DEFAULT_IBS.keys())) -def test_schema_probe_missing_single_ibs_field_disables_publisher( - publisher_mod, missing_field -): +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(publisher_mod) + pub, _ = _build_schema_probe_publisher() original_publisher = pub.fpm_publisher stat = _build_fake_stat() @@ -386,11 +365,11 @@ def test_schema_probe_missing_single_ibs_field_disables_publisher( original_publisher.shutdown.assert_called_once() -def test_schema_probe_missing_ibs_dict_disables_publisher_legacy_trtllm(publisher_mod): +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(publisher_mod) + pub, _ = _build_schema_probe_publisher() original_publisher = pub.fpm_publisher pub._check_fpm_schema({"iterLatencyMS": 10.0, "attentionDpRank": 0}) @@ -400,11 +379,11 @@ def test_schema_probe_missing_ibs_dict_disables_publisher_legacy_trtllm(publishe original_publisher.shutdown.assert_called_once() -def test_schema_probe_ibs_not_a_dict_disables_publisher(publisher_mod): +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(publisher_mod) + pub, _ = _build_schema_probe_publisher() original_publisher = pub.fpm_publisher pub._check_fpm_schema({"inflightBatchingStats": None}) @@ -413,11 +392,10 @@ def test_schema_probe_ibs_not_a_dict_disables_publisher(publisher_mod): original_publisher.shutdown.assert_called_once() -def test_schema_probe_noop_when_fpm_publisher_already_none(publisher_mod): - """If initialization failed and fpm_publisher is already None, the probe - must not blow up and must still flip _fpm_schema_checked so we do not - re-enter.""" - pub, _ = _build_schema_probe_publisher(publisher_mod, fpm_publisher_mock=None) +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()) @@ -426,11 +404,11 @@ def test_schema_probe_noop_when_fpm_publisher_already_none(publisher_mod): assert pub.fpm_publisher is None -def test_schema_probe_shutdown_exception_still_disables_publisher(publisher_mod): +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(publisher_mod) + pub, _ = _build_schema_probe_publisher() pub.fpm_publisher.shutdown.side_effect = RuntimeError("tokio runtime gone") stat = _build_fake_stat() @@ -441,11 +419,11 @@ def test_schema_probe_shutdown_exception_still_disables_publisher(publisher_mod) assert pub._fpm_schema_checked is True -def test_handle_stat_probe_gate_fires_once_and_skips_subsequent_stats(publisher_mod): +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(publisher_mod) + pub, _ = _build_schema_probe_publisher() original_publisher = pub.fpm_publisher if pub.fpm_publisher is not None and not pub._fpm_schema_checked: @@ -461,10 +439,12 @@ def test_handle_stat_probe_gate_fires_once_and_skips_subsequent_stats(publisher_ original_publisher.shutdown.assert_not_called() -def test_schema_probe_field_list_matches_default_ibs_set(publisher_mod): +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()) diff --git a/docs/backends/trtllm/README.md b/docs/backends/trtllm/README.md index 2e439f4c6cdf..6b36aae2951e 100644 --- a/docs/backends/trtllm/README.md +++ b/docs/backends/trtllm/README.md @@ -86,6 +86,5 @@ You can deploy TensorRT-LLM with Dynamo on Kubernetes using a `DynamoGraphDeploy - **[Examples](trtllm-examples.md)**: All deployment patterns with launch scripts - **[KV Cache Transfer](trtllm-kv-cache-transfer.md)**: KV cache transfer methods for disaggregated serving - **[Observability](trtllm-observability.md)**: Metrics and monitoring -- **[Attention-DP FPM Metrics Repro](trtllm-attention-dp-fpm-repro.md)**: Local build, launch, and metrics capture for attention-DP FPM validation - **[Multinode Examples](multinode/trtllm-multinode-examples.md)**: Multi-node deployment with SLURM - **[Deploying TensorRT-LLM with Dynamo on Kubernetes](https://github.com/ai-dynamo/dynamo/tree/main/examples/backends/trtllm/deploy/README.md)**: Kubernetes deployment guide From 23c8a4284c9fae03f123c65ae3609774cd54d587 Mon Sep 17 00:00:00 2001 From: Indrajit Bhosale Date: Wed, 6 May 2026 19:09:44 -0700 Subject: [PATCH 3/4] Emit per rank stats Signed-off-by: Indrajit Bhosale --- components/src/dynamo/trtllm/publisher.py | 40 ++++++---- .../trtllm/tests/test_trtllm_fpm_publisher.py | 79 +++++++++++++++++-- 2 files changed, 96 insertions(+), 23 deletions(-) diff --git a/components/src/dynamo/trtllm/publisher.py b/components/src/dynamo/trtllm/publisher.py index ed064c5fbbe0..dad9419d1666 100644 --- a/components/src/dynamo/trtllm/publisher.py +++ b/components/src/dynamo/trtllm/publisher.py @@ -400,18 +400,18 @@ def __init__( # Needed by the events and metrics publishers self.metrics_publisher: Optional[WorkerMetricsPublisher] = None - # FPM is emitted as one logical channel per TRT-LLM engine. Under - # attention-DP, TRT-LLM rank 0 emits an engine-wide aggregate; ranks - # 1..N-1 have no independent planner-visible worker identity. + # FPM is emitted as one logical channel per TRT-LLM attention-DP rank. + # TRT-LLM tags each IterationStats row with attentionDpRank, which is + # forwarded as Dynamo's dp_rank for planner-visible per-rank load. 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 + self.kv_event_publishers: Optional[Dict[int, KvEventPublisher]] = ( + None # One per attention_dp_rank + ) self.zmq_kv_event_publisher = None # ZMQ publisher for consolidator self.publish_kv_cache_events_thread: Optional[ManagedThread] = None self.publish_stats_thread: Optional[ManagedThread] = None @@ -452,17 +452,18 @@ def initialize(self) -> None: lambda _: logging.debug("metrics publisher endpoint created") ) - # Setup the ForwardPassMetrics publisher. Always use a single FPM - # channel: non-attention-DP has one logical rank, and attention-DP - # reports one rank-0 engine aggregate. If rank 0 stalls, the whole - # TRT-LLM engine stalls, so FPM emission is intentionally all-or-nothing. + # Setup the ForwardPassMetrics publisher with one internal channel per + # attention-DP rank. Non-attention-DP engines report size 1. Under + # attention-DP, TRT-LLM emits one IterationStats row per rank and + # Dynamo forwards attentionDpRank as the FPM dp_rank. try: + fpm_dp_size = max(1, int(self.attention_dp_size or 1)) self.fpm_publisher = FpmDirectPublisher( endpoint=self.endpoint, worker_id=str(self.worker_id), - dp_size=1, + dp_size=fpm_dp_size, ) - logging.info("FpmDirectPublisher initialized with dp_size=1") + logging.info(f"FpmDirectPublisher initialized with dp_size={fpm_dp_size}") except RuntimeError as e: # PyO3 surfaces all FpmDirectPublisher::new failures as # PyRuntimeError (Endpoint missing, tokio runtime missing, @@ -654,10 +655,11 @@ def handle_stat(stat): # Publish ForwardPassMetrics. TRT-LLM tags each stat dict with # top-level attentionDpRank inside BaseWorker._stats_serializer. - # Under attention-DP, the companion TRT-LLM change emits one - # rank-0 engine aggregate. The FPM source fields live nested under - # stat["inflightBatchingStats"] (camelCase from NLOHMANN - # serialization). Variance fields are not yet computed in + # Under attention-DP, TRT-LLM emits one row per rank. Scheduled + # fields are rank-local, while engine-global queued fields are + # naturally nonzero only on rank 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: @@ -699,8 +701,12 @@ def handle_stat(stat): ) + int(ibs.get("numQueuedGenKvTokens", 0)) # iterLatencyMS is ms; the Rust snapshot expects seconds. wall_time_secs = float(stat.get("iterLatencyMS", 0.0)) / 1000.0 + attention_dp_rank = stat.get("attentionDpRank") + dp_rank = ( + int(attention_dp_rank) if attention_dp_rank is not None else 0 + ) self.fpm_publisher.publish( - dp_rank=int(stat.get("attentionDpRank", 0)), + dp_rank=dp_rank, 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, 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 47d856cb89f6..ee5d01f8eb84 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py @@ -11,7 +11,10 @@ * 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. + unchanged; missing/null key defaults to 0. + * Attention-DP fanout emits one FPM per attentionDpRank; queued fields are + forwarded from the stat row as-is, so they remain nonzero only on the + rank-0 row produced by TRT-LLM. * 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 @@ -89,8 +92,10 @@ def _invoke_handler(stat, fpm_publisher): queued_sum_decode_kv_tokens = int(ibs.get("numPausedKvTokens", 0)) + int( ibs.get("numQueuedGenKvTokens", 0) ) + attention_dp_rank = stat.get("attentionDpRank") + dp_rank = int(attention_dp_rank) if attention_dp_rank is not None else 0 fpm_publisher.publish( - dp_rank=int(stat.get("attentionDpRank", 0)), + dp_rank=dp_rank, 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)), @@ -178,6 +183,62 @@ def test_handle_stat_routes_per_attention_dp_rank(): assert call.kwargs["scheduled_sum_prefill_tokens"] == 100 * (i + 1) +def test_handle_stat_attention_dp_fanout_preserves_rank0_queue_only(): + """TRT-LLM emits one stat row per attention-DP rank. Dynamo forwards each + row as its own FPM and does not smear rank-0 queued fields onto rank 1.""" + fpm = MagicMock() + rank0 = _build_fake_stat( + attentionDpRank=0, + iterLatencyMS=2147.591, + ibs_overrides={ + "numContextRequests": 1, + "numCtxTokens": 15, + "numCtxKvTokens": 1, + "numGenRequests": 0, + "numGenKvTokens": 0, + "numPausedRequests": 0, + "numPausedKvTokens": 0, + "numQueuedContextRequests": 8664, + "numQueuedCtxTokens": 155952, + "numQueuedGenRequests": 0, + "numQueuedGenKvTokens": 0, + }, + ) + rank1 = _build_fake_stat( + attentionDpRank=1, + iterLatencyMS=2147.591, + ibs_overrides={ + "numContextRequests": 0, + "numCtxTokens": 0, + "numCtxKvTokens": 0, + "numGenRequests": 2, + "numGenKvTokens": 33, + "numPausedRequests": 0, + "numPausedKvTokens": 0, + "numQueuedContextRequests": 0, + "numQueuedCtxTokens": 0, + "numQueuedGenRequests": 0, + "numQueuedGenKvTokens": 0, + }, + ) + + _invoke_handler(rank0, fpm) + _invoke_handler(rank1, fpm) + + rank0_call, rank1_call = fpm.publish.call_args_list + assert rank0_call.kwargs["dp_rank"] == 0 + assert rank0_call.kwargs["queued_num_prefill_requests"] == 8664 + assert rank0_call.kwargs["queued_sum_prefill_tokens"] == 155952 + assert rank0_call.kwargs["scheduled_num_decode_requests"] == 0 + + assert rank1_call.kwargs["dp_rank"] == 1 + assert rank1_call.kwargs["queued_num_prefill_requests"] == 0 + assert rank1_call.kwargs["queued_sum_prefill_tokens"] == 0 + assert rank1_call.kwargs["scheduled_num_decode_requests"] == 2 + assert rank1_call.kwargs["scheduled_sum_decode_kv_tokens"] == 33 + assert rank0_call.kwargs["wall_time_secs"] == rank1_call.kwargs["wall_time_secs"] + + def test_handle_stat_missing_attention_dp_rank_defaults_zero(): fpm = MagicMock() stat = _build_fake_stat() @@ -186,6 +247,12 @@ def test_handle_stat_missing_attention_dp_rank_defaults_zero(): assert fpm.publish.call_args.kwargs["dp_rank"] == 0 +def test_handle_stat_null_attention_dp_rank_defaults_zero(): + fpm = MagicMock() + _invoke_handler(_build_fake_stat(attentionDpRank=None), 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), @@ -304,7 +371,7 @@ def test_publisher_initialize_constructs_fpm_direct_publisher_when_fpm_enabled( def test_publisher_initializes_fpm_publisher_under_attention_dp(monkeypatch): """Under attention-DP (attention_dp_size > 1), Publisher.initialize() - still constructs the single engine-wide FpmDirectPublisher channel.""" + constructs one FpmDirectPublisher channel per attention-DP rank.""" pub, _publisher_mod, fake_fpm_cls = _build_publisher_stub( monkeypatch, attention_dp_size=4, fpm_enabled=False ) @@ -312,7 +379,7 @@ def test_publisher_initializes_fpm_publisher_under_attention_dp(monkeypatch): 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 kwargs["dp_size"] == 4 assert pub.fpm_publisher is not None @@ -393,8 +460,8 @@ def test_schema_probe_ibs_not_a_dict_disables_publisher(): 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.""" + """If initialization already disabled fpm_publisher, 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 From b3ef42cbddb8588c95cb4f172bfb12bd3b9e3098 Mon Sep 17 00:00:00 2001 From: Indrajit Bhosale Date: Wed, 6 May 2026 19:10:06 -0700 Subject: [PATCH 4/4] Emit per rank stats Signed-off-by: Indrajit Bhosale --- components/src/dynamo/trtllm/publisher.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/components/src/dynamo/trtllm/publisher.py b/components/src/dynamo/trtllm/publisher.py index dad9419d1666..281d6c1d23f5 100644 --- a/components/src/dynamo/trtllm/publisher.py +++ b/components/src/dynamo/trtllm/publisher.py @@ -409,9 +409,9 @@ def __init__( # 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 - ) + self.kv_event_publishers: Optional[ + Dict[int, KvEventPublisher] + ] = None # One per attention_dp_rank self.zmq_kv_event_publisher = None # ZMQ publisher for consolidator self.publish_kv_cache_events_thread: Optional[ManagedThread] = None self.publish_stats_thread: Optional[ManagedThread] = None