diff --git a/.agents/contributor-skills/build-and-dependency/SKILL.md b/.agents/contributor-skills/build-and-dependency/SKILL.md index 980aa1efeca..2868963efa6 100644 --- a/.agents/contributor-skills/build-and-dependency/SKILL.md +++ b/.agents/contributor-skills/build-and-dependency/SKILL.md @@ -26,10 +26,11 @@ docker buildx build -f docker/Dockerfile \ Skip optional backends to reduce build time: ```bash -# Skip vLLM and SGLang +# Skip vLLM, SGLang, and TRT-LLM docker buildx build -f docker/Dockerfile \ --build-arg SKIP_VLLM_BUILD=1 \ --build-arg SKIP_SGLANG_BUILD=1 \ + --build-arg SKIP_TRTLLM_BUILD=1 \ --tag nemo-rl:latest . ``` diff --git a/.github/workflows/_build_container.yml b/.github/workflows/_build_container.yml index 4b6e527e062..0bd1d93dd57 100644 --- a/.github/workflows/_build_container.yml +++ b/.github/workflows/_build_container.yml @@ -35,6 +35,16 @@ on: default: "" description: Additional Docker build contexts. type: string + trtllm-ccache-tag: + required: false + default: "" + description: Prefix for the TRT-LLM ccache image tag; the target architecture and a hash of the effective base image are appended automatically. Empty disables the ccache. + type: string + trtllm-wheel-cache-tag: + required: false + default: "" + description: Prefix for the TRT-LLM wheel cache image tag; the target architecture and a hash of the effective base image are appended automatically. Empty disables the cache. + type: string dockerfile: required: true description: Path to the Dockerfile. @@ -83,6 +93,142 @@ jobs: - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 + - name: Resolve TRT-LLM cache configuration + id: trtllm_cache + shell: bash + env: + BUILD_ARGS: ${{ inputs.build-args }} + CCACHE_TAG_PREFIX: ${{ inputs.trtllm-ccache-tag }} + WHEEL_TAG_PREFIX: ${{ inputs.trtllm-wheel-cache-tag }} + DOCKERFILE: ${{ inputs.dockerfile }} + PLATFORM: ${{ inputs.platform }} + run: | + set -euo pipefail + + # Match Dockerfile semantics: any non-empty value skips TRT-LLM. + SKIP_VALUE=$(printf '%s\n' "$BUILD_ARGS" | sed -nE 's/^[[:space:]]*SKIP_TRTLLM_BUILD[[:space:]]*=(.*)$/\1/p' | tail -1) + SKIP_VALUE=$(printf '%s' "$SKIP_VALUE" | sed -E 's/^[[:space:]]+//; s/[[:space:]]+$//') + if [[ -n "$SKIP_VALUE" ]]; then + echo "TRT-LLM build is disabled by SKIP_TRTLLM_BUILD" + echo "enabled=false" >> "$GITHUB_OUTPUT" + exit 0 + fi + echo "enabled=true" >> "$GITHUB_OUTPUT" + + # No cache configuration is needed when both image prefixes are empty. + if [[ -z "$CCACHE_TAG_PREFIX" && -z "$WHEEL_TAG_PREFIX" ]]; then + exit 0 + fi + + BASE_IMAGE=$(printf '%s\n' "$BUILD_ARGS" | sed -nE 's/^[[:space:]]*BASE_IMAGE[[:space:]]*=[[:space:]]*(.+)[[:space:]]*$/\1/p' | tail -1) + if [[ -z "$BASE_IMAGE" ]]; then + BASE_IMAGE=$(sed -nE 's/^[[:space:]]*ARG[[:space:]]+BASE_IMAGE[[:space:]]*=[[:space:]]*(.+)[[:space:]]*$/\1/p' "$DOCKERFILE" | tail -1) + fi + if [[ -z "$BASE_IMAGE" ]]; then + echo "Could not resolve BASE_IMAGE from build args or $DOCKERFILE" >&2 + exit 1 + fi + + ARCH="${PLATFORM#*/}" + ARCH="${ARCH%%/*}" + if [[ -z "$ARCH" || "$PLATFORM" != */* || "$PLATFORM" == *,* ]]; then + echo "Expected a single OS/architecture platform, got: $PLATFORM" >&2 + exit 1 + fi + + CACHE_KEY=$(printf '%s' "$BASE_IMAGE" | sha256sum | cut -c1-12) + { + echo "arch=$ARCH" + if [[ -n "$CCACHE_TAG_PREFIX" ]]; then + echo "ccache_tag=${CCACHE_TAG_PREFIX}-${ARCH}-${CACHE_KEY}" + fi + if [[ -n "$WHEEL_TAG_PREFIX" ]]; then + echo "wheel_tag=${WHEEL_TAG_PREFIX}-${ARCH}-${CACHE_KEY}" + fi + } >> "$GITHUB_OUTPUT" + echo "TRT-LLM caches use $BASE_IMAGE on $ARCH (key $CACHE_KEY)" + + - name: Resolve TRT-LLM cache seeds + id: trtllm_cache_seed + if: steps.trtllm_cache.outputs.ccache_tag != '' || steps.trtllm_cache.outputs.wheel_tag != '' + shell: bash + env: + CCACHE_TAG: ${{ steps.trtllm_cache.outputs.ccache_tag }} + WHEEL_TAG: ${{ steps.trtllm_cache.outputs.wheel_tag }} + run: | + set -euo pipefail + + resolve_seed() { + local name="$1" + local tag="$2" + local context="$3" + local description="$4" + [[ -n "$tag" ]] || return 0 + + local ref="$REGISTRY/$IMAGE_NAME:$tag" + if docker buildx imagetools inspect "$ref" >/dev/null 2>&1; then + echo "Using TRT-LLM $description seed $ref" + { + echo "${name}_ref=$ref" + echo "${name}_context=${context}=docker-image://$ref" + } >> "$GITHUB_OUTPUT" + else + echo "No TRT-LLM $description image found at $ref; starting with an empty cache" + fi + } + + resolve_seed ccache "$CCACHE_TAG" trtllm-ccache-seed ccache + resolve_seed wheel "$WHEEL_TAG" trtllm-wheel-cache-seed "wheel cache" + + - name: Configure TRT-LLM cache exports + id: trtllm_cache_export + if: steps.trtllm_cache.outputs.ccache_tag != '' || steps.trtllm_cache.outputs.wheel_tag != '' + shell: bash + env: + CCACHE_SEED_REF: ${{ steps.trtllm_cache_seed.outputs.ccache_ref }} + CCACHE_TAG: ${{ steps.trtllm_cache.outputs.ccache_tag }} + WHEEL_SEED_REF: ${{ steps.trtllm_cache_seed.outputs.wheel_ref }} + WHEEL_TAG: ${{ steps.trtllm_cache.outputs.wheel_tag }} + PLATFORM: ${{ inputs.platform }} + run: | + set -euo pipefail + + configure_export() { + local name="$1" + local seed_ref="$2" + local output_context="$3" + local description="$4" + local max_incremental_layers=16 + local cache_layout=incremental-v1 + local export_mode=full + + if [[ -n "$seed_ref" ]]; then + local image_metadata + image_metadata=$(docker buildx imagetools inspect "$seed_ref" --format "{{len (index .Image \"$PLATFORM\").RootFS.DiffIDs}} {{index (index .Image \"$PLATFORM\").Config.Labels \"com.nvidia.nemo-rl.cache-layout\"}}" 2>/dev/null || true) + if [[ -z "$image_metadata" ]]; then + image_metadata=$(docker buildx imagetools inspect "$seed_ref" --format '{{len .Image.RootFS.DiffIDs}} {{index .Image.Config.Labels "com.nvidia.nemo-rl.cache-layout"}}' 2>/dev/null || true) + fi + + local layer_count existing_layout + read -r layer_count existing_layout <<< "$image_metadata" + if [[ "$existing_layout" == "$cache_layout" && "$layer_count" =~ ^[0-9]+$ ]] && (( layer_count < max_incremental_layers )); then + export_mode=incremental + echo "Appending an incremental $description layer to $seed_ref ($layer_count existing layers)" + echo "${name}_output_context=${output_context}=docker-image://$seed_ref" >> "$GITHUB_OUTPUT" + else + echo "Compacting TRT-LLM $description into a full snapshot (layout: ${existing_layout:-legacy}, layers: ${layer_count:-unknown})" + fi + fi + echo "${name}_mode=$export_mode" >> "$GITHUB_OUTPUT" + } + + if [[ -n "$CCACHE_TAG" ]]; then + configure_export ccache "$CCACHE_SEED_REF" trtllm-ccache-output-base ccache + fi + if [[ -n "$WHEEL_TAG" ]]; then + configure_export wheel "$WHEEL_SEED_REF" trtllm-wheel-cache-output-base "wheel cache" + fi + - name: Get recently merged PR cache refs id: recent_pr_cache_refs uses: actions/github-script@v8 @@ -160,7 +306,10 @@ jobs: push: true context: . platforms: ${{ inputs.platform }} - build-contexts: ${{ inputs.build-contexts }} + build-contexts: | + ${{ inputs.build-contexts }} + ${{ steps.trtllm_cache_seed.outputs.ccache_context }} + ${{ steps.trtllm_cache_seed.outputs.wheel_context }} build-args: ${{ inputs.build-args }} cache-from: | ${{ steps.build_meta.outputs.cache-from }} @@ -170,3 +319,121 @@ jobs: tags: | ${{ steps.build_meta.outputs.tags }} target: ${{ inputs.target }} + + - name: Probe TRT-LLM incremental cache exports + id: trtllm_cache_export_status + if: ${{ always() && !cancelled() && steps.trtllm_cache.outputs.enabled == 'true' && steps.trtllm_cache_export.outcome == 'success' && (steps.trtllm_cache_export.outputs.ccache_mode == 'incremental' || steps.trtllm_cache_export.outputs.wheel_mode == 'incremental') }} + # This is an optional churn optimization. If the probe fails, leave its + # outputs unset so the cache update steps below retain their old, + # fail-open behavior and push. + continue-on-error: true + uses: docker/build-push-action@v5 + with: + file: ${{ inputs.dockerfile }} + push: false + context: . + platforms: ${{ inputs.platform }} + build-contexts: | + ${{ inputs.build-contexts }} + ${{ steps.trtllm_cache_seed.outputs.ccache_context }} + ${{ steps.trtllm_cache_seed.outputs.wheel_context }} + build-args: | + ${{ inputs.build-args }} + CCACHE_EXPORT_NONCE=${{ github.run_id }}-${{ github.run_attempt }} + CCACHE_EXPORT_MODE=${{ steps.trtllm_cache_export.outputs.ccache_mode || 'full' }} + WHEEL_CACHE_EXPORT_NONCE=${{ github.run_id }}-${{ github.run_attempt }} + WHEEL_CACHE_EXPORT_MODE=${{ steps.trtllm_cache_export.outputs.wheel_mode || 'full' }} + cache-from: | + ${{ steps.build_meta.outputs.cache-from }} + ${{ steps.recent_pr_cache_refs.outputs.cache-from }} + no-cache: false + provenance: false + outputs: type=local,dest=${{ runner.temp }}/trtllm-cache-export-status-${{ github.run_id }}-${{ github.run_attempt }} + target: trtllm-cache-export-status + + - name: Check TRT-LLM incremental cache deltas + id: trtllm_cache_delta + if: steps.trtllm_cache_export_status.outcome == 'success' + continue-on-error: true + shell: bash + env: + STATUS_DIR: ${{ runner.temp }}/trtllm-cache-export-status-${{ github.run_id }}-${{ github.run_attempt }} + CCACHE_MODE: ${{ steps.trtllm_cache_export.outputs.ccache_mode }} + WHEEL_MODE: ${{ steps.trtllm_cache_export.outputs.wheel_mode }} + run: | + set -euo pipefail + + resolve_delta() { + local mode="$1" + local status_file="$2" + if [[ "$mode" != "incremental" ]]; then + echo true + return + fi + [[ -f "$status_file" ]] || return 1 + local has_delta + has_delta=$(<"$status_file") + [[ "$has_delta" == "true" || "$has_delta" == "false" ]] || return 1 + echo "$has_delta" + } + + CCACHE_HAS_DELTA=$(resolve_delta "$CCACHE_MODE" "$STATUS_DIR/ccache-has-delta") + WHEEL_HAS_DELTA=$(resolve_delta "$WHEEL_MODE" "$STATUS_DIR/wheel-has-delta") + { + echo "ccache_has_delta=$CCACHE_HAS_DELTA" + echo "wheel_has_delta=$WHEEL_HAS_DELTA" + } >> "$GITHUB_OUTPUT" + echo "TRT-LLM incremental cache deltas: ccache=$CCACHE_HAS_DELTA, wheel=$WHEEL_HAS_DELTA" + + - name: Update TRT-LLM ccache image + # Cache mounts survive a failed RUN in the current Buildx builder. Run + # this step after both successful and failed image builds so completed + # compiler outputs are preserved for the next retry. + if: ${{ always() && !cancelled() && steps.trtllm_cache.outputs.enabled == 'true' && steps.trtllm_cache.outputs.ccache_tag != '' && steps.trtllm_cache_export.outcome == 'success' && steps.trtllm_cache_delta.outputs.ccache_has_delta != 'false' }} + uses: docker/build-push-action@v5 + with: + file: ${{ inputs.dockerfile }} + push: true + context: . + platforms: ${{ inputs.platform }} + build-contexts: | + ${{ inputs.build-contexts }} + ${{ steps.trtllm_cache_seed.outputs.ccache_context }} + ${{ steps.trtllm_cache_export.outputs.ccache_output_context }} + build-args: | + ${{ inputs.build-args }} + CCACHE_EXPORT_NONCE=${{ github.run_id }}-${{ github.run_attempt }} + CCACHE_EXPORT_MODE=${{ steps.trtllm_cache_export.outputs.ccache_mode }} + cache-from: | + ${{ steps.build_meta.outputs.cache-from }} + ${{ steps.recent_pr_cache_refs.outputs.cache-from }} + no-cache: false + provenance: false + tags: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:${{ steps.trtllm_cache.outputs.ccache_tag }} + target: trtllm-ccache + + - name: Update TRT-LLM wheel cache image + # Persist the custom wheel even when the release build fails, so a + # retry can reuse completed compilation work. + if: ${{ always() && !cancelled() && steps.trtllm_cache.outputs.enabled == 'true' && steps.trtllm_cache.outputs.wheel_tag != '' && steps.trtllm_cache_export.outcome == 'success' && steps.trtllm_cache_delta.outputs.wheel_has_delta != 'false' }} + uses: docker/build-push-action@v5 + with: + file: ${{ inputs.dockerfile }} + push: true + context: . + platforms: ${{ inputs.platform }} + build-contexts: | + ${{ inputs.build-contexts }} + ${{ steps.trtllm_cache_seed.outputs.wheel_context }} + ${{ steps.trtllm_cache_export.outputs.wheel_output_context }} + build-args: | + ${{ inputs.build-args }} + WHEEL_CACHE_EXPORT_NONCE=${{ github.run_id }}-${{ github.run_attempt }} + WHEEL_CACHE_EXPORT_MODE=${{ steps.trtllm_cache_export.outputs.wheel_mode }} + cache-from: | + ${{ steps.build_meta.outputs.cache-from }} + ${{ steps.recent_pr_cache_refs.outputs.cache-from }} + no-cache: false + provenance: false + tags: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:${{ steps.trtllm_cache.outputs.wheel_tag }} + target: trtllm-wheel-cache diff --git a/.github/workflows/cicd-main.yml b/.github/workflows/cicd-main.yml index 63ead502048..e1ec80a0c81 100644 --- a/.github/workflows/cicd-main.yml +++ b/.github/workflows/cicd-main.yml @@ -359,8 +359,11 @@ jobs: build-contexts: | nemo-rl=. ${{ vars.UV_BUILD_CACHE == 'enabled' && format('uv-cache-seed=docker-image://{0}/{1}:uv-cache', needs.org-member-pre-flight.outputs.registry, vars.CI_CONTAINER_NAME) || '' }} + trtllm-ccache-tag: ${{ vars.TRTLLM_BUILD_CACHE == 'enabled' && 'trtllm-ccache' || '' }} + trtllm-wheel-cache-tag: ${{ vars.TRTLLM_BUILD_CACHE == 'enabled' && 'trtllm-wheel-cache' || '' }} build-args: | MAX_JOBS=4 + TRTLLM_BUILD_JOBS=24 NEMO_RL_COMMIT=${{ needs.pre-flight.outputs.test_sha }} build-container-gb200: @@ -398,8 +401,11 @@ jobs: build-contexts: | nemo-rl=. ${{ vars.UV_BUILD_CACHE == 'enabled' && format('uv-cache-seed=docker-image://{0}/{1}:uv-cache', needs.gb200-config.outputs.registry, vars.CI_CONTAINER_NAME) || '' }} + trtllm-ccache-tag: ${{ vars.TRTLLM_BUILD_CACHE == 'enabled' && 'trtllm-ccache' || '' }} + trtllm-wheel-cache-tag: ${{ vars.TRTLLM_BUILD_CACHE == 'enabled' && 'trtllm-wheel-cache' || '' }} build-args: | MAX_JOBS=4 + TRTLLM_BUILD_JOBS=8 NEMO_RL_COMMIT=${{ needs.pre-flight.outputs.test_sha }} update-uv-cache: @@ -522,6 +528,8 @@ jobs: runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }} - script: L0_Unit_Tests_Sglang runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }} + - script: L0_Unit_Tests_Trtllm + runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }} - script: L0_Unit_Tests_Mcore runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }} - script: L0_Unit_Tests_Mcore_Policy_1 @@ -582,6 +590,9 @@ jobs: - name: Checkout uses: actions/checkout@v6 - name: main + # Lfast reuses the main image. Until TRT-LLM lands on main, running this + # shard would make prefetch_venvs compile the cp313 wheel from source. + if: ${{ needs.pre-flight.outputs.test_level != 'Lfast' || matrix.script != 'L0_Unit_Tests_Trtllm' }} uses: ./.github/actions/test-template env: HF_TOKEN: ${{ secrets.HF_TOKEN }} @@ -706,6 +717,8 @@ jobs: runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }} - script: L1_Functional_Tests_SGLang runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }} + - script: L1_Functional_Tests_Trtllm + runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }} - script: L1_Functional_Tests_Gym runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }} - script: L1_Functional_Tests_GRPO_1 diff --git a/3rdparty/TensorRT-LLM-workspace/_backend.py b/3rdparty/TensorRT-LLM-workspace/_backend.py new file mode 100644 index 00000000000..fe217ef0e16 --- /dev/null +++ b/3rdparty/TensorRT-LLM-workspace/_backend.py @@ -0,0 +1,258 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# 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. + +"""Custom PEP 517 build backend for TensorRT-LLM. + +Two hooks implement the two-phase build: + + prepare_metadata_for_build_wheel + Returns static dist-info fast — no GPU / compilation required. + Called by `uv lock` and by `uv sync` before deciding whether to build. + + build_wheel + Invokes tools/build-custom-trtllm.sh (≈60 min on GB200). + Called by `uv sync --extra trtllm` when the wheel is not yet cached. + +The package is declared as no-build-isolation-package in the root +pyproject.toml so this backend runs inside the main venv and has access +to torch, ninja, cmake, and the CUDA toolkit. + +Build coordinates (git URL / ref) are read from the workspace's +``[tool.trtllm]`` table in pyproject.toml. +""" + +from __future__ import annotations + +import hashlib +import os +import platform +import shutil +import subprocess +import sys +import tomllib +from pathlib import Path + +# --------------------------------------------------------------------------- +# Constants — must stay in sync with 3rdparty/TensorRT-LLM-workspace/pyproject.toml +# --------------------------------------------------------------------------- +_HERE = Path(__file__).parent.resolve() +_PYPROJECT = _HERE / "pyproject.toml" + +with _PYPROJECT.open("rb") as _f: + _META = tomllib.load(_f) + +VERSION: str = _META["project"]["version"] +NAME: str = _META["project"]["name"].replace("-", "_") # tensorrt_llm +DIST_NAME: str = _META["project"]["name"] # tensorrt-llm +REQUIRES: list[str] = _META["project"].get("dependencies", []) + +# Fork URL + commit ref to build. The [tool.trtllm] table in this same +# pyproject.toml is the sole source of truth. This backend reads it here, folds +# it into the wheel cache key, and passes it to tools/build-custom-trtllm.sh as +# argv (the script takes no defaults). The CI ccache image is scoped separately +# by base image and architecture so it can span ref and uv-version changes. +# There is no env-var override — to build a different fork/ref, edit +# [tool.trtllm]. +_TRTLLM: dict[str, str] = _META["tool"]["trtllm"] +TRTLLM_URL: str = _TRTLLM["url"] +TRTLLM_REF: str = _TRTLLM["ref"] + + +def _wheel_platform_tag() -> str: + """Return the real wheel tag for the current interpreter, e.g. cp313-cp313-linux_aarch64. + + Used only for the cache key — NOT for prepare_metadata_for_build_wheel, which + must report py3-none-any so that uv lock succeeds on both x86_64 and aarch64. + """ + py = f"cp{sys.version_info.major}{sys.version_info.minor}" + machine = platform.machine() # aarch64 | x86_64 + return f"{py}-{py}-linux_{machine}" + + +# py3-none-any is intentional: prepare_metadata_for_build_wheel is called by +# uv lock, which resolves for both x86_64 and aarch64. A platform-specific tag +# here would make tensorrt-llm appear incompatible with one of the two arches +# and break the lock. The real platform tag is used only inside _wheel_cache_dir. +_METADATA_WHEEL_TAG = "py3-none-any" + + +# Default SM arch list passed to build_wheel.py's -a flag. MUST stay in sync +# with tools/build-custom-trtllm.sh, which reads BUILD_CUSTOM_TRTLLM_ARCH and +# falls back to this same default. Folded into the wheel cache key below so +# editing the arch list forces a rebuild instead of reusing a stale wheel. +_DEFAULT_ARCH = "90-real;100-real" + + +def _build_input_tag(arch: str) -> str: + """Build-affecting inputs (beyond url/ref/version/platform) for the cache key. + + The compiled wheel depends on the SM arch list and the torch/CUDA toolchain + it links against, so a change to any of these — without a git_ref bump — + would otherwise silently reuse a stale cached wheel. torch is imported + lazily so prepare_metadata_for_build_wheel (called under ``uv lock`` without + torch) never triggers it. + """ + # Import lazily because metadata-only hooks do not need this heavy dependency. + import torch # noqa: PLC0415 + + toolchain = f"torch{torch.__version__},cuda{torch.version.cuda}" + return f"arch={arch}|{toolchain}" + + +def _wheel_cache_dir(base: str, git_url: str, git_ref: str, build_inputs: str) -> Path: + """Return a per-(url, ref, version, platform, build-inputs) cache subdir. + + Using a content-addressed subdir means different commits never collide, + and a stale wheel from a previous ref is never accidentally reused. + The cache key uses the real platform tag (not py3-none-any) so aarch64 and + x86_64 wheels built in separate Docker runs never overwrite each other, and + ``build_inputs`` (arch list + toolchain) so editing a build-affecting input + without bumping git_ref still forces a rebuild. + """ + key = hashlib.sha256( + f"{git_url}|{git_ref}|{VERSION}|{_wheel_platform_tag()}|{build_inputs}".encode() + ).hexdigest()[:16] + return Path(base) / key + + +# --------------------------------------------------------------------------- +# PEP 517 hooks +# --------------------------------------------------------------------------- + + +def get_requires_for_build_wheel(config_settings=None): + """No isolated-build requirements; deps come from the main venv (no-build-isolation).""" + return [] + + +def prepare_metadata_for_build_wheel(metadata_directory, config_settings=None): + """Write minimal .dist-info without compiling anything. + + This is the fast path used by ``uv lock`` and ``uv sync``'s preflight + metadata check. It must not require CUDA or take significant time. + """ + dist_info_name = f"{NAME}-{VERSION}.dist-info" + dist_info = Path(metadata_directory) / dist_info_name + dist_info.mkdir(parents=True, exist_ok=True) + + requires_lines = "\n".join(f"Requires-Dist: {r}" for r in REQUIRES) + (dist_info / "METADATA").write_text( + f"Metadata-Version: 2.1\n" + f"Name: {DIST_NAME}\n" + f"Version: {VERSION}\n" + f"{requires_lines}\n", + encoding="utf-8", + ) + (dist_info / "WHEEL").write_text( + "Wheel-Version: 1.0\n" + f"Generator: TensorRT-LLM-workspace-backend\n" + "Root-Is-Purelib: false\n" + f"Tag: {_METADATA_WHEEL_TAG}\n", + encoding="utf-8", + ) + return dist_info_name + + +def build_wheel(wheel_directory, config_settings=None, metadata_directory=None): + """Build the real TRT-LLM wheel by running tools/build-custom-trtllm.sh. + + The script compiles TensorRT-LLM (≈60 min) and copies the resulting + ``tensorrt_llm-*.whl`` into *wheel_directory*. + """ + repo_root = (_HERE / "../..").resolve() + script = repo_root / "tools" / "build-custom-trtllm.sh" + if not script.exists(): + raise FileNotFoundError(f"Build script not found: {script}") + + env = os.environ.copy() + git_url = TRTLLM_URL + git_ref = TRTLLM_REF + # NOTE: when bumping the ref (in the [tool.trtllm] table of this pyproject.toml), + # re-sync the Requires-Dist list in [project].dependencies of the same file. + # uv resolves this package's runtime deps from that hand-curated static list + # (surfaced by prepare_metadata_for_build_wheel), NOT from the built wheel's + # METADATA — so a ref bump can silently drop or miss real deps. Regenerate: + # git -C checkout && cat requirements.txt + # then reconcile [project].dependencies against it (dropping build-only pins). + + # SM arch list — single source of truth for both the cache key (below) and + # the build script (which reads BUILD_CUSTOM_TRTLLM_ARCH). Exporting it into + # env guarantees the script compiles exactly what the cache key was keyed on. + arch = env.get("BUILD_CUSTOM_TRTLLM_ARCH", _DEFAULT_ARCH) + env["BUILD_CUSTOM_TRTLLM_ARCH"] = arch + + # Our own cache keyed by (git_url, git_ref, version, platform_tag, + # build_inputs=arch+toolchain). + # uv's built-in build cache misses across venvs for no-build-isolation + # packages because its cache key incorporates the build-environment hash. + # We bypass that by always building into TRTLLM_WHEEL_CACHE_DIR (a stable + # path that persists across all venv sync calls in the same Docker build), + # then copying the result into wheel_directory for uv to consume. + cache_base = env.get("TRTLLM_WHEEL_CACHE_DIR", "/opt/trtllm_wheels") + cache_dir = _wheel_cache_dir(cache_base, git_url, git_ref, _build_input_tag(arch)) + + wheel = max(cache_dir.glob("tensorrt_llm-*.whl"), default=None) + if wheel is None: + if env.get("TRTLLM_REQUIRE_CACHED_WHEEL") == "1": + raise RuntimeError( + "TRT-LLM cached wheel is required but was not found at " + f"{cache_dir}. Refusing to compile TRT-LLM because " + "TRTLLM_REQUIRE_CACHED_WHEEL=1." + ) + + # Build directly into cache_dir so later venv syncs can reuse the wheel. + cache_dir.mkdir(parents=True, exist_ok=True) + env["WHEEL_OUTPUT_DIR"] = str(cache_dir) + venv_bin = str(Path(sys.executable).parent) + env["PATH"] = f"{venv_bin}:{env.get('PATH', os.defpath)}" + subprocess.run( + ["bash", str(script), git_url, git_ref], + check=True, + env=env, + cwd=str(repo_root), + ) + wheel = max(cache_dir.glob("tensorrt_llm-*.whl"), default=None) + if wheel is None: + raise RuntimeError( + f"No tensorrt_llm-*.whl found in {cache_dir} after build. " + "Check the build-custom-trtllm.sh output above for errors." + ) + print(f"[trtllm-backend] Wheel built and cached to: {cache_dir}", flush=True) + else: + print(f"[trtllm-backend] Cache hit — reusing wheel: {wheel}", flush=True) + + # BuildKit cache mounts are not included in the resulting image. Allow the + # release build to mirror only this wheel's content-addressed cache entry + # into a persistent image path for later `uv run --extra trtllm` calls. + mirror_base = env.get("TRTLLM_WHEEL_CACHE_MIRROR_DIR") + if mirror_base: + mirror_dir = _wheel_cache_dir( + mirror_base, git_url, git_ref, _build_input_tag(arch) + ) + mirror_dir.mkdir(parents=True, exist_ok=True) + mirror = mirror_dir / wheel.name + if wheel.resolve() != mirror.resolve(): + shutil.copy2(wheel, mirror) + print(f"[trtllm-backend] Mirrored cached wheel to: {mirror}", flush=True) + + destination = Path(wheel_directory) / wheel.name + shutil.copy2(wheel, destination) + return destination.name + + +def build_sdist(sdist_directory, config_settings=None): + raise NotImplementedError( + "TRT-LLM workspace wrapper does not support sdist builds. " + "Use `uv sync --extra trtllm` to build the wheel." + ) diff --git a/3rdparty/TensorRT-LLM-workspace/pyproject.toml b/3rdparty/TensorRT-LLM-workspace/pyproject.toml new file mode 100644 index 00000000000..5751291addc --- /dev/null +++ b/3rdparty/TensorRT-LLM-workspace/pyproject.toml @@ -0,0 +1,121 @@ +# Workspace wrapper for TensorRT-LLM. +# +# This package exposes tensorrt-llm as a path source so that: +# - `uv lock` can resolve the dep graph without building the wheel +# (prepare_metadata_for_build_wheel returns static metadata fast). +# - `uv sync --extra trtllm` triggers the actual ~60-min wheel build +# (build_wheel in _backend.py calls tools/build-custom-trtllm.sh). +# +# Declared as no-build-isolation-package in the root pyproject.toml so the +# build runs in the main venv context and has access to torch / CUDA headers. +# +# Sourced from requirements.txt on github.com/shuyixiong/TensorRT-LLM@nemorl. +# Intentional omissions from requirements.txt: +# - datasets==3.1.0 → conflicts with nemorl base datasets>=4.0.0 +# - setuptools<80 → conflicts with nemorl override setuptools>=80.10.2 +# - llguidance==0.7.29 → conflicts with nemorl override llguidance>=1.3.0,<1.4.0 +# - triton==3.6.0 → nemorl base manages triton via pytorch-cu130 index +# - evaluate → pulls in conflicting datasets==3.1.0 transitively +# - build / ninja / meson / patchelf → build-time only, not runtime deps +# - cuda-core → included in cuda-python>=13 +# - graphviz / polygraphy / mcp → optional dev/test tools + +[project] +name = "tensorrt-llm" +version = "1.3.0rc21" +requires-python = ">=3.13" +dependencies = [ + "tensorrt~=10.16.1", + "mpi4py", + "numpy>=2.0.0,<2.4", + "torch>=2.11.0", + "torchvision", + "nvidia-modelopt>=0.37.0", + "nvidia-nccl-cu13>=2.28.9,<=2.29.7", + "nvidia-cuda-nvrtc", + "nvidia-cutlass-dsl[cu13]==4.5.0", + "nvidia-ml-py>=13", + "cuda-python>=13", + "flash-attn-4==4.0.0b11", + "flashinfer-python==0.6.12", + "torchao>=0.14.1,<0.16.0", + "apache-tvm-ffi==0.1.6", + "torch-c-dlpack-ext==0.1.3", + "xdsl>=0.59.0", + "cuda-tile>=1.0.1", + "nvidia-cuda-tileiras>=13.1,<13.2", + "transformers==5.5.4", + "accelerate>=1.7.0", + "sentencepiece>=0.1.99", + "mistral-common>=1.10.0", + "tiktoken", + "diffusers>=0.37.1", + "peft>=0.18.1", + "optimum", + "ftfy", + "xgrammar>=0.1.32", + "lark", + "partial_json_parser", + "onnx>=1.21.0", + "onnx_graphsurgeon>=0.5.2", + "fastapi>=0.120.1", + "starlette>=0.49.1", + "uvicorn", + "pydantic>=2.9.1", + "pydantic-settings[yaml]", + "python-multipart", + "prometheus_client", + "prometheus-fastapi-instrumentator>=8.0.2", + "openai-harmony==0.0.4", + "etcd-sdk-python==0.0.7", + "smg-grpc-proto>=0.4.2", + "cache-dit>=1.3.5", + "pillow", + "opencv-python-headless", + "soundfile", + "librosa", + "h5py==3.12.1", + "pandas", + "numexpr", + "mpmath>=1.3.0", + "colored", + "psutil", + "omegaconf", + "pyzmq", + "einops", + "aenum", + "strenum", + "llguidance>=1.3.0,<1.4.0", + "backoff", + "jsonschema", + "click>=8.3.1", + "click_option_group", + "ordered-set", + "msgpack", + "blake3", + "llist", + "pulp", + "blobfile", + "nvtx", + "plotly", + "matplotlib", + "openai", +] + +[build-system] +# No extra build requirements here — build deps (torch, ninja, cmake, etc.) +# come from the main venv via no-build-isolation-package. +requires = [] +build-backend = "_backend" +backend-path = ["."] + +[tool.trtllm] +# Single source of truth for which TensorRT-LLM fork/commit the custom wheel +# builds. _backend.py reads these, folds them into the wheel cache key, and +# passes them to tools/build-custom-trtllm.sh as its required argv (the script +# itself takes no defaults). The wheel cache is ref-scoped, while the CI ccache +# image is scoped by base image and architecture so valid compiler objects can +# be reused across ref and uv-version changes. When bumping `ref`, also re-sync +# [project].dependencies against the new requirements.txt. +url = "https://github.com/NVIDIA/TensorRT-LLM.git" +ref = "bf2ef86f9a2652132b11773d4041e292c553c142" # pragma: allowlist secret diff --git a/docker/Dockerfile b/docker/Dockerfile index 7f85701f20c..35dabf3c5a0 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -21,6 +21,7 @@ # Optional build args to skip vLLM or SGLang dependencies: # --build-arg SKIP_VLLM_BUILD=1 # Skip vLLM dependencies # --build-arg SKIP_SGLANG_BUILD=1 # Skip SGLang dependencies +# --build-arg SKIP_TRTLLM_BUILD=1 # Skip TRT-LLM dependencies ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.03-cuda13.2-devel-ubuntu24.04 FROM scratch AS nemo-rl @@ -30,6 +31,20 @@ ADD --keep-git-dir=true https://github.com/NVIDIA-NeMo/RL.git#${NRL_GIT_REF} / # Empty default; CI overrides with --build-context to inject pre-compiled wheels FROM scratch AS uv-cache-seed +# Empty default; CI overrides with --build-context to restore custom TRT-LLM +# wheels across builds. Normal downloads use the shared uv cache. +FROM scratch AS trtllm-wheel-cache-seed + +# Empty default; CI overrides with --build-context to restore the TRT-LLM +# compiler cache from a registry image. +FROM scratch AS trtllm-ccache-seed + +# Empty defaults for incremental cache export. CI overrides these contexts +# with the current cache images when appending a delta layer, and leaves them +# empty when periodically compacting the cache into a full snapshot. +FROM scratch AS trtllm-wheel-cache-output-base +FROM scratch AS trtllm-ccache-output-base + # Default custom-setup stage: installs apptainer. # Override with: --build-context custom-setup= FROM scratch AS custom-setup @@ -55,6 +70,7 @@ apt-get install -y --no-install-recommends \ wget \ less \ vim \ + ccache \ # Nsight apt install -y --no-install-recommends gnupg @@ -79,7 +95,7 @@ EOF # CMake (for CUDA extension builds: transformer-engine, mamba-ssm, etc.) RUN GITHUB_ARTIFACTORY=github.com \ - && CMAKE_VERSION=3.31.1 \ + && CMAKE_VERSION=4.0.3 \ && ARCH=$(uname -m) \ && CMAKE_INSTALLER="cmake-${CMAKE_VERSION}-linux-${ARCH}" \ && curl --retry 3 --retry-delay 2 -fsSL -o "${CMAKE_INSTALLER}.tar.gz" \ @@ -124,11 +140,16 @@ ARG BUILD_CUSTOM_VLLM_PRECOMPILED_WHEEL_LOCATION ARG BUILD_CUSTOM_FLASHINFER ARG BUILD_CUSTOM_FLASHINFER_URL ARG BUILD_CUSTOM_FLASHINFER_REF -# Skip building vLLM or SGLang dependencies (set to any non-empty value to skip) + +# Skip building vLLM, SGLang, or TRT-LLM dependencies (set to any non-empty value to skip) ARG SKIP_VLLM_BUILD ARG SKIP_SGLANG_BUILD +ARG SKIP_TRTLLM_BUILD ARG BASE_IMAGE ARG UV_VERSION +ARG TARGETARCH +ARG TRTLLM_CCACHE_MAXSIZE=20G +ARG TRTLLM_BUILD_JOBS=24 ENV UV_PROJECT_ENVIRONMENT=/opt/nemo_rl_venv ENV UV_LINK_MODE=copy @@ -143,6 +164,17 @@ ENV CUDA_HOME=/usr/local/cuda ## builds that include (e.g., deep_gemm) can find them. ENV CPLUS_INCLUDE_PATH=/usr/local/cuda/include/cccl +# Point TE at the pip-installed cuDNN instead of the system one. +# The container's system cuDNN may differ from what pip installed; TE prioritizes +# system libraries by default, causing a version mismatch crash at runtime. +# CUDNN_HOME makes TE's Python code find the pip version first; LD_LIBRARY_PATH +# makes the dynamic linker resolve cuDNN sub-libraries from pip when loading +# libtransformer_engine.so. +# Verify with: python -c "import transformer_engine.pytorch as te; print(te.get_cudnn_version())" +ENV CUDNN_HOME=/opt/nemo_rl_venv/lib/python3.13/site-packages/nvidia/cudnn +ENV CUDNN_PATH=/opt/nemo_rl_venv/lib/python3.13/site-packages/nvidia/cudnn +ENV LD_LIBRARY_PATH="/opt/nemo_rl_venv/lib/python3.13/site-packages/nvidia/cudnn/lib:${LD_LIBRARY_PATH}" + # First copy only the dependency files COPY --from=nemo-rl pyproject.toml uv.lock ./ # Copy in the top level __init__.py/package_info.py since build-custom-vllm.sh needs the nemo_rl package to exist. @@ -150,7 +182,10 @@ COPY --from=nemo-rl nemo_rl/__init__.py nemo_rl/package_info.py ./nemo_rl/ COPY --from=nemo-rl tools/build-custom-vllm.sh ./tools/build-custom-vllm.sh COPY --from=nemo-rl tools/build-custom-flashinfer.sh ./tools/build-custom-flashinfer.sh COPY --from=nemo-rl --link research/ ./research/ -COPY --from=nemo-rl --link 3rdparty/ ./3rdparty/ +COPY --from=nemo-rl --link \ + --exclude=TensorRT-LLM-workspace \ + --exclude=TensorRT-LLM-workspace/** \ + 3rdparty/ ./3rdparty/ RUN --mount=type=ssh \ --mount=type=bind,from=uv-cache-seed,source=.,target=/tmp/uv-cache-seed <<"EOF" bash -exu @@ -173,24 +208,26 @@ fi if [[ -n "${BUILD_CUSTOM_FLASHINFER:-}" ]]; then bash tools/build-custom-flashinfer.sh ${BUILD_CUSTOM_FLASHINFER_URL:-} ${BUILD_CUSTOM_FLASHINFER_REF:-} fi -# uv sync has a more reliable resolver than simple uv pip install which can fail +# The TRT-LLM workspace is copied later. --frozen consumes the committed lock +# without evaluating that absent local source; the late layer runs uv lock +# --check with the complete workspace before starting the TRT-LLM build. # Sync each training + inference backend one at a time (since they may conflict) # to warm the uv cache, then at the end just sync the default dependencies. # Do everything in one layer to prevent large layers. # The venv is symlinked to avoid bloating the layer size -UV_LINK_MODE=hardlink uv sync --locked --no-install-project +UV_LINK_MODE=hardlink uv sync --frozen --no-install-project if [[ -z "${SKIP_VLLM_BUILD:-}" ]]; then - UV_LINK_MODE=hardlink uv sync --locked --extra vllm --no-install-project + UV_LINK_MODE=hardlink uv sync --frozen --extra vllm --no-install-project fi if [[ -z "${SKIP_SGLANG_BUILD:-}" ]]; then - UV_LINK_MODE=hardlink uv sync --locked --extra sglang --no-install-project + UV_LINK_MODE=hardlink uv sync --frozen --extra sglang --no-install-project fi -uv sync --link-mode symlink --locked --extra mcore --no-install-project -uv sync --link-mode symlink --locked --extra automodel --no-install-project -uv sync --link-mode symlink --locked --extra modelopt --no-install-project -uv sync --link-mode symlink --locked --all-groups --no-install-project +uv sync --link-mode symlink --frozen --extra mcore --no-install-project +uv sync --link-mode symlink --frozen --extra automodel --no-install-project +uv sync --link-mode symlink --frozen --extra modelopt --no-install-project +uv sync --link-mode symlink --frozen --all-groups --no-install-project # Remove the aiohttp in this uv cache dir to fully address CVE GHSA-mqqc-3gqh-h2x8 # The ray install will include the older aiohttp version in its cache @@ -199,17 +236,80 @@ find /root/.cache/uv -type d -path "*ray/_private/runtime_env/agent/thirdparty_f echo "${CACHE_KEY}" > /root/.cache/uv/.cache-key EOF +# Keep the frequently changing TRT-LLM build inputs after the other backend +# warmups so TRT-only changes can reuse expensive wheels such as Transformer +# Engine from the preceding layer. +COPY --from=nemo-rl tools/build-custom-trtllm.sh ./tools/build-custom-trtllm.sh +COPY --from=nemo-rl --link 3rdparty/TensorRT-LLM-workspace/ ./3rdparty/TensorRT-LLM-workspace/ + +RUN --mount=type=ssh \ + --mount=type=bind,from=trtllm-ccache-seed,source=.,target=/tmp/trtllm-ccache-seed \ + --mount=type=cache,id=trtllm-wheel-cache-${TARGETARCH}-${BASE_IMAGE},from=trtllm-wheel-cache-seed,source=.,target=/root/.cache/trtllm-wheels,sharing=locked \ + --mount=type=cache,id=trtllm-ccache-${TARGETARCH}-${BASE_IMAGE},target=/root/.cache/ccache,sharing=locked <<"EOF" bash -exu +# Lock/pyproject consistency is verified LATE here, for the whole image. The +# dependency syncs above run --frozen (read-only on uv.lock: they neither update +# nor validate it), so this uv lock --check is the ONLY place a stale lock is +# caught. It can't move earlier: it needs the full workspace, but the tensorrt-llm +# path-source member isn't copied until the COPY just above (kept late so a +# TRT-LLM ref bump doesn't bust the cached TE/backend build). +# WARNING: removing this check (or this layer) without also reverting --frozen +# above would let an out-of-date lock ship silently. Keep it before the SKIP exit. +uv lock --check + +if [[ -n "${SKIP_TRTLLM_BUILD:-}" ]]; then + echo "SKIP_TRTLLM_BUILD is set; skipping the isolated TRT-LLM build" + exit 0 +fi + +# These settings are build-only and intentionally do not persist in the final +# NeMo-RL image. COMPILERCHECK=content is robust to compiler timestamp changes +# between ephemeral CI runners. +export CCACHE_DIR=/root/.cache/ccache +export CCACHE_COMPILERCHECK=content +export CCACHE_COMPRESS=true +export CCACHE_MAXSIZE="${TRTLLM_CCACHE_MAXSIZE}" + +# Cache entries are content-addressed, so keep newer builder-local entries when +# merging the registry seed. +rsync -a --ignore-existing /tmp/trtllm-ccache-seed/ "${CCACHE_DIR}/" + +# uv uses the shared image cache. Only the custom wheel needs a cache mount +# because uv can miss no-build-isolation wheels across build environments. +echo "TRT-LLM cache state before build:" +du -sh /root/.cache/uv /root/.cache/trtllm-wheels + +TRTLLM_SYNC_LOG=$(mktemp /tmp/trtllm-sync.XXXXXX.log) +set +e +UV_CACHE_DIR=/root/.cache/uv \ + TRTLLM_WHEEL_CACHE_DIR=/root/.cache/trtllm-wheels \ + uv sync --verbose --link-mode symlink --locked --extra trtllm --no-install-project \ + 2>&1 | tee "$TRTLLM_SYNC_LOG" \ + | awk '/\[TRTLLM_CCACHE\]|Ninja progress:/ { print; fflush() }' +TRTLLM_SYNC_STATUS=$? +set -e +if [[ "${TRTLLM_SYNC_STATUS}" -ne 0 ]]; then + echo "[ERROR] TRT-LLM uv sync failed with exit code ${TRTLLM_SYNC_STATUS}." >&2 + echo "[ERROR] Last 200 lines of captured uv output:" >&2 + tail -n 200 "$TRTLLM_SYNC_LOG" >&2 + rm -f "$TRTLLM_SYNC_LOG" + exit "${TRTLLM_SYNC_STATUS}" +fi +rm -f "$TRTLLM_SYNC_LOG" + +echo "TRT-LLM cache state after build:" +du -sh /root/.cache/uv /root/.cache/trtllm-wheels + +# Restore the intended default environment in the same layer so the transient +# TRT-LLM installation does not add a large intermediate venv layer. +uv sync --link-mode symlink --locked --all-groups --no-install-project + +# The final sync can repopulate the shared uv cache, so repeat the security +# cleanup performed in the preceding dependency layer. +find /root/.cache/uv -type d -path "*ray/_private/runtime_env/agent/thirdparty_files/aiohttp*" -exec rm -rf {} + +EOF + ENV PATH="/opt/nemo_rl_venv/bin:$PATH" ENV NEMO_RL_VENV_DIR=/opt/ray_venvs -# Point TE at the pip-installed cuDNN (nvidia-cudnn-cu12) instead of the system one. -# The container's system cuDNN (/lib/x86_64-linux-gnu/libcudnn*.so.9) may be a different -# version than what pip installed. TE prioritizes system libraries by default, causing a -# version mismatch crash (e.g. system 9.10.1 vs pip 9.19.0). CUDNN_HOME makes TE's Python -# code find the pip version first, and LD_LIBRARY_PATH makes the dynamic linker resolve -# cuDNN sub-libraries from pip when loading libtransformer_engine.so. -ENV CUDNN_HOME=/opt/nemo_rl_venv/lib/python3.13/site-packages/nvidia/cudnn -ENV LD_LIBRARY_PATH="/opt/nemo_rl_venv/lib/python3.13/site-packages/nvidia/cudnn/lib:${LD_LIBRARY_PATH}" -# Verify with: python -c "import transformer_engine.pytorch as te; print(te.get_cudnn_version())" # Custom setup layer (override with: --build-context custom-setup= --build-arg CUSTOM_SETUP_FNAME=