fix: TRTLLM container -copy NVIDIA torch lib into uv venv location

[dmitry-tokarev-nv]

Overview:

Details:

Where should the reviewer start?

Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)

• closes GitHub issue: #xxx

Summary by CodeRabbit

• New Features

  • None

• Refactor

  • Reworked TensorRT-LLM installation flow with architecture-aware dependencies and streamlined wheel sourcing.
  • Aligned runtime to use a dedicated virtual environment for PyTorch-related packages.

• Chores

  • Tightened dependency installs using no-cache and cleanup to avoid conflicts.
  • Updated benchmarks installation for consistency.
  • Added temporary diagnostic package listings to aid build verification.

• Bug Fixes

  • Reduced environment inconsistencies by relocating Python packages into the runtime virtual environment.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[dmitry-tokarev-nv]

/ok to test 918326d

[coderabbitai]

Walkthrough

Refactors container/Dockerfile.trtllm to rebase Python packages into a venv, adjust TensorRT-LLM installation with custom indices and conditional Triton, add uninstall steps for conflicting packages, install local wheelhouse artifacts, and insert uv pip list checkpoints. Runtime stage mirrors venv packaging and relocations.

Changes

Cohort / File(s)

Summary of Changes

TRT-LLM container build flow container/Dockerfile.trtllm

Tighten installs with --no-cache; uninstall networkx/packaging/torch/triton pre-NVIDIA stages; relocate PyTorch ecosystem from system dist-packages to venv site-packages; copy libs to /usr/local/lib; add multiple uv pip list checkpoints; refactor TRT-LLM install using TENSORRTLLM_INDEX_URL and wheel var; install cuda-python>=12,<13; conditionally install triton==3.3.1 on amd64; install local wheelhouse (ai_dynamo*, nixl*); update runtime stage to rely on venv packages and add checkpoints; switch benchmarks install to --no-cache.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  participant Builder as Docker Build
  participant SysPy as System Python (dist-packages)
  participant Venv as /opt/dynamo/venv (site-packages)
  participant Index as TRT-LLM Index / Extra Wheel
  participant Wheelhouse as Local Wheelhouse
  participant Runtime as Runtime Image

  Builder->>SysPy: Install base deps (no-cache)
  Builder->>SysPy: Uninstall networkx, packaging, torch, triton
  note right of Builder: uv pip list (debug)

  Builder->>Venv: Copy torch/torchvision/triton/... from dist-packages
  Builder->>/usr/local/lib: Copy lib*
  note over Venv: Consolidate runtime to venv

  Builder->>Index: Install cuda-python (>=12,<13)
  Builder->>Index: Install TRT-LLM via extra-index/wheel var
  alt amd64
    Builder->>Index: Install triton==3.3.1
  end
  note right of Builder: uv pip list (debug)

  Builder->>Wheelhouse: Install ai_dynamo*, ai_dynamo_runtime*, nixl* (no-cache)
  note right of Builder: uv pip list (debug)

  Builder-->>Runtime: Copy venv and packaged libs
  note over Runtime: Use venv-based dependencies

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

• fix: triton ver lock fix and trtllm env var declaration #2300 — Also updates the TRT-LLM Dockerfile and pins triton==3.3.1 for amd64.
• fix: use wheel files for installation in trtllm build #2372 — Adjusts trtllm Dockerfile to install local ai-dynamo/nixl wheels, similar wheel flow.
• build: Add TensorRT-LLM to optional dependency and corresponding instructions #2113 — Modifies TensorRT-LLM container setup and installation pipeline akin to this refactor.

Poem

In a burrow of bits I hop with delight,
Venvs aglow in the Docker-night.
Triton tamed, wheels aligned,
Pip lists twinkle, neatly timed.
Torch and friends, in one snug den—
Carrots cached? No—fresh again! 🥕✨

Tip

🔌 Remote MCP (Model Context Protocol) integration is now available!

Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

CodeRabbit Commands (Invoked using PR/Issue comments)

Type @coderabbitai help to get the list of available commands.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

Status, Documentation and Community

• Visit our Status Page to check the current availability of CodeRabbit.
• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[coderabbitai]

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

