Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
32 changes: 32 additions & 0 deletions docs/source/installation/build-from-source.md
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,38 @@ Key flags used above:
| `--fast_build` | Skip compiling some kernels to speed up compilation -- for development only |
| `--cpp_only` | Build only the C++ runtime library, without Python bindings |

### Building from a checkout on a network filesystem

Network filesystems (Lustre, NFS, GPFS) handle large streaming I/O well but are slow for metadata-heavy workloads. A full build creates a very large number of small files (CMake state, object files, downloaded dependencies, the build virtual environment, wheel staging), so keeping that state in the checkout makes builds on such filesystems far slower than necessary.

Pass `--build_root` (or set the `TRTLLM_BUILD_ROOT` environment variable) to keep all high-churn build state on fast local storage while the checkout stays on shared storage:

```bash
python3 scripts/build_wheel.py --build_root /tmp/trtllm-build --use_ccache -a "90-real" --skip_building_wheel --linking_install_binary
```

With `--build_root <dir>` set, the following default under `<dir>` instead of the checkout:

| State | Location under `<dir>` | Individual override |
|-------|------------------------|---------------------|
| CMake build directory (objects, `_deps` downloads, conan output) | `cpp-build*` | `--build_dir` |
| Build virtual environment | `venv-<python version>` | run inside an activated venv, or `--no-venv` |
| Wheel staging tree and `*.egg-info` | `wheel-staging` | `TRTLLM_WHEEL_STAGING_DIR` |
| ccache directory (with `--use_ccache`) | `ccache` | `CCACHE_DIR` |
| Intermediate extension-module objects | `kv_cache_manager_v2-temp` | — |

Conan's `cpp/CMakeUserPresets.json` convenience file is also skipped in this mode, since it would reference the (possibly ephemeral) out-of-tree build directory.

Only final artifacts are still written into the checkout: `tensorrt_llm/libs`, `tensorrt_llm/include`, Python bindings and stubs, generated FMHA kernel sources, and the `.whl` output directory (`--dist_dir`).

Related knobs for shared-storage workflows:

- `CCACHE_DIR`: point at persistent storage so compile results survive container or job restarts even when `<dir>` is ephemeral (for example, node-local `/tmp`).
- `CONAN_HOME`: conan's download cache defaults to `~/.conan2`; relocate it if your home directory is small or slow.
- `--use-3rdparty-cache`: cache FetchContent git clones as bare repos under `TRTLLM_FETCHCONTENT_CACHE` (defaults to `3rdparty/.cache_3rdparty`), avoiding repeated full clones after a clean.

Plain local-disk builds are unaffected: without `--build_root`, all paths behave as before.

### Python-only build (no C++ compilation)

If you only need to modify Python code, you can skip C++ compilation entirely by reusing precompiled binaries:
Expand Down
129 changes: 104 additions & 25 deletions scripts/build_wheel.py
Original file line number Diff line number Diff line change
Expand Up @@ -95,10 +95,13 @@ def get_source_dir():
return get_project_dir() / "cpp"


def get_build_dir(build_dir, build_type):
def get_build_dir(build_dir, build_type, build_root=None):
if build_dir is None:
build_dir = get_source_dir() / ("build" if build_type == "Release" else
f"build_{build_type}")
dir_name = "build" if build_type == "Release" else f"build_{build_type}"
if build_root is not None:
build_dir = Path(build_root).resolve() / f"cpp-{dir_name}"
else:
build_dir = get_source_dir() / dir_name
Comment thread
brnguyen2 marked this conversation as resolved.
else:
build_dir = Path(build_dir).resolve()
return build_dir
Expand Down Expand Up @@ -135,10 +138,9 @@ def sysconfig_scheme(override_vars=None):
return {key: value.format(**vars_) for key, value in scheme.items()}


