[https://nvbugs/5680911][fix] Remove @cache decorator to enhance CI stability for unit tests using single process mode by zheyuf · Pull Request #10730 · NVIDIA/TensorRT-LLM

zheyuf · 2026-01-16T00:39:22Z

Summary by CodeRabbit

Tests
- Updated speculative decoding test infrastructure with improved test isolation and per-test fixture enforcement.
- Added skip markers for tests with known issues.
- Enhanced spec gate testing with better validation of acceptance thresholds and gate behavior.
Chores
- Removed caching from utility function to ensure fresh computation.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

Description

This PR has the same content with a previous PR: #10098. The only difference is that the current PR targets release/1.2, while the previous PR targets main (changes on main currently will not be reflected to release/1.2, so I open a new PR).

This PR contains a one-line change that requires review from trt-llm-llmapi-devs (remove the @cache decorator). This PR enhance the CI stability for tests using single process mode. Moreover, it fixes three NVBugs and waive one test.

Problem
The @cache decorator on enable_worker_single_process_for_tp1() causes the function to return a stale cached result instead of re-reading os.environ on each call. This breaks directory-level CI test runs where different unit tests need different process modes.

Root Cause
When CI runs directory tests like:python -m pytest unittest/_torch/speculative -k not eagle3
The first test (test_draft_len_schedule.py::test_correctness_across_batch_sizes[ngram-schedule0]) runs without setting TLLM_WORKER_USE_SINGLE_PROCESS, so enable_worker_single_process_for_tp1() caches False.

Later tests like test_dynamic_spec_decode.py::test_dynamic_spec_decode[True] set the env var to True via monkeypatch.setenv(), but enable_worker_single_process_for_tp1() always returns the cached False instead of reading the updated env var.

Impact
This creates a particularly insidious bug: test_dynamic_spec_decode requires single-process mode. When the test silently runs in multi-process mode due to the stale cache, the mock function inside the unit test has no effect—the test passes without actually validating the intended behavior. In CI directory runs, it becomes a no-op that doesn't catch regressions.

Fix
Remove the @cache decorator. The performance impact is negligible since enable_worker_single_process_for_tp1() is only called once per LLM initialization (not in the inference hot path), and os.environ.get() is a simple dictionary lookup.

NVBugs Fixed in this PR

Additional Changes
Waived test_dynamic_spec_decode (https://nvbugspro.nvidia.com/bug/5758449), which is buggy: Due to the @cache decorator, this test has been silently running in multi-process mode, making it never raise alarms in CI. Now that the decorator is removed and the test would run correctly in single-process mode, it needs to be fixed to actually work as intended.

Test Coverage

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

Details

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

For guidance on mapping tests to stage names, see docs/source/reference/ci-overview.md
and the scripts/test_to_stage_mapping.py helper.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

zheyuf · 2026-01-16T00:53:45Z

Hi @Superjomn, would you mind reviewing and approving this PR? It contains the exact same content as the previous one you reviewed here: #10098

The only difference is that this PR targets release/1.2 instead of main. Since changes to main are not currently being merged into the release/1.2 branch, I’ve opened this separate PR to ensure the fix is included in 1.2.

Would appreciate your help.

coderabbitai · 2026-01-16T00:56:23Z

📝 Walkthrough

Walkthrough

The changes remove caching from a utility function and refactor speculative decoding tests to use function-scoped single-worker enforcement, update test skip lists, and enhance the spec gate test with acceptance window mocking and assertions for low-acceptance scenarios.

Changes

Cohort / File(s)	Summary
Caching removal `tensorrt_llm/llmapi/utils.py`	Removed `@cache` decorator from `enable_worker_single_process_for_tp1` function; minor import adjustment from functools. Function now recomputes on each invocation rather than caching results.
Test skip list updates `tests/integration/test_lists/waives.txt`	Replaced `test_draft_len_schedule_functionality` and `test_spec_gate_e2e` entries with `test_dynamic_spec_decode` marked as SKIP.
Test fixture refactoring `tests/unittest/_torch/speculative/test_draft_len_schedule.py`	Converted `enforce_single_worker` fixture from module-scoped direct environment manipulation to function-scoped with `monkeypatch.setenv`. Updated `test_draft_len_schedule_functionality` to accept the fixture as a parameter.
Test skip marker `tests/unittest/_torch/speculative/test_dynamic_spec_decode.py`	Added `@pytest.mark.skip` decorator with reference to issue tracker on `test_dynamic_spec_decode`.
Spec gate test enhancement `tests/unittest/_torch/speculative/test_spec_gate.py`	Refactored `test_spec_gate_e2e` to use `enforce_single_worker` fixture, added acceptance window/threshold parameters to EagleDecodingConfig, mocked `SpeculationGate.record_avg_decoded` to simulate low-acceptance scenario, added assertions verifying gate disables speculation when acceptance drops below threshold, wrapped execution in try/finally for shutdown, included logging for gate trigger progress.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 40.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically identifies the main change: removing the `@cache` decorator to enhance CI stability for single-process mode unit tests.
Linked Issues check	✅ Passed	The PR links three NVBugs (5698292, 5710045, 5680911) in the title and description with clear explanation of the connection.
Out of Scope Changes check	✅ Passed	All changes are directly related to fixing the `@cache` decorator issue: one-line fix in utils.py, test fixture updates for single-process enforcement, test adjustments, and a test waiver for a previously masked bug.
Description check	✅ Passed	The PR description comprehensively explains the problem, root cause, impact, fix, and additional changes.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

📝 Generate docstrings

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

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

zheyuf · 2026-01-16T00:57:47Z

/bot run

tensorrt-cicd · 2026-01-16T01:03:52Z

PR_Github #32194 [ run ] triggered by Bot. Commit: f86ddcc

tensorrt-cicd · 2026-01-16T03:36:21Z

PR_Github #32194 [ run ] completed with state SUCCESS. Commit: f86ddcc
/LLM/release-1.2/L0_MergeRequest_PR pipeline #6 completed with status: 'FAILURE'

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

zheyuf · 2026-01-16T04:06:53Z

/bot run

zheyuf · 2026-01-16T04:30:03Z

/bot run

tensorrt-cicd · 2026-01-16T04:35:43Z

PR_Github #32231 [ run ] triggered by Bot. Commit: f86ddcc

zheyuf · 2026-01-17T00:11:31Z

Hi @QiJune, @schetlur-nv mentioned you've been focusing on CI improvements. Since this PR may enhance CI stability in some scenarios, we’ve added you as a reviewer to keep you informed.

zheyuf · 2026-01-17T00:21:10Z

/bot run

tensorrt-cicd · 2026-01-17T00:26:51Z

PR_Github #32367 [ run ] triggered by Bot. Commit: cb2df3a

tensorrt-cicd · 2026-01-17T04:04:07Z

PR_Github #32367 [ run ] completed with state SUCCESS. Commit: cb2df3a
/LLM/release-1.2/L0_MergeRequest_PR pipeline #34 completed with status: 'FAILURE'

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

zheyuf · 2026-01-19T21:36:16Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-01-19T21:42:42Z

PR_Github #32611 [ run ] triggered by Bot. Commit: 32db66b

tensorrt-cicd · 2026-01-19T22:24:28Z

PR_Github #32611 [ run ] completed with state FAILURE. Commit: 32db66b
/LLM/release-1.2/L0_MergeRequest_PR pipeline #60 completed with status: 'FAILURE'

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

zheyuf · 2026-01-19T22:25:40Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-01-19T22:32:38Z

PR_Github #32613 [ run ] triggered by Bot. Commit: 32db66b

tensorrt-cicd · 2026-01-20T03:38:01Z

PR_Github #32613 [ run ] completed with state SUCCESS. Commit: 32db66b
/LLM/release-1.2/L0_MergeRequest_PR pipeline #61 completed with status: 'FAILURE'

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

zheyuf · 2026-01-20T18:47:35Z

/bot run --disable-fail-fast

zheyuf · 2026-01-21T03:25:51Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-01-21T03:31:54Z

PR_Github #32851 [ run ] triggered by Bot. Commit: 3f6be68

tensorrt-cicd · 2026-01-21T09:18:32Z

PR_Github #32851 [ run ] completed with state SUCCESS. Commit: 3f6be68
/LLM/release-1.2/L0_MergeRequest_PR pipeline #100 completed with status: 'FAILURE'

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

zheyuf · 2026-01-21T17:15:13Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-01-21T17:23:39Z

PR_Github #32985 [ run ] triggered by Bot. Commit: 21797d8

tensorrt-cicd · 2026-01-21T23:44:24Z

PR_Github #32985 [ run ] completed with state SUCCESS. Commit: 21797d8
/LLM/release-1.2/L0_MergeRequest_PR pipeline #132 completed with status: 'FAILURE'

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>

zheyuf · 2026-01-22T02:02:56Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-01-22T02:08:48Z

PR_Github #33042 [ run ] triggered by Bot. Commit: e285711

tensorrt-cicd · 2026-01-22T08:15:27Z

PR_Github #33042 [ run ] completed with state SUCCESS. Commit: e285711
/LLM/release-1.2/L0_MergeRequest_PR pipeline #143 completed with status: 'FAILURE'

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>

zheyuf · 2026-01-22T18:51:54Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-01-22T18:58:56Z

PR_Github #33214 [ run ] triggered by Bot. Commit: 696d7df

tensorrt-cicd · 2026-01-23T00:30:08Z

PR_Github #33214 [ run ] completed with state SUCCESS. Commit: 696d7df
/LLM/release-1.2/L0_MergeRequest_PR pipeline #178 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

…tability for unit tests using single process mode (NVIDIA#10730) Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>

…tability for unit tests using single process mode (NVIDIA#10730) Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com> Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>

…tability for unit tests using single process mode (#10730) Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com> Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>

zheyuf changed the title ~~[https://nvbugs/5680911][fix] Remove @cache decorator from enable_worker_single_process_for_tp1() to fix env var staleness~~ [https://nvbugs/5680911][fix] Remove @cache decorator to enhance CI stability for unit tests using single process mode Jan 16, 2026

zheyuf requested a review from Superjomn January 16, 2026 00:53

zheyuf marked this pull request as ready for review January 16, 2026 00:54

zheyuf requested a review from a team as a code owner January 16, 2026 00:54

schetlur-nv approved these changes Jan 16, 2026

View reviewed changes

schetlur-nv assigned QiJune and unassigned QiJune Jan 16, 2026

schetlur-nv requested a review from QiJune January 16, 2026 23:54

zheyuf force-pushed the fix_1.2 branch from f86ddcc to cb2df3a Compare January 17, 2026 00:20

zheyuf force-pushed the fix_1.2 branch from cb2df3a to 32db66b Compare January 19, 2026 21:35

zheyuf force-pushed the fix_1.2 branch from 32db66b to c350c2f Compare January 20, 2026 18:46

zheyuf force-pushed the fix_1.2 branch from 3f6be68 to 21797d8 Compare January 21, 2026 17:14

Fix CI stability for single process mode tests.

e285711

Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>

zheyuf force-pushed the fix_1.2 branch from 21797d8 to e285711 Compare January 22, 2026 02:02

Merge branch 'release/1.2' into fix_1.2

696d7df

Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>

zheyuf merged commit 27feb0c into NVIDIA:release/1.2 Jan 23, 2026
5 checks passed

Uh oh!

Conversation

zheyuf commented Jan 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Description

Test Coverage

PR Checklist

GitHub Bot Help

kill

skip

reuse-pipeline

Uh oh!

zheyuf commented Jan 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

coderabbitai Bot commented Jan 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

Uh oh!

zheyuf commented Jan 16, 2026

Uh oh!

tensorrt-cicd commented Jan 16, 2026

Uh oh!

tensorrt-cicd commented Jan 16, 2026

Uh oh!

zheyuf commented Jan 16, 2026

Uh oh!

zheyuf commented Jan 16, 2026

Uh oh!

tensorrt-cicd commented Jan 16, 2026

Uh oh!

zheyuf commented Jan 17, 2026

Uh oh!

zheyuf commented Jan 17, 2026

Uh oh!

tensorrt-cicd commented Jan 17, 2026

Uh oh!

tensorrt-cicd commented Jan 17, 2026

Uh oh!

zheyuf commented Jan 19, 2026

Uh oh!

tensorrt-cicd commented Jan 19, 2026

Uh oh!

tensorrt-cicd commented Jan 19, 2026

Uh oh!

zheyuf commented Jan 19, 2026

Uh oh!

tensorrt-cicd commented Jan 19, 2026

Uh oh!

tensorrt-cicd commented Jan 20, 2026

Uh oh!

zheyuf commented Jan 20, 2026

Uh oh!

zheyuf commented Jan 21, 2026

Uh oh!

tensorrt-cicd commented Jan 21, 2026

Uh oh!

tensorrt-cicd commented Jan 21, 2026

Uh oh!

zheyuf commented Jan 21, 2026

Uh oh!

tensorrt-cicd commented Jan 21, 2026

Uh oh!

tensorrt-cicd commented Jan 21, 2026

Uh oh!

zheyuf commented Jan 22, 2026

Uh oh!

tensorrt-cicd commented Jan 22, 2026

Uh oh!

tensorrt-cicd commented Jan 22, 2026

Uh oh!

zheyuf commented Jan 22, 2026

Uh oh!

tensorrt-cicd commented Jan 22, 2026

Uh oh!

tensorrt-cicd commented Jan 23, 2026

Uh oh!

zheyuf commented Jan 16, 2026 •

edited

Loading

zheyuf commented Jan 16, 2026 •

edited

Loading

coderabbitai Bot commented Jan 16, 2026 •

edited

Loading