diff --git a/docs/source/installation/build-from-source.md b/docs/source/installation/build-from-source.md index a6d66c7f3d58..30b0d6b0fc73 100644 --- a/docs/source/installation/build-from-source.md +++ b/docs/source/installation/build-from-source.md @@ -119,6 +119,8 @@ Related knobs for shared-storage workflows: Plain local-disk builds are unaffected: without `--build_root`, all paths behave as before. +Artifact copy-back into the checkout (`tensorrt_llm/include`, the `deep_gemm`/`deep_ep`/`flash_mla` Python trees) is incremental: directory trees are populated once via a streamed tar pipeline and then kept in sync by size/mtime comparison, so a rebuild rewrites only what changed instead of re-copying ~10k files onto the (possibly network) filesystem. + #### Out-of-tree wheel builds (read-only checkout) For CI or ephemeral-node workflows that only need a wheel, add `--out-of-tree` to guarantee the checkout is never written — it can even be mounted read-only: diff --git a/scripts/build_wheel.py b/scripts/build_wheel.py index 4a9f9ee02f21..44aeb405a604 100755 --- a/scripts/build_wheel.py +++ b/scripts/build_wheel.py @@ -28,7 +28,8 @@ from multiprocessing import cpu_count from pathlib import Path from shutil import copy, copytree, rmtree -from subprocess import DEVNULL, CalledProcessError, check_output, run +from subprocess import (DEVNULL, PIPE, CalledProcessError, Popen, check_output, + run) from typing import Optional, Sequence try: @@ -579,6 +580,107 @@ def build_kv_cache_manager_v2(project_dir, print("-- Done building kv_cache_manager_v2.") +def _tar_pipe_copy(src: Path, dst: Path) -> bool: + """Populate dst from src as one streamed tar pipeline. + + A single reader/writer pair with kernel-buffered pipe I/O is much faster + than per-file copies on network filesystems. Dereferences symlinks (-h) + and preserves mtimes, matching copytree(symlinks=False) + copystat. + Returns False if tar is unavailable or fails, so callers can fall back. + """ + tar_bin = shutil.which("tar") + if tar_bin is None: + return False + dst.mkdir(parents=True, exist_ok=True) + # posix (pax) format keeps sub-second mtimes; the gnu default truncates + # to whole seconds, which would defeat sync_tree's mtime comparison. + # + # Chain the producer and consumer tars directly through an OS pipe rather + # than a shell string. Checking both return codes reports a failing + # producer (e.g. an unreadable source file) that a shell pipeline would + # mask behind the consumer's exit status, without depending on a bash that + # supports `set -o pipefail`; it also avoids shell quoting entirely. + producer = Popen( + [tar_bin, "--format=posix", "-C", + str(src), "-chf", "-", "."], + stdout=PIPE) + consumer = Popen([tar_bin, "-C", str(dst), "-xf", "-"], + stdin=producer.stdout) + # Close our copy of the write end so the consumer sees EOF when the + # producer exits (and the producer gets SIGPIPE if the consumer dies). + producer.stdout.close() + consumer.wait() + producer.wait() + return producer.returncode == 0 and consumer.returncode == 0 + + +def sync_tree(src, dst, exclude: Sequence[str] = ()): + """Mirror the src directory into dst, touching only what changed. + + Replaces the rmtree+copytree pattern for artifact copy-back: files are + rewritten only when size or mtime differs and entries missing from src + are deleted, so incremental rebuilds cause almost no I/O on the + destination (which may be a slow network filesystem). A missing dst is + populated via a streamed tar pipeline instead of per-file copies. + Symlinks are dereferenced like copytree(symlinks=False); mtimes are + preserved so the next sync can compare against them. exclude lists + fnmatch patterns for entry names to skip. + """ + import fnmatch + + src = Path(src).resolve() + dst = Path(dst) + + def excluded(name): + return any(fnmatch.fnmatch(name, pat) for pat in exclude) + + if dst.is_symlink(): + dst.unlink() + elif dst.exists() and src == dst.resolve(): + return + + if not dst.exists(): + if not exclude and _tar_pipe_copy(src, dst): + return + copytree(src, + dst, + symlinks=False, + ignore=shutil.ignore_patterns(*exclude) if exclude else None) + return + + for root, dirs, files in os.walk(src, followlinks=True): + rel = Path(root).relative_to(src) + dirs[:] = [d for d in dirs if not excluded(d)] + files = [f for f in files if not excluded(f)] + dst_root = dst / rel + if dst_root.exists() and not dst_root.is_dir(): + dst_root.unlink() + dst_root.mkdir(exist_ok=True) + keep = set(dirs) | set(files) + for stale in os.listdir(dst_root): + if stale not in keep: + stale_path = dst_root / stale + if stale_path.is_dir() and not stale_path.is_symlink(): + rmtree(stale_path) + else: + stale_path.unlink() + for name in files: + src_file = Path(root) / name + dst_file = dst_root / name + try: + src_stat = src_file.stat() + dst_stat = dst_file.stat() + if (src_stat.st_size == dst_stat.st_size + and abs(src_stat.st_mtime - dst_stat.st_mtime) < 1e-3): + continue + except OSError: + pass + if dst_file.is_dir() and not dst_file.is_symlink(): + rmtree(dst_file) + # copy2: mtime must survive for the next sync's comparison. + shutil.copy2(src_file, dst_file) + + def stage_python_package(project_dir: Path, staging_dir: Path) -> None: """Copy the sources setup.py packages into an out-of-tree staging project. @@ -587,15 +689,13 @@ def stage_python_package(project_dir: Path, staging_dir: Path) -> None: """ print(f"-- Staging python package sources into {staging_dir} ...") staging_dir.mkdir(parents=True, exist_ok=True) - ignore = shutil.ignore_patterns("__pycache__", "*.pyc") # examples: setup.py's root-level find_packages() ships the # examples.configs.database package from it. for tree in ("tensorrt_llm", "triton_kernels", "examples", "3rdparty/MSA/python/fmha_sm100"): - dst = staging_dir / tree - if dst.exists(): - rmtree(dst) - copytree(project_dir / tree, dst, symlinks=False, ignore=ignore) + sync_tree(project_dir / tree, + staging_dir / tree, + exclude=("__pycache__", "*.pyc")) top_level_files = [ "setup.py", "pyproject.toml", "requirements.txt", "requirements-dev.txt", "constraints.txt", "LICENSE", "README.md" @@ -937,8 +1037,9 @@ def main(*, include_dir = pkg_dir / "include" if lib_dir.exists(): clear_folder(lib_dir) - if include_dir.exists(): - clear_folder(include_dir) + # include_dir is not cleared: its subtrees are synced with deletion of + # extraneous entries (sync_tree) or guarded by generation stamps, so + # incremental rebuilds skip the ~10k-file rewrite of the include tree. # Remove auto-generated attributions file from previous builds auto_attr_file = wheel_project_dir / "ATTRIBUTIONS.md" if auto_attr_file.exists(): @@ -972,21 +1073,7 @@ def safe_copy(src, dst): install_file = safe_copy - # Wrapper for copytree that checks if source and destination are the same - def safe_copytree(src, dst, dirs_exist_ok=True): - """Copy tree, but skip if source and destination resolve to the same directory.""" - src_path = Path(src).resolve() - dst_path = Path(dst) - if dst_path.is_symlink(): - dst_path.unlink() - elif src_path == dst_path.resolve(): - # Source and destination are the same, skip copying - return - if dst_path.exists() and dirs_exist_ok: - rmtree(dst_path) - copytree(src_path, dst_path, dirs_exist_ok=dirs_exist_ok) - - install_tree = safe_copytree + install_tree = sync_tree if skip_building_wheel and linking_install_binary: def symlink_remove_dst(src, dst): @@ -1000,10 +1087,12 @@ def symlink_remove_dst(src, dst): install_file = symlink_remove_dst - def symlink_remove_dst_tree(src, dst, dirs_exist_ok=True): + def symlink_remove_dst_tree(src, dst): src = os.path.abspath(src) dst = os.path.abspath(dst) - if dirs_exist_ok and os.path.lexists(dst): + if os.path.isdir(dst) and not os.path.islink(dst): + rmtree(dst) # left behind by a previous copy-mode build + elif os.path.lexists(dst): os.remove(dst) os.symlink(src, dst) @@ -1012,8 +1101,7 @@ def symlink_remove_dst_tree(src, dst, dirs_exist_ok=True): lib_dir.mkdir(parents=True, exist_ok=True) include_dir.mkdir(parents=True, exist_ok=True) install_tree(get_source_dir() / "include" / "tensorrt_llm" / "deep_gemm", - include_dir / "deep_gemm", - dirs_exist_ok=True) + include_dir / "deep_gemm") # Copy FMHA kernel generation headers for JIT compilation fmha_build_dir = build_dir / "tensorrt_llm" / "kernels" / "trtllmGenKernels" / "fmha" @@ -1148,7 +1236,7 @@ def copy_resolving_symlink(src_path, dst_path): ucx_dir = lib_dir / "ucx" if ucx_dir.exists(): clear_folder(ucx_dir) - install_tree("/usr/local/ucx/lib", ucx_dir, dirs_exist_ok=True) + install_tree("/usr/local/ucx/lib", ucx_dir) build_run( f"find {ucx_dir} -type f -name '*.so*' -exec patchelf --set-rpath \'$ORIGIN:$ORIGIN/ucx:$ORIGIN/../\' {{}} \\;" ) @@ -1170,7 +1258,7 @@ def copy_resolving_symlink(src_path, dst_path): nixl_lib_path = "/opt/nvidia/nvda_nixl/lib/aarch64-linux-gnu" if not os.path.exists(nixl_lib_path): nixl_lib_path = "/opt/nvidia/nvda_nixl/lib64" - install_tree(nixl_lib_path, nixl_dir, dirs_exist_ok=True) + install_tree(nixl_lib_path, nixl_dir) build_run( f"find {nixl_dir} -type f -name '*.so*' -exec patchelf --set-rpath \'$ORIGIN:$ORIGIN/plugins:$ORIGIN/../:$ORIGIN/../ucx/:$ORIGIN/../../ucx/\' {{}} \\;" ) @@ -1201,20 +1289,11 @@ def copy_resolving_symlink(src_path, dst_path): install_file(build_dir / "tensorrt_llm/runtime/utils/libpg_utils.so", lib_dir / "libpg_utils.so") + # deep_ep/deep_gemm are synced in place below (sync_tree removes stale + # entries); deep_ep is deleted explicitly when this build does not + # produce it. deep_ep_dir = pkg_dir / "deep_ep" - if deep_ep_dir.is_symlink(): - deep_ep_dir.unlink() - elif deep_ep_dir.is_dir(): - clear_folder(deep_ep_dir) - deep_ep_dir.rmdir() - - # Handle deep_gemm installation deep_gemm_dir = pkg_dir / "deep_gemm" - if deep_gemm_dir.is_symlink(): - deep_gemm_dir.unlink() - elif deep_gemm_dir.is_dir(): - clear_folder(deep_gemm_dir) - deep_gemm_dir.rmdir() scripts_dir = pkg_dir / "scripts" if scripts_dir.exists(): @@ -1241,13 +1320,17 @@ def get_binding_lib(subdirectory, name): with (build_dir / "tensorrt_llm" / "deep_ep" / "cuda_architectures.txt").open() as f: deep_ep_cuda_architectures = f.read().strip().strip(";") + if not deep_ep_cuda_architectures and deep_ep_dir.exists(): + if deep_ep_dir.is_symlink(): + deep_ep_dir.unlink() + else: + rmtree(deep_ep_dir) if deep_ep_cuda_architectures: install_file(get_binding_lib("deep_ep", "deep_ep_cpp_tllm"), pkg_dir) - install_tree(build_dir / "tensorrt_llm" / "deep_ep" / "python" / - "deep_ep", - deep_ep_dir, - dirs_exist_ok=True) + install_tree( + build_dir / "tensorrt_llm" / "deep_ep" / "python" / "deep_ep", + deep_ep_dir) (lib_dir / "nvshmem").mkdir(exist_ok=True) install_file( build_dir / "tensorrt_llm/deep_ep/nvshmem-build/License.txt", @@ -1263,10 +1346,9 @@ def get_binding_lib(subdirectory, name): install_file(get_binding_lib("deep_gemm", "deep_gemm_cpp_tllm"), pkg_dir) - install_tree(build_dir / "tensorrt_llm" / "deep_gemm" / "python" / - "deep_gemm", - deep_gemm_dir, - dirs_exist_ok=True) + install_tree( + build_dir / "tensorrt_llm" / "deep_gemm" / "python" / "deep_gemm", + deep_gemm_dir) with (build_dir / "tensorrt_llm" / "flash_mla" / "cuda_architectures.txt").open() as f: @@ -1274,10 +1356,9 @@ def get_binding_lib(subdirectory, name): if flash_mla_cuda_architectures: install_file(get_binding_lib("flash_mla", "flash_mla_cpp_tllm"), pkg_dir) - install_tree(build_dir / "tensorrt_llm" / "flash_mla" / "python" / - "flash_mla", - pkg_dir / "flash_mla", - dirs_exist_ok=True) + install_tree( + build_dir / "tensorrt_llm" / "flash_mla" / "python" / + "flash_mla", pkg_dir / "flash_mla") if not skip_stubs: with working_directory(pkg_dir): diff --git a/tests/unittest/scripts/test_build_wheel_copy_back.py b/tests/unittest/scripts/test_build_wheel_copy_back.py new file mode 100644 index 000000000000..8d1449405911 --- /dev/null +++ b/tests/unittest/scripts/test_build_wheel_copy_back.py @@ -0,0 +1,196 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Equivalence tests for build_wheel.py's sync_tree copy-back. + +sync_tree replaced an rmtree+copytree copy-back of build artifacts into the +wheel. These tests pin the on-disk result of sync_tree to that of the old +path across cold populate and incremental re-sync, so the optimization can +never silently drop, stale, or corrupt a file relative to a clean recopy. + +Small synthetic trees only (no build required), so this runs in well under a +second on a CPU node. +""" + +from __future__ import annotations + +import importlib.util +import os +import shutil +import stat +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.cpu_only + +REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent +SCRIPT_PATH = REPO_ROOT / "scripts" / "build_wheel.py" + + +@pytest.fixture(scope="module") +def sync_tree(): + spec = importlib.util.spec_from_file_location("build_wheel", SCRIPT_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.sync_tree + + +def old_copy(src, dst, exclude=()): + """Replicate the pre-sync_tree copy-back. + + Wipe dst, then copytree dereferencing symlinks. ``exclude`` mirrors the + old shutil.ignore_patterns call. + """ + src = Path(src).resolve() + dst = Path(dst) + if dst.is_symlink(): + dst.unlink() + elif dst.exists(): + shutil.rmtree(dst) + ignore = shutil.ignore_patterns(*exclude) if exclude else None + shutil.copytree(src, dst, symlinks=False, ignore=ignore) + + +def snapshot(root): + """Map relpath -> (kind, content-or-None, perm-bits). + + A surviving symlink is recorded as its own kind so that a deref mismatch + shows up as a diff. + """ + root = Path(root) + out = {} + for dirpath, dirnames, filenames in os.walk(root): + for name in dirnames: + p = Path(dirpath) / name + rel = str(p.relative_to(root)) + if p.is_symlink(): + out[rel] = ("dirsymlink", None, None) + else: + out[rel] = ("dir", None, stat.S_IMODE(p.stat().st_mode)) + for name in filenames: + p = Path(dirpath) / name + rel = str(p.relative_to(root)) + if p.is_symlink(): + out[rel] = ("filesymlink", None, None) + else: + out[rel] = ("file", p.read_bytes(), stat.S_IMODE(p.stat().st_mode)) + return out + + +def assert_trees_equal(new_dir, old_dir): + new_snap, old_snap = snapshot(new_dir), snapshot(old_dir) + assert new_snap == old_snap + + +def build_tree(base): + """Build a small tree covering where the two copy paths could differ. + + Nested dirs, an empty dir, an executable file, a symlink to a file, a + symlink to a dir, and __pycache__/*.pyc that the stage path excludes. + """ + base = Path(base) + (base / "pkg/sub").mkdir(parents=True) + (base / "pkg/__init__.py").write_bytes(b"x = 1\n") + (base / "pkg/sub/mod.py").write_bytes(b"def f():\n return 42\n") + (base / "pkg/data.bin").write_bytes(bytes(range(256))) + run_sh = base / "pkg/run.sh" + run_sh.write_bytes(b"#!/bin/sh\necho hi\n") + run_sh.chmod(0o755) + (base / "pkg/emptydir").mkdir() + (base / "pkg/link_to_init.py").symlink_to(base / "pkg/__init__.py") + (base / "pkg/link_to_sub").symlink_to(base / "pkg/sub") + (base / "pkg/__pycache__").mkdir() + (base / "pkg/__pycache__/mod.cpython-312.pyc").write_bytes(b"\x00\x01") + (base / "pkg/stale.pyc").write_bytes(b"\x00") + + +def test_cold_populate_matches_copytree(sync_tree, tmp_path): + """Cold sync (streamed tar populate) == old clean copytree.""" + src = tmp_path / "src" + src.mkdir() + build_tree(src) + new_dir, old_dir = tmp_path / "new", tmp_path / "old" + sync_tree(src, new_dir) + old_copy(src, old_dir) + assert_trees_equal(new_dir, old_dir) + + +def test_exclude_matches_ignore_patterns(sync_tree, tmp_path): + """sync_tree(exclude=...) drops the same entries as ignore_patterns.""" + src = tmp_path / "src" + src.mkdir() + build_tree(src) + exclude = ("__pycache__", "*.pyc") + new_dir, old_dir = tmp_path / "new", tmp_path / "old" + sync_tree(src, new_dir, exclude=exclude) + old_copy(src, old_dir, exclude=exclude) + assert_trees_equal(new_dir, old_dir) + # exclusion actually happened + assert not (new_dir / "pkg/__pycache__").exists() + assert not (new_dir / "pkg/stale.pyc").exists() + + +def test_incremental_converges_to_fresh_copy(sync_tree, tmp_path): + """Incremental re-sync after mutation must match a fresh full recopy. + + Covers changed/added/deleted files and file<->dir type swaps. + """ + src = tmp_path / "src" + src.mkdir() + build_tree(src) + new_dir = tmp_path / "new" + sync_tree(src, new_dir) # initial populate + + # mutate the source + (src / "pkg/sub/mod.py").write_bytes(b"def f():\n return 99\n") # change + (src / "pkg/data.bin").unlink() # delete + (src / "pkg/added.txt").write_bytes(b"new\n") # add + (src / "pkg/run.sh").unlink() # file -> dir + (src / "pkg/run.sh").mkdir() + (src / "pkg/run.sh/inner.txt").write_bytes(b"now a dir\n") + (src / "pkg/emptydir").rmdir() # dir -> file + (src / "pkg/emptydir").write_bytes(b"now a file\n") + + sync_tree(src, new_dir) # incremental re-sync onto existing dst + old_dir = tmp_path / "old" + old_copy(src, old_dir) # ground truth + assert_trees_equal(new_dir, old_dir) + + +def test_warm_noop_is_stable(sync_tree, tmp_path): + """Re-syncing an unchanged source must not corrupt the destination.""" + src = tmp_path / "src" + src.mkdir() + build_tree(src) + new_dir = tmp_path / "new" + sync_tree(src, new_dir) + sync_tree(src, new_dir) # second sync, no source change + old_dir = tmp_path / "old" + old_copy(src, old_dir) + assert_trees_equal(new_dir, old_dir) + + +def test_same_src_dst_is_noop(sync_tree, tmp_path): + """sync_tree onto itself must not wipe the tree. + + Guards the case where source and destination resolve to the same path. + """ + src = tmp_path / "src" + src.mkdir() + build_tree(src) + before = snapshot(src) + sync_tree(src, src) + assert snapshot(src) == before