diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md index 2d0a9f0cf64f..97dcc5baaabd 100644 --- a/docs/source/llm-api/index.md +++ b/docs/source/llm-api/index.md @@ -130,6 +130,50 @@ The following tips typically assist new LLM API users who are familiar with othe This limitation is applicable for multi-GPU inference only. +### Unified compiled-artifact cache + +Set `TRTLLM_CACHE_DIR` before importing TensorRT-LLM to place its compiled +artifact caches under a single root. TensorRT-LLM supplies the following +defaults. Except for the isolation behavior described below, an individually +configured environment variable takes precedence. + +| Environment variable | Default under `TRTLLM_CACHE_DIR` | +|----------------------|------------------------------------| +| `TLLM_AUTOTUNER_CACHE_PATH` | `autotuner/cache.json` | +| `TORCHINDUCTOR_CACHE_DIR` | `inductor` | +| `TRITON_CACHE_DIR` | `triton` | +| `TORCH_EXTENSIONS_DIR` | `torch_extensions` | +| `FLASHINFER_WORKSPACE_BASE` | `flashinfer` | +| `CUTE_DSL_CACHE_DIR` | `cute_dsl` | +| `DG_JIT_CACHE_DIR` | `deep_gemm` | +| `TRTLLM_DG_CACHE_DIR` | `trtllm_deep_gemm` | +| `CUDA_CACHE_PATH` | `cuda` | + +The directories are created by their respective consumers when needed. +TensorRT-LLM does not add cache locking, cleanup, or per-rank isolation for +these unified defaults (with the exceptions described below). The +configured values are forwarded to dynamically spawned MPI workers. + +Configuring `TRTLLM_DG_CACHE_DIR`, either directly or through +`TRTLLM_CACHE_DIR`, also defaults `TRTLLM_DG_JIT_DUMP_CUBIN` to `1` so its +NVRTC-generated cubins are persisted. An explicitly configured dump setting +still takes precedence. See the +[DeepGEMM environment variables](source:examples/models/core/deepseek_v3/README.md#deepgemm) +for details. + +Note that, when `TRTLLM_CACHE_DIR` is set, the FlashInfer workspace isolation +mechanism described below still applies. This will prevent race conditions +during cache read/writes, at the cost of limiting cache sharing. To disable +workspace isolation, either set `TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS=0`, +or set `FLASHINFER_WORKSPACE_BASE` to a different path from the one assigned +under `TRTLLM_CACHE_DIR`. + +Ray workers treat `DG_JIT_CACHE_DIR` equal to +`/deep_gemm` as unified-cache-managed and isolate it by rank +and GPU, even when that value was set explicitly. Other explicit values are +preserved. Set `TRTLLM_DEEP_GEMM_CACHE_PER_PROCESS=0` to preserve and share the +unified path. + ### FlashInfer JIT workspaces for MPI workers `trtllm-llmapi-launch` ranks and dynamically spawned `MpiPoolSession` workers @@ -151,7 +195,7 @@ process-unique temporary workspace that is removed, artifacts included, when the process exits. An explicitly configured `FLASHINFER_WORKSPACE_BASE` takes precedence in both -launch modes. An explicitly configured `FLASHINFER_CUBIN_DIR` is propagated +launch modes, while setting `TRTLLM_CACHE_DIR` does not. An explicitly configured `FLASHINFER_CUBIN_DIR` is propagated unchanged to every rank and is the way to share downloaded artifacts between ranks; populate it before the multi-rank launch, either by installing the `flashinfer-cubin` package matching `flashinfer-python` or with a single-rank diff --git a/examples/models/core/deepseek_v3/README.md b/examples/models/core/deepseek_v3/README.md index 205078dc028b..8dc58e1a653a 100644 --- a/examples/models/core/deepseek_v3/README.md +++ b/examples/models/core/deepseek_v3/README.md @@ -697,6 +697,7 @@ DeepGEMM-related behavior can be controlled by the following environment variabl | `TRTLLM_DG_JIT_DEBUG` | When set to `1`, enable JIT debugging. | | `TRTLLM_DG_JIT_USE_NVCC` | When set to `1`, use NVCC instead of NVRTC to compile the kernel, which has slightly better performance but requires CUDA Toolkit (>=12.3) and longer compilation time.| | `TRTLLM_DG_JIT_DUMP_CUBIN` | When set to `1`, dump the cubin file. This is only effective with NVRTC since NVCC will always dump the cubin file. NVRTC-based JIT will store the generated kernels in memory by default. If you want to persist the kernels across multiple runs, you can either use this variable or use NVCC. | +| `TRTLLM_DG_CACHE_DIR` | Sets the root directory for TensorRT-LLM DeepGEMM JIT artifacts. Temporary files are stored under `tmp/`, and compiled kernels are stored under `cache/` when `TRTLLM_DG_JIT_DUMP_CUBIN=1` or NVCC is used. | #### MOE GEMM Optimization diff --git a/tensorrt_llm/__init__.py b/tensorrt_llm/__init__.py index 516e31ad2c52..e2ba5f70fcc3 100644 --- a/tensorrt_llm/__init__.py +++ b/tensorrt_llm/__init__.py @@ -23,8 +23,9 @@ from ._bootstrap import _init, _prepare_environment -# Phase 1: DLL search path, Python-library preload and vendored triton_kernels -# precedence. Must run before torch and before any TensorRT-LLM shared object. +# Phase 1: cache environment, DLL search path, Python-library preload and +# vendored triton_kernels precedence. Must run before torch and before any +# TensorRT-LLM shared object. _prepare_environment() # The package's public surface is loaded lazily (PEP 562): importing diff --git a/tensorrt_llm/_bootstrap.py b/tensorrt_llm/_bootstrap.py index af9ed206a6f5..33e0f2b56fc1 100644 --- a/tensorrt_llm/_bootstrap.py +++ b/tensorrt_llm/_bootstrap.py @@ -16,10 +16,10 @@ Two phases, in this order, both driven from ``tensorrt_llm/__init__.py``: -1. :func:`_prepare_environment` -- process environment defaults, DLL search - path, Python-library preload and vendored ``triton_kernels`` precedence. It - must run *before* ``torch`` and before any TensorRT-LLM shared object is - loaded. +1. :func:`_prepare_environment` -- process environment and compiled-artifact + cache defaults, DLL search path, Python-library preload and vendored + ``triton_kernels`` precedence. It must run *before* ``torch`` and before + any TensorRT-LLM shared object is loaded. 2. :func:`_init` -- custom-op library loading and MPI initialization. It runs after the package's own imports have completed. @@ -45,6 +45,48 @@ _inited = False +_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" +_UNIFIED_CACHE_ENV_VARS = ( + "TLLM_AUTOTUNER_CACHE_PATH", + "TORCHINDUCTOR_CACHE_DIR", + "TRITON_CACHE_DIR", + "TORCH_EXTENSIONS_DIR", + "FLASHINFER_WORKSPACE_BASE", + "CUTE_DSL_CACHE_DIR", + "DG_JIT_CACHE_DIR", + _TRTLLM_DG_CACHE_ENV, + "CUDA_CACHE_PATH", +) + + +def _unified_cache_defaults(cache_root: str) -> dict[str, str]: + """Return cache environment defaults rooted at ``cache_root``.""" + cache_root = os.path.expanduser(cache_root) + return { + "TLLM_AUTOTUNER_CACHE_PATH": os.path.join(cache_root, "autotuner", "cache.json"), + "TORCHINDUCTOR_CACHE_DIR": os.path.join(cache_root, "inductor"), + "TRITON_CACHE_DIR": os.path.join(cache_root, "triton"), + "TORCH_EXTENSIONS_DIR": os.path.join(cache_root, "torch_extensions"), + "FLASHINFER_WORKSPACE_BASE": os.path.join(cache_root, "flashinfer"), + "CUTE_DSL_CACHE_DIR": os.path.join(cache_root, "cute_dsl"), + "DG_JIT_CACHE_DIR": os.path.join(cache_root, "deep_gemm"), + _TRTLLM_DG_CACHE_ENV: os.path.join(cache_root, "trtllm_deep_gemm"), + "CUDA_CACHE_PATH": os.path.join(cache_root, "cuda"), + } + + +def _setup_unified_cache() -> None: + """Configure compiled-artifact cache defaults before framework imports.""" + cache_root = os.environ.get(_UNIFIED_CACHE_ROOT_ENV) + if cache_root: + for name, path in _unified_cache_defaults(cache_root).items(): + os.environ.setdefault(name, path) + + if os.environ.get(_TRTLLM_DG_CACHE_ENV): + os.environ.setdefault(_TRTLLM_DG_DUMP_CUBIN_ENV, "1") + def _set_numexpr_thread_default() -> None: # TRT-LLM uses NumExpr only for scalar checkpoint formulas. Avoid creating @@ -120,6 +162,7 @@ def _setup_vendored_triton_kernels(): def _prepare_environment() -> None: """Phase 1: environment and library preparation, before the Torch import.""" _set_numexpr_thread_default() + _setup_unified_cache() _add_trt_llm_dll_directory() _preload_python_lib() _setup_vendored_triton_kernels() diff --git a/tensorrt_llm/executor/ray/gpu_worker.py b/tensorrt_llm/executor/ray/gpu_worker.py index c6c8653a2031..0be1ddd88ade 100644 --- a/tensorrt_llm/executor/ray/gpu_worker.py +++ b/tensorrt_llm/executor/ray/gpu_worker.py @@ -15,7 +15,6 @@ import gc import importlib import os -import tempfile from functools import wraps from pathlib import Path from queue import Queue @@ -27,7 +26,8 @@ from tensorrt_llm._torch.utils import get_device_uuid from tensorrt_llm._torch.virtual_memory import (materialize_with_tag, release_with_tag) -from tensorrt_llm.executor.ray.utils import control_action_decorator +from tensorrt_llm.executor.ray.utils import (_configure_deep_gemm_cache, + control_action_decorator) from ... import TorchLlmArgs from ...llmapi.llm_args import BaseLlmArgs, ExecutorMemoryType @@ -73,9 +73,7 @@ def __init__(self, worker_cls, worker_kwargs, world_size, rank): self.gpu = int(ray.get_gpu_ids()[0]) self.local_gpu = self.physical_to_local_id(self.gpu) - # Per-worker DeepGemm JIT cache to avoid rename race across co-located workers - os.environ["DG_JIT_CACHE_DIR"] = os.path.join( - tempfile.gettempdir(), f"deep_gemm_rank{rank}_gpu{self.gpu}") + _configure_deep_gemm_cache(rank, self.gpu) torch.cuda.set_device(self.local_gpu) diff --git a/tensorrt_llm/executor/ray/utils.py b/tensorrt_llm/executor/ray/utils.py index 8df771ece793..f481077533d2 100644 --- a/tensorrt_llm/executor/ray/utils.py +++ b/tensorrt_llm/executor/ray/utils.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools +import os +import tempfile from contextlib import contextmanager from typing import Callable, Optional @@ -22,6 +24,29 @@ import tensorrt_llm.executor.ray.stub as ray +def _configure_deep_gemm_cache(rank: int, gpu: int) -> None: + if os.environ.get("TRTLLM_DEEP_GEMM_CACHE_PER_PROCESS", "1") == "0": + return + + cache_dir = os.environ.get("TRTLLM_CACHE_DIR") + if (cache_dir and os.environ.get("DG_JIT_CACHE_DIR") == os.path.join( + os.path.expanduser(cache_dir), "deep_gemm")): + from tensorrt_llm.logger import logger + + logger.warning_once( + "TRTLLM_CACHE_DIR keeps DeepGEMM isolation enabled; set " + "TRTLLM_DEEP_GEMM_CACHE_PER_PROCESS=0 for better cache " + "reuse at the risk of concurrent writes.", + key="deep_gemm_unified_cache_isolation") + os.environ["DG_JIT_CACHE_DIR"] = os.path.join( + os.environ["DG_JIT_CACHE_DIR"], f"deep_gemm_rank{rank}_gpu{gpu}") + else: + os.environ.setdefault( + "DG_JIT_CACHE_DIR", + os.path.join(tempfile.gettempdir(), + f"deep_gemm_rank{rank}_gpu{gpu}")) + + @contextmanager def unwrap_ray_errors(): try: diff --git a/tensorrt_llm/llmapi/mpi_session.py b/tensorrt_llm/llmapi/mpi_session.py index c1fb41c42346..41b8c45d7f76 100644 --- a/tensorrt_llm/llmapi/mpi_session.py +++ b/tensorrt_llm/llmapi/mpi_session.py @@ -17,6 +17,7 @@ import zmq +from tensorrt_llm._bootstrap import _UNIFIED_CACHE_ENV_VARS from tensorrt_llm.bindings.BuildInfo import ENABLE_MULTI_DEVICE from tensorrt_llm.logger import logger @@ -549,14 +550,20 @@ def _start_mpi_pool(self): env = { key: value for key, value in os.environ.items() - if key.startswith("TRTLLM") or key.startswith("TLLM") or key in ( - "FLASHINFER_WORKSPACE_BASE", "FLASHINFER_CUBIN_DIR", - "MASTER_ADDR", "MASTER_PORT") + if key.startswith("TRTLLM") or key.startswith("TLLM") or key in + _UNIFIED_CACHE_ENV_VARS or key in ("FLASHINFER_CUBIN_DIR", + "MASTER_ADDR", "MASTER_PORT") } workspace_managed = env.get(_FLASHINFER_WORKSPACE_MANAGED_ENV) == "1" env.update(self._env_overrides) explicit_workspace_override = (_FLASHINFER_WORKSPACE_ENV in self._env_overrides) + workspace_root = _FLASHINFER_WORKSPACE_ROOT + if cache_root := env.get("TRTLLM_CACHE_DIR"): + workspace_root = os.path.join(os.path.expanduser(cache_root), + "flashinfer") + workspace_managed |= env.get( + _FLASHINFER_WORKSPACE_ENV) == workspace_root isolate_workspace = ( (self.n_workers > 1 or workspace_managed) and env.get("TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS", "1") != "0" @@ -564,12 +571,17 @@ def _start_mpi_pool(self): (workspace_managed and not explicit_workspace_override))) if isolate_workspace: env.pop(_FLASHINFER_WORKSPACE_ENV, None) + if cache_root: + logger.warning_once( + "TRTLLM_CACHE_DIR keeps FlashInfer isolation enabled; " + "set TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS=0 for better " + "cache reuse at the risk of concurrent writes.", + key="flashinfer_unified_cache_isolation") elif explicit_workspace_override: # The override is user-owned, including in any further nested pool. env.pop(_FLASHINFER_WORKSPACE_MANAGED_ENV, None) - python_args = ([ - "-c", _FLASHINFER_WORKER_BOOTSTRAP, _FLASHINFER_WORKSPACE_ROOT - ] if isolate_workspace else None) + python_args = (["-c", _FLASHINFER_WORKER_BOOTSTRAP, workspace_root] + if isolate_workspace else None) self.mpi_pool = MPIPoolExecutor(max_workers=self.n_workers, path=sys.path, env=env, diff --git a/tensorrt_llm/llmapi/trtllm-llmapi-launch b/tensorrt_llm/llmapi/trtllm-llmapi-launch index e9323d93f3ba..0242ccf64086 100755 --- a/tensorrt_llm/llmapi/trtllm-llmapi-launch +++ b/tensorrt_llm/llmapi/trtllm-llmapi-launch @@ -84,15 +84,38 @@ PY # which include compiler inputs, inside the per-rank workspace. A caller-set # value is preserved. See docs/source/llm-api/index.md. function configure_flashinfer_workspace { + local workspace_root="${HOME:+$HOME/.cache/tensorrt_llm/flashinfer}" + if [ -n "$TRTLLM_CACHE_DIR" ]; then + local cache_root + if ! cache_root=$(python3 -S -c \ + 'import os, sys; print(os.path.expanduser(sys.argv[1]))' "$TRTLLM_CACHE_DIR"); then + log_stderr "Failed to expand TRTLLM_CACHE_DIR=$TRTLLM_CACHE_DIR" + use_unique_flashinfer_workspace + return + fi + workspace_root="${cache_root%/}/flashinfer" + fi + + if [ -n "$FLASHINFER_WORKSPACE_BASE" ] && \ + [ "$FLASHINFER_WORKSPACE_BASE" != "$workspace_root" ]; then + return + fi + # PMI_RANK or PMI_ID still proves that this process is launcher-managed # when PMI_SIZE is unavailable. Isolate even rank 0 in that case so rank-0 # processes from concurrent jobs cannot share FlashInfer's default workspace. - if [ -n "$FLASHINFER_WORKSPACE_BASE" ] || \ - { [ "$tllm_mpi_size" -le 1 ] && [ -z "$PMI_RANK" ] && [ -z "$PMI_ID" ]; } || \ + if { [ "$tllm_mpi_size" -le 1 ] && [ -z "$PMI_RANK" ] && [ -z "$PMI_ID" ]; } || \ [ "${TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS:-1}" = "0" ]; then + if [ -n "$TRTLLM_CACHE_DIR" ]; then + export FLASHINFER_WORKSPACE_BASE="$workspace_root" + fi return fi + if [ -n "$TRTLLM_CACHE_DIR" ]; then + log_stderr "Warning: FlashInfer isolation remains enabled; set TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS=0 for better cache reuse at the risk of concurrent writes." + fi + # Do not overwrite a descriptor inherited from the launcher caller. A # process-unique workspace still provides safe isolation in this rare case. if { true >&200; } 2>/dev/null; then @@ -101,13 +124,11 @@ function configure_flashinfer_workspace { return fi - if [ -z "$HOME" ]; then + if [ -z "$workspace_root" ]; then log_stderr "HOME is unset; persistent FlashInfer JIT workspace isolation is unavailable" use_unique_flashinfer_workspace return fi - - local workspace_root="$HOME/.cache/tensorrt_llm/flashinfer" if ! mkdir -p "$workspace_root"; then log_stderr "Failed to create the persistent FlashInfer JIT workspace root: $workspace_root" use_unique_flashinfer_workspace diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 4b1f1e4efab7..a60f907517e9 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -44,6 +44,7 @@ l0_cpu: - unittest/_torch/distributed - unittest/_torch/executor - unittest/test_startup.py + - unittest/test_unified_cache.py - unittest/_torch/disaggregation - unittest/_torch/peft - unittest/_torch/memory diff --git a/tests/unittest/llmapi/test_mpi_session.py b/tests/unittest/llmapi/test_mpi_session.py index 3f94a2e97c93..58ffd85a993d 100644 --- a/tests/unittest/llmapi/test_mpi_session.py +++ b/tests/unittest/llmapi/test_mpi_session.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import os +import shlex import subprocess # nosec B404 import sys import threading @@ -283,14 +284,16 @@ def _launcher_env(tmp_path: Path, home: str) -> dict: """Environment for a rank-0 launcher-managed run of ``trtllm-llmapi-launch``. ``python3`` and ``openssl`` are stubbed so the launcher can hand out an IPC - address and an HMAC key without importing tensorrt_llm; the stubbed - ``python3 -S`` used for the workspace lock exits 0, so the persistent slot - is treated as acquired. + address and an HMAC key without importing tensorrt_llm. Stdlib-only + ``python3 -S`` calls use the test interpreter. """ stub_bin = tmp_path / "bin" stub_bin.mkdir(exist_ok=True) python_stub = stub_bin / "python3" python_stub.write_text("#!/bin/sh\n" + "if [ \"$1\" = \"-S\" ]; then\n" + f" exec {shlex.quote(sys.executable)} \"$@\"\n" + "fi\n" "if [ \"$1\" = \"-c\" ]; then\n" " echo ipc:///tmp/trtllm-pmi-workspace-test\n" "fi\n") @@ -378,6 +381,77 @@ def test_llmapi_launch_temporary_fallback_leaves_cubin_dir_unset( assert list(tmpdir.iterdir()) == [] +@pytest.mark.cpu_only +def test_llmapi_launch_respects_unified_cache_root(tmp_path: Path) -> None: + stub_bin = tmp_path / "bin" + stub_bin.mkdir() + python_stub = stub_bin / "python3" + python_stub.write_text("#!/bin/sh\n" + "if [ \"$1\" = \"-S\" ]; then\n" + f" exec {shlex.quote(sys.executable)} \"$@\"\n" + "fi\n" + "if [ \"$1\" = \"-c\" ]; then\n" + " echo ipc:///tmp/trtllm-unified-cache-test\n" + "fi\n") + python_stub.chmod(0o755) + openssl_stub = stub_bin / "openssl" + openssl_stub.write_text("#!/bin/sh\nprintf '%064d\\n' 0\n") + openssl_stub.chmod(0o755) + + env = os.environ.copy() + for name in ( + "SLURM_NTASKS", + "SLURM_PROCID", + "OMPI_COMM_WORLD_SIZE", + "OMPI_COMM_WORLD_RANK", + "PMI_ID", + "FLASHINFER_WORKSPACE_BASE", + "FLASHINFER_CUBIN_DIR", + "TRTLLM_FLASHINFER_WORKSPACE_MANAGED", + "TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS", + ): + env.pop(name, None) + cache_root = tmp_path / "home" / "unified" + env["PMI_RANK"] = "0" + env["HOME"] = str(tmp_path / "home") + env["TRTLLM_CACHE_DIR"] = "~/unified" + env["FLASHINFER_WORKSPACE_BASE"] = str(cache_root / "flashinfer") + env["PATH"] = f"{stub_bin}{os.pathsep}{env['PATH']}" + + launcher = Path(__file__).parents[ + 3] / "tensorrt_llm" / "llmapi" / "trtllm-llmapi-launch" + result = subprocess.run( # nosec B603 + ["bash", str(launcher), "/usr/bin/env"], + check=True, + capture_output=True, + env=env, + text=True, + timeout=10, + ) + + workspace = cache_root / "flashinfer" / "rank-0" + assert f"FLASHINFER_WORKSPACE_BASE={workspace}" in result.stdout + assert "TRTLLM_FLASHINFER_WORKSPACE_MANAGED=1" in result.stdout + assert "better cache reuse" in result.stderr + + +@pytest.mark.cpu_only +def test_llmapi_launch_matches_bootstrap_expansion_with_trailing_home( + tmp_path: Path) -> None: + home = tmp_path / "home" + home.mkdir() + cache_root = home / "unified" + env = _launcher_env(tmp_path, f"{home}{os.sep}") + env["TRTLLM_CACHE_DIR"] = "~/unified" + env["FLASHINFER_WORKSPACE_BASE"] = str(cache_root / "flashinfer") + + result = _run_launcher_env(env) + + workspace = cache_root / "flashinfer" / "rank-0" + assert f"FLASHINFER_WORKSPACE_BASE={workspace}" in result.stdout + assert "TRTLLM_FLASHINFER_WORKSPACE_MANAGED=1" in result.stdout + + @pytest.mark.cpu_only def test_llmapi_launch_aborts_when_no_workspace_is_available( tmp_path: Path) -> None: diff --git a/tests/unittest/test_unified_cache.py b/tests/unittest/test_unified_cache.py new file mode 100644 index 000000000000..91eaa5eadffe --- /dev/null +++ b/tests/unittest/test_unified_cache.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from tensorrt_llm import _bootstrap +from tensorrt_llm.llmapi import mpi_session + +pytestmark = pytest.mark.cpu_only + + +@pytest.fixture(autouse=True) +def _clear_cache_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv(_bootstrap._UNIFIED_CACHE_ROOT_ENV, raising=False) + monkeypatch.delenv(_bootstrap._TRTLLM_DG_DUMP_CUBIN_ENV, raising=False) + monkeypatch.delenv("TRTLLM_DEEP_GEMM_CACHE_PER_PROCESS", raising=False) + monkeypatch.delenv("TRTLLM_FLASHINFER_WORKSPACE_MANAGED", raising=False) + monkeypatch.delenv("TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS", raising=False) + for name in _bootstrap._UNIFIED_CACHE_ENV_VARS: + monkeypatch.delenv(name, raising=False) + + +def test_unified_cache_is_disabled_without_root() -> None: + _bootstrap._setup_unified_cache() + + assert all(name not in os.environ for name in _bootstrap._UNIFIED_CACHE_ENV_VARS) + assert _bootstrap._TRTLLM_DG_DUMP_CUBIN_ENV not in os.environ + + +def test_unified_cache_enables_trtllm_deep_gemm_cubin_dump( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + cache_root = tmp_path / "unified" + monkeypatch.setenv(_bootstrap._UNIFIED_CACHE_ROOT_ENV, str(cache_root)) + + _bootstrap._setup_unified_cache() + + assert os.environ[_bootstrap._TRTLLM_DG_CACHE_ENV] == str(cache_root / "trtllm_deep_gemm") + assert os.environ[_bootstrap._TRTLLM_DG_DUMP_CUBIN_ENV] == "1" + + +def test_explicit_trtllm_deep_gemm_cache_enables_cubin_dump( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setenv(_bootstrap._TRTLLM_DG_CACHE_ENV, str(tmp_path / "deep_gemm")) + + _bootstrap._setup_unified_cache() + + assert os.environ[_bootstrap._TRTLLM_DG_DUMP_CUBIN_ENV] == "1" + + +def test_explicit_trtllm_deep_gemm_dump_setting_takes_precedence( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setenv(_bootstrap._TRTLLM_DG_CACHE_ENV, str(tmp_path / "deep_gemm")) + monkeypatch.setenv(_bootstrap._TRTLLM_DG_DUMP_CUBIN_ENV, "0") + + _bootstrap._setup_unified_cache() + + assert os.environ[_bootstrap._TRTLLM_DG_DUMP_CUBIN_ENV] == "0" + + +def test_unified_cache_respects_individual_overrides( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setenv(_bootstrap._UNIFIED_CACHE_ROOT_ENV, str(tmp_path / "unified")) + for name in _bootstrap._UNIFIED_CACHE_ENV_VARS: + monkeypatch.setenv(name, f"/explicit/{name.lower()}") + + _bootstrap._setup_unified_cache() + + for name in _bootstrap._UNIFIED_CACHE_ENV_VARS: + assert os.environ[name] == f"/explicit/{name.lower()}" + + +@pytest.mark.parametrize( + "cache_override,isolate,expected", + [ + (None, True, "isolated"), + ("unified", True, "isolated"), + (None, False, "unified"), + ("explicit", True, "explicit"), + ], +) +def test_ray_deep_gemm_cache_configuration( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + cache_override: str | None, + isolate: bool, + expected: str, +) -> None: + from tensorrt_llm.executor.ray.utils import _configure_deep_gemm_cache + + cache_root = tmp_path / "unified" + cache_dir = cache_root / "deep_gemm" + monkeypatch.setenv(_bootstrap._UNIFIED_CACHE_ROOT_ENV, str(cache_root)) + if cache_override: + override = cache_dir if cache_override == "unified" else tmp_path / "explicit" + monkeypatch.setenv("DG_JIT_CACHE_DIR", str(override)) + if not isolate: + monkeypatch.setenv("TRTLLM_DEEP_GEMM_CACHE_PER_PROCESS", "0") + + _bootstrap._setup_unified_cache() + _configure_deep_gemm_cache(rank=2, gpu=3) + + expected_cache = { + "isolated": cache_dir / "deep_gemm_rank2_gpu3", + "unified": cache_dir, + "explicit": tmp_path / "explicit", + }[expected] + assert os.environ["DG_JIT_CACHE_DIR"] == str(expected_cache) + + +def test_prepare_environment_configures_cache_first(monkeypatch: pytest.MonkeyPatch) -> None: + calls = [] + monkeypatch.setattr(_bootstrap, "_setup_unified_cache", lambda: calls.append("cache")) + monkeypatch.setattr(_bootstrap, "_add_trt_llm_dll_directory", lambda: calls.append("dll")) + monkeypatch.setattr(_bootstrap, "_preload_python_lib", lambda: calls.append("python")) + monkeypatch.setattr( + _bootstrap, "_setup_vendored_triton_kernels", lambda: calls.append("triton") + ) + + _bootstrap._prepare_environment() + + assert calls == ["cache", "dll", "python", "triton"] + + +def test_mpi_pool_environment_forwards_unified_cache_variables( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, object] = {} + + class FakeMpiPoolExecutor: + def __init__(self, **kwargs: object) -> None: + captured.update(kwargs) + + monkeypatch.setenv(_bootstrap._UNIFIED_CACHE_ROOT_ENV, "/cache") + for name in _bootstrap._UNIFIED_CACHE_ENV_VARS: + monkeypatch.setenv(name, f"/cache/{name.lower()}") + monkeypatch.setenv(_bootstrap._TRTLLM_DG_DUMP_CUBIN_ENV, "1") + monkeypatch.setenv("UNRELATED_CACHE_DIR", "/not-forwarded") + monkeypatch.setattr(mpi_session, "MPIPoolExecutor", FakeMpiPoolExecutor, raising=False) + session = SimpleNamespace(mpi_pool=None, n_workers=1, _env_overrides={}) + + mpi_session.MpiPoolSession._start_mpi_pool(session) + + worker_env = captured["env"] + assert isinstance(worker_env, dict) + assert worker_env[_bootstrap._UNIFIED_CACHE_ROOT_ENV] == "/cache" + assert worker_env[_bootstrap._TRTLLM_DG_DUMP_CUBIN_ENV] == "1" + assert all(worker_env[name] == os.environ[name] for name in _bootstrap._UNIFIED_CACHE_ENV_VARS) + assert "UNRELATED_CACHE_DIR" not in worker_env + + +@pytest.mark.parametrize( + "use_unified_cache,isolate,explicit_workspace", + [ + (False, True, False), + (True, True, False), + (True, False, False), + (True, True, True), + ], +) +def test_mpi_pool_flashinfer_isolation_respects_unified_cache_configuration( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + use_unified_cache: bool, + isolate: bool, + explicit_workspace: bool, +) -> None: + captured: dict[str, object] = {} + + class FakeMpiPoolExecutor: + def __init__(self, **kwargs: object) -> None: + captured.update(kwargs) + + cache_root = tmp_path / "unified" + if use_unified_cache: + monkeypatch.setenv(_bootstrap._UNIFIED_CACHE_ROOT_ENV, str(cache_root)) + if not isolate: + monkeypatch.setenv("TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS", "0") + if explicit_workspace: + monkeypatch.setenv("FLASHINFER_WORKSPACE_BASE", str(tmp_path / "explicit-flashinfer")) + monkeypatch.setattr(mpi_session, "MPIPoolExecutor", FakeMpiPoolExecutor, raising=False) + _bootstrap._setup_unified_cache() + session = SimpleNamespace(mpi_pool=None, n_workers=2, _env_overrides={}) + + mpi_session.MpiPoolSession._start_mpi_pool(session) + + worker_env = captured["env"] + assert isinstance(worker_env, dict) + if use_unified_cache: + workspace_root = str(cache_root / "flashinfer") + else: + workspace_root = mpi_session._FLASHINFER_WORKSPACE_ROOT + if isolate and not explicit_workspace: + assert captured["python_args"] == [ + "-c", + mpi_session._FLASHINFER_WORKER_BOOTSTRAP, + workspace_root, + ] + assert "FLASHINFER_WORKSPACE_BASE" not in worker_env + else: + expected_workspace = ( + str(tmp_path / "explicit-flashinfer") if explicit_workspace else workspace_root + ) + assert captured["python_args"] is None + assert worker_env["FLASHINFER_WORKSPACE_BASE"] == expected_workspace