Skip to content

[https://nvbugs/6930163][fix] Shut down Inductor compile workers in KV pool rebalance accuracy test #20450

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
trtllm-agent wants to merge 2 commits into from
Closed

[https://nvbugs/6930163][fix] Shut down Inductor compile workers in KV pool rebalance accuracy test #20450

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
eea624f
[https://nvbugs/6930163][fix] Shut down Inductor compile workers in K…
trtllm-agent Oct 16, 2026
445d383
Merge branch 'main' into repair-bot-bug6930163
thorjohnsen Oct 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 31 additions & 6 deletions tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,11 @@
tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py
"""

import threading

import pytest
import torch._inductor.config as inductor_config
from torch._inductor.async_compile import shutdown_compile_workers

@thorjohnsen thorjohnsen Oct 16, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.


from tensorrt_llm import LLM
from tensorrt_llm.llmapi import KvCacheConfig, SamplingParams
Expand Down Expand Up @@ -59,6 +63,22 @@ def _inject_pool_ratio_mismatch(llm: LLM, *, skew: float = 2.0) -> None:
_introspection.force_rebalance_precondition(kv_cache_manager.impl, skew=skew)


def _shutdown_inductor_compile_workers() -> None:

@thorjohnsen thorjohnsen Oct 16, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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/.

"""Stop Inductor's compile-worker pool and wait for its threads to exit.

The sampler's logprob fusions are ``torch.compile``d, so the first request
with ``logprobs`` lazily starts Inductor's process-global compile pool.
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()

@thorjohnsen thorjohnsen Oct 16, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

# The pool's health monitor and quiesce timer only notice the shutdown at

@thorjohnsen thorjohnsen Oct 16, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

# their next poll; the timer polls every quiesce_async_compile_time / 2.
for thread in threading.enumerate():

@thorjohnsen thorjohnsen Oct 16, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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).

if thread.name.startswith("InductorSubproc") or thread.name == "subproc_worker_timer":
thread.join(timeout=inductor_config.quiesce_async_compile_time)

@thorjohnsen thorjohnsen Oct 16, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.



# --------------------------------------------------------------------------- #
# Test
# --------------------------------------------------------------------------- #
Expand Down Expand Up @@ -232,13 +252,18 @@ def test_rebalance_matches_baseline(self, disable_overlap, monkeypatch):
# can reach .engine on the client side.
monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1")

baseline = _generate_tokens(
model_path=self.MODEL_PATH, disable_overlap=disable_overlap, enable_rebalance=False
)
# pytest-threadleak checks before fixture teardown, so the compile pool
# has to be shut down inside the test body.
try:

@thorjohnsen thorjohnsen Oct 16, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

baseline = _generate_tokens(
model_path=self.MODEL_PATH, disable_overlap=disable_overlap, enable_rebalance=False
)

treated = _generate_tokens(
model_path=self.MODEL_PATH, disable_overlap=disable_overlap, enable_rebalance=True
)
treated = _generate_tokens(
model_path=self.MODEL_PATH, disable_overlap=disable_overlap, enable_rebalance=True
)
finally:

@thorjohnsen thorjohnsen Oct 16, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

_shutdown_inductor_compile_workers()

assert len(baseline) == len(treated) == len(_PROMPTS)
for i, (b, t) in enumerate(zip(baseline, treated)):
Expand Down
2 changes: 0 additions & 2 deletions tests/integration/test_lists/waives.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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)

@thorjohnsen thorjohnsen Oct 16, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This part is worth keeping. With #20294 merged, both waivers can come out on their own.

accuracy/test_kv_pool_rebalance_accuracy.py::TestKvPoolRebalanceAccuracy::test_rebalance_matches_baseline[overlap] SKIP (https://nvbugs/6930163)
accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[ep4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6402058)
accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp2pp2-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6198774)
accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6198774)
Expand Down