ci: split vllm venv COPY into parallel layers for faster pulls (#5494)

Signed-off-by: alec-flowers <aflowers@nvidia.com>
Signed-off-by: Alec <35311602+alec-flowers@users.noreply.github.com>

.dockerignore

@@ -41,4 +41,38 @@
 **/*safetensors
 container/Dockerfile*
 .venv
-.venv-docs
\ No newline at end of file
+.venv-docs
+
+# Python
+__pycache__/
+*.pyc
+*.pyo
+*.egg-info/
+build/
+dist/
+*.egg
+
+# Rust
+target/
+
+# IDE/Editor
+.idea/
+*.swp
+*.swo
+.vscode/
+.build/
+
+# Tests & Linters
+.pytest_cache/
+.coverage
+htmlcov/
+.tox/
+.ruff_cache/
+.mypy_cache/
+
+# Docs
+docs/_build/
+
+# AI
+.cursor/
+.claude/
\ No newline at end of file

container/Dockerfile.vllm

-# syntax=docker/dockerfile:1.10.0
+# syntax=docker/dockerfile:1.10.0-labs
 # SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
@@ -637,6 +637,7 @@ RUN userdel -r ubuntu > /dev/null 2>&1 || true \
 
 ARG ARCH_ALT
 ARG PYTHON_VERSION
+ENV PYTHON_VERSION=${PYTHON_VERSION}
 
 # Install Python, build-essential and python3-dev as apt dependencies
 RUN apt-get update && \
@@ -674,10 +675,31 @@ ENV NIXL_PREFIX=/opt/nvidia/nvda_nixl
 ENV NIXL_LIB_DIR=$NIXL_PREFIX/lib/${ARCH_ALT}-linux-gnu
 ENV NIXL_PLUGIN_DIR=$NIXL_LIB_DIR/plugins
 
+# Site-packages path derived from PYTHON_VERSION ARG
+ARG SITE_PACKAGES=${VIRTUAL_ENV}/lib/python${PYTHON_VERSION}/site-packages
+
 ### VIRTUAL ENVIRONMENT SETUP ###
-# Copy entire virtual environment from framework container with correct ownership
-# Pattern: COPY --chmod=775 <path>; chmod g+w <path> done later as root because COPY --chmod only affects <path>/*, not <path>
-COPY --chmod=775 --chown=dynamo:0 --from=framework ${VIRTUAL_ENV} ${VIRTUAL_ENV}
+# Copy virtual environment from framework container, splitting large packages into separate layers
+# to enable parallel downloads. Pattern: COPY --chmod=775 <path>; chmod g+w <path> done later as
+# root because COPY --chmod only affects <path>/*, not <path>
+#
+# Layer sizes (uncompressed): nvidia=4.5GB, flashinfer_jit_cache=4.1GB, torch=2.1GB,
+#                             vllm=1.2GB, triton=592MB, flashinfer_cubin=437MB
+COPY --chmod=775 --chown=dynamo:0 --from=framework ${SITE_PACKAGES}/nvidia ${SITE_PACKAGES}/nvidia
+COPY --chmod=775 --chown=dynamo:0 --from=framework ${SITE_PACKAGES}/flashinfer_jit_cache ${SITE_PACKAGES}/flashinfer_jit_cache
+COPY --chmod=775 --chown=dynamo:0 --from=framework ${SITE_PACKAGES}/torch ${SITE_PACKAGES}/torch
+COPY --chmod=775 --chown=dynamo:0 --from=framework ${SITE_PACKAGES}/vllm ${SITE_PACKAGES}/vllm
+COPY --chmod=775 --chown=dynamo:0 --from=framework ${SITE_PACKAGES}/triton ${SITE_PACKAGES}/triton
+COPY --chmod=775 --chown=dynamo:0 --from=framework ${SITE_PACKAGES}/flashinfer_cubin ${SITE_PACKAGES}/flashinfer_cubin
+# Remaining packages and venv structure (bin/, include/, share/, etc.)
+COPY --chmod=775 --chown=dynamo:0 --from=framework \
+    --exclude=lib/python*/site-packages/nvidia \
+    --exclude=lib/python*/site-packages/flashinfer_jit_cache \
+    --exclude=lib/python*/site-packages/torch \
+    --exclude=lib/python*/site-packages/vllm \
+    --exclude=lib/python*/site-packages/triton \
+    --exclude=lib/python*/site-packages/flashinfer_cubin \
+    ${VIRTUAL_ENV} ${VIRTUAL_ENV}
 
 # Copy vllm with correct ownership (read-only, no group-write needed)
 COPY --chown=dynamo:0 --from=framework /opt/vllm /opt/vllm

container/dev/50-framework-paths.sh

@@ -34,7 +34,11 @@ ARCH_ALT=$(uname -m | sed 's/aarch64/aarch64/;s/x86_64/x86_64/')
     export LD_LIBRARY_PATH="/usr/lib/${ARCH_ALT}-linux-gnu/nvshmem/13:${LD_LIBRARY_PATH}"
 
 # PyTorch libraries (TRT-LLM)
-PYTHON_VERSION=${PYTHON_VERSION:-3.12}
+# PYTHON_VERSION should be set via ENV in container; fail early if missing
+if [ -z "${PYTHON_VERSION}" ]; then
+    echo "WARNING: PYTHON_VERSION not set, defaulting to 3.12" >&2
+    PYTHON_VERSION=3.12
+fi
 [ -d "/opt/dynamo/venv/lib/python${PYTHON_VERSION}/site-packages/torch/lib" ] && \
     export LD_LIBRARY_PATH="/opt/dynamo/venv/lib/python${PYTHON_VERSION}/site-packages/torch/lib:${LD_LIBRARY_PATH}"
 [ -d "/opt/dynamo/venv/lib/python${PYTHON_VERSION}/site-packages/torch_tensorrt/lib" ] && \

container/dev/Dockerfile.dev

@@ -324,8 +324,8 @@ RUN if [ "${FRAMEWORK}" = "sglang" ]; then \
         # Copy all packages from runtime stage system site-packages into venv
         # This includes ai-dynamo-runtime, kubernetes, and all other dependencies
         # Use --no-preserve=mode so copied files inherit umask 002 (group-writable)
-        cp -r --no-preserve=mode /usr/local/lib/python3.12/dist-packages/* \
-              /opt/dynamo/venv/lib/python3.12/site-packages/; \
+        cp -r --no-preserve=mode /usr/local/lib/python${PYTHON_VERSION}/dist-packages/* \
+              /opt/dynamo/venv/lib/python${PYTHON_VERSION}/site-packages/; \
         # Ensure `uv` is available on PATH for subsequent `uv pip ...` steps.
         cp /tmp/uv-binary /opt/dynamo/venv/bin/uv && \
         chmod +x /opt/dynamo/venv/bin/uv && \