[None][feat] Add a link mode to the precompiled editable-install path

docs/source/installation/build-from-source.md

@@ -139,3 +139,13 @@ TRTLLM_USE_PRECOMPILED=1 pip install -e .
 ```
 
 This downloads a precompiled wheel matching the version in `tensorrt_llm/version.py` and extracts its compiled libraries into your working directory. Override the version with `TRTLLM_USE_PRECOMPILED=x.y.z` or specify a custom URL/path with `TRTLLM_PRECOMPILED_LOCATION`.
+
+#### Sharing one build tree between checkouts
+
+`TRTLLM_PRECOMPILED_LOCATION` also accepts a local directory in git-clone layout, which lets a second checkout reuse a build you already have instead of downloading a wheel. By default the compiled artifacts are copied into the checkout, so each one holds its own multi-gigabyte copy. Add `TRTLLM_PRECOMPILED_LINK=1` to symlink them instead:
+
+```bash
+TRTLLM_PRECOMPILED_LINK=1 TRTLLM_PRECOMPILED_LOCATION=/path/to/built/checkout pip install -e .
+```
+
+Use this when several checkouts (for example git worktrees carrying Python-only changes) share a single built tree: nothing is duplicated, and rebuilding the shared tree updates every checkout at once. An existing `3rdparty/fmha_sm100` symlink is left in place rather than replaced, so a checkout that already shares a build tree keeps its links. The flag only applies to a local directory; it is an error to combine it with a wheel or a URL. All checkouts must stay on a commit whose C++ sources match the shared build, since the reused artifacts are not rebuilt. When the source is a local directory, the install prints a warning if the two checkouts are on different commits and any of `cpp/`, `3rdparty/`, `setup.py`, `scripts/build_wheel.py` or `requirements.txt` differ between them, naming the two commits and the first few differing files. It is only a warning, and it applies to both copy and link mode: the artifacts are equally stale either way. Like the rest of this path's output it comes from `setup.py`, which pip shows only with `pip install -v`. If an import fails afterwards with a message about rebuilding, that skew is the first thing to check.

setup.py

@@ -250,6 +250,73 @@ def should_skip_precompiled_package_data(filename: str) -> bool:
         source_owned_package_data_prefixes)
 
 
+def warn_on_build_skew(precompiled_location: str) -> None:
+    """Warn when the source checkout differs from this one where it matters.
+
+    The precompiled artifacts are reused as they are, never rebuilt, so any
+    difference in what feeds the native build makes them stale. This is
+    advisory: it never fails the install, since the two checkouts are often
+    meant to differ (that is the point of reusing a build) and only some of
+    those differences matter.
+    """
+    import subprocess
+
+    NATIVE_BUILD_INPUTS = [
+        "cpp/", "3rdparty/", "setup.py", "scripts/build_wheel.py",
+        "requirements.txt"
+    ]
+
+    def head_of(checkout: str) -> str | None:
+        try:
+            done = subprocess.run(["git", "-C", checkout, "rev-parse", "HEAD"],
+                                  capture_output=True,
+                                  text=True,
+                                  check=True)
+        except (OSError, subprocess.SubprocessError):
+            return None
+        return done.stdout.strip() or None
+
+    source_head = head_of(precompiled_location)
+    current_head = head_of(".")
+    if source_head is None or current_head is None:
+        print("Cannot check for build skew: one of the two checkouts is not a "
+              "git repository. Make sure the precompiled artifacts were built "
+              "from these sources.")
+        return
+    if source_head == current_head:
+        return
+
+    stale = ("Import errors mentioning 'rebuild and install' after this may be "
+             "ABI skew; rebuild, or pick a precompiled source that matches.")
+    try:
+        # Both revisions are reachable here when the two checkouts are
+        # worktrees of one clone, which is the case this is meant to catch.
+        done = subprocess.run(
+            ["git", "diff", "--name-only", source_head, current_head, "--"] +
+            NATIVE_BUILD_INPUTS,
+            capture_output=True,
+            text=True,
+            check=True)
+    except (OSError, subprocess.SubprocessError):
+        print(
+            f"WARNING: the precompiled artifacts come from {source_head[:12]} "
+            f"but this checkout is {current_head[:12]}, and the difference "
+            f"could not be inspected from here. {stale}")
+        return
+
+    changed = done.stdout.split()
+    if not changed:
+        return
+    shown = ", ".join(changed[:3])
+    if len(changed) > 3:
+        shown += ", ..."
+    count = f"{len(changed)} file" + ("s" if len(changed) > 1 else "")
+    print(
+        f"WARNING: the precompiled artifacts come from {source_head[:12]} but "
+        f"this checkout is {current_head[:12]}; {count} feeding the "
+        f"native build changed ({shown}). {stale}")
+
+
 def extract_from_precompiled(precompiled_location: str, package_data: list[str],
                              workspace: str) -> None:
     """Extract package data (binaries and other materials) from a precompiled wheel or local directory to the working directory.
@@ -259,6 +326,9 @@ def extract_from_precompiled(precompiled_location: str, package_data: list[str],
     - Local directory (git clone structure): e.g., /home/dev/TensorRT-LLM
     - Local wheel file: e.g., /path/to/tensorrt_llm-*.whl
     - Remote URL: Downloads and extracts from URL (wheel or tar.gz)
+
+    With TRTLLM_PRECOMPILED_LINK=1 a local directory is symlinked instead of
+    copied, so several checkouts can share one build tree.
     """
     import fnmatch
     import shutil
@@ -268,12 +338,22 @@ def extract_from_precompiled(precompiled_location: str, package_data: list[str],
 
     from setuptools.errors import SetupError
 
+    # Only a local directory can be linked; a wheel or a URL has no build tree
+    # to point at.
+    link_artifacts = os.getenv("TRTLLM_PRECOMPILED_LINK", "0") not in ("", "0")
+    if link_artifacts and not os.path.isdir(precompiled_location):
+        raise SetupError(
+            "TRTLLM_PRECOMPILED_LINK=1 requires TRTLLM_PRECOMPILED_LOCATION to "
+            "be a local directory in git-clone layout, but got "
+            f"{precompiled_location}.")
+
     # Handle local directory (assuming repo structure)
     if os.path.isdir(precompiled_location):
         precompiled_location = os.path.abspath(precompiled_location)
         print(
             f"Using local directory as precompiled source: {precompiled_location}"
         )
+        warn_on_build_skew(precompiled_location)
         source_tensorrt_llm = os.path.join(precompiled_location, "tensorrt_llm")
         if not os.path.isdir(source_tensorrt_llm):
             raise SetupError(
@@ -317,8 +397,14 @@ def extract_from_precompiled(precompiled_location: str, package_data: list[str],
                 dst_dir = os.path.dirname(dst_file)
                 if dst_dir:
                     os.makedirs(dst_dir, exist_ok=True)
-                print(f"Copying {rel_path} from local directory.")
-                shutil.copy2(src_file, dst_file)
+                if link_artifacts:
+                    if os.path.lexists(dst_file):
+                        os.unlink(dst_file)
+                    print(f"Linking {rel_path} from local directory.")
+                    os.symlink(os.path.abspath(src_file), dst_file)
+                else:
+                    print(f"Copying {rel_path} from local directory.")
+                    shutil.copy2(src_file, dst_file)
 
         source_fmha = os.path.join(precompiled_location, "3rdparty",
                                    "fmha_sm100")
@@ -328,12 +414,25 @@ def extract_from_precompiled(precompiled_location: str, package_data: list[str],
                 "packaging and does not contain 3rdparty/fmha_sm100. Use a "
                 "precompiled source built with MSA packaging support.")
         dst_fmha = os.path.join("3rdparty", "fmha_sm100")
-        print(f"Copying fmha_sm100 from local directory: {source_fmha}")
-        if os.path.islink(dst_fmha):
-            os.unlink(dst_fmha)
-        elif os.path.isdir(dst_fmha):
-            shutil.rmtree(dst_fmha)
-        shutil.copytree(source_fmha, dst_fmha)
+        if link_artifacts:
+            if os.path.islink(dst_fmha):
+                # The checkout already shares a build tree; replacing the link
+                # would undo that.
+                print(f"Keeping existing fmha_sm100 symlink: {dst_fmha}")
+            else:
+                if os.path.isdir(dst_fmha):
+                    shutil.rmtree(dst_fmha)
+                # copytree() creates the parent below; os.symlink() does not.
+                os.makedirs(os.path.dirname(dst_fmha), exist_ok=True)
+                print(f"Linking fmha_sm100 from local directory: {source_fmha}")
+                os.symlink(source_fmha, dst_fmha)
+        else:
+            print(f"Copying fmha_sm100 from local directory: {source_fmha}")
+            if os.path.islink(dst_fmha):
+                os.unlink(dst_fmha)
+            elif os.path.isdir(dst_fmha):
+                shutil.rmtree(dst_fmha)
+            shutil.copytree(source_fmha, dst_fmha)
         return
 
     # Handle local file or remote URL

tests/unittest/others/test_precompiled_link_mode.py

@@ -0,0 +1,208 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""``TRTLLM_PRECOMPILED_LINK`` in ``setup.py``'s ``extract_from_precompiled``.
+
+The local-directory source normally copies the compiled artifacts into the
+checkout. ``TRTLLM_PRECOMPILED_LINK=1`` symlinks them instead so several
+checkouts can share one build tree.
+
+``setup.py`` cannot be imported (module scope runs ``setup()``), so the two
+functions under test are pulled out of its AST and executed on their own. They
+are self-contained: every import they need is inside the function body, apart
+from ``os``.
+"""
+
+import ast
+import os
+from pathlib import Path
+
+import pytest
+
+_SETUP_PY = Path(__file__).resolve().parents[3] / "setup.py"
+_WANTED = ("should_skip_precompiled_package_data", "warn_on_build_skew", "extract_from_precompiled")
+
+
+def _load(name):
+    assert _SETUP_PY.is_file(), _SETUP_PY
+    tree = ast.parse(_SETUP_PY.read_text(encoding="utf-8"))
+    wanted = [
+        node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in _WANTED
+    ]
+    assert {node.name for node in wanted} == set(_WANTED)
+    namespace = {"os": os}
+    exec(compile(ast.Module(body=wanted, type_ignores=[]), str(_SETUP_PY), "exec"), namespace)
+    return namespace[name]
+
+
+@pytest.fixture(scope="module")
+def extract_from_precompiled():
+    return _load("extract_from_precompiled")
+
+
+@pytest.fixture(scope="module")
+def warn_on_build_skew():
+    return _load("warn_on_build_skew")
+
+
+@pytest.fixture
+def build_tree(tmp_path, monkeypatch):
+    """A source checkout holding compiled artifacts, and an empty destination."""
+    source = tmp_path / "built-checkout"
+    (source / "tensorrt_llm" / "libs").mkdir(parents=True)
+    (source / "tensorrt_llm" / "libs" / "libtensorrt_llm.so").write_text("so")
+    (source / "3rdparty" / "fmha_sm100").mkdir(parents=True)
+    (source / "3rdparty" / "fmha_sm100" / "__init__.py").write_text("")
+
+    destination = tmp_path / "checkout"
+    destination.mkdir()
+    monkeypatch.chdir(destination)
+    return source
+
+
+PACKAGE_DATA = ["libs/*.so"]
+
+
+def _run(extract, source, workspace, link):
+    previous = os.environ.get("TRTLLM_PRECOMPILED_LINK")
+    os.environ["TRTLLM_PRECOMPILED_LINK"] = "1" if link else "0"
+    try:
+        extract(str(source), PACKAGE_DATA, str(workspace))
+    finally:
+        if previous is None:
+            os.environ.pop("TRTLLM_PRECOMPILED_LINK", None)
+        else:
+            os.environ["TRTLLM_PRECOMPILED_LINK"] = previous
+
+
+def test_link_mode_symlinks_the_artifacts(extract_from_precompiled, build_tree, tmp_path):
+    _run(extract_from_precompiled, build_tree, tmp_path, link=True)
+
+    lib = Path("tensorrt_llm/libs/libtensorrt_llm.so")
+    assert lib.is_symlink()
+    assert Path(os.readlink(lib)) == (build_tree / "tensorrt_llm" / "libs" / "libtensorrt_llm.so")
+    fmha = Path("3rdparty/fmha_sm100")
+    assert fmha.is_symlink()
+    assert Path(os.readlink(fmha)) == build_tree / "3rdparty" / "fmha_sm100"
+
+
+def test_link_mode_keeps_an_existing_fmha_symlink(extract_from_precompiled, build_tree, tmp_path):
+    """A checkout already sharing a build tree keeps the link it has."""
+    other = tmp_path / "other-fmha"
+    other.mkdir()
+    Path("3rdparty").mkdir()
+    os.symlink(other, "3rdparty/fmha_sm100")
+
+    _run(extract_from_precompiled, build_tree, tmp_path, link=True)
+
+    assert Path(os.readlink("3rdparty/fmha_sm100")) == other
+
+
+def test_link_mode_replaces_a_stale_artifact(extract_from_precompiled, build_tree, tmp_path):
+    """A real file left by an earlier copy-mode install is not kept."""
+    Path("tensorrt_llm/libs").mkdir(parents=True)
+    Path("tensorrt_llm/libs/libtensorrt_llm.so").write_text("stale")
+
+    _run(extract_from_precompiled, build_tree, tmp_path, link=True)
+
+    assert Path("tensorrt_llm/libs/libtensorrt_llm.so").is_symlink()
+
+
+def test_copy_mode_is_unchanged(extract_from_precompiled, build_tree, tmp_path):
+    _run(extract_from_precompiled, build_tree, tmp_path, link=False)
+
+    lib = Path("tensorrt_llm/libs/libtensorrt_llm.so")
+    assert lib.is_file() and not lib.is_symlink()
+    assert lib.read_text() == "so"
+    fmha = Path("3rdparty/fmha_sm100")
+    assert fmha.is_dir() and not fmha.is_symlink()
+    assert (fmha / "__init__.py").is_file()
+
+
+def test_link_mode_rejects_a_wheel(extract_from_precompiled, build_tree, tmp_path):
+    """There is no build tree to point at, so fail instead of copying."""
+    from setuptools.errors import SetupError
+
+    wheel = tmp_path / "tensorrt_llm-0.0.0.whl"
+    wheel.write_text("")
+
+    with pytest.raises(SetupError, match="TRTLLM_PRECOMPILED_LINK"):
+        _run(extract_from_precompiled, wheel, tmp_path, link=True)
+
+
+# --------------------------------------------------------------------------
+# Build-skew warning
+# --------------------------------------------------------------------------
+def _fake_git(monkeypatch, heads, diff_output="", diff_fails=False):
+    """Stub ``git rev-parse`` / ``git diff`` with canned output."""
+    import subprocess
+
+    class Done:
+        def __init__(self, stdout):
+            self.stdout = stdout
+
+    def run(argv, **kwargs):
+        if "rev-parse" in argv:
+            return Done(heads[argv[argv.index("-C") + 1]] + "\n")
+        assert argv[1] == "diff", argv
+        if diff_fails:
+            raise subprocess.CalledProcessError(1, argv)
+        return Done(diff_output)
+
+    monkeypatch.setattr(subprocess, "run", run)
+
+
+def test_no_warning_when_the_checkouts_match(warn_on_build_skew, monkeypatch, capfd):
+    _fake_git(monkeypatch, {"/src": "a" * 40, ".": "a" * 40})
+
+    warn_on_build_skew("/src")
+
+    assert capfd.readouterr().out == ""
+
+
+def test_no_warning_when_nothing_native_differs(warn_on_build_skew, monkeypatch, capfd):
+    """The two checkouts are meant to differ; only the native inputs matter."""
+    _fake_git(monkeypatch, {"/src": "a" * 40, ".": "b" * 40}, diff_output="")
+
+    warn_on_build_skew("/src")
+
+    assert capfd.readouterr().out == ""
+
+
+def test_warns_when_native_inputs_differ(warn_on_build_skew, monkeypatch, capfd):
+    _fake_git(
+        monkeypatch,
+        {"/src": "a" * 40, ".": "b" * 40},
+        diff_output="cpp/one.cu\ncpp/two.cu\nsetup.py\n3rdparty/x\n",
+    )
+
+    warn_on_build_skew("/src")
+
+    out = capfd.readouterr().out
+    assert "WARNING" in out
+    assert "aaaaaaaaaaaa" in out and "bbbbbbbbbbbb" in out
+    assert "4 files feeding" in out
+    assert "cpp/one.cu" in out and "..." in out
+    assert "ABI skew" in out
+
+
+def test_warns_generically_when_the_diff_cannot_be_taken(warn_on_build_skew, monkeypatch, capfd):
+    """Two unrelated clones do not share an object store."""
+    _fake_git(monkeypatch, {"/src": "a" * 40, ".": "b" * 40}, diff_fails=True)
+
+    warn_on_build_skew("/src")
+
+    out = capfd.readouterr().out
+    assert "WARNING" in out and "could not be inspected" in out
+
+
+def test_skew_check_is_skipped_outside_a_repo(warn_on_build_skew, monkeypatch, capfd):
+    import subprocess
+
+    def run(argv, **kwargs):
+        raise subprocess.CalledProcessError(128, argv)
+
+    monkeypatch.setattr(subprocess, "run", run)
+
+    warn_on_build_skew("/src")
+
+    assert "Cannot check for build skew" in capfd.readouterr().out