Repository navigation
[https://nvbugs/6930163][fix] Shut down Inductor compile workers in KV pool rebalance accuracy test - #20450
[https://nvbugs/6930163][fix] Shut down Inductor compile workers in KV pool rebalance accuracy test#20450trtllm-agent wants to merge 2 commits into
Conversation
…V pool rebalance test The logprobs=2 sampling params route the in-process PyExecutor through the torch.compile'd logprob fusions, which lazily start Inductor's compile-worker pool. Its daemon threads outlived the test and failed the integration thread-leak check. Shut the pool down inside the test body and join the threads it was told to stop; remove the nvbugs/6930163 waivers. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
|
Important Review skippedReview was skipped as selected files did not have any reviewable changes. ⚙️ Run configuration
You can disable this status message by setting the Use the checkbox below for a quick retry:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
WalkthroughThe KV pool rebalance accuracy test now shuts down Inductor compile workers after its generation runs. The waiver list no longer includes the ChangesKV pool rebalance accuracy test
Priority: ⬇️ Low Estimated code review effort: 2 (Simple) | ~10 minutes Change: Bug fix Suggested reviewers: Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 1 files. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Title check The title clearly identifies NVBugs 6930163, uses the required [fix] type, and summarizes the main change: shutting down Inductor compile workers in the KV pool rebalance accuracy test. Description check The description explains the root cause, fix, affected tests, bug link, and test results. It is mostly complete, although it uses Summary and Test plan headings instead of the template's Description a… Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
|
/bot run --disable-fail-fast |
|
/bot run --disable-fail-fast |
|
PR_Github #78487 [ run ] triggered by Bot. Commit: |
|
PR_Github #78488 [ run ] triggered by Bot. Commit: |
|
PR_Github #78487 [ run ] completed with state |
thorjohnsen
left a comment
thorjohnsen
left a comment
There was a problem hiding this comment.
#20294 (merged 2026-10-16, 6544bcd91a) added InductorSubproc|InductorSubprocHealth|subproc_worker_timer to threadleak_exclude in tests/integration/defs/pytest.ini. That fixes nvbugs/6930163 for the whole suite. This branch already has main merged in, so the new try/finally and _shutdown_inductor_compile_workers() now wait for threads the leak check ignores. The repro in the bug was run at 9ec1bf3e58, before #20294 landed.
Suggestion: reduce this PR to the two waives.txt removals, and run both parametrizations on the H100 PyTorch stages to confirm they pass. Inline comments below cover issues in the test-body change, in case it stays.
| ) | ||
| # pytest-threadleak checks before fixture teardown, so the compile pool | ||
| # has to be shut down inside the test body. | ||
| try: |
There was a problem hiding this comment.
With #20294 merged, this is no longer needed: the threads it shuts down are already in the integration threadleak_exclude. I'd drop the try/finally and the helper, and keep only the waiver removal.
| treated = _generate_tokens( | ||
| model_path=self.MODEL_PATH, disable_overlap=disable_overlap, enable_rebalance=True | ||
| ) | ||
| finally: |
There was a problem hiding this comment.
This finally can hide the real failure. shutdown_compile_workers() goes through SubprocPool.shutdown(), which calls process.wait(300) and catches only OSError. If _generate_tokens raises (e.g. the rebalance assert or a CUDA error) while the compile sidecar is stuck, this blocks for up to 5 min, then raises TimeoutExpired. The report shows that instead of the original error.
| With the PyExecutor in-process, that pool's daemon threads would otherwise | ||
| outlive the test and trip the integration suite's thread-leak check. | ||
| """ | ||
| shutdown_compile_workers() |
There was a problem hiding this comment.
This only covers the SubprocPool threads. shutdown_compile_workers() doesn't touch AsyncCompile.pool(), whose threads are named ThreadPoolExecutor-N_M, or the manager threads used with TORCHINDUCTOR_WORKER_START=fork/spawn. Neither is matched here or excluded in the integration pytest.ini (the unittest one excludes ThreadPoolExecutor-\d+_\d+). The waivers are removed on the assumption that this cleanup is complete, so that's a gap.
| outlive the test and trip the integration suite's thread-leak check. | ||
| """ | ||
| shutdown_compile_workers() | ||
| # The pool's health monitor and quiesce timer only notice the shutdown at |
There was a problem hiding this comment.
Possible race with the quiesce timer: it checks exit, releases its lock, then calls pool.quiesce() → _send(). If shutdown() sets running=False in between, _send raises RuntimeError('Attempting to use a closed pool') in a background thread. pytest then reports a PytestUnhandledThreadExceptionWarning, which is an error under -W error.
| shutdown_compile_workers() | ||
| # The pool's health monitor and quiesce timer only notice the shutdown at | ||
| # their next poll; the timer polls every quiesce_async_compile_time / 2. | ||
| for thread in threading.enumerate(): |
There was a problem hiding this comment.
subproc_worker_timer only rechecks exit after sleep(quiesce_async_compile_time / 2), 30 s by default, so this loop can sit idle for up to ~30 s per parametrization. The joins also run one after another, each with a 60 s timeout. _disable_inductor_compile_worker_quiesce in visual_gen_test_utils.py avoids creating the timer at all (inductor_config.quiesce_async_compile_pool = False).
| # their next poll; the timer polls every quiesce_async_compile_time / 2. | ||
| for thread in threading.enumerate(): | ||
| if thread.name.startswith("InductorSubproc") or thread.name == "subproc_worker_timer": | ||
| thread.join(timeout=inductor_config.quiesce_async_compile_time) |
There was a problem hiding this comment.
The result of the join is never checked. If a thread is still alive after the timeout, this returns silently and any later leak report has nothing tying it back here. Consider assert not thread.is_alive(), thread.name or at least a warning.
| _introspection.force_rebalance_precondition(kv_cache_manager.impl, skew=skew) | ||
|
|
||
|
|
||
| def _shutdown_inductor_compile_workers() -> None: |
There was a problem hiding this comment.
This duplicates the shutdown-and-join in tests/integration/defs/examples/visual_gen/visual_gen_test_utils.py::_cleanup_cuda, and the two already differ (that one joins only InductorSubprocHealth, with a 10 s timeout). If this logic is needed at all, it should be one shared helper or fixture under tests/integration/defs/.
|
|
||
| import pytest | ||
| import torch._inductor.config as inductor_config | ||
| from torch._inductor.async_compile import shutdown_compile_workers |
There was a problem hiding this comment.
Module-level imports of private torch internals (torch._inductor.config, torch._inductor.async_compile.shutdown_compile_workers). If a torch upgrade moves them, collection fails and both parametrizations error out. If the threads are renamed, the join silently does nothing. If kept, import lazily inside the helper.
| @@ -1,7 +1,5 @@ | |||
| accuracy/test_disaggregated_serving.py::TestQwen3_30B_A3B::test_mixed_ctx_gen_model[ctxpp2gentp2] SKIP (https://nvbugs/5748664) | |||
| accuracy/test_kimi3.py::TestKimiK3DSpark::test_gsm8k_tep8 SKIP (https://nvbugs/6766812) | |||
| accuracy/test_kv_pool_rebalance_accuracy.py::TestKvPoolRebalanceAccuracy::test_rebalance_matches_baseline[no_overlap] SKIP (https://nvbugs/6930163) | |||
There was a problem hiding this comment.
This part is worth keeping. With #20294 merged, both waivers can come out on their own.
|
I ran the failing test locally with and without #20294, for both parametrizations. #20294 fixes the leak by itself, without the test-body changes in this PR. Setup: 1x H100, each parametrization run on its own in a fresh pytest process,
In every failing run, the leak check was the only failure; the rebalance and baseline outputs matched. Caveat: the Suggestion: reduce this PR to the two |
|
Closing this PR since it has been superseded. |
|
PR_Github #78488 [ run ] completed with state
|
Summary
logprobs=2(4f442e5) runs the in-process executor throughtorch.compiled logprob functions; the compile-worker threads they start (InductorSubproc, InductorSubprocHealth, subproc_worker_timer) outlive the test and fail the integration thread-leak check.shutdown_compile_workers()and joins the signalled pool threads before the leak check; the two nvbugs/6930163 waivers are removed; criteria, logprobs and threadleak settings are unchanged.bash .repair-bot/run_repro.sh 10Test plan
Links
Reproduction comparison
Only the ToT run produced a result, so the two cannot be compared.
Dev Engineer Review
The test now shuts down PyTorch Inductor compile workers in a
finallyblock, including when either generation run raises. It joins matching worker and timer threads for a bounded interval. The join does not verify that threads exited after the timeout, so persistent workers could still trigger the thread-leak check.The change is limited to test cleanup and removal of the two NVBUG 6930163 waivers. It does not change test criteria,
logprobs, or thread-leak settings.QA Engineer Review
The parameterized
test_rebalance_matches_baselinetest still covers theno_overlapandoverlapcases. Its assertions check that rebalance changes the pool ratio and that generated tokens match the baseline, except for an allowed near-tie divergence. Both cases remain listed in thel0_h100.ymlCI list andllm_function_core.txtmanual-QA list. Their NVBUG 6930163 waivers were removed.The supplied test report says the failure reproduced on ToT in three runs. It also reports checks on the same GPU type and related tests, but does not provide results that establish the fix passed. Coverage verdict: needs follow-up to confirm the fix on the affected GPU and confirm the thread-leak check passes.
Per-File QA Perspective
tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py: Covers KV pool rebalance output accuracy for both overlap modes. Cleanup runs after success or failure; verify that Inductor workers exit before the thread-leak check.tests/integration/test_lists/waives.txt: Removes theno_overlapandoverlapwaivers for NVBUG 6930163. Both test cases remain listed for CI and manual QA.