[https://nvbugs/6971730][fix] Keep Inductor's async-compile pool alive (#21007)

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

Author: yuxianq

tensorrt_llm/_bootstrap.py

@@ -45,6 +45,7 @@
 
 _inited = False
 
+_INDUCTOR_QUIESCE_POOL_ENV = "TORCHINDUCTOR_QUIESCE_ASYNC_COMPILE_POOL"
 _UNIFIED_CACHE_ROOT_ENV = "TRTLLM_CACHE_DIR"
 _TRTLLM_DG_CACHE_ENV = "TRTLLM_DG_CACHE_DIR"
 _TRTLLM_DG_DUMP_CUBIN_ENV = "TRTLLM_DG_JIT_DUMP_CUBIN"
@@ -94,6 +95,14 @@ def _set_numexpr_thread_default() -> None:
     os.environ.setdefault("NUMEXPR_NUM_THREADS", "1")
 
 
+def _keep_inductor_compile_pool_alive() -> None:
+    # Inductor shuts down its async-compile worker pool after a short idle
+    # period. A compile submitted while that shutdown is in flight can wait
+    # forever, and a rank stuck there stalls every collective of its group
+    # (e.g. a torch.compile'd helper recompiling inside a TP forward).
+    os.environ.setdefault(_INDUCTOR_QUIESCE_POOL_ENV, "0")
+
+
 def _add_trt_llm_dll_directory():
     on_windows = platform.system() == "Windows"
     if on_windows:
@@ -162,6 +171,7 @@ def _setup_vendored_triton_kernels():
 def _prepare_environment() -> None:
     """Phase 1: environment and library preparation, before the Torch import."""
     _set_numexpr_thread_default()
+    _keep_inductor_compile_pool_alive()
     _setup_unified_cache()
     _add_trt_llm_dll_directory()
     _preload_python_lib()

tests/unittest/test_unified_cache.py

@@ -129,6 +129,33 @@ def test_prepare_environment_configures_cache_first(monkeypatch: pytest.MonkeyPa
     assert calls == ["cache", "dll", "python", "triton"]
 
 
+def test_prepare_environment_keeps_inductor_compile_pool_alive(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.delenv(_bootstrap._INDUCTOR_QUIESCE_POOL_ENV, raising=False)
+    for name in (
+        "_setup_unified_cache",
+        "_add_trt_llm_dll_directory",
+        "_preload_python_lib",
+        "_setup_vendored_triton_kernels",
+    ):
+        monkeypatch.setattr(_bootstrap, name, lambda: None)
+
+    _bootstrap._prepare_environment()
+
+    assert os.environ[_bootstrap._INDUCTOR_QUIESCE_POOL_ENV] == "0"
+
+
+def test_explicit_inductor_compile_pool_quiesce_setting_takes_precedence(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv(_bootstrap._INDUCTOR_QUIESCE_POOL_ENV, "1")
+
+    _bootstrap._keep_inductor_compile_pool_alive()
+
+    assert os.environ[_bootstrap._INDUCTOR_QUIESCE_POOL_ENV] == "1"
+
+
 def test_mpi_pool_environment_forwards_unified_cache_variables(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None: