[None][infra] Reduce Docker image layer count in Dockerfile.multi

Replace COPY+RUN patterns with --mount=source bind mounts for install
scripts, constraints.txt, and OSS attribution script. Consolidate
release stage COPYs and RUNs into single layers using bind mounts
and cp. Use pip cache mounts instead of manual mkdir/rm of cache dirs.

This reduces the TRT-LLM-contributed layers by ~17 (from ~24 to ~7),
leaving more headroom under Docker's 128-layer overlay2 limit for
downstream consumers.

Signed-off-by: Jan Bernlöhr <jbernloehr@nvidia.com>
Signed-off-by: jbernloehr <jbernloehr@nvidia.com>

Author: janbernloehr

docker/Dockerfile.multi

@@ -18,10 +18,7 @@ ARG BASH_ENV="/etc/bash.bashrc"
 ENV BASH_ENV=${BASH_ENV}
 
 ARG GITHUB_MIRROR=""
-RUN echo "Using GitHub mirror: $GITHUB_MIRROR"
-
 ARG PYTHON_VERSION="3.12.3"
-RUN echo "Using Python version: $PYTHON_VERSION"
 
 SHELL ["/bin/bash", "-c"]
 
@@ -31,62 +28,36 @@ FROM base AS devel
 # NB: PyTorch requires this to be < 1.0
 ENV PYTORCH_ALLOC_CONF="garbage_collection_threshold:0.99999"
 
-# Copy all installation scripts at once to reduce layers
-COPY docker/common/install.sh \
-     docker/common/install_base.sh \
-     docker/common/install_cmake.sh \
-     docker/common/install_ccache.sh \
-     docker/common/install_cuda_toolkit.sh \
-     docker/common/install_tensorrt.sh \
-     docker/common/install_polygraphy.sh \
-     docker/common/install_mpi4py.sh \
-     docker/common/install_pytorch.sh \
-     docker/common/install_ucx.sh \
-     docker/common/install_nixl.sh \
-     docker/common/install_etcd.sh \
-     ./
-
 ARG TRT_VER
 ARG CUDA_VER
 ARG CUDNN_VER
 ARG NCCL_VER
 ARG CUBLAS_VER
 ARG TORCH_INSTALL_TYPE="skip"
-RUN GITHUB_MIRROR=${GITHUB_MIRROR} \
+RUN --mount=type=cache,target=/root/.cache/pip \
+    --mount=source=docker/common,target=/tmp/docker/common \
+    --mount=source=constraints.txt,target=/tmp/constraints.txt \
+    GITHUB_MIRROR=${GITHUB_MIRROR} \
     PYTHON_VERSION=${PYTHON_VERSION} \
     TRT_VER=${TRT_VER} \
     CUDA_VER=${CUDA_VER} \
     CUDNN_VER=${CUDNN_VER} \
     NCCL_VER=${NCCL_VER} \
     CUBLAS_VER=${CUBLAS_VER} \
     TORCH_INSTALL_TYPE=${TORCH_INSTALL_TYPE} \
-    bash ./install.sh --base --cmake --ccache --cuda_toolkit --tensorrt --polygraphy --mpi4py --pytorch --opencv && \
-    rm install_base.sh && \
-    rm install_cmake.sh && \
-    rm install_ccache.sh && \
-    rm install_cuda_toolkit.sh && \
-    rm install_tensorrt.sh && \
-    rm install_polygraphy.sh && \
-    rm install_mpi4py.sh && \
-    rm install_pytorch.sh && \
-    rm install.sh
-
-# Copy and install dependencies from constraints.txt
-COPY constraints.txt /tmp/constraints.txt
-RUN pip3 install --no-cache-dir -r /tmp/constraints.txt && rm /tmp/constraints.txt
-
-# Remove nbconvert to avoid https://github.com/advisories/GHSA-xm59-rqc7-hhvf in the base NGC PyTorch image.
-RUN pip3 uninstall -y nbconvert || true
+    echo "Using GitHub mirror: $GITHUB_MIRROR" && \
+    echo "Using Python version: $PYTHON_VERSION" && \
+    bash /tmp/docker/common/install.sh --base --cmake --ccache --cuda_toolkit --tensorrt --polygraphy --mpi4py --pytorch --opencv && \
+    pip3 install --no-cache-dir -r /tmp/constraints.txt && \
+    # Remove nbconvert to avoid https://github.com/advisories/GHSA-xm59-rqc7-hhvf in the base NGC PyTorch image.
+    (pip3 uninstall -y nbconvert || true)
 
 # Install UCX, NIXL, etcd
-# TODO: Combine these into the main install.sh script
-RUN GITHUB_MIRROR=${GITHUB_MIRROR} bash ./install_ucx.sh && \
-    GITHUB_MIRROR=${GITHUB_MIRROR} bash ./install_nixl.sh && \
-    bash ./install_etcd.sh && \
-    rm install_ucx.sh && \
-    rm install_nixl.sh && \
-    rm install_etcd.sh && \
-    rm -rf /root/.cache/pip && \
+RUN --mount=type=cache,target=/root/.cache/pip \
+    --mount=source=docker/common,target=/tmp/docker/common \
+    GITHUB_MIRROR=${GITHUB_MIRROR} bash /tmp/docker/common/install_ucx.sh && \
+    GITHUB_MIRROR=${GITHUB_MIRROR} bash /tmp/docker/common/install_nixl.sh && \
+    bash /tmp/docker/common/install_etcd.sh && \
     rm -rf /root/.cache/uv/archive-v0 && \
     # WAR against https://github.com/advisories/GHSA-58pv-8j8x-9vj2
     rm -rf /usr/local/lib/python3.12/dist-packages/setuptools/_vendor/jaraco.context-5.3.0.dist-info && \
@@ -96,9 +67,8 @@ RUN GITHUB_MIRROR=${GITHUB_MIRROR} bash ./install_ucx.sh && \
 # Generate OSS attribution file for devel image
 ARG TRT_LLM_VER
 ARG TARGETARCH
-COPY scripts/generate_container_oss_attribution.sh /tmp/generate_container_oss_attribution.sh
-RUN bash /tmp/generate_container_oss_attribution.sh "devel" "${TRT_LLM_VER}" "${TARGETARCH}" && \
-    rm /tmp/generate_container_oss_attribution.sh
+RUN --mount=source=scripts/generate_container_oss_attribution.sh,target=/tmp/generate_container_oss_attribution.sh \
+    bash /tmp/generate_container_oss_attribution.sh "devel" "${TRT_LLM_VER}" "${TARGETARCH}"
 
 FROM ${TRITON_IMAGE}:${TRITON_BASE_TAG} AS triton
 
@@ -111,20 +81,14 @@ COPY --from=triton /opt/tritonserver/include /opt/tritonserver/include
 COPY --from=triton /opt/tritonserver/bin /opt/tritonserver/bin
 COPY --from=triton /opt/tritonserver/caches /opt/tritonserver/caches
 
-# Copy all installation scripts at once to reduce layers
-COPY docker/common/install_triton.sh \
-     docker/common/install_mooncake.sh \
-     ./
-
-# Install Mooncake, after triton handles boost requirement
-RUN GITHUB_MIRROR=${GITHUB_MIRROR} bash ./install_triton.sh && \
+# Install Triton and Mooncake
+RUN --mount=source=docker/common,target=/tmp/docker/common \
+    GITHUB_MIRROR=${GITHUB_MIRROR} bash /tmp/docker/common/install_triton.sh && \
     if [ -f /etc/redhat-release ]; then \
         echo "Rocky8 detected, skipping mooncake installation"; \
     else \
-        bash ./install_mooncake.sh; \
-    fi && \
-    rm install_triton.sh && \
-    rm install_mooncake.sh
+        bash /tmp/docker/common/install_mooncake.sh; \
+    fi
 
 FROM ${DEVEL_IMAGE} AS wheel
 WORKDIR /src/tensorrt_llm
