copy NVIDIA torch lib into uv venv location

Author: dmitry-tokarev-nv

container/Dockerfile.trtllm

@@ -408,12 +408,14 @@ RUN uv venv $VIRTUAL_ENV --python 3.12 && \
 # Common dependencies
 # TODO: Remove extra install and use pyproject.toml to define all dependencies
 RUN --mount=type=bind,source=./container/deps/requirements.txt,target=/tmp/requirements.txt \
-    uv pip install --requirement /tmp/requirements.txt
+    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
 
 # Install test dependencies
 # TODO: Remove this once we have a functional CI image built on top of the runtime image
 RUN --mount=type=bind,source=./container/deps/requirements.test.txt,target=/tmp/requirements.txt \
-    uv pip install --requirement /tmp/requirements.txt
+    uv pip install --no-cache --requirement /tmp/requirements.txt
 
 # Copy CUDA toolkit components needed for nvcc, cudafe, cicc etc.
 COPY --from=build /usr/local/cuda/bin/nvcc /usr/local/cuda/bin/nvcc
@@ -435,29 +437,31 @@ ARG SYMPY_VER=1.14.0
 ARG PACKAGING_VER=23.2
 ARG FLASH_ATTN_VER=2.7.3
 ARG MPMATH_VER=1.3.0
+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 /usr/local/lib/python3.12/dist-packages/torch
-COPY --from=build /usr/local/lib/python3.12/dist-packages/torch-${TORCH_VER}.dist-info /usr/local/lib/python3.12/dist-packages/torch-${TORCH_VER}.dist-info
-COPY --from=build /usr/local/lib/python3.12/dist-packages/torchgen /usr/local/lib/python3.12/dist-packages/torchgen
-COPY --from=build /usr/local/lib/python3.12/dist-packages/torchvision /usr/local/lib/python3.12/dist-packages/torchvision
-COPY --from=build /usr/local/lib/python3.12/dist-packages/torchvision-${TORCHVISION_VER}.dist-info /usr/local/lib/python3.12/dist-packages/torchvision-${TORCHVISION_VER}.dist-info
-COPY --from=build /usr/local/lib/python3.12/dist-packages/torchvision.libs /usr/local/lib/python3.12/dist-packages/torchvision.libs
-COPY --from=build /usr/local/lib/python3.12/dist-packages/setuptools /usr/local/lib/python3.12/dist-packages/setuptools
-COPY --from=build /usr/local/lib/python3.12/dist-packages/setuptools-${SETUPTOOLS_VER}.dist-info /usr/local/lib/python3.12/dist-packages/setuptools-${SETUPTOOLS_VER}.dist-info
-COPY --from=build /usr/local/lib/python3.12/dist-packages/functorch /usr/local/lib/python3.12/dist-packages/functorch
-COPY --from=build /usr/local/lib/python3.12/dist-packages/triton /usr/local/lib/python3.12/dist-packages/triton
-COPY --from=build /usr/local/lib/python3.12/dist-packages/pytorch_triton-${PYTORCH_TRITON_VER}.dist-info /usr/local/lib/python3.12/dist-packages/pytorch_triton-${PYTORCH_TRITON_VER}.dist-info
-COPY --from=build /usr/local/lib/python3.12/dist-packages/jinja2 /usr/local/lib/python3.12/dist-packages/jinja2
-COPY --from=build /usr/local/lib/python3.12/dist-packages/jinja2-${JINJA2_VER}.dist-info /usr/local/lib/python3.12/dist-packages/jinja2-${JINJA2_VER}.dist-info
-COPY --from=build /usr/local/lib/python3.12/dist-packages/networkx /usr/local/lib/python3.12/dist-packages/networkx
-COPY --from=build /usr/local/lib/python3.12/dist-packages/networkx-${NETWORKX_VER}.dist-info /usr/local/lib/python3.12/dist-packages/networkx-${NETWORKX_VER}.dist-info
-COPY --from=build /usr/local/lib/python3.12/dist-packages/sympy /usr/local/lib/python3.12/dist-packages/sympy
-COPY --from=build /usr/local/lib/python3.12/dist-packages/sympy-${SYMPY_VER}.dist-info /usr/local/lib/python3.12/dist-packages/sympy-${SYMPY_VER}.dist-info
-COPY --from=build /usr/local/lib/python3.12/dist-packages/packaging /usr/local/lib/python3.12/dist-packages/packaging
-COPY --from=build /usr/local/lib/python3.12/dist-packages/packaging-${PACKAGING_VER}.dist-info /usr/local/lib/python3.12/dist-packages/packaging-${PACKAGING_VER}.dist-info
-COPY --from=build /usr/local/lib/python3.12/dist-packages/flash_attn /usr/local/lib/python3.12/dist-packages/flash_attn
-COPY --from=build /usr/local/lib/python3.12/dist-packages/flash_attn-${FLASH_ATTN_VER}.dist-info /usr/local/lib/python3.12/dist-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 /usr/local/lib/python3.12/dist-packages/
+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
 
 
 # Install TensorRT-LLM (same as in build stage)
@@ -478,19 +482,22 @@ COPY --from=dev /workspace/target/release/metrics /usr/local/bin/metrics
 # NOTE: locking cuda-python version to <13 to avoid breaks with tensorrt-llm 1.0.0rc4. This
 #       can be removed after https://github.com/NVIDIA/TensorRT-LLM/pull/6703 is merged
 #       we upgrade to a published pip wheel containing this change.
-RUN uv pip install "cuda-python>=12,<13" && \
-    uv pip install --extra-index-url "${TENSORRTLLM_INDEX_URL}" "${TENSORRTLLM_PIP_WHEEL}" && \
-    if [ "$ARCH" = "amd64" ]; then \
-        pip install "triton==3.3.1"; \
-    fi; \
-    uv pip install /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}" && \
+    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
 
 # Copy benchmarks, backends and tests for CI
 # TODO: Remove this once we have a functional CI image built on top of the runtime image
 COPY tests /workspace/tests
 COPY benchmarks /workspace/benchmarks
 COPY components/backends/trtllm /workspace/components/backends/trtllm
-RUN uv pip install /workspace/benchmarks
+RUN uv pip install --no-cache /workspace/benchmarks
+RUN uv pip list # TODO remove this line
 
 # Copy files for legal compliance
 COPY ATTRIBUTION* LICENSE /workspace/