container/Dockerfile.trtllm (1)

135-151: Missing trtllm_wheel build stage causes COPY to fail

The COPY --from=trtllm_wheel . /trtllm_wheel/ instruction in container/Dockerfile.trtllm refers to a stage alias that isn’t defined anywhere, so the build will always error on that line. You must either:

• Add a multistage build stage named trtllm_wheel before line 135, for example:

  FROM python:3.10-slim AS trtllm_wheel
  # …install/build the TRTLLM wheel here…

• Or guard the COPY behind your existing HAS_TRTLLM_CONTEXT check so that it’s only executed if the wheel stage was actually built.

Please update container/Dockerfile.trtllm accordingly to prevent build failures.

🧹 Nitpick comments (2)

container/Dockerfile.trtllm (2)

411-414: Avoid uninstalling widely used pure-Python deps; make uninstall idempotent and non-interactive

Uninstalling packaging and networkx is unnecessary and can create a mixed “copied” vs “installed” state that confuses tooling. If the goal is only to prefer NVIDIA’s torch/triton, restrict the uninstall to those packages and make it non-interactive and idempotent.

Apply this diff:

-    uv pip install --no-cache --requirement /tmp/requirements.txt && \
-    echo "uninstall (networkx packaging torch triton) as we will use NVIDIA's versions later" && \
-    uv pip uninstall networkx packaging torch triton
+    uv pip install --no-cache --requirement /tmp/requirements.txt && \
+    echo "uninstalling (torch triton) to prefer NGC-provided builds" && \
+    uv pip uninstall -y torch triton || true

---

440-440: Drop or gate uv pip list debug hooks to keep layers lean

These no-op runs add build time and layers. Either remove them or guard behind a DEBUG arg.

Apply this diff at each listed line:

-RUN uv pip list # TODO remove this line
+RUN if [ "${DEBUG_PIP_LIST:-0}" = "1" ]; then uv pip list; fi

Add this near the top of the runtime stage (outside this hunk):

# at the start of the Runtime stage
ARG DEBUG_PIP_LIST=0

Also applies to: 464-464, 485-485, 492-492, 500-501

📜 Review details

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

📥 Commits

Reviewing files that changed from the base of the PR and between fb15e20 and 918326d.

📒 Files selected for processing (1)

• container/Dockerfile.trtllm (3 hunks)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)

• GitHub Check: Build and Test - dynamo

[coderabbitai]

🛠️ Refactor suggestion

Brittle wholesale copy from system dist-packages into venv; reduce surface to essentials and avoid ABI drift

Copying system-wide libs and many third-party packages (setuptools, packaging, networkx, sympy, jinja2, triton) into the venv is fragile:

• Hard-coded dist-info versions will break when the base image updates.
• /usr/local/lib/lib* can clobber runtime CUDA/OS libs and cause subtle ABI issues.
• Mixing copied files with uv-managed packages makes upgrades, uninstalls, and debugging harder.

Prefer copying only the PyTorch stack you explicitly need (torch/torchgen/torchvision/functorch and flash_attn artifacts) and let uv manage generic Python deps. Also avoid copying triton from NGC if you intend to pin a specific triton for TRT-LLM later.

Slim the COPY set like this:

-RUN uv pip list # TODO remove this line
-COPY --from=build /usr/local/lib/lib* /usr/local/lib/
 COPY --from=build /usr/local/lib/python3.12/dist-packages/torch /opt/dynamo/venv/lib/python3.12/site-packages/torch
 COPY --from=build /usr/local/lib/python3.12/dist-packages/torch-${TORCH_VER}.dist-info /opt/dynamo/venv/lib/python3.12/site-packages/torch-${TORCH_VER}.dist-info
 COPY --from=build /usr/local/lib/python3.12/dist-packages/torchgen /opt/dynamo/venv/lib/python3.12/site-packages/torchgen
 COPY --from=build /usr/local/lib/python3.12/dist-packages/torchvision /opt/dynamo/venv/lib/python3.12/site-packages/torchvision
 COPY --from=build /usr/local/lib/python3.12/dist-packages/torchvision-${TORCHVISION_VER}.dist-info /opt/dynamo/venv/lib/python3.12/site-packages/torchvision-${TORCHVISION_VER}.dist-info
 COPY --from=build /usr/local/lib/python3.12/dist-packages/torchvision.libs /opt/dynamo/venv/lib/python3.12/site-packages/torchvision.libs
