[TRTLLM-14628][perf] Incremental, streamed artifact copy-back into the checkout

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:

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,21 +1346,19 @@ 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:
             flash_mla_cuda_architectures = f.read().strip().strip(";")
         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):