feat: Atn-dp fpm upstream integration

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

components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py

@@ -302,18 +302,18 @@ 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)."""
+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, attention_dp_size=4, fpm_enabled=False
     )
     pub.initialize()
-    fake_fpm_cls.assert_not_called()
-    assert pub.fpm_publisher is None
+    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
 
 
 # ---------------------------------------------------------------------------