-COPY --from=build /usr/local/lib/python3.12/dist-packages/setuptools /opt/dynamo/venv/lib/python3.12/site-packages/setuptools
-COPY --from=build /usr/local/lib/python3.12/dist-packages/setuptools-${SETUPTOOLS_VER}.dist-info /opt/dynamo/venv/lib/python3.12/site-packages/setuptools-${SETUPTOOLS_VER}.dist-info
 COPY --from=build /usr/local/lib/python3.12/dist-packages/functorch /opt/dynamo/venv/lib/python3.12/site-packages/functorch
-COPY --from=build /usr/local/lib/python3.12/dist-packages/triton /opt/dynamo/venv/lib/python3.12/site-packages/triton
-COPY --from=build /usr/local/lib/python3.12/dist-packages/pytorch_triton-${PYTORCH_TRITON_VER}.dist-info /opt/dynamo/venv/lib/python3.12/site-packages/pytorch_triton-${PYTORCH_TRITON_VER}.dist-info
-COPY --from=build /usr/local/lib/python3.12/dist-packages/jinja2 /opt/dynamo/venv/lib/python3.12/site-packages/jinja2
-COPY --from=build /usr/local/lib/python3.12/dist-packages/jinja2-${JINJA2_VER}.dist-info /opt/dynamo/venv/lib/python3.12/site-packages/jinja2-${JINJA2_VER}.dist-info
-COPY --from=build /usr/local/lib/python3.12/dist-packages/networkx /opt/dynamo/venv/lib/python3.12/site-packages/networkx
-COPY --from=build /usr/local/lib/python3.12/dist-packages/networkx-${NETWORKX_VER}.dist-info /opt/dynamo/venv/lib/python3.12/site-packages/networkx-${NETWORKX_VER}.dist-info
-COPY --from=build /usr/local/lib/python3.12/dist-packages/sympy /opt/dynamo/venv/lib/python3.12/site-packages/sympy
-COPY --from=build /usr/local/lib/python3.12/dist-packages/sympy-${SYMPY_VER}.dist-info /opt/dynamo/venv/lib/python3.12/site-packages/sympy-${SYMPY_VER}.dist-info
-COPY --from=build /usr/local/lib/python3.12/dist-packages/packaging /opt/dynamo/venv/lib/python3.12/site-packages/packaging
-COPY --from=build /usr/local/lib/python3.12/dist-packages/packaging-${PACKAGING_VER}.dist-info /opt/dynamo/venv/lib/python3.12/site-packages/packaging-${PACKAGING_VER}.dist-info
 COPY --from=build /usr/local/lib/python3.12/dist-packages/flash_attn /opt/dynamo/venv/lib/python3.12/site-packages/flash_attn
 COPY --from=build /usr/local/lib/python3.12/dist-packages/flash_attn-${FLASH_ATTN_VER}.dist-info /opt/dynamo/venv/lib/python3.12/site-packages/flash_attn-${FLASH_ATTN_VER}.dist-info
 COPY --from=build /usr/local/lib/python3.12/dist-packages/flash_attn_2_cuda.cpython-312-*-linux-gnu.so /opt/dynamo/venv/lib/python3.12/site-packages/
-RUN uv pip list # TODO remove this line

If you want to future-proof the dist-info names, switch to a single RUN cp -a inside the runtime stage that uses wildcard globs (e.g., torch*, torchvision*, flash_attn*) rather than hard-coded versions.

Committable suggestion skipped: line range outside the PR's diff.

🤖 Prompt for AI Agents