@@ -148,43 +112,45 @@ RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=${CCACH
 
 FROM ${DEVEL_IMAGE} AS release
 
-# Create a cache directory for pip
-RUN mkdir -p /root/.cache/pip
-
 WORKDIR /app/tensorrt_llm
 RUN --mount=type=cache,target=/root/.cache/pip --mount=type=bind,from=wheel,source=/src/tensorrt_llm/build,target=/tmp/wheel \
     pip install /tmp/wheel/tensorrt_llm*.whl
 
-COPY README.md ./
-COPY --from=wheel /src/tensorrt_llm/build/tensorrt_llm*.whl ./
-COPY docs docs
-COPY cpp/include include
-
-RUN ln -sv $(python3 -c 'import site; print(f"{site.getsitepackages()[0]}/tensorrt_llm/bin")') bin && \
-    test -f bin/executorWorker && \
-    ln -sv $(python3 -c 'import site; print(f"{site.getsitepackages()[0]}/tensorrt_llm/libs")') lib && \
-    test -f lib/libnvinfer_plugin_tensorrt_llm.so && \
-    echo "/app/tensorrt_llm/lib" > /etc/ld.so.conf.d/tensorrt_llm.conf && \
-    ldconfig && \
-    ! ( ldd -v bin/executorWorker | grep tensorrt_llm | grep -q "not found" )
-
-ARG SRC_DIR=/src/tensorrt_llm
-COPY --from=wheel ${SRC_DIR}/benchmarks benchmarks
-ARG CPP_BUILD_DIR=${SRC_DIR}/cpp/build
-COPY --from=wheel \
-     ${CPP_BUILD_DIR}/benchmarks/bertBenchmark \
-     ${CPP_BUILD_DIR}/benchmarks/gptManagerBenchmark \
-     ${CPP_BUILD_DIR}/benchmarks/disaggServerBenchmark \
-     benchmarks/cpp/
-
-COPY examples examples
-RUN chmod -R a+w examples && \
+RUN --mount=source=README.md,target=/tmp/ctx/README.md \
+    --mount=source=docs,target=/tmp/ctx/docs \
+    --mount=source=cpp/include,target=/tmp/ctx/include \
+    --mount=source=examples,target=/tmp/ctx/examples \
+    --mount=type=bind,from=wheel,source=/src/tensorrt_llm/build,target=/tmp/wheel \
+    --mount=type=bind,from=wheel,source=/src/tensorrt_llm/benchmarks,target=/tmp/benchmarks \
+    --mount=type=bind,from=wheel,source=/src/tensorrt_llm/cpp/build/benchmarks,target=/tmp/cpp_benchmarks \
+    # Copy build context files
+    cp /tmp/ctx/README.md ./ && \
+    cp -r /tmp/ctx/docs ./docs && \
+    cp -r /tmp/ctx/include ./include && \
+    cp -r /tmp/ctx/examples ./examples && \
+    chmod -R a+w examples && \
+    # Copy wheel stage outputs
+    cp /tmp/wheel/tensorrt_llm*.whl ./ && \
+    cp -r /tmp/benchmarks ./benchmarks && \
+    mkdir -p benchmarks/cpp && \
+    cp /tmp/cpp_benchmarks/bertBenchmark \
+       /tmp/cpp_benchmarks/gptManagerBenchmark \
+       /tmp/cpp_benchmarks/disaggServerBenchmark \
+       benchmarks/cpp/ && \
     rm -v \
       benchmarks/cpp/bertBenchmark.cpp \
       benchmarks/cpp/gptManagerBenchmark.cpp \
       benchmarks/cpp/disaggServerBenchmark.cpp \
       benchmarks/cpp/CMakeLists.txt && \
-    rm -rf /root/.cache/pip && \
+    # Create symlinks to installed package binaries and libraries
+    ln -sv $(python3 -c 'import site; print(f"{site.getsitepackages()[0]}/tensorrt_llm/bin")') bin && \
+    test -f bin/executorWorker && \
+    ln -sv $(python3 -c 'import site; print(f"{site.getsitepackages()[0]}/tensorrt_llm/libs")') lib && \
+    test -f lib/libnvinfer_plugin_tensorrt_llm.so && \
+    echo "/app/tensorrt_llm/lib" > /etc/ld.so.conf.d/tensorrt_llm.conf && \
+    ldconfig && \
+    ! ( ldd -v bin/executorWorker | grep tensorrt_llm | grep -q "not found" ) && \
+    # Clean up uv cache and CVE workarounds
     rm -rf /root/.cache/uv/archive-v0 && \
     # WAR against https://github.com/advisories/GHSA-58pv-8j8x-9vj2
     rm -rf /usr/local/lib/python3.12/dist-packages/setuptools/_vendor/jaraco.context-5.3.0.dist-info && \
@@ -198,8 +164,8 @@ ENV TRT_LLM_GIT_COMMIT=${GIT_COMMIT} \
     TRT_LLM_VERSION=${TRT_LLM_VER}
 
 # Generate OSS attribution file for release image
-COPY scripts/generate_container_oss_attribution.sh /tmp/generate_container_oss_attribution.sh
-RUN bash /tmp/generate_container_oss_attribution.sh "release" "${TRT_LLM_VER}" "${TARGETARCH}" && rm /tmp/generate_container_oss_attribution.sh
+RUN --mount=source=scripts/generate_container_oss_attribution.sh,target=/tmp/generate_container_oss_attribution.sh \
+    bash /tmp/generate_container_oss_attribution.sh "release" "${TRT_LLM_VER}" "${TARGETARCH}"
 
 FROM wheel AS tritonbuild