def create_venv(project_dir: Path):
def create_venv(venv_prefix: Path):
py_major = sys.version_info.major
py_minor = sys.version_info.minor
venv_prefix = project_dir / f".venv-{py_major}.{py_minor}"
print(
f"-- Using virtual environment at: {venv_prefix} (Python {py_major}.{py_minor})"
)
Expand All @@ -162,13 +164,16 @@ def create_venv(project_dir: Path):
def setup_venv(project_dir: Path,
requirements_file: Path,
no_venv: bool,
yes: bool = False) -> tuple[Path, Path]:
yes: bool = False,
build_root: Optional[Path] = None) -> tuple[Path, Path]:
"""Creates/updates a venv and installs requirements.

Args:
project_dir: The root directory of the project.
requirements_file: Path to the requirements file.
no_venv: Use current Python environment as is.
build_root: Directory for out-of-tree build state; when set, the venv
is created there instead of inside the checkout.

Returns:
Tuple[Path, Path]: Paths to the python and conan executables in the venv.
Expand All @@ -178,7 +183,12 @@ def setup_venv(project_dir: Path,
print(f"-- {reason}, using environment {sys.prefix} as is.")
venv_prefix = Path(sys.prefix)
else:
venv_prefix = create_venv(project_dir)
py_version = f"{sys.version_info.major}.{sys.version_info.minor}"
if build_root is not None:
venv_prefix = build_root / f"venv-{py_version}"
else:
venv_prefix = project_dir / f".venv-{py_version}"
venv_prefix = create_venv(venv_prefix)

scheme = sysconfig_scheme({'base': venv_prefix})
# Determine venv executable paths
Expand Down Expand Up @@ -484,11 +494,20 @@ def generate_python_stubs_windows(venv_python: Path, pkg_dir: Path,
exit(1)


def build_kv_cache_manager_v2(project_dir, venv_python, use_mypyc=False):
def build_kv_cache_manager_v2(project_dir,
venv_python,
use_mypyc=False,
build_root=None):
print("-- Building kv_cache_manager_v2...")
kv_cache_mgr_dir = project_dir / "tensorrt_llm/runtime/kv_cache_manager_v2"
runtime_dir = project_dir / "tensorrt_llm/runtime"

# The produced .so files always land in-place (they are final artifacts);
# only the intermediate object files are redirected out of the checkout.
build_temp_arg = ""
if build_root is not None:
build_temp_arg = f' --build-temp "{build_root / "kv_cache_manager_v2-temp"}"'

# Clean up any existing mypyc artifacts in runtime directory to prevent stale inclusion
# when switching from --mypyc to standard build
if not use_mypyc:
Expand All @@ -504,16 +523,18 @@ def build_kv_cache_manager_v2(project_dir, venv_python, use_mypyc=False):
# Build rawref
print("-- Building kv_cache_manager_v2 rawref extension...", end=" ")
rawref_dir = kv_cache_mgr_dir / "rawref"
build_run(f'"{venv_python}" setup.py build_ext --inplace', cwd=rawref_dir)
build_run(f'"{venv_python}" setup.py build_ext --inplace{build_temp_arg}',
cwd=rawref_dir)
print("Done")

if use_mypyc:
# Build mypyc
print("-- Building kv_cache_manager_v2 mypyc extensions...", end=" ")
# setup_mypyc.py is in kv_cache_manager_v2 but executed from runtime dir
setup_mypyc = kv_cache_mgr_dir / "setup_mypyc.py"
build_run(f'"{venv_python}" "{setup_mypyc}" build_ext --inplace',
cwd=runtime_dir)
build_run(
f'"{venv_python}" "{setup_mypyc}" build_ext --inplace{build_temp_arg}',
cwd=runtime_dir)

# Verify that the shared library was generated
if not list(runtime_dir.glob("*__mypyc*.so")):
Expand All @@ -527,16 +548,17 @@ def build_kv_cache_manager_v2(project_dir, venv_python, use_mypyc=False):
def main(*,
build_type: str = "Release",
generator: str = "",
build_dir: Path = None,
dist_dir: Path = None,
cuda_architectures: str = None,
job_count: int = None,
build_root: Optional[Path] = None,
build_dir: Optional[Path] = None,
dist_dir: Optional[Path] = None,
cuda_architectures: Optional[str] = None,
job_count: Optional[int] = None,
extra_cmake_vars: Sequence[str] = tuple(),
extra_make_targets: str = "",
nccl_root: str = None,
nixl_root: str = None,
mooncake_root: str = None,
internal_cutlass_kernels_root: str = None,
nccl_root: Optional[str] = None,
nixl_root: Optional[str] = None,
mooncake_root: Optional[str] = None,
internal_cutlass_kernels_root: Optional[str] = None,
clean: bool = False,
clean_wheel: bool = False,
configure_cmake: bool = False,
Expand Down Expand Up @@ -565,6 +587,23 @@ def main(*,

project_dir = get_project_dir()
apply_version_override(project_dir, version_override)

# Out-of-tree build state: everything metadata-heavy (venv, wheel
# staging, ccache, intermediate objects) goes under build_root, keeping
# the checkout free of high-churn I/O (important on network filesystems).
# Resolve before chdir so a relative path stays anchored to the caller's
# working directory.
if build_root is None and os.environ.get("TRTLLM_BUILD_ROOT"):
build_root = Path(os.environ["TRTLLM_BUILD_ROOT"])
if build_root is not None:
build_root = build_root.resolve()
build_root.mkdir(parents=True, exist_ok=True)
print(f"-- Out-of-tree build state under: {build_root}")
# setup.py redirects the setuptools staging tree and *.egg-info
# to this directory; an explicit env var set by the user wins.
os.environ.setdefault("TRTLLM_WHEEL_STAGING_DIR",
str(build_root / "wheel-staging"))
Comment thread
brnguyen2 marked this conversation as resolved.

os.chdir(project_dir)

# Get all submodules and check their folder exists. If not,
Expand All @@ -584,7 +623,8 @@ def main(*,
venv_python, venv_conan = setup_venv(project_dir,
project_dir / requirements_filename,
no_venv,
yes=yes)
yes=yes,
build_root=build_root)

if cuda_architectures is not None:
if "70-real" in cuda_architectures:
Expand Down Expand Up @@ -652,14 +692,22 @@ def main(*,
raise RuntimeError("Mooncake is not supported on Windows.")
cmake_def_args.append(f"-DMOONCAKE_ROOT={mooncake_root}")

build_dir = get_build_dir(build_dir, build_type)
build_dir = get_build_dir(build_dir, build_type, build_root)
first_build = not Path(build_dir, "CMakeFiles").exists()

if clean and build_dir.exists():
clear_folder(build_dir) # Keep the folder in case it is mounted.
build_dir.mkdir(parents=True, exist_ok=True)

if use_ccache:
if build_root is not None and "CCACHE_DIR" not in os.environ:
# Default the cache next to the rest of the out-of-tree build
# state. Point CCACHE_DIR at persistent storage instead to keep
# compile results across ephemeral nodes/containers.
ccache_dir = build_root / "ccache"
ccache_dir.mkdir(parents=True, exist_ok=True)
os.environ["CCACHE_DIR"] = str(ccache_dir)
print(f"-- ccache directory: {ccache_dir}")
cmake_def_args.append(
f"-DCMAKE_CXX_COMPILER_LAUNCHER=ccache -DCMAKE_CUDA_COMPILER_LAUNCHER=ccache"
)
Expand Down Expand Up @@ -730,8 +778,15 @@ def main(*,

with working_directory(build_dir):
if clean or first_build or configure_cmake or configure_only:
# Conan writes a CMakeUserPresets.json convenience file next to
# cpp/CMakeLists.txt; with out-of-tree build state it would be
# the only build file left in the checkout (and would point at a
# possibly ephemeral location), so skip generating it.
conan_extra_args = (
" -c tools.cmake.cmaketoolchain:user_presets=False"
if build_root is not None else "")
build_run(
f"\"{venv_conan}\" install --build=missing --no-remote --output-folder={build_dir}/conan -s 'build_type={build_type}' {source_dir}"
f"\"{venv_conan}\" install --build=missing --no-remote --output-folder={build_dir}/conan -s 'build_type={build_type}'{conan_extra_args} {source_dir}"
)
cmake_def_args.append(
f"-DCMAKE_TOOLCHAIN_FILE={build_dir}/conan/conan_toolchain.cmake"
Expand Down Expand Up @@ -1128,7 +1183,10 @@ def get_binding_lib(subdirectory, name):
nixl_root is not None or mooncake_root is not None,
binding_lib_file_name)

build_kv_cache_manager_v2(project_dir, venv_python, use_mypyc=mypyc)
build_kv_cache_manager_v2(project_dir,
venv_python,
use_mypyc=mypyc,
build_root=build_root)

if not skip_building_wheel:
if dist_dir is None:
Expand All @@ -1146,6 +1204,17 @@ def get_binding_lib(subdirectory, name):
# This breaks the Windows CI/CD pipeline when building
# and validating python changes in the whl.
clear_folder(dist_dir)
# Without --build_root the setuptools staging tree (build_base)
# lives at project_dir/build == dist_dir, so the clear above
# already wipes it. With --build_root it moves under
# TRTLLM_WHEEL_STAGING_DIR, so clearing dist_dir alone would
# leave stale copies of deleted package files there to be
# re-packed into the next "clean" wheel. Clear it too.
staging_dir = os.environ.get("TRTLLM_WHEEL_STAGING_DIR")
if staging_dir:
staging_build = Path(staging_dir) / "build"
if staging_build.exists():
clear_folder(staging_build)

extra_wheel_build_args = os.getenv("EXTRA_WHEEL_BUILD_ARGS", "")
plat_name_arg = ""
Expand Down Expand Up @@ -1299,12 +1368,22 @@ def add_arguments(parser: ArgumentParser):
help=
"Directory containing internal_cutlass_kernels sources. If specified, the internal_cutlass_kernels and NVRTC wrapper libraries will be built from source."
)
parser.add_argument(
"--build_root",
type=Path,
help=
"Directory for all out-of-tree build state (also via TRTLLM_BUILD_ROOT env var). "
"When set, the CMake build dir, build venv, wheel staging, intermediate "
"objects and (with --use_ccache) the ccache directory default under this "
"directory instead of the checkout. Point it at fast local storage (e.g. "
"/tmp) when the checkout lives on a network filesystem. Individual "
"options like --build_dir and CCACHE_DIR still override their piece.")
parser.add_argument(
"--build_dir",
type=Path,
help=
"Directory where C++ sources are built (default: cpp/build or cpp/build_<build_type>)"
)
"Directory where C++ sources are built (default: cpp/build or cpp/build_<build_type>, "
"or <build_root>/cpp-build* when --build_root is set)")
parser.add_argument(
"--dist_dir",
type=Path,
Expand Down
26 changes: 26 additions & 0 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -442,6 +442,31 @@ def extract_from_precompiled(precompiled_location: str, package_data: list[str],
msa_package_dir = {"fmha_sm100": "3rdparty/MSA/python/fmha_sm100"}
packages += ["fmha_sm100"]


def get_build_state_options():
"""Optionally redirect setuptools build state out of the source tree.

When TRTLLM_WHEEL_STAGING_DIR is set (e.g. by scripts/build_wheel.py
--build_root), the setuptools staging tree (build/) and *.egg-info are
written there instead of into the checkout. This keeps metadata-heavy
churn off slow network filesystems; behavior is unchanged when unset.
"""
staging_dir = os.environ.get("TRTLLM_WHEEL_STAGING_DIR")
if not staging_dir:
return {}
staging = Path(staging_dir)
egg_base = staging / "egg-info"
egg_base.mkdir(parents=True, exist_ok=True)
return {
"build": {
"build_base": str(staging / "build")
},
"egg_info": {
"egg_base": str(egg_base)
},
}


# https://setuptools.pypa.io/en/latest/references/keywords.html
setup(
name='tensorrt_llm',
Expand All @@ -466,6 +491,7 @@ def extract_from_precompiled(precompiled_location: str, package_data: list[str],
"Programming Language :: Python :: 3.12",
],
distclass=BinaryDistribution,
options=get_build_state_options(),
license="Apache License 2.0",
keywords="nvidia tensorrt deeplearning inference",
package_data={
Expand Down
Loading