In container/Dockerfile.trtllm around lines 441-464, the PR currently copies a
large set of system dist-packages and /usr/local/lib/lib* into the venv using
hard-coded dist-info names, which is brittle and risks ABI drift; replace those
COPY lines so you only bring the PyTorch runtime artifacts the image actually
needs (torch, torchgen, torchvision, functorch, flash_attn and the compiled .so
artifacts) and stop copying generic packages (setuptools, networkx, sympy,
packaging, jinja2, triton, etc.) or system libs. Implement this by removing the
/usr/local/lib/lib* COPY and the hard-coded dist-info COPYs, and instead either
(A) use wildcarded COPY/RUN cp -a inside the runtime stage to copy torch*
torchvision* functorch flash_attn and their associated .so files and dist-info
globs, or (B) only copy explicit package directories for the minimal PyTorch set
and let uv handle generic Python deps; also avoid copying triton from NGC unless
you intend to pin it, and remove the temporary "uv pip list" RUN line.

[coderabbitai]

⚠️ Potential issue

Parity with build stage: ensure compatible triton for TRT-LLM on amd64

In build you pin triton==3.3.1 for amd64, but in runtime you don’t. Given TRT-LLM’s sensitivity to triton versions, this can lead to import/runtime failures if the copied PyTorch triton doesn’t match. Align runtime with build by installing triton==3.3.1 on amd64 after TRT-LLM.

Apply this diff:

-RUN uv pip list # TODO remove this line
-RUN uv pip install --no-cache "cuda-python>=12,<13" && \
-    uv pip install --no-cache --extra-index-url "${TENSORRTLLM_INDEX_URL}" "${TENSORRTLLM_PIP_WHEEL}" && \
-    uv pip install --no-cache \
+RUN uv pip install --no-cache "cuda-python>=12,<13" && \
+    uv pip install --no-cache --extra-index-url "${TENSORRTLLM_INDEX_URL}" "${TENSORRTLLM_PIP_WHEEL}" && \
+    if [ "$ARCH" = "amd64" ]; then uv pip install --no-cache "triton==3.3.1"; fi && \
+    uv pip install --no-cache \
         /workspace/wheelhouse/ai_dynamo_runtime*cp312*.whl \
         /workspace/wheelhouse/ai_dynamo*any.whl \
         /workspace/wheelhouse/nixl*.whl
-RUN uv pip list # TODO remove this line

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	RUN uv pip list # TODO remove this line
	RUN uv pip install --no-cache "cuda-python>=12,<13" && \
	uv pip install --no-cache --extra-index-url "${TENSORRTLLM_INDEX_URL}" "${TENSORRTLLM_PIP_WHEEL}" && \
	uv pip install --no-cache \
	/workspace/wheelhouse/ai_dynamo_runtime*cp312*.whl \
	/workspace/wheelhouse/ai_dynamo*any.whl \
	/workspace/wheelhouse/nixl*.whl
	RUN uv pip list # TODO remove this line
	RUN uv pip install --no-cache "cuda-python>=12,<13" && \
	uv pip install --no-cache --extra-index-url "${TENSORRTLLM_INDEX_URL}" "${TENSORRTLLM_PIP_WHEEL}" && \
	if [ "$ARCH" = "amd64" ]; then uv pip install --no-cache "triton==3.3.1"; fi && \
	uv pip install --no-cache \
	/workspace/wheelhouse/ai_dynamo_runtime*cp312*.whl \
	/workspace/wheelhouse/ai_dynamo*any.whl \
	/workspace/wheelhouse/nixl*.whl

🤖 Prompt for AI Agents

In container/Dockerfile.trtllm around lines 485 to 493, the runtime stage does
not pin triton while the build stage pins triton==3.3.1 for amd64, which can
cause import/runtime failures; add a conditional install right after the TRT-LLM
pip installs that, when building for amd64, runs pip install --no-cache
triton==3.3.1 (using the same uv pip invocation pattern and environment) so the
runtime triton version matches the build; implement the condition using the
existing architecture/build-arg mechanism in the Dockerfile (e.g., check
TARGETARCH or similar) and keep the command order so triton is installed after
the TRT-LLM wheels are installed.

[github-actions]

This PR is stale because it has been open 30 days with no activity. Remove stale label or comment or this will be closed in 5 days.