[None][feat] Add startup self-benchmarking with batched prefill - #16191
Open
venkywonka wants to merge 12 commits into
Open
[None][feat] Add startup self-benchmarking with batched prefill#16191venkywonka wants to merge 12 commits into
venkywonka wants to merge 12 commits into
Conversation
Signed-off-by: Sachal Malick <s@chal.ai>
Signed-off-by: Sachal Malick <s@chal.ai>
Signed-off-by: Sachal Malick <s@chal.ai>
Signed-off-by: Sachal Malick <s@chal.ai>
Signed-off-by: Sachal Malick <s@chal.ai>
Signed-off-by: Sachal Malick <s@chal.ai>
Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>
venkywonka
added
the
api-compatible
Contributor
📝 WalkthroughWalkthroughAdds configurable startup self-benchmarking to PyExecutor, including synthetic prefill/decode workloads, tensor-parallel planning, metrics capture, cache validation, terminal JSON output, request handling, configuration validation, documentation, and extensive tests. ChangesSelf-benchmark configuration
Estimated code review effort: 5 (Critical) | ~90+ minutes Sequence Diagram(s)sequenceDiagram
participant PyExecutor
participant SelfBenchmark
participant KVCacheManager
participant JSONOutput
PyExecutor->>SelfBenchmark: synchronize plan and outcome
SelfBenchmark->>KVCacheManager: allocate synthetic decode KV entries
SelfBenchmark->>PyExecutor: inject synthetic requests
PyExecutor->>SelfBenchmark: report scheduled and executed batches
SelfBenchmark->>JSONOutput: write benchmark results
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py (1)
282-296: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
.get("enable_self_benchmark", True)default can mask a missing-kwarg regression.If
create_py_executor_instanceever forgot to passenable_self_benchmarkfor one of the two calls,kwargs.get(..., True)would silently reportTrue, and since the expected second-call value is alsoTrue, the assertion at Line 387-390 would still pass without ever confirming the kwarg was actually threaded through. Consider using a sentinel default (e.g.kwargs.get("enable_self_benchmark", "MISSING")) so an omitted kwarg fails the assertion instead of being masked.🔧 Proposed fix
executor_call_records.append( { - "enable_self_benchmark": kwargs.get("enable_self_benchmark", True), + "enable_self_benchmark": kwargs.get("enable_self_benchmark", "MISSING"), "executor_uses_canonical_args": kwargs["llm_args"] is llm_args, "model_engine_uses_canonical_args": kwargs["model_engine"].llm_args is llm_args, } )🤖 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 `@tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py` around lines 282 - 296, The `.get("enable_self_benchmark", True)` call in the _create_py_executor_instance function provides a default value of True which can mask a missing-kwarg regression. If the enable_self_benchmark kwarg is ever forgotten in the create_py_executor_instance calls, the default True would still match the expected True value, allowing the assertion to pass undetected. Replace the default value True with a sentinel value such as "MISSING" so that if the kwarg is omitted, the recorded value will be "MISSING" instead of True, causing the assertion at lines 387-390 to fail and properly detect the regression.
🤖 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
`@tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py`:
- Around line 282-296: The `.get("enable_self_benchmark", True)` call in the
_create_py_executor_instance function provides a default value of True which can
mask a missing-kwarg regression. If the enable_self_benchmark kwarg is ever
forgotten in the create_py_executor_instance calls, the default True would still
match the expected True value, allowing the assertion to pass undetected.
Replace the default value True with a sentinel value such as "MISSING" so that
if the kwarg is omitted, the recorded value will be "MISSING" instead of True,
causing the assertion at lines 387-390 to fail and properly detect the
regression.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: a9a96856-88b8-4fd8-987e-bfc09deca3af
📒 Files selected for processing (15)
docs/source/commands/trtllm-serve/trtllm-serve.rsttensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/llm_request.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.pytensorrt_llm/_torch/pyexecutor/request_utils.pytensorrt_llm/_torch/pyexecutor/self_benchmark.pytensorrt_llm/llmapi/__init__.pytensorrt_llm/llmapi/llm_args.pytests/integration/test_lists/test-db/l0_a10.ymltests/unittest/_torch/executor/test_dual_pool_kv_cache.pytests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.pytests/unittest/_torch/executor/test_self_benchmark.pytests/unittest/api_stability/references/llm.yamltests/unittest/llmapi/test_llm_args.py
Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>
Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>
Collaborator
Author
|
@coderabbitai full review |
Contributor
✅ Action performedFull review finished. |
…anks Replace counter-driven startup self-benchmark completion with exact per-request lifecycle tracking, scheduler-authoritative admission planning, rank-symmetric TP consensus, and additive schema-version-2 diagnostics. SelfBenchmark owns immutable BenchmarkCase, mutable BenchmarkTrial, frozen BenchmarkTrialResult, exact per-request schedule signatures, drain/skip/abort lifecycle, coverage, and schema-v2 terminal artifacts (complete, invalid, aborted, interrupted). PyExecutor performs ordered tp_allreduce(MAX) outcome consensus, startup planner consensus, scheduler filtering with overlap-safe draining, execution/termination/failure callbacks, and post-resource-release notification. llm_request and request_utils carry the trial marker through request conversion and broadcast with a deletion-only ownership boundary; when self_benchmark_config is absent the normal serving path reaches no benchmark helper, protected-ID set, marker scan, or collective beyond direct None guards. Validated on 4x NVIDIA L40S with a golden real-GPU campaign (all_passed=true): disabled-mode A/B non-regression (median throughput ratio 1.00027), single-GPU aggregate exactness, TP2 deferred-request, TP2 partial-admission and forced-chunking skip, TP2 rank-asymmetric consensus skip and construction- failure abort, and disaggregated TP2+TP2 over NIXL/UCX, plus a 273-test focused CPU suite. Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>
…ive alias Quality-only cleanup of the startup self-benchmark implementation with no behavior change: - remove dead helpers _num_context_requests/_num_decode_requests and the write-only fields BenchmarkTrial.scheduled_total_batch_size, RequestScheduleSignature.request_id, and SelfBenchmark._planner_synchronized - fold the sync_active property (a pure alias) into active - hoist the duplicated attn_metadata/kv_cache_dtype_byte_size getattr above the benchmark_active branch in the overlap loop to match _process_previous_batch Focused suite green on L40S: test_self_benchmark.py 90 passed; the request_utils/benchmark_disagg/llm_args unit tests 282 passed. Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>
Collaborator
Author
|
@coderabbitai full review |
Contributor
✅ Action performedFull review finished. |
Collaborator
Author
|
@coderabbitai call sequence diagram |
Contributor
|
Tip For best results, initiate chat on the files or code changes. 🧩 Analysis chain🏁 Script executed: #!/bin/bash
set -e
FILE=$(fd -a py_executor.py -p tensorrt_llm/_torch/pyexecutor)
echo "== py_executor.py outline =="
ast-grep outline "$FILE" --items all --type function,class --match "Executor|benchmark|fetch|process_iter|handle_responses"
echo "== self_benchmark.py outline =="
FILE2=$(fd -a self_benchmark.py -p tensorrt_llm/_torch/pyexecutor)
ast-grep outline "$FILE2" --items all
echo "== _fetch_and_enqueue_requests =="
rg -n "_fetch_and_enqueue_requests" -A 40 "$FILE" | head -100
echo "== _fetch_and_activate_new_requests (self-benchmark part) =="
rg -n "make_prefill_requests|make_decode_requests" -B5 -A15 "$FILE"
echo "== _process_iter_stats (self-benchmark part) =="
rg -n "_process_iter_stats" -A 40 "$FILE" | head -80
echo "== _handle_responses (self-benchmark part) =="
rg -n "is_self_benchmark_request" -B5 -A15 "$FILE"Length of output: 25245 🏁 Script executed: #!/bin/bash
set -e
FILE=$(fd -a py_executor.py -p tensorrt_llm/_torch/pyexecutor)
echo "== _sync_self_benchmark_plan / _sync_self_benchmark_outcome =="
rg -n "_sync_self_benchmark_plan|_sync_self_benchmark_outcome|start_trial|apply_planner_consensus|apply_global_outcome|local_outcome|local_planner_record" "$FILE" -B2 -A20 | head -220
echo "== make_prefill_requests / make_decode_requests bodies =="
FILE2=$(fd -a self_benchmark.py -p tensorrt_llm/_torch/pyexecutor)
rg -n "def make_prefill_requests" -A 40 "$FILE2"
rg -n "def make_decode_requests" -A 40 "$FILE2"
rg -n "def observe_scheduled_batch|def observe_executed_batch|def observe_finished_requests" -A 25 "$FILE2"Length of output: 15136
sequenceDiagram
autonumber
participant Init as PyExecutor.__init__
participant PE as PyExecutor (per-iter loop)
participant SB as SelfBenchmark
participant Dist as tp_rank Dist collectives
participant KV as KVCacheManager
participant WQ as WaitingQueue
participant Perf as PerfManager
Note over Init,SB: Startup planning (once, at construction)
Init->>SB: _create_self_benchmark(enable_self_benchmark)
Init->>PE: _sync_self_benchmark_plan()
PE->>SB: local_planner_record(tp_rank)
PE->>Dist: tp_allgather(local_record)
PE->>SB: apply_planner_consensus(records)
Note over SB: builds grid of BenchmarkCase (prefill/prefill_seed/decode/warmup)
loop Executor iteration (while benchmark.active)
PE->>PE: _prepare_and_schedule_batch()
PE->>PE: _fetch_and_activate_new_requests()
PE->>SB: make_prefill_requests(active_requests, waiting_queue)
alt next case is prefill/warmup/prefill_seed
SB->>SB: _peek_next_case() / _make_prefill_request()
SB->>SB: executor_request_to_llm_request(...)
SB->>SB: _mark_benchmark_request(llm_request, trial_id)
SB->>SB: start_trial(case, requests)
end
PE->>SB: make_decode_requests(active_requests, waiting_queue)
alt next case is decode
SB->>KV: add_dummy_requests(request_ids, token_nums, is_gen=True, ...)
SB->>SB: start_trial(case, requests)
end
PE->>PE: active_requests.extend(benchmark_requests)
PE->>PE: _schedule()
PE->>SB: observe_scheduled_batch(scheduled_batch)
SB->>SB: match scheduled ids vs trial.expected_request_ids
PE->>PE: _sync_self_benchmark_outcome()
PE->>SB: local_outcome()
PE->>Dist: tp_allreduce(local_outcome, op=MAX)
PE->>SB: apply_global_outcome(outcome, reason, origin_rank, drain_target)
PE->>PE: _drain_self_benchmark_requests(protected_ids)
PE->>PE: _terminate_request(skipped/aborted trial requests)
PE->>PE: _forward_step() / model execution
PE->>PE: _process_iter_stats(...)
PE->>SB: observe_executed_batch(scheduled_batch, stats)
alt batch shape mismatch
SB->>SB: _request_local_transition(SKIP, reason)
else full expected batch executed
SB->>SB: record stats into current BenchmarkTrial
SB-->>SB: _record_cache_hit_validation (for prefill w/ kv_read_tokens)
end
PE->>PE: _handle_responses()
loop for each is_self_benchmark_request
PE->>Perf: append_step_metrics(request, iter_counter, batch_token_time)
PE->>SB: should_terminate(request)
alt should terminate
PE->>PE: _do_terminate_request(request)
PE->>SB: observe_terminated_requests([request])
end
end
PE->>SB: observe_finished_requests(requests)
end
Note over SB: _peek_next_case() returns None → grid exhausted
SB->>SB: _finish_complete() / _finish_aborted() / _finish_interrupted()
SB->>SB: _atomic_write_json(...) via write_output()
Key points captured in the diagram:
Let me know if you'd like a narrower diagram (e.g. just the prefill-injection path, or just the outcome-consensus/draining path). |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
tensorrt_llm/llmapi/llm_args.py (1)
90-91: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueDefault output path collision risk for concurrent self-benchmark runs.
_DEFAULT_SELF_BENCHMARK_OUTPUT_PATHresolves to a fixed filename in the shared OS temp directory. Two concurrenttrtllm-serve/LLM()instances on the same host that both enableself_benchmark_configwithout an explicitoutput_pathwill write to the same file and can clobber each other's results. This is intentional/tested (test_self_benchmark_config_default_output_path_uses_tempdir), so not blocking, but worth calling out for multi-instance/multi-node deployments given the feature is "prototype" status.Also applies to: 2725-2728
🤖 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 `@tensorrt_llm/llmapi/llm_args.py` around lines 90 - 91, Make the default self-benchmark output path unique per LLM/self-benchmark instance instead of using the fixed filename from _DEFAULT_SELF_BENCHMARK_OUTPUT_PATH. Update the default resolution near the self_benchmark_config handling so concurrent runs cannot overwrite one another, while preserving explicitly configured output_path values and the existing temp-directory requirement.tensorrt_llm/_torch/pyexecutor/self_benchmark.py (1)
53-79: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftType the planner record and KV-manager interfaces
Rawdictpayloads still flow throughlocal_planner_record(),_merge_planner_records(),_merge_axis_plans(), and_build_admission_record(). AddTypedDicts for the planner/admission records and a smallProtocolfor the KV-cache manager methods used here so schema mismatches surface earlier.🤖 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 `@tensorrt_llm/_torch/pyexecutor/self_benchmark.py` around lines 53 - 79, Define TypedDicts for planner records, axis/admission records, and KV snapshots, then apply them throughout local_planner_record(), _merge_planner_records(), _merge_axis_plans(), and _build_admission_record() instead of raw dict annotations. Add a small KV-cache-manager Protocol declaring the methods these functions use, and type the relevant parameters and attributes against it so payload and interface mismatches are detected statically.Source: Coding guidelines
🤖 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.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/self_benchmark.py`:
- Around line 1007-1039: Update the terminal artifact handling in
_write_terminal and its callers _finish_complete, _finish_aborted, and
_finish_interrupted so a final write failure is not silently swallowed while the
running sentinel remains. Propagate the write exception to the run lifecycle or
explicitly remove the stale output artifact and mark the benchmark failed,
ensuring readers cannot continue interpreting a terminal run as in progress.
- Around line 201-228: Narrow the exception handler around _invalidate_output()
and _build_cases() from Exception to the recoverable planning exception types
already used by the consensus path: KeyError, TypeError, ValueError, and
RuntimeError. Apply the same restriction to prefill/decode request-construction
handlers, while preserving the existing broad abort handling at the consensus
boundary so unexpected defects propagate with their traceback.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/self_benchmark.py`:
- Around line 53-79: Define TypedDicts for planner records, axis/admission
records, and KV snapshots, then apply them throughout local_planner_record(),
_merge_planner_records(), _merge_axis_plans(), and _build_admission_record()
instead of raw dict annotations. Add a small KV-cache-manager Protocol declaring
the methods these functions use, and type the relevant parameters and attributes
against it so payload and interface mismatches are detected statically.
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 90-91: Make the default self-benchmark output path unique per
LLM/self-benchmark instance instead of using the fixed filename from
_DEFAULT_SELF_BENCHMARK_OUTPUT_PATH. Update the default resolution near the
self_benchmark_config handling so concurrent runs cannot overwrite one another,
while preserving explicitly configured output_path values and the existing
temp-directory requirement.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 411f0a07-3832-489e-a31e-8335aba3f34e
📒 Files selected for processing (17)
docs/source/commands/trtllm-serve/trtllm-serve.rsttensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.pytensorrt_llm/_torch/pyexecutor/request_utils.pytensorrt_llm/_torch/pyexecutor/self_benchmark.pytensorrt_llm/llmapi/__init__.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/integration/test_lists/test-db/l0_a10.ymltests/unittest/_torch/executor/test_benchmark_disagg.pytests/unittest/_torch/executor/test_dual_pool_kv_cache.pytests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.pytests/unittest/_torch/executor/test_request_utils.pytests/unittest/_torch/executor/test_self_benchmark.pytests/unittest/api_stability/references/llm.yamltests/unittest/llmapi/test_llm_args.py
💤 Files with no reviewable changes (1)
- tensorrt_llm/_torch/pyexecutor/request_utils.py
Comment on lines
+201
to
+228
| initialization_stage = "output_invalidation" | ||
| try: | ||
| # Invalidate any stale results from a previous run up front, | ||
| # so a crash mid-sweep cannot leave a prior run looking valid. | ||
| self._invalidate_output() | ||
| initialization_stage = "case_planning" | ||
| self._cases = self._build_cases() | ||
| except Exception as exc: | ||
| if initialization_stage == "output_invalidation": | ||
| # The synchronized abort cannot publish on a path that | ||
| # already failed. Keep this rank alive through consensus; | ||
| # writable peers still publish their terminal artifacts. | ||
| self._output_path = None | ||
| self._planner_error = { | ||
| "type": type(exc).__name__, | ||
| "reason": str(exc), | ||
| } | ||
| self._cases = [] | ||
| self._case_plans = {} | ||
| self._planner_events = [ | ||
| { | ||
| "event": "planner_error", | ||
| "stage": initialization_stage, | ||
| "error_type": type(exc).__name__, | ||
| "reason": str(exc), | ||
| } | ||
| ] | ||
| logger.error("Self-benchmark planning failed on this rank: %s", exc) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
path = Path("tensorrt_llm/_torch/pyexecutor/self_benchmark.py")
lines = path.read_text().splitlines()
# Print the relevant regions with line numbers.
ranges = [(180, 250), (680, 790)]
for start, end in ranges:
print(f"\n### {path}:{start}-{end}")
for i in range(start, min(end, len(lines)) + 1):
print(f"{i:4d}: {lines[i-1]}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 9227
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
import re
path = Path("tensorrt_llm/_torch/pyexecutor/self_benchmark.py")
text = path.read_text()
symbols = [
"_invalidate_output",
"_build_cases",
"_make_prefill_request",
"executor_request_to_llm_request",
"add_dummy_requests",
]
for sym in symbols:
m = re.search(rf"^\s*def\s+{re.escape(sym)}\b.*?:", text, re.M)
if not m:
m = re.search(rf"^\s*class\s+{re.escape(sym)}\b.*?:", text, re.M)
print(f"{sym}: {'found' if m else 'not found in file'}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 333
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate the relevant helpers and call sites.
rg -n --hidden --glob '!**/.git/**' \
'_build_cases|_make_prefill_request|executor_request_to_llm_request|add_dummy_requests|observe_failed_requests|start_trial|_invalidate_output' \
tensorrt_llm
# Show the surrounding definitions for the request-builder helpers.
python3 - <<'PY'
from pathlib import Path
import re
paths = [Path("tensorrt_llm/_torch/pyexecutor/self_benchmark.py")]
for path in paths:
text = path.read_text()
for name in ["_make_prefill_request", "make_prefill_requests", "make_decode_requests", "_invalidate_output"]:
m = re.search(rf"^\s*def\s+{re.escape(name)}\b.*?:", text, re.M)
if not m:
continue
start = text.rfind("\n", 0, m.start()) + 1
lines = text.splitlines()
line_no = text[:m.start()].count("\n") + 1
print(f"\n### {path}:{line_no}")
for i in range(max(1, line_no - 25), min(len(lines), line_no + 80) + 1):
print(f"{i:4d}: {lines[i-1]}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 18189
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
import re
targets = {
"tensorrt_llm/_torch/pyexecutor/self_benchmark.py": [
"_build_cases",
"make_prefill_requests",
"make_decode_requests",
"_make_prefill_request",
"_invalidate_output",
],
"tensorrt_llm/_torch/pyexecutor/llm_request.py": [
"executor_request_to_llm_request",
],
"tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py": [
"add_dummy_requests",
],
}
for file, names in targets.items():
path = Path(file)
if not path.exists():
continue
text = path.read_text().splitlines()
print(f"\n### {file}")
for name in names:
for idx, line in enumerate(text, start=1):
if re.match(rf"^\s*(def|class)\s+{re.escape(name)}\b", line):
start = max(1, idx - 25)
end = min(len(text), idx + 120)
print(f"\n--- {file}:{idx} {name} ---")
for i in range(start, end + 1):
print(f"{i:4d}: {text[i-1]}")
break
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
path = Path("tensorrt_llm/_torch/pyexecutor/self_benchmark.py")
lines = path.read_text().splitlines()
# Print the distributed-planning call site and the exception-heavy helpers.
ranges = [(280, 330), (520, 590), (1040, 1075), (2589, 2675)]
for start, end in ranges:
print(f"\n### {path}:{start}-{end}")
for i in range(start, min(end, len(lines)) + 1):
print(f"{i:4d}: {lines[i-1]}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 7886
Narrow the planner/request exception handlers.
The consensus path already limits planning failures to KeyError/TypeError/ValueError/RuntimeError; mirror that here instead of catching Exception, and let unexpected defects surface with their traceback. The same applies to prefill/decode request construction—keep only the recoverable failures and leave the broad abort path at the consensus boundary.
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 208-208: Do not catch blind exception: Exception
(BLE001)
🤖 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 `@tensorrt_llm/_torch/pyexecutor/self_benchmark.py` around lines 201 - 228,
Narrow the exception handler around _invalidate_output() and _build_cases() from
Exception to the recoverable planning exception types already used by the
consensus path: KeyError, TypeError, ValueError, and RuntimeError. Apply the
same restriction to prefill/decode request-construction handlers, while
preserving the existing broad abort handling at the consensus boundary so
unexpected defects propagate with their traceback.
Sources: Coding guidelines, Linters/SAST tools
Comment on lines
+1007
to
+1039
| except Exception as exc: | ||
| self._artifact_write_error = { | ||
| "stage": "terminal", | ||
| "type": type(exc).__name__, | ||
| "reason": str(exc), | ||
| } | ||
| logger.error( | ||
| "Failed to write self-benchmark terminal artifact %s: %s", | ||
| self._output_path, | ||
| exc, | ||
| ) | ||
| return | ||
| logger.info( | ||
| "Self-benchmark terminal artifact written to %s (%s, valid=%s)", | ||
| self._output_path, | ||
| status, | ||
| valid, | ||
| ) | ||
|
|
||
| def _finish_complete(self) -> None: | ||
| self._done = True | ||
| self._run_state = RunState.COMPLETE | ||
| self._write_terminal("complete") | ||
|
|
||
| def _finish_aborted(self) -> None: | ||
| self._done = True | ||
| self._run_state = RunState.ABORTED | ||
| self._write_terminal("aborted") | ||
|
|
||
| def _finish_interrupted(self) -> None: | ||
| self._done = True | ||
| self._run_state = RunState.INTERRUPTED | ||
| self._write_terminal("interrupted") |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file='tensorrt_llm/_torch/pyexecutor/self_benchmark.py'
echo '=== outline ==='
ast-grep outline "$file" --view expanded || true
echo
echo '=== relevant line ranges ==='
python3 - <<'PY'
from pathlib import Path
path = Path("tensorrt_llm/_torch/pyexecutor/self_benchmark.py")
lines = path.read_text().splitlines()
for start, end in [(930, 1085), (1, 220)]:
print(f"\n--- {path}:{start}-{end} ---")
for i in range(start, min(end, len(lines)) + 1):
print(f"{i:4d}: {lines[i-1]}")
PY
echo
echo '=== search usages ==='
rg -n "_artifact_write_error|running sentinel|terminal artifact|_finish_complete|_finish_aborted|_finish_interrupted|_write_terminal" tensorrt_llm/_torch/pyexecutor -S || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 27974
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n "_artifact_write_error|running|invalidate_output|_write_terminal|self_benchmark" tensorrt_llm/_torch/pyexecutor -SRepository: NVIDIA/TensorRT-LLM
Length of output: 10963
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file='tensorrt_llm/_torch/pyexecutor/self_benchmark.py'
python3 - <<'PY'
from pathlib import Path
path = Path("tensorrt_llm/_torch/pyexecutor/self_benchmark.py")
lines = path.read_text().splitlines()
for start, end in [(780, 825), (1888, 1955), (3130, 3235), (4200, 4315), (4538, 4710)]:
print(f"\n--- {path}:{start}-{end} ---")
for i in range(start, min(end, len(lines)) + 1):
print(f"{i:4d}: {lines[i-1]}")
PY
echo
echo '=== search consumers / references ==='
rg -n "_artifact_write_error|_output_path|_invalidate_output\(|_write_terminal\(|RunState\.COMPLETE|RunState\.ABORTED|RunState\.INTERRUPTED|active\)" tensorrt_llm/_torch/pyexecutor -SRepository: NVIDIA/TensorRT-LLM
Length of output: 8950
Propagate terminal-artifact write failures
At tensorrt_llm/_torch/pyexecutor/self_benchmark.py:968-1039, a failed final write leaves the init-time "running" sentinel in place, so readers keep seeing the benchmark as in progress even after _run_state is terminal. Bubble the error up, or clear the output path and mark the run failed, instead of swallowing it here.
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 1007-1007: Do not catch blind exception: Exception
(BLE001)
🤖 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 `@tensorrt_llm/_torch/pyexecutor/self_benchmark.py` around lines 1007 - 1039,
Update the terminal artifact handling in _write_terminal and its callers
_finish_complete, _finish_aborted, and _finish_interrupted so a final write
failure is not silently swallowed while the running sentinel remains. Propagate
the write exception to the run lifecycle or explicitly remove the stale output
artifact and mark the benchmark failed, ensuring readers cannot continue
interpreting a terminal run as in progress.
venkywonka
commented
Jul 13, 2026
| virtual_memory_pools: Optional[dict] = None, | ||
| execution_stream: Optional[torch.cuda.Stream] = None, | ||
| dwdp_manager: Optional[DwdpManager] = None, | ||
| enable_self_benchmark: bool = True, |
Collaborator
Author
There was a problem hiding this comment.
enable_self_benchmark=True is a "don't-forbid" veto, not an enable � self-benchmark still requires a user-set self_benchmark_config; the flag is only flipped to False to suppress benchmarking on the KV-cache-estimation warmup pass.
This should be made explicit to avoid miscommunication.
…self-benchmark JSON The startup self-benchmark serializes each trial/skipped case under a `case`/`case_type` shape. Dynamo's TRT-LLM self-benchmark consumer (ai-dynamo/dynamo NVIDIA#10541) keys results by `point`/`point_type` (matching the vLLM instrumented-scheduler payload) and builds ForwardPassMetrics from them. Emit a `point` projection alongside `case` (and a `skipped_reason` alias) so the emitted JSON is consumable by that normalizer without changing the engine's internal representation. Signed-off-by: venkywonka <23023424+venkywonka@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
2 participants
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.
Description
This carries forward @sachalmalick's startup self-benchmarking implementation and adds a prefill batch-size dimension, keeping TensorRT-LLM aligned with the related Dynamo/vLLM work in ai-dynamo/dynamo#7779.
The startup benchmark covers prefill input length, prefill batch size, cached-KV reads, decode context length, and decode batch size. Prefill points are capped by both the scheduler sequence limit and token budget. Every batch member receives deterministic per-rank request metadata and a distinct cache salt, and a point is recorded only after telemetry observes the full batch. The change preserves TP-rank lockstep, KV-read seed/measurement pairing, output-schema behavior, and disaggregated context/generation startup.
The PR also adds the nested
SelfBenchmarkConfigsurface, serving documentation, API-stability updates, CI enrollment, and focused unit coverage for the complete startup self-benchmarking feature.Validation
(ISL, batch) = (1,1), (1,4), (64,1). The corresponding(numContextRequests, numCtxTokens)telemetry was exactly(1,1), (4,4), (1,64).status=complete,valid=true,world_size=2, andtp_size=2, with identical prefill points including batch 4. A normal 2-token generation also completed after the startup sweep.numGenRequests=4at context lengths 1 and 63. NIXL over UCX initialized on every worker rank, and the proxy returned an 8-token OpenAI completion.On the tested node, NCCL's P2P/CUMEM path also hung in a minimal one-float control all-reduce, independent of TensorRT-LLM and this PR. Setting
NCCL_P2P_DISABLE=1selected the shared-memory transport; the raw two-rank control then returned the expected reduced value 3.0, and both TP=2 and disaggregated validations passed.PR Checklist
api-compatible.Summary by CodeRabbit
trtllm-servesetup.