fix: synchronize self-benchmark capacity across DP ranks

[liyuanzhe1991]

Overview:

Self-benchmarking with attention-DP > 1 intermittently deadlocks at startup: each DP rank probes its local KV-cache capacity independently, and any run-to-run variation (block-pool rounding, allocator state) makes ranks generate different benchmark grids. The lockstep synchronization protocol then waits forever on points that only a subset of ranks generated. In our GLM-5.2-NVFP4 / 4×B200 campaigns, decode benchmarks with dp=4 failed 3 out of 4 attempts before this fix.

Details:

• Each rank builds a local capacity envelope (_bench_make_local_capacity: max_model_len, max_num_running_reqs, max_num_scheduled_tokens, usable KV blocks) and exchanges it through the existing benchmark synchronizer before grid generation.
• Grid construction reads capacity exclusively through _bench_capacity_limit() / _bench_grid_usable_blocks(), which return the minimum across ranks — every rank derives byte-identical grids by construction.
• _bench_decode_feasible_max_batch_size() now computes from the synchronized envelope, so both call sites (the native decode grid and explicit benchmark points) agree across ranks.
• Single-rank (dp=1) behavior is unchanged: the synchronized envelope degenerates to the local one.

Validated on 4-rank and 8-rank attention-DP topologies (GLM-5.2-NVFP4, B200, vLLM 0.25.1): previously flaky decode cells pass deterministically; full prefill+decode campaigns (5 466 and 6 508 grid points) completed with zero lockstep timeouts.

Where should the reviewer start?

• components/src/dynamo/vllm/instrumented_scheduler.py:
  • _BenchmarkCapacityEnvelope, _bench_make_local_capacity, _bench_capacity_limit, _bench_grid_usable_blocks
  • the capacity exchange in the _BenchmarkSynchronizer handshake
  • _bench_decode_feasible_max_batch_size (now envelope-based)
• components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py — capacity-divergence regression tests

Related Issues

🚫 This PR is NOT linked to an issue:

• Confirmed — no related issue

Summary by CodeRabbit

• New Features

  • Improved distributed benchmark setup by negotiating shared capacity across ranks.
  • Benchmark grids now adapt to the common capacity available across participating ranks.
  • Benchmark results include the negotiated common capacity when available.

• Bug Fixes

  • Added validation to detect capacity and benchmark-grid mismatches before warmup.
  • Improved KV-cache capacity accounting for prefill and decode benchmark scenarios.

• Tests

  • Added regression coverage for capacity negotiation, grid synchronization, and mismatch handling.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

👋 Hi liyuanzhe1991! Thank you for contributing to ai-dynamo/dynamo.

Just a reminder: The NVIDIA Test Github Validation CI runs an essential subset of the testing framework to quickly catch errors.Your PR reviewers may elect to test the changes comprehensively before approving your changes.

🚀

[coderabbitai]

Walkthrough

Changes

Benchmark capacity coordination

Layer / File(s)	Summary
Capacity envelope and rank synchronization components/src/dynamo/vllm/instrumented_scheduler.py, components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py	Capacity reports are validated and reduced to common minimum limits across attention-DP ranks. Grid digests, expected point counts, and missing phases are synchronized, with regression tests for mismatch errors.
Benchmark grid construction and lifecycle state components/src/dynamo/vllm/instrumented_scheduler.py, components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py	Benchmark state tracks local and negotiated capacity, grid construction performs synchronization, results include common capacity, and existing grid tests install capacity preflight stubs.
Capacity-aware prefill and decode feasibility components/src/dynamo/vllm/instrumented_scheduler.py, components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py	Prefill and decode generation use negotiated limits and grid-aware usable-block accounting, with regression tests verifying identical grids across differing rank capacities.

Estimated code review effort: 4 (Complex) | ~60 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: synchronizing self-benchmark capacity across DP ranks.
Description check	✅ Passed	The description follows the template with Overview, Details, reviewer guidance, and a completed Related Issues section.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py (1)

552-699: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add a @pytest.mark.timeout(...) to the threaded synchronizer tests.

These three tests spawn a follower thread and rely on ZMQ socket polling/waits between ranks. If coordination ever wedges, follower.join(timeout=2) returns while the socket-blocked thread lingers; a bounded per-test timeout keeps CI from hanging on regressions. As per path instructions: "Add @pytest.mark.timeout(...) for tests that may exceed 30s or contain polling/sleeps/waits."

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py` around
lines 552 - 699, Add an appropriate `@pytest.mark.timeout`(...) decorator to each
threaded synchronizer test:
test_benchmark_synchronizer_negotiates_minimum_capacity_and_grid,
test_benchmark_synchronizer_rejects_capacity_invariant_mismatch, and
test_benchmark_synchronizer_rejects_grid_mismatch_before_warmup. Keep the
existing thread joins and test behavior unchanged.

Source: Path instructions

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py`:
- Around line 552-699: Add an appropriate `@pytest.mark.timeout`(...) decorator to
each threaded synchronizer test:
test_benchmark_synchronizer_negotiates_minimum_capacity_and_grid,
test_benchmark_synchronizer_rejects_capacity_invariant_mismatch, and
test_benchmark_synchronizer_rejects_grid_mismatch_before_warmup. Keep the
existing thread joins and test behavior unchanged.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 75d54ebb-7b81-43c1-b341-42f7548b71af

📥 Commits

Reviewing files that changed from the base of the PR and between 116a1cb and 06ab021.

📒 Files selected for processing (2)

• components/src/dynamo/vllm/instrumented_scheduler.py
• components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py

[tedzhouhk]

Three blocking findings from the agentic review are noted inline.

[dynamo-ops]

/ok to test 79717ca

[dynamo-ops]

/ok to test 077a2c5

[dynamo-ops]

/ok to test 1f16d7e

[dynamo-ops]

/ok to test bc00233

[tedzhouhk]

Verified the three previously raised P1 concerns against the current head: decode capture invariants are negotiated correctly, explicit prefill metadata uses the common limit, and the unused locals are removed. The added regression coverage matches the fixes.