feat(trtllm): strict first-stat FPM schema probe; gate off on missing fields - #8396
Draft
nv-yna wants to merge 5 commits into
Draft
feat(trtllm): strict first-stat FPM schema probe; gate off on missing fields#8396nv-yna wants to merge 5 commits into
nv-yna wants to merge 5 commits into
Conversation
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 <nv-yna@users.noreply.github.com>
Signed-off-by: Yuewei Na <nv-yna@users.noreply.github.com>
…r 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 <nv-yna@users.noreply.github.com>
… fields Closes the "old TRT-LLM" planner-poison gap left by ai-dynamo#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 ai-dynamo#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 <nv-yna@users.noreply.github.com>
Contributor
|
👋 Hi nv-yna! Thank you for contributing to ai-dynamo/dynamo. Just a reminder: The 🚀 |
github-actions
Bot
added
external-contribution
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 <nv-yna@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Stacked on top of #8356 — this PR assumes that PR lands first. The only new commit here is
038e668. Once #8356 merges, this PR's diff narrows to just that commit.Closes the "old TRT-LLM" planner-poison gap that the MVP doesn't cover.
Problem
handle_stat's.publish(...)call incomponents/src/dynamo/trtllm/publisher.pyreads every requiredIterationStatsfield viastat.get(field, 0). On a TRT-LLM version that predates the companion TensorRT-LLM#13199, every field defaults to0and the emittedForwardPassSnapshotbecomes byte-identical to the Rust-side idle heartbeat. The Planner reads that as "worker is idle" and can scale down a worker that is actually under real load.This is the same class of bug as the attention-DP ranks 1..N-1 case — which #8356 explicitly gates via
fpm_enabled = attention_dp_size == 1— but ungated for the old-TRT-LLM path.Fix
One-shot strict schema probe on the first
IterationStatsdelivered tohandle_stat:self.fpm_publisher.shutdown(),self.fpm_publisher = None.if self.fpm_publisher is not None:guard then short-circuits all further FPM emission for the lifetime of the worker, matching the attention-DP gate's "zero messages on the topic" contract.Fail-safe: if
shutdown()itself raises, the publisher is stillNone'd out — the goal is to suppress emission, not succeed at shutdown.Required fields (strict probe — all must be present)
scheduledNumPrefillRequestsscheduledSumPrefillTokensscheduledSumPrefillKvTokensscheduledNumDecodeRequestsscheduledSumDecodeKvTokensqueuedNumPrefillRequestsqueuedSumPrefillTokensqueuedNumDecodeRequestsqueuedSumDecodeKvTokensiterLatencyMSandattentionDpRankare intentionally not probed — both predate #13199 and their presence isn't a signal about whether the companion PR is applied.Changes
components/src/dynamo/trtllm/publisher.py:_FPM_REQUIRED_STAT_FIELDS(9 fields).self._fpm_schema_checked: bool = FalseinPublisher.__init__.Publisher._check_fpm_schema(stat)method.handle_statahead of the existing FPM publish branch.components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py:No Rust changes. No API surface changes.
Test plan