diff --git a/docs/source/installation/build-from-source.md b/docs/source/installation/build-from-source.md index 4e656d420499..3178db219a3e 100644 --- a/docs/source/installation/build-from-source.md +++ b/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. diff --git a/setup.py b/setup.py index 205735a4daae..b589c6163213 100644 --- a/setup.py +++ b/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,31 @@ 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}.") + # Linking a checkout onto itself would unlink the real artifacts and + # replace them with symlinks that point to themselves, destroying the + # build tree this mode is meant to share. + if link_artifacts and os.path.realpath( + precompiled_location) == os.path.realpath("."): + raise SetupError( + "TRTLLM_PRECOMPILED_LINK=1 needs a source checkout separate from " + "this one, but TRTLLM_PRECOMPILED_LOCATION resolves to the current " + "directory.") + # 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 +406,19 @@ 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.") + # Drop a stale symlink from a prior link-mode run so + # copy2 writes a real file instead of following the link + # into the shared build tree. + if os.path.islink(dst_file): + os.unlink(dst_file) + shutil.copy2(src_file, dst_file) source_fmha = os.path.join(precompiled_location, "3rdparty", "fmha_sm100") @@ -328,12 +428,30 @@ 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) and os.path.realpath( + dst_fmha) == os.path.realpath(source_fmha): + # Already points at this source; leave the shared link alone. + print(f"Keeping existing fmha_sm100 symlink: {dst_fmha}") + else: + # A stale link (pointing at a different source) or a real + # directory: replace it so fmha_sm100 tracks the same source + # as the other linked artifacts. + if os.path.islink(dst_fmha): + os.unlink(dst_fmha) + elif 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 diff --git a/tests/unittest/others/test_precompiled_link_mode.py b/tests/unittest/others/test_precompiled_link_mode.py new file mode 100644 index 000000000000..dd72ffbb7f2a --- /dev/null +++ b/tests/unittest/others/test_precompiled_link_mode.py @@ -0,0 +1,230 @@ +# 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 + +pytestmark = pytest.mark.cpu_only + +_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 link already pointing at this source is left untouched.""" + Path("3rdparty").mkdir() + os.symlink(build_tree / "3rdparty" / "fmha_sm100", "3rdparty/fmha_sm100") + + _run(extract_from_precompiled, build_tree, tmp_path, link=True) + + assert Path(os.readlink("3rdparty/fmha_sm100")) == build_tree / "3rdparty" / "fmha_sm100" + + +def test_link_mode_relinks_fmha_when_the_source_changed( + extract_from_precompiled, build_tree, tmp_path +): + """A link to a different source is repointed, not kept stale.""" + 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")) == build_tree / "3rdparty" / "fmha_sm100" + + +def test_link_mode_rejects_the_current_checkout(extract_from_precompiled, build_tree, tmp_path): + """Linking a checkout onto itself would destroy its own artifacts.""" + from setuptools.errors import SetupError + + with pytest.raises(SetupError, match="current directory"): + _run(extract_from_precompiled, Path.cwd(), tmp_path, link=True) + + +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