ci: OPS-2136: Setup caching for docker build (#4535)

Signed-off-by: Dillon Cullinan <dcullinan@nvidia.com>
Signed-off-by: Dillon Cullinan <dillon-cullinan@users.noreply.github.com>
Co-authored-by: Dillon Cullinan <dillon-cullinan@users.noreply.github.com>

.github/actions/docker-build/action.yml

@@ -16,9 +16,6 @@ inputs:
   image_tag:
     description: 'Custom image tag (optional, defaults to framework:latest)'
     required: false
-  ngc_ci_access_token:
-    description: 'NGC CI Access Token'
-    required: false
   ci_token:
     description: 'CI Token'
     required: false
@@ -67,20 +64,14 @@ runs:
     - name: Set up Docker Buildx
       uses: docker/setup-buildx-action@e468171a9de216ec08956ac3ada2f0791b6bd435 #v3.11.1
       with:
-        driver: docker
+        driver: docker-container
         # Enable BuildKit for enhanced metadata
         buildkitd-flags: --debug
-    - name: Login to ECR
+    - name: Cleanup
+      if: always()
       shell: bash
-      env:
-        ECR_HOSTNAME: ${{ inputs.aws_account_id }}.dkr.ecr.${{ inputs.aws_default_region }}.amazonaws.com
       run: |
-        aws ecr get-login-password --region ${{ inputs.aws_default_region }} | docker login --username AWS --password-stdin ${ECR_HOSTNAME}
-    - name: Login to NGC
-      if: github.event.pull_request.head.repo.full_name == github.repository || github.event_name == 'push'
-      uses: ./.github/actions/docker-login
-      with:
-        ngc_ci_access_token: ${{ inputs.ngc_ci_access_token }}
+        docker system prune -af
     - name: Build image
       id: build
       shell: bash
@@ -91,9 +82,12 @@ runs:
         AWS_ACCESS_KEY_ID: ${{ inputs.aws_access_key_id }}
         AWS_SECRET_ACCESS_KEY: ${{ inputs.aws_secret_access_key }}
         PLATFORM: ${{ inputs.platform }}
+        ECR_HOSTNAME: ${{ inputs.aws_account_id }}.dkr.ecr.${{ inputs.aws_default_region }}.amazonaws.com
         GITHUB_RUN_ID: ${{ github.run_id }}
         GITHUB_JOB: ${{ github.job }}
+        GITHUB_REF_NAME: ${{ github.ref_name }}
       run: |
+        set -x
         # Determine image tag
         if [ -n "${{ inputs.image_tag }}" ]; then
           IMAGE_TAG="${{ inputs.image_tag }}"
@@ -113,18 +107,28 @@ runs:
         echo "📝 Build log will be saved to: ${BUILD_LOG_FILE}"
 
         # Collect optional overrides provided by the workflow
+        # Set base cache args and set --cache-to if this is a main commit
         EXTRA_ARGS=""
+        EXTRA_ARGS="--cache-to type=inline "
+        EXTRA_ARGS+="--cache-from type=registry,ref=${ECR_HOSTNAME}/ai-dynamo/dynamo:${{ inputs.framework }}-${PLATFORM##*/}-cache "
+        EXTRA_ARGS+="--cache-from type=registry,ref=${ECR_HOSTNAME}/ai-dynamo/dynamo:main-${{ inputs.framework }}-${PLATFORM##*/} "
+        if [[ "$GITHUB_REF_NAME" == "main" ]]; then
+          EXTRA_ARGS+="--cache-to type=registry,ref=${ECR_HOSTNAME}/ai-dynamo/dynamo:${{ inputs.framework }}-${PLATFORM##*/}-cache,mode=max "
+        fi
+
+        echo "$EXTRA_ARGS"
+        # Collect optional overrides provided by the workflow
         if [ -n "${{ inputs.base_image_tag }}" ]; then
-          EXTRA_ARGS+=" --base-image-tag ${{ inputs.base_image_tag }}"
+          EXTRA_ARGS+="--base-image-tag ${{ inputs.base_image_tag }} "
         fi
         if [ -n "${{ inputs.runtime_image_tag }}" ]; then
-          EXTRA_ARGS+=" --build-arg RUNTIME_IMAGE_TAG=${{ inputs.runtime_image_tag }}"
+          EXTRA_ARGS+="--build-arg RUNTIME_IMAGE_TAG=${{ inputs.runtime_image_tag }} "
         fi
         if [ -n "${{ inputs.cuda_version }}" ]; then
-          EXTRA_ARGS+=" --build-arg CUDA_VERSION=${{ inputs.cuda_version }}"
+          EXTRA_ARGS+="--build-arg CUDA_VERSION=${{ inputs.cuda_version }} "
         fi
         if [ -n "${{ inputs.torch_backend }}" ]; then
-          EXTRA_ARGS+=" --build-arg TORCH_BACKEND=${{ inputs.torch_backend }}"
+          EXTRA_ARGS+="--build-arg TORCH_BACKEND=${{ inputs.torch_backend }} "
         fi
         if [ -n "${{ inputs.dynamo_base_image }}" ]; then
           EXTRA_ARGS+=" --dynamo-base-image ${{ inputs.dynamo_base_image }}"
@@ -250,8 +254,7 @@ runs:
         chmod +x .github/scripts/parse_buildkit_output.py
 
         # Check for build logs and build stage arguments dynamically
-        BASE_BUILD_LOG="build-logs/base-image-build.log"
-        FRAMEWORK_BUILD_LOG="build-logs/framework-${FRAMEWORK_LOWER}-build.log"
+        BUILD_LOG="build-logs/single-stage-build.log"
 
         # Path to container metadata created in previous step
         CONTAINER_METADATA="build-metrics/metrics-${{ inputs.framework }}-${PLATFORM_ARCH}-${WORKFLOW_ID}-${JOB_ID}.json"
@@ -264,18 +267,11 @@ runs:
         # Build stage arguments dynamically based on which logs exist
         STAGE_ARGS=()
 
-        if [ -f "$BASE_BUILD_LOG" ]; then
-          echo "  ✓ Found base image log: ${BASE_BUILD_LOG}"
-          STAGE_ARGS+=("base:${BASE_BUILD_LOG}")
-        else
-          echo "  ℹ️  No base image log found"
-        fi
-
-        if [ -f "$FRAMEWORK_BUILD_LOG" ]; then
-          echo "  ✓ Found framework log: ${FRAMEWORK_BUILD_LOG}"
-          STAGE_ARGS+=("runtime:${FRAMEWORK_BUILD_LOG}")
+        if [ -f "$BUILD_LOG" ]; then
+          echo "  ✓ Found base image log: ${BUILD_LOG}"
+          STAGE_ARGS+=("runtime:${BUILD_LOG}")
         else
-          echo "  ℹ️  No framework log found"
+          echo "  ℹ️  No image log found"
         fi
 
         # Check for any additional stage logs (e.g., build-logs/stage3-*.log)

.github/actions/docker-tag-push/action.yml

@@ -8,6 +8,11 @@ inputs:
   push_tags:
     description: 'Target Name:Tag (newline-separated list for multiple tags)'
     required: true
+  # There isn't a clean way to have an additional tag that is conditional
+  # Adding this to handle this use-case (we want multiple tags for main builds)
+  conditional_tag:
+    description: 'Optional tag for conditionals'
+    required: false
   aws_push:
     description: 'Push to AWS Boolean'
     required: false
@@ -22,21 +27,9 @@ inputs:
   aws_default_region:
     description: 'AWS Default Region'
     required: false
-  aws_access_key_id:
-    description: 'AWS Access Key ID'
-    required: false
-  aws_secret_access_key:
-    description: 'AWS Secret Access Key'
-    required: false
   azure_acr_hostname:
     description: 'Azure ACR hostname'
     required: false
-  azure_acr_user:
-    description: 'Azure ACR user'
-    required: false
-  azure_acr_password:
-    description: 'Azure ACR password'
-    required: false
 
 outputs:
   image_tags:
@@ -48,16 +41,20 @@ runs:
   steps:
     - name: Set up Docker Buildx
       uses: docker/setup-buildx-action@v3
-
     - name: ECR Tag and Push
       shell: bash
       if: ${{ inputs.aws_push == 'true' }}
       env:
         LOCAL_IMAGE: ${{ inputs.local_image }}
         PUSH_TAGS: ${{ inputs.push_tags }}
+        CONDITIONAL_TAG: ${{ inputs.conditional_tag }}
         ECR_HOSTNAME: ${{ inputs.aws_account_id }}.dkr.ecr.${{ inputs.aws_default_region }}.amazonaws.com
       run: |
         set -euo pipefail
+        if [[ ${CONDITIONAL_TAG} != '' ]]; then
+            docker tag ${LOCAL_IMAGE} ${ECR_HOSTNAME}/${CONDITIONAL_TAG}
+            docker push ${ECR_HOSTNAME}/${CONDITIONAL_TAG}
+        fi
         while IFS= read -r TAG; do
           if [ -z "$TAG" ]; then
             continue
@@ -66,7 +63,6 @@ runs:
           docker tag "${LOCAL_IMAGE}" "${ECR_HOSTNAME}/${TAG}"
           docker push "${ECR_HOSTNAME}/${TAG}"
         done <<< "$PUSH_TAGS"
-
     - name: ACR Tag and Push
       shell: bash
       if: ${{ inputs.azure_push == 'true' }}

.github/workflows/container-validation-backends.yml

@@ -69,11 +69,10 @@ jobs:
         uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955  # v4.3.0
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v3
-        with:
-          driver: docker
-      - name: Login to ECR
+      - name: Docker Login
         uses: ./.github/actions/docker-login
         with:
+          ngc_ci_access_token: ${{ secrets.NGC_CI_ACCESS_TOKEN }}
           aws_default_region: ${{ secrets.AWS_DEFAULT_REGION }}
           aws_account_id: ${{ secrets.AWS_ACCOUNT_ID }}
           azure_acr_hostname: ${{ secrets.AZURE_ACR_HOSTNAME }}
@@ -93,7 +92,6 @@ jobs:
         run: |
           cd deploy/cloud/operator
           docker build --target tester --progress=plain --build-arg DOCKER_PROXY=${ECR_HOSTNAME}/dockerhub/ .
-
       - name: Set up Go
         uses: actions/setup-go@44694675825211faa026b3c33043df3e48a5fa00 # v6.0.0
         with:
@@ -125,11 +123,7 @@ jobs:
           push_tags: ai-dynamo/dynamo:${{ github.sha }}-operator-${{ matrix.platform.arch }}
           aws_push: 'false'
           azure_push: 'true'
-          aws_account_id: ${{ secrets.AWS_ACCOUNT_ID }}
-          aws_default_region: ${{ secrets.AWS_DEFAULT_REGION }}
           azure_acr_hostname: ${{ secrets.AZURE_ACR_HOSTNAME }}
-          azure_acr_user: ${{ secrets.AZURE_ACR_USER }}
-          azure_acr_password: ${{ secrets.AZURE_ACR_PASSWORD }}
 
   vllm:
     needs: changed-files
@@ -149,6 +143,15 @@ jobs:
           echo ${K8S_NODE_NAME}
       - name: Checkout code
         uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955  # v4.3.0
+      - name: Docker Login
+        uses: ./.github/actions/docker-login
+        with:
+          ngc_ci_access_token: ${{ secrets.NGC_CI_ACCESS_TOKEN }}
+          aws_default_region: ${{ secrets.AWS_DEFAULT_REGION }}
+          aws_account_id: ${{ secrets.AWS_ACCOUNT_ID }}
+          azure_acr_hostname: ${{ secrets.AZURE_ACR_HOSTNAME }}
+          azure_acr_user: ${{ secrets.AZURE_ACR_USER }}
+          azure_acr_password: ${{ secrets.AZURE_ACR_PASSWORD }}
       - name: Build Container
         id: build-image
         uses: ./.github/actions/docker-build
@@ -160,7 +163,6 @@ jobs:
           runtime_image_tag: ${{ matrix.platform.arch == 'arm64' && '12.9.0-runtime-ubuntu24.04' || '' }}
           cuda_version: ${{ matrix.platform.arch == 'arm64' && '129' || '' }}
           torch_backend: ${{ matrix.platform.arch == 'arm64' && 'cu129' || '' }}
-          ngc_ci_access_token: ${{ secrets.NGC_CI_ACCESS_TOKEN }}
           ci_token: ${{ secrets.CI_TOKEN }}
           aws_default_region: ${{ secrets.AWS_DEFAULT_REGION }}
           sccache_s3_bucket: ${{ secrets.SCCACHE_S3_BUCKET }}
@@ -179,15 +181,12 @@ jobs:
         with:
           local_image: ${{ steps.build-image.outputs.image_tag }}
           push_tags: ai-dynamo/dynamo:${{ github.sha }}-vllm-${{ matrix.platform.arch }}
-          # OPS-1145: Switch aws_push to true
-          aws_push: 'false'
+          conditional_tag: ${{ github.ref_name == 'main' && format('ai-dynamo/dynamo:main-vllm-{0}', matrix.platform.arch) || '' }}
+          aws_push: 'true'
           azure_push: 'true'
           aws_account_id: ${{ secrets.AWS_ACCOUNT_ID }}
           aws_default_region: ${{ secrets.AWS_DEFAULT_REGION }}
           azure_acr_hostname: ${{ secrets.AZURE_ACR_HOSTNAME }}
-          azure_acr_user: ${{ secrets.AZURE_ACR_USER }}
-          azure_acr_password: ${{ secrets.AZURE_ACR_PASSWORD }}
-
       - name: Run tests
         if: ${{ matrix.platform.arch != 'arm64' }}
         uses: ./.github/actions/pytest
@@ -216,7 +215,15 @@ jobs:
           echo ${K8S_NODE_NAME}
       - name: Checkout repository
         uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955  # v4.3.0
-
+      - name: Docker Login
+        uses: ./.github/actions/docker-login
+        with:
+          ngc_ci_access_token: ${{ secrets.NGC_CI_ACCESS_TOKEN }}
+          aws_default_region: ${{ secrets.AWS_DEFAULT_REGION }}
+          aws_account_id: ${{ secrets.AWS_ACCOUNT_ID }}
+          azure_acr_hostname: ${{ secrets.AZURE_ACR_HOSTNAME }}
+          azure_acr_user: ${{ secrets.AZURE_ACR_USER }}
+          azure_acr_password: ${{ secrets.AZURE_ACR_PASSWORD }}
       - name: Build Container
         id: build-image
         uses: ./.github/actions/docker-build
@@ -224,35 +231,23 @@ jobs:
           framework: sglang
           target: runtime
           platform: 'linux/${{ matrix.platform.arch }}'
-          ngc_ci_access_token: ${{ secrets.NGC_CI_ACCESS_TOKEN }}
           ci_token: ${{ secrets.CI_TOKEN }}
           aws_default_region: ${{ secrets.AWS_DEFAULT_REGION }}
           sccache_s3_bucket: ${{ secrets.SCCACHE_S3_BUCKET }}
           aws_account_id: ${{ secrets.AWS_ACCOUNT_ID }}
           aws_access_key_id: ${{ secrets.AWS_ACCESS_KEY_ID }}
           aws_secret_access_key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
-
-      - name: Login to Container Registries
-        uses: ./.github/actions/docker-login
-        with:
-          azure_acr_hostname: ${{ secrets.AZURE_ACR_HOSTNAME }}
-          azure_acr_user: ${{ secrets.AZURE_ACR_USER }}
-          azure_acr_password: ${{ secrets.AZURE_ACR_PASSWORD }}
-          ngc_ci_access_token: ${{ secrets.NGC_CI_ACCESS_TOKEN }}
       - name: Docker Tag and Push
         uses: ./.github/actions/docker-tag-push
         with:
           local_image: ${{ steps.build-image.outputs.image_tag }}
           push_tags: ai-dynamo/dynamo:${{ github.sha }}-sglang-${{ matrix.platform.arch }}
-          # OPS-1145: Switch aws_push to true
-          aws_push: 'false'
+          conditional_tag: ${{ github.ref_name == 'main' && format('ai-dynamo/dynamo:main-sglang-{0}', matrix.platform.arch) || '' }}
+          aws_push: 'true'
           azure_push: 'true'
           aws_account_id: ${{ secrets.AWS_ACCOUNT_ID }}
           aws_default_region: ${{ secrets.AWS_DEFAULT_REGION }}
           azure_acr_hostname: ${{ secrets.AZURE_ACR_HOSTNAME }}
-          azure_acr_user: ${{ secrets.AZURE_ACR_USER }}
-          azure_acr_password: ${{ secrets.AZURE_ACR_PASSWORD }}
-
       - name: Run tests
         if: ${{ matrix.platform.arch != 'arm64' }}
         uses: ./.github/actions/pytest
@@ -281,7 +276,15 @@ jobs:
           echo ${K8S_NODE_NAME}
       - name: Checkout code
         uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955  # v4.3.0
-
+      - name: Docker Login
+        uses: ./.github/actions/docker-login
+        with:
+          ngc_ci_access_token: ${{ secrets.NGC_CI_ACCESS_TOKEN }}
+          aws_default_region: ${{ secrets.AWS_DEFAULT_REGION }}
+          aws_account_id: ${{ secrets.AWS_ACCOUNT_ID }}
+          azure_acr_hostname: ${{ secrets.AZURE_ACR_HOSTNAME }}
+          azure_acr_user: ${{ secrets.AZURE_ACR_USER }}
+          azure_acr_password: ${{ secrets.AZURE_ACR_PASSWORD }}
       - name: Build Container
         id: build-image
         uses: ./.github/actions/docker-build
@@ -289,35 +292,23 @@ jobs:
           framework: trtllm
           target: runtime
           platform: 'linux/${{ matrix.platform.arch }}'
-          ngc_ci_access_token: ${{ secrets.NGC_CI_ACCESS_TOKEN }}
           ci_token: ${{ secrets.CI_TOKEN }}
           aws_default_region: ${{ secrets.AWS_DEFAULT_REGION }}
           sccache_s3_bucket: ${{ secrets.SCCACHE_S3_BUCKET }}
           aws_account_id: ${{ secrets.AWS_ACCOUNT_ID }}
           aws_access_key_id: ${{ secrets.AWS_ACCESS_KEY_ID }}
           aws_secret_access_key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
-
-      - name: Login to Container Registries
-        uses: ./.github/actions/docker-login
-        with:
-          azure_acr_hostname: ${{ secrets.AZURE_ACR_HOSTNAME }}
-          azure_acr_user: ${{ secrets.AZURE_ACR_USER }}
-          azure_acr_password: ${{ secrets.AZURE_ACR_PASSWORD }}
-          ngc_ci_access_token: ${{ secrets.NGC_CI_ACCESS_TOKEN }}
       - name: Docker Tag and Push
         uses: ./.github/actions/docker-tag-push
         with:
           local_image: ${{ steps.build-image.outputs.image_tag }}
           push_tags: ai-dynamo/dynamo:${{ github.sha }}-trtllm-${{ matrix.platform.arch }}
-          # OPS-1145: Switch aws_push to true
-          aws_push: 'false'
+          conditional_tag: ${{ github.ref_name == 'main' && format('ai-dynamo/dynamo:main-trtllm-{0}', matrix.platform.arch) || '' }}
+          aws_push: 'true'
           azure_push: 'true'
           aws_account_id: ${{ secrets.AWS_ACCOUNT_ID }}
           aws_default_region: ${{ secrets.AWS_DEFAULT_REGION }}
           azure_acr_hostname: ${{ secrets.AZURE_ACR_HOSTNAME }}
-          azure_acr_user: ${{ secrets.AZURE_ACR_USER }}
-          azure_acr_password: ${{ secrets.AZURE_ACR_PASSWORD }}
-
       - name: Run tests
         if: ${{ matrix.platform.arch != 'arm64' }}
         uses: ./.github/actions/pytest

container/Dockerfile

 # syntax=docker/dockerfile:1.10.0
 # SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
+#
+# NOTE FOR dynamo_base AND wheel_builder STAGES:
+#
+# All changes to dynamo_base and wheel_builder stages should be replicated across
+# Dockerfile and Dockerfile.<framework> images.:
+#   - Dockerfile
+#   - Dockerfile.vllm
+#   - Dockerfile.sglang
+#   - Dockerfile.trtllm
+# This duplication was introduced purposely to quickly enable Docker layer caching and
+# deduplication. Please ensure these stages stay in sync until the duplication can be
+# addressed.
 
 ##################################
 ########## Build Arguments ########
@@ -36,9 +48,9 @@ ARG SCCACHE_BUCKET=""
 ARG SCCACHE_REGION=""
 
 # NIXL configuration
-ARG NIXL_UCX_REF=v1.19.0
-ARG NIXL_REF=0.7.1
-ARG NIXL_GDRCOPY_REF=v2.5.1
+ARG NIXL_UCX_REF
+ARG NIXL_REF
+ARG NIXL_GDRCOPY_REF
 
 ##################################
 ########## Base Image ############
@@ -202,10 +214,9 @@ ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET}} \
 RUN --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
     --mount=type=secret,id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY \
     export SCCACHE_S3_KEY_PREFIX="${SCCACHE_S3_KEY_PREFIX:-${ARCH}}" && \
-    CC=${USE_SCCACHE:+sccache gcc} && \
-    CXX=${USE_SCCACHE:+sccache g++} && \
-    export CC=${CC} && \
-    export CXX=${CXX} && \
+    export CMAKE_C_COMPILER_LAUNCHER="sccache" && \
+    export CMAKE_CXX_COMPILER_LAUNCHER="sccache" && \
+    export CMAKE_CUDA_COMPILER_LAUNCHER="sccache" && \
     cd /usr/local/src && \
      git clone https://github.com/openucx/ucx.git && \
      cd ucx && 			     \
@@ -236,6 +247,8 @@ RUN --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
 RUN --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
     --mount=type=secret,id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY \
     export SCCACHE_S3_KEY_PREFIX="${SCCACHE_S3_KEY_PREFIX:-${ARCH}}" && \
+    export CC=$(which gcc) && \
+    export CXX=$(which g++) && \
     source ${VIRTUAL_ENV}/bin/activate && \
     git clone --depth 1 --branch ${NIXL_REF} "https://github.com/ai-dynamo/nixl.git" && \
     cd nixl && \
@@ -261,6 +274,8 @@ RUN echo "$NIXL_LIB_DIR" > /etc/ld.so.conf.d/nixl.conf && \
 RUN --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
     --mount=type=secret,id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY \
     export SCCACHE_S3_KEY_PREFIX="${SCCACHE_S3_KEY_PREFIX:-${ARCH}}" && \
+    export CC=$(which gcc) && \
+    export CXX=$(which g++) && \
     cd /workspace/nixl && \
     uv build . --out-dir /opt/dynamo/dist/nixl --python $PYTHON_VERSION
 

container/Dockerfile.sglang

@@ -2,6 +2,18 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
+# NOTE FOR dynamo_base AND wheel_builder STAGES:
+#
+# All changes to dynamo_base and wheel_builder stages should be replicated across
+# Dockerfile and Dockerfile.<framework> images.:
+#   - Dockerfile
+#   - Dockerfile.vllm
+#   - Dockerfile.sglang
+#   - Dockerfile.trtllm
+# This duplication was introduced purposely to quickly enable Docker layer caching and
+# deduplication. Please ensure these stages stay in sync until the duplication can be
+# addressed.
+#
 # Throughout this file, we make certain paths group-writable because this allows
 # both the dynamo user (UID 1000) and Dev Container users (UID != 1000) to work
 # properly without needing slow chown -R operations (which can add 2-10 extra
@@ -26,19 +38,286 @@ ARG BASE_IMAGE_TAG
 ARG FRAMEWORK_IMAGE
 ARG FRAMEWORK_IMAGE_TAG
 ARG PYTHON_VERSION
+ARG ENABLE_KVBM
+ARG ENABLE_MEDIA_NIXL
+ARG CARGO_BUILD_JOBS
 ARG CUDA_VERSION
 
 ARG ARCH=amd64
 ARG ARCH_ALT=x86_64
-ARG CARGO_BUILD_JOBS
 
 # sccache configuration - inherit from base build
 ARG USE_SCCACHE
 ARG SCCACHE_BUCKET=""
 ARG SCCACHE_REGION=""
 
-ARG DYNAMO_BASE_IMAGE="dynamo:latest-none"
-FROM ${DYNAMO_BASE_IMAGE} AS dynamo_base
+# NIXL configuration
+ARG NIXL_UCX_REF
+ARG NIXL_REF
+ARG NIXL_GDRCOPY_REF
+
+##################################
+########## Base Image ############
+##################################
+
+FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG} AS dynamo_base
+
+ARG ARCH
+ARG ARCH_ALT
+
+USER root
+WORKDIR /opt/dynamo
+
+# Install uv package manager
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
+
+# Install NATS server
+ENV NATS_VERSION="v2.10.28"
+RUN --mount=type=cache,target=/var/cache/apt \
+    wget --tries=3 --waitretry=5 https://github.com/nats-io/nats-server/releases/download/${NATS_VERSION}/nats-server-${NATS_VERSION}-${ARCH}.deb && \
+    dpkg -i nats-server-${NATS_VERSION}-${ARCH}.deb && rm nats-server-${NATS_VERSION}-${ARCH}.deb
+
+# Install etcd
+ENV ETCD_VERSION="v3.5.21"
+RUN wget --tries=3 --waitretry=5 https://github.com/etcd-io/etcd/releases/download/$ETCD_VERSION/etcd-$ETCD_VERSION-linux-${ARCH}.tar.gz -O /tmp/etcd.tar.gz && \
+    mkdir -p /usr/local/bin/etcd && \
+    tar -xvf /tmp/etcd.tar.gz -C /usr/local/bin/etcd --strip-components=1 && \
+    rm /tmp/etcd.tar.gz
+ENV PATH=/usr/local/bin/etcd/:$PATH
+
+# Rust Setup
+# Rust environment setup
+ENV RUSTUP_HOME=/usr/local/rustup \
+    CARGO_HOME=/usr/local/cargo \
+    PATH=/usr/local/cargo/bin:$PATH \
+    RUST_VERSION=1.90.0
+
+# Define Rust target based on ARCH_ALT ARG
+ARG RUSTARCH=${ARCH_ALT}-unknown-linux-gnu
+
+# Install Rust
+RUN wget --tries=3 --waitretry=5 "https://static.rust-lang.org/rustup/archive/1.28.1/${RUSTARCH}/rustup-init" && \
+    chmod +x rustup-init && \
+    ./rustup-init -y --no-modify-path --profile minimal --default-toolchain $RUST_VERSION --default-host ${RUSTARCH} && \
+    rm rustup-init && \
+    chmod -R a+w $RUSTUP_HOME $CARGO_HOME
+
+
+##################################
+##### Wheel Build Image ##########
+##################################
+
+# Redeclare ARCH_ALT ARG so it's available for interpolation in the FROM instruction
+ARG ARCH_ALT
+
+FROM quay.io/pypa/manylinux_2_28_${ARCH_ALT} AS wheel_builder
+
+# Redeclare ARGs for this stage
+ARG ARCH
+ARG ARCH_ALT
+ARG CARGO_BUILD_JOBS
+
+WORKDIR /workspace
+
+# Copy CUDA from base stage
+COPY --from=dynamo_base /usr/local/cuda /usr/local/cuda
+COPY --from=dynamo_base /etc/ld.so.conf.d/hpcx.conf /etc/ld.so.conf.d/hpcx.conf
+
+# Set environment variables first so they can be used in COPY commands
+ENV CARGO_BUILD_JOBS=${CARGO_BUILD_JOBS:-16} \
+    RUSTUP_HOME=/usr/local/rustup \
+    CARGO_HOME=/usr/local/cargo \
+    CARGO_TARGET_DIR=/opt/dynamo/target \
+    PATH=/usr/local/cargo/bin:$PATH
+
+# Copy artifacts from base stage
+COPY --from=dynamo_base $RUSTUP_HOME $RUSTUP_HOME
+COPY --from=dynamo_base $CARGO_HOME $CARGO_HOME
+# Install system dependencies
+RUN yum groupinstall -y 'Development Tools' &&  \
+    dnf install -y almalinux-release-synergy && \
+    dnf config-manager --set-enabled powertools && \
+    dnf install -y \
+        # Build tools
+        cmake \
+        ninja-build \
+        clang-devel \
+        gcc-c++ \
+        flex \
+        wget \
+        # Kernel module build dependencies
+        dkms \
+        # Protobuf support
+        protobuf-compiler \
+        # RDMA/InfiniBand support (required for UCX build with --with-verbs)
+        libibverbs \
+        libibverbs-devel \
+        rdma-core \
+        rdma-core-devel \
+        libibumad \
+        libibumad-devel \
+        librdmacm-devel \
+        numactl-devel
+
+# Ensure a modern protoc is available (required for --experimental_allow_proto3_optional)
+RUN set -eux; \
+    PROTOC_VERSION=25.3; \
+    case "${ARCH_ALT}" in \
+      x86_64) PROTOC_ZIP="protoc-${PROTOC_VERSION}-linux-x86_64.zip" ;; \
+      aarch64) PROTOC_ZIP="protoc-${PROTOC_VERSION}-linux-aarch_64.zip" ;; \
+      *) echo "Unsupported architecture: ${ARCH_ALT}" >&2; exit 1 ;; \
+    esac; \
+    wget --tries=3 --waitretry=5 -O /tmp/protoc.zip "https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOC_VERSION}/${PROTOC_ZIP}"; \
+    rm -f /usr/local/bin/protoc /usr/bin/protoc; \
+    unzip -o /tmp/protoc.zip -d /usr/local bin/protoc include/*; \
+    chmod +x /usr/local/bin/protoc; \
+    ln -s /usr/local/bin/protoc /usr/bin/protoc; \
+    protoc --version
+
+# Point build tools explicitly at the modern protoc
+ENV PROTOC=/usr/local/bin/protoc
+
+ENV CUDA_PATH=/usr/local/cuda \
+    PATH=/usr/local/cuda/bin:$PATH \
+    LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/lib:/usr/local/lib64:$LD_LIBRARY_PATH \
+    NVIDIA_DRIVER_CAPABILITIES=video,compute,utility
+
+# Create virtual environment for building wheels
+ARG PYTHON_VERSION
+ENV VIRTUAL_ENV=/workspace/.venv
+RUN uv venv ${VIRTUAL_ENV} --python $PYTHON_VERSION && \
+    uv pip install --upgrade meson pybind11 patchelf maturin[patchelf]
+
+ARG NIXL_UCX_REF
+ARG NIXL_REF
+ARG NIXL_GDRCOPY_REF
+
+# Build and install gdrcopy
+RUN git clone --depth 1 --branch ${NIXL_GDRCOPY_REF} https://github.com/NVIDIA/gdrcopy.git && \
+    cd gdrcopy/packages && \
+    CUDA=/usr/local/cuda ./build-rpm-packages.sh && \
+    rpm -Uvh gdrcopy-kmod-*.el8.noarch.rpm && \
+    rpm -Uvh gdrcopy-*.el8.${ARCH_ALT}.rpm && \
+    rpm -Uvh gdrcopy-devel-*.el8.noarch.rpm
+
+# Install SCCACHE if requested
+ARG USE_SCCACHE
+ARG SCCACHE_BUCKET
+ARG SCCACHE_REGION
+COPY container/use-sccache.sh /tmp/use-sccache.sh
+RUN if [ "$USE_SCCACHE" = "true" ]; then \
+        /tmp/use-sccache.sh install; \
+    fi
+
+# Set SCCACHE environment variables
+ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET}} \
+    SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION}} \
+    RUSTC_WRAPPER=${USE_SCCACHE:+sccache}
+
+# Build and install UCX
+RUN --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
+    --mount=type=secret,id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY \
+    export SCCACHE_S3_KEY_PREFIX="${SCCACHE_S3_KEY_PREFIX:-${ARCH}}" && \
+    export CMAKE_C_COMPILER_LAUNCHER="sccache" && \
+    export CMAKE_CXX_COMPILER_LAUNCHER="sccache" && \
+    export CMAKE_CUDA_COMPILER_LAUNCHER="sccache" && \
+    cd /usr/local/src && \
+     git clone https://github.com/openucx/ucx.git && \
+     cd ucx && 			     \
+     git checkout $NIXL_UCX_REF &&	 \
+     ./autogen.sh &&      \
+     ./contrib/configure-release    \
+        --prefix=/usr/local/ucx     \
+        --enable-shared             \
+        --disable-static            \
+        --disable-doxygen-doc       \
+        --enable-optimizations      \
+        --enable-cma                \
+        --enable-devel-headers      \
+        --with-cuda=/usr/local/cuda \
+        --with-verbs                \
+        --with-dm                   \
+        --with-gdrcopy=/usr/local   \
+        --with-efa                  \
+        --enable-mt &&              \
+     make -j &&                      \
+     make -j install-strip &&        \
+     /tmp/use-sccache.sh show-stats "UCX" && \
+     echo "/usr/local/ucx/lib" > /etc/ld.so.conf.d/ucx.conf && \
+     echo "/usr/local/ucx/lib/ucx" >> /etc/ld.so.conf.d/ucx.conf && \
+     ldconfig
+
+# build and install nixl
+RUN --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
+    --mount=type=secret,id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY \
+    export SCCACHE_S3_KEY_PREFIX="${SCCACHE_S3_KEY_PREFIX:-${ARCH}}" && \
+    export CC=$(which gcc) && \
+    export CXX=$(which g++) && \
+    source ${VIRTUAL_ENV}/bin/activate && \
+    git clone --depth 1 --branch ${NIXL_REF} "https://github.com/ai-dynamo/nixl.git" && \
+    cd nixl && \
+    mkdir build && \
+    meson setup build/ --prefix=/opt/nvidia/nvda_nixl --buildtype=release \
+    -Dcudapath_lib="/usr/local/cuda/lib64" \
+    -Dcudapath_inc="/usr/local/cuda/include" \
+    -Ducx_path="/usr/local/ucx" && \
+    cd build && \
+    ninja && \
+    ninja install && \
+    /tmp/use-sccache.sh show-stats "NIXL"
+
+ENV NIXL_LIB_DIR=/opt/nvidia/nvda_nixl/lib64  \
+    NIXL_PLUGIN_DIR=/opt/nvidia/nvda_nixl/lib64/plugins \
+    NIXL_PREFIX=/opt/nvidia/nvda_nixl
+ENV LD_LIBRARY_PATH=${NIXL_LIB_DIR}:${NIXL_PLUGIN_DIR}:/usr/local/ucx/lib:/usr/local/ucx/lib/ucx:${LD_LIBRARY_PATH}
+
+RUN echo "$NIXL_LIB_DIR" > /etc/ld.so.conf.d/nixl.conf && \
+    echo "$NIXL_PLUGIN_DIR" >> /etc/ld.so.conf.d/nixl.conf && \
+    ldconfig
+
+RUN --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
+    --mount=type=secret,id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY \
+    export SCCACHE_S3_KEY_PREFIX="${SCCACHE_S3_KEY_PREFIX:-${ARCH}}" && \
+    export CC=$(which gcc) && \
+    export CXX=$(which g++) && \
+    cd /workspace/nixl && \
+    uv build . --out-dir /opt/dynamo/dist/nixl --python $PYTHON_VERSION
+
+# Copy source code (order matters for layer caching)
+COPY pyproject.toml README.md LICENSE Cargo.toml Cargo.lock rust-toolchain.toml hatch_build.py /opt/dynamo/
+COPY launch/ /opt/dynamo/launch/
+COPY lib/ /opt/dynamo/lib/
+COPY components/ /opt/dynamo/components/
+
+# Build dynamo wheels
+ARG ENABLE_KVBM
+RUN --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
+    --mount=type=secret,id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY \
+    export SCCACHE_S3_KEY_PREFIX=${SCCACHE_S3_KEY_PREFIX:-${ARCH}} && \
+    export CMAKE_C_COMPILER_LAUNCHER="sccache" && \
+    export CMAKE_CXX_COMPILER_LAUNCHER="sccache" && \
+    export RUSTC_WRAPPER="sccache" && \
+    source ${VIRTUAL_ENV}/bin/activate && \
+    cd /opt/dynamo && \
+    uv build --wheel --out-dir /opt/dynamo/dist && \
+    cd /opt/dynamo/lib/bindings/python && \
+    if [ "$ENABLE_MEDIA_NIXL" = "true" ]; then \
+        maturin build --release --features dynamo-llm/media-nixl --out /opt/dynamo/dist; \
+    else \
+        maturin build --release --out /opt/dynamo/dist; \
+    fi && \
+    if [ "$ENABLE_KVBM" = "true" ]; then \
+        cd /opt/dynamo/lib/bindings/kvbm && \
+        maturin build --release --out target/wheels && \
+        auditwheel repair \
+            --exclude libnixl.so \
+            --exclude libnixl_build.so \
+            --exclude libnixl_common.so \
+            --plat manylinux_2_28_${ARCH_ALT} \
+            --wheel-dir /opt/dynamo/dist \
+            target/wheels/*.whl; \
+    fi && \
+    /tmp/use-sccache.sh show-stats "Dynamo"
 
 ########################################################
 ########## Framework Development Image ################
@@ -187,11 +466,7 @@ fi
 # Set environment variables - they'll be empty strings if USE_SCCACHE=false
 ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET}} \
     SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION}} \
-    SCCACHE_S3_KEY_PREFIX=${USE_SCCACHE:+${ARCH}} \
-    RUSTC_WRAPPER=${USE_SCCACHE:+sccache} \
-    CMAKE_C_COMPILER_LAUNCHER=${USE_SCCACHE:+sccache} \
-    CMAKE_CXX_COMPILER_LAUNCHER=${USE_SCCACHE:+sccache} \
-    CMAKE_CUDA_COMPILER_LAUNCHER=${USE_SCCACHE:+sccache}
+    SCCACHE_S3_KEY_PREFIX=${USE_SCCACHE:+${ARCH}}
 
 WORKDIR /sgl-workspace
 
@@ -260,6 +535,9 @@ RUN --mount=type=cache,target=/var/cache/curl,uid=1000,gid=0 \
 RUN --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
     --mount=type=secret,id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY \
     export SCCACHE_S3_KEY_PREFIX=${SCCACHE_S3_KEY_PREFIX:-${ARCH}} && \
+    export CC=$(which gcc) && \
+    export CXX=$(which g++) && \
+    export CUDA_CXX=$(which nvcc) && \
     cd /sgl-workspace/nvshmem && \
     if [ "$GRACE_BLACKWELL" = true ]; then CUDA_ARCH="90;100;120"; else CUDA_ARCH="90"; fi && \
     NVSHMEM_SHMEM_SUPPORT=0 \
@@ -279,6 +557,9 @@ RUN --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
 RUN --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
     --mount=type=secret,id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY \
     export SCCACHE_S3_KEY_PREFIX=${SCCACHE_S3_KEY_PREFIX:-${ARCH}} && \
+    export CC=$(which gcc) && \
+    export CXX=$(which g++) && \
+    export CUDA_CXX=$(which nvcc) && \
     cd /sgl-workspace/nvshmem && \
     if [ "$GRACE_BLACKWELL" = true ]; then CUDA_ARCH="90;100;120"; else CUDA_ARCH="90"; fi && \
     NVSHMEM_SHMEM_SUPPORT=0 \
@@ -357,17 +638,20 @@ ${NIXL_PLUGIN_DIR}:\
 /usr/local/nvidia/lib64:\
 ${LD_LIBRARY_PATH}
 
-# Copy NATS and ETCD from dynamo_base, and UCX/NIXL
+# Copy NATS and ETCD from dynamo_base, and UCX/NIXL from wheel_builder
 COPY --from=dynamo_base /usr/bin/nats-server /usr/bin/nats-server
 COPY --from=dynamo_base /usr/local/bin/etcd/ /usr/local/bin/etcd/
-COPY --from=dynamo_base /usr/local/ucx /usr/local/ucx
-COPY --from=dynamo_base $NIXL_PREFIX $NIXL_PREFIX
+COPY --from=wheel_builder /usr/local/ucx /usr/local/ucx
+COPY --chown=dynamo: --from=wheel_builder $NIXL_PREFIX $NIXL_PREFIX
+COPY --chown=dynamo: --from=wheel_builder /opt/nvidia/nvda_nixl/lib64/. ${NIXL_LIB_DIR}/
+COPY --chown=dynamo: --from=wheel_builder /opt/dynamo/dist/nixl/ /opt/dynamo/wheelhouse/nixl/
+COPY --chown=dynamo: --from=wheel_builder /workspace/nixl/build/src/bindings/python/nixl-meta/nixl-*.whl /opt/dynamo/wheelhouse/nixl/
 ENV PATH=/usr/local/bin/etcd/:/usr/local/cuda/nvvm/bin:${HOME}/.local/bin:$PATH
 
 # Install Dynamo wheels from dynamo_base wheelhouse
 # 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 benchmarks/ /opt/dynamo/benchmarks/
-COPY --chmod=775 --chown=dynamo:0 --from=dynamo_base /opt/dynamo/wheelhouse/ /opt/dynamo/wheelhouse/
+COPY --chmod=775 --chown=dynamo:0 --from=wheel_builder /opt/dynamo/dist/*.whl /opt/dynamo/wheelhouse/
 RUN python3 -m pip install \
     /opt/dynamo/wheelhouse/ai_dynamo_runtime*.whl \
     /opt/dynamo/wheelhouse/ai_dynamo*any.whl \

container/Dockerfile.trtllm

 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
+# NOTE FOR dynamo_base AND wheel_builder STAGES:
+#
+# All changes to dynamo_base and wheel_builder stages should be replicated across
+# Dockerfile and Dockerfile.<framework> images.:
+#   - Dockerfile
+#   - Dockerfile.vllm
+#   - Dockerfile.sglang
+#   - Dockerfile.trtllm
+# This duplication was introduced purposely to quickly enable Docker layer caching and
+# deduplication. Please ensure these stages stay in sync until the duplication can be
+# addressed.
+#
 # Throughout this file, we make certain paths group-writable because this allows
 # both the dynamo user (UID 1000) and Dev Container users (UID != 1000) to work
 # properly without needing slow chown -R operations (which can add 2-10 extra
@@ -23,6 +35,8 @@ ARG BASE_IMAGE_TAG
 
 ARG PYTHON_VERSION
 ARG ENABLE_KVBM
+ARG ENABLE_MEDIA_NIXL
+ARG CARGO_BUILD_JOBS
 
 ARG PYTORCH_BASE_IMAGE="nvcr.io/nvidia/pytorch"
 ARG PYTORCH_BASE_IMAGE_TAG="25.10-py3"
@@ -35,6 +49,16 @@ ARG TENSORRTLLM_PIP_WHEEL="tensorrt-llm"
 ARG TENSORRTLLM_INDEX_URL="https://pypi.nvidia.com/"
 ARG GITHUB_TRTLLM_COMMIT
 
+# SCCACHE configuration
+ARG USE_SCCACHE
+ARG SCCACHE_BUCKET=""
+ARG SCCACHE_REGION=""
+
+# NIXL configuration
+ARG NIXL_UCX_REF
+ARG NIXL_REF
+ARG NIXL_GDRCOPY_REF
+
 # Define general architecture ARGs for supporting both x86 and aarch64 builds.
 #   ARCH: Used for package suffixes (e.g., amd64, arm64)
 #   ARCH_ALT: Used for Rust targets, manylinux suffix (e.g., x86_64, aarch64)
@@ -50,12 +74,272 @@ ARG GITHUB_TRTLLM_COMMIT
 ARG ARCH=amd64
 ARG ARCH_ALT=x86_64
 
-ARG DYNAMO_BASE_IMAGE="dynamo:latest-none"
-FROM ${DYNAMO_BASE_IMAGE} AS dynamo_base
-
 # Copy artifacts from NGC PyTorch image
 FROM ${PYTORCH_BASE_IMAGE}:${PYTORCH_BASE_IMAGE_TAG} AS pytorch_base
 
+##################################
+########## Base Image ############
+##################################
+
+FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG} AS dynamo_base
+
+ARG ARCH
+ARG ARCH_ALT
+
+USER root
+WORKDIR /opt/dynamo
+
+# Install uv package manager
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
+
+# Install NATS server
+ENV NATS_VERSION="v2.10.28"
+RUN --mount=type=cache,target=/var/cache/apt \
+    wget --tries=3 --waitretry=5 https://github.com/nats-io/nats-server/releases/download/${NATS_VERSION}/nats-server-${NATS_VERSION}-${ARCH}.deb && \
+    dpkg -i nats-server-${NATS_VERSION}-${ARCH}.deb && rm nats-server-${NATS_VERSION}-${ARCH}.deb
+
+# Install etcd
+ENV ETCD_VERSION="v3.5.21"
+RUN wget --tries=3 --waitretry=5 https://github.com/etcd-io/etcd/releases/download/$ETCD_VERSION/etcd-$ETCD_VERSION-linux-${ARCH}.tar.gz -O /tmp/etcd.tar.gz && \
+    mkdir -p /usr/local/bin/etcd && \
+    tar -xvf /tmp/etcd.tar.gz -C /usr/local/bin/etcd --strip-components=1 && \
+    rm /tmp/etcd.tar.gz
+ENV PATH=/usr/local/bin/etcd/:$PATH
+
+# Rust Setup
+# Rust environment setup
+ENV RUSTUP_HOME=/usr/local/rustup \
+    CARGO_HOME=/usr/local/cargo \
+    PATH=/usr/local/cargo/bin:$PATH \
+    RUST_VERSION=1.90.0
+
+# Define Rust target based on ARCH_ALT ARG
+ARG RUSTARCH=${ARCH_ALT}-unknown-linux-gnu
+
+# Install Rust
+RUN wget --tries=3 --waitretry=5 "https://static.rust-lang.org/rustup/archive/1.28.1/${RUSTARCH}/rustup-init" && \
+    chmod +x rustup-init && \
+    ./rustup-init -y --no-modify-path --profile minimal --default-toolchain $RUST_VERSION --default-host ${RUSTARCH} && \
+    rm rustup-init && \
+    chmod -R a+w $RUSTUP_HOME $CARGO_HOME
+
+
+##################################
+##### Wheel Build Image ##########
+##################################
+
+# Redeclare ARCH_ALT ARG so it's available for interpolation in the FROM instruction
+ARG ARCH_ALT
+
+FROM quay.io/pypa/manylinux_2_28_${ARCH_ALT} AS wheel_builder
+
+# Redeclare ARGs for this stage
+ARG ARCH
+ARG ARCH_ALT
+ARG CARGO_BUILD_JOBS
+
+WORKDIR /workspace
+
+# Copy CUDA from base stage
+COPY --from=dynamo_base /usr/local/cuda /usr/local/cuda
+COPY --from=dynamo_base /etc/ld.so.conf.d/hpcx.conf /etc/ld.so.conf.d/hpcx.conf
+
+# Set environment variables first so they can be used in COPY commands
+ENV CARGO_BUILD_JOBS=${CARGO_BUILD_JOBS:-16} \
+    RUSTUP_HOME=/usr/local/rustup \
+    CARGO_HOME=/usr/local/cargo \
+    CARGO_TARGET_DIR=/opt/dynamo/target \
+    PATH=/usr/local/cargo/bin:$PATH
+
+# Copy artifacts from base stage
+COPY --from=dynamo_base $RUSTUP_HOME $RUSTUP_HOME
+COPY --from=dynamo_base $CARGO_HOME $CARGO_HOME
+# Install system dependencies
+RUN yum groupinstall -y 'Development Tools' &&  \
+    dnf install -y almalinux-release-synergy && \
+    dnf config-manager --set-enabled powertools && \
+    dnf install -y \
+        # Build tools
+        cmake \
+        ninja-build \
+        clang-devel \
+        gcc-c++ \
+        flex \
+        wget \
+        # Kernel module build dependencies
+        dkms \
+        # Protobuf support
+        protobuf-compiler \
+        # RDMA/InfiniBand support (required for UCX build with --with-verbs)
+        libibverbs \
+        libibverbs-devel \
+        rdma-core \
+        rdma-core-devel \
+        libibumad \
+        libibumad-devel \
+        librdmacm-devel \
+        numactl-devel
+
+# Ensure a modern protoc is available (required for --experimental_allow_proto3_optional)
+RUN set -eux; \
+    PROTOC_VERSION=25.3; \
+    case "${ARCH_ALT}" in \
+      x86_64) PROTOC_ZIP="protoc-${PROTOC_VERSION}-linux-x86_64.zip" ;; \
+      aarch64) PROTOC_ZIP="protoc-${PROTOC_VERSION}-linux-aarch_64.zip" ;; \
+      *) echo "Unsupported architecture: ${ARCH_ALT}" >&2; exit 1 ;; \
+    esac; \
+    wget --tries=3 --waitretry=5 -O /tmp/protoc.zip "https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOC_VERSION}/${PROTOC_ZIP}"; \
+    rm -f /usr/local/bin/protoc /usr/bin/protoc; \
+    unzip -o /tmp/protoc.zip -d /usr/local bin/protoc include/*; \
+    chmod +x /usr/local/bin/protoc; \
+    ln -s /usr/local/bin/protoc /usr/bin/protoc; \
+    protoc --version
+
+# Point build tools explicitly at the modern protoc
+ENV PROTOC=/usr/local/bin/protoc
+
+ENV CUDA_PATH=/usr/local/cuda \
+    PATH=/usr/local/cuda/bin:$PATH \
+    LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/lib:/usr/local/lib64:$LD_LIBRARY_PATH \
+    NVIDIA_DRIVER_CAPABILITIES=video,compute,utility
+
+# Create virtual environment for building wheels
+ARG PYTHON_VERSION
+ENV VIRTUAL_ENV=/workspace/.venv
+RUN uv venv ${VIRTUAL_ENV} --python $PYTHON_VERSION && \
+    uv pip install --upgrade meson pybind11 patchelf maturin[patchelf]
+
+ARG NIXL_UCX_REF
+ARG NIXL_REF
+ARG NIXL_GDRCOPY_REF
+
+# Build and install gdrcopy
+RUN git clone --depth 1 --branch ${NIXL_GDRCOPY_REF} https://github.com/NVIDIA/gdrcopy.git && \
+    cd gdrcopy/packages && \
+    CUDA=/usr/local/cuda ./build-rpm-packages.sh && \
+    rpm -Uvh gdrcopy-kmod-*.el8.noarch.rpm && \
+    rpm -Uvh gdrcopy-*.el8.${ARCH_ALT}.rpm && \
+    rpm -Uvh gdrcopy-devel-*.el8.noarch.rpm
+
+# Install SCCACHE if requested
+ARG USE_SCCACHE
+ARG SCCACHE_BUCKET
+ARG SCCACHE_REGION
+COPY container/use-sccache.sh /tmp/use-sccache.sh
+RUN if [ "$USE_SCCACHE" = "true" ]; then \
+        /tmp/use-sccache.sh install; \
+    fi
+
+# Set SCCACHE environment variables
+ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET}} \
+    SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION}} \
+    RUSTC_WRAPPER=${USE_SCCACHE:+sccache}
+
+# Build and install UCX
+RUN --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
+    --mount=type=secret,id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY \
+    export SCCACHE_S3_KEY_PREFIX="${SCCACHE_S3_KEY_PREFIX:-${ARCH}}" && \
+    export CMAKE_C_COMPILER_LAUNCHER="sccache" && \
+    export CMAKE_CXX_COMPILER_LAUNCHER="sccache" && \
+    export CMAKE_CUDA_COMPILER_LAUNCHER="sccache" && \
+    cd /usr/local/src && \
+     git clone https://github.com/openucx/ucx.git && \
+     cd ucx && 			     \
+     git checkout $NIXL_UCX_REF &&	 \
+     ./autogen.sh &&      \
+     ./contrib/configure-release    \
+        --prefix=/usr/local/ucx     \
+        --enable-shared             \
+        --disable-static            \
+        --disable-doxygen-doc       \
+        --enable-optimizations      \
+        --enable-cma                \
+        --enable-devel-headers      \
+        --with-cuda=/usr/local/cuda \
+        --with-verbs                \
+        --with-dm                   \
+        --with-gdrcopy=/usr/local   \
+        --with-efa                  \
+        --enable-mt &&              \
+     make -j &&                      \
+     make -j install-strip &&        \
+     /tmp/use-sccache.sh show-stats "UCX" && \
+     echo "/usr/local/ucx/lib" > /etc/ld.so.conf.d/ucx.conf && \
+     echo "/usr/local/ucx/lib/ucx" >> /etc/ld.so.conf.d/ucx.conf && \
+     ldconfig
+
+# build and install nixl
+RUN --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
+    --mount=type=secret,id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY \
+    export SCCACHE_S3_KEY_PREFIX="${SCCACHE_S3_KEY_PREFIX:-${ARCH}}" && \
+    export CC=$(which gcc) && \
+    export CXX=$(which g++) && \
+    source ${VIRTUAL_ENV}/bin/activate && \
+    git clone --depth 1 --branch ${NIXL_REF} "https://github.com/ai-dynamo/nixl.git" && \
+    cd nixl && \
+    mkdir build && \
+    meson setup build/ --prefix=/opt/nvidia/nvda_nixl --buildtype=release \
+    -Dcudapath_lib="/usr/local/cuda/lib64" \
+    -Dcudapath_inc="/usr/local/cuda/include" \
+    -Ducx_path="/usr/local/ucx" && \
+    cd build && \
+    ninja && \
+    ninja install && \
+    /tmp/use-sccache.sh show-stats "NIXL"
+
+ENV NIXL_LIB_DIR=/opt/nvidia/nvda_nixl/lib64  \
+    NIXL_PLUGIN_DIR=/opt/nvidia/nvda_nixl/lib64/plugins \
+    NIXL_PREFIX=/opt/nvidia/nvda_nixl
+ENV LD_LIBRARY_PATH=${NIXL_LIB_DIR}:${NIXL_PLUGIN_DIR}:/usr/local/ucx/lib:/usr/local/ucx/lib/ucx:${LD_LIBRARY_PATH}
+
+RUN echo "$NIXL_LIB_DIR" > /etc/ld.so.conf.d/nixl.conf && \
+    echo "$NIXL_PLUGIN_DIR" >> /etc/ld.so.conf.d/nixl.conf && \
+    ldconfig
+
+RUN --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
+    --mount=type=secret,id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY \
+    export CC=$(which gcc) && \
+    export CXX=$(which g++) && \
+    export SCCACHE_S3_KEY_PREFIX="${SCCACHE_S3_KEY_PREFIX:-${ARCH}}" && \
+    cd /workspace/nixl && \
+    uv build . --out-dir /opt/dynamo/dist/nixl --python $PYTHON_VERSION
+
+# Copy source code (order matters for layer caching)
+COPY pyproject.toml README.md LICENSE Cargo.toml Cargo.lock rust-toolchain.toml hatch_build.py /opt/dynamo/
+COPY launch/ /opt/dynamo/launch/
+COPY lib/ /opt/dynamo/lib/
+COPY components/ /opt/dynamo/components/
+
+# Build dynamo wheels
+ARG ENABLE_KVBM
+RUN --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
+    --mount=type=secret,id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY \
+    export SCCACHE_S3_KEY_PREFIX=${SCCACHE_S3_KEY_PREFIX:-${ARCH}} && \
+    export CMAKE_C_COMPILER_LAUNCHER="sccache" && \
+    export CMAKE_CXX_COMPILER_LAUNCHER="sccache" && \
+    export RUSTC_WRAPPER="sccache" && \
+    source ${VIRTUAL_ENV}/bin/activate && \
+    cd /opt/dynamo && \
+    uv build --wheel --out-dir /opt/dynamo/dist && \
+    cd /opt/dynamo/lib/bindings/python && \
+    if [ "$ENABLE_MEDIA_NIXL" = "true" ]; then \
+        maturin build --release --features dynamo-llm/media-nixl --out /opt/dynamo/dist; \
+    else \
+        maturin build --release --out /opt/dynamo/dist; \
+    fi && \
+    if [ "$ENABLE_KVBM" = "true" ]; then \
+        cd /opt/dynamo/lib/bindings/kvbm && \
+        maturin build --release --out target/wheels && \
+        auditwheel repair \
+            --exclude libnixl.so \
+            --exclude libnixl_build.so \
+            --exclude libnixl_common.so \
+            --plat manylinux_2_28_${ARCH_ALT} \
+            --wheel-dir /opt/dynamo/dist \
+            target/wheels/*.whl; \
+    fi && \
+    /tmp/use-sccache.sh show-stats "Dynamo"
+
 ##################################################
 ########## Framework Builder Stage ##############
 ##################################################
@@ -329,8 +613,11 @@ COPY --chmod=775 --chown=dynamo:0 --from=framework ${VIRTUAL_ENV} ${VIRTUAL_ENV}
 # Copy UCX from framework image as plugin for NIXL
 # Copy NIXL source from framework image
 # Copy dynamo wheels for gitlab artifacts (read-only, no group-write needed)
-COPY --chown=dynamo:0 --from=dynamo_base /usr/local/ucx /usr/local/ucx
-COPY --chown=dynamo:0 --from=dynamo_base $NIXL_PREFIX $NIXL_PREFIX
+COPY --chown=dynamo: --from=wheel_builder /usr/local/ucx /usr/local/ucx
+COPY --chown=dynamo: --from=wheel_builder $NIXL_PREFIX $NIXL_PREFIX
+COPY --chown=dynamo: --from=wheel_builder /opt/nvidia/nvda_nixl/lib64/. ${NIXL_LIB_DIR}/
+COPY --chown=dynamo: --from=wheel_builder /opt/dynamo/dist/nixl/ /opt/dynamo/wheelhouse/nixl/
+COPY --chown=dynamo: --from=wheel_builder /workspace/nixl/build/src/bindings/python/nixl-meta/nixl-*.whl /opt/dynamo/wheelhouse/nixl/
 
 ENV TENSORRT_LIB_DIR=/usr/local/tensorrt/targets/${ARCH_ALT}-linux-gnu/lib
 ENV PATH="/usr/local/ucx/bin:${VIRTUAL_ENV}/bin:/opt/hpcx/ompi/bin:/usr/local/bin/etcd/:/usr/local/cuda/bin:/usr/local/cuda/nvvm/bin:$PATH"
@@ -353,7 +640,7 @@ COPY --chmod=775 --chown=dynamo:0 benchmarks/ /workspace/benchmarks/
 # Install dynamo, NIXL, and dynamo-specific dependencies
 # Pattern: COPY --chmod=775 <path>; chmod g+w <path> done later as root because COPY --chmod only affects <path>/*, not <path>
 ARG ENABLE_KVBM
-COPY --chmod=775 --chown=dynamo:0 --from=dynamo_base /opt/dynamo/wheelhouse/ /opt/dynamo/wheelhouse/
+COPY --chmod=775 --chown=dynamo:0 --from=wheel_builder /opt/dynamo/dist/*.whl /opt/dynamo/wheelhouse/
 RUN uv pip install \
       --no-cache \
       /opt/dynamo/wheelhouse/ai_dynamo_runtime*.whl \

container/Dockerfile.vllm

@@ -2,6 +2,18 @@
 # SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
+# NOTE FOR dynamo_base AND wheel_builder STAGES:
+#
+# All changes to dynamo_base and wheel_builder stages should be replicated across
+# Dockerfile and Dockerfile.<framework> images.:
+#   - Dockerfile
+#   - Dockerfile.vllm
+#   - Dockerfile.sglang
+#   - Dockerfile.trtllm
+# This duplication was introduced purposely to quickly enable Docker layer caching and
+# deduplication. Please ensure these stages stay in sync until the duplication can be
+# addressed.
+#
 # Throughout this file, we make certain paths group-writable because this allows
 # both the dynamo user (UID 1000) and Dev Container users (UID != 1000) to work
 # properly without needing slow chown -R operations (which can add 2-10 extra
@@ -17,13 +29,43 @@
 # 2. COPY --chmod=775       - Sets permissions on copied children (not destination)
 # 3. chmod g+w (no -R)      - Fixes destination dirs only (milliseconds vs minutes)
 
-# This section contains build arguments that are common and shared with
-# the plain Dockerfile, so they should NOT have a default. The source of truth is from build.sh.
+##################################
+########## Build Arguments ########
+##################################
+
+# This section contains build arguments that are common and shared across various
+# Dockerfile.<frameworks>, so they should NOT have a default. The source of truth is from build.sh.
+
 ARG BASE_IMAGE
 ARG BASE_IMAGE_TAG
 
 ARG PYTHON_VERSION
 ARG ENABLE_KVBM
+ARG ENABLE_MEDIA_NIXL
+ARG CARGO_BUILD_JOBS
+
+# Define general architecture ARGs for supporting both x86 and aarch64 builds.
+#   ARCH: Used for package suffixes (e.g., amd64, arm64)
+#   ARCH_ALT: Used for Rust targets, manylinux suffix (e.g., x86_64, aarch64)
+#
+# Default values are for x86/amd64:
+#   --build-arg ARCH=amd64 --build-arg ARCH_ALT=x86_64
+#
+# For arm64/aarch64, build with:
+#   --build-arg ARCH=arm64 --build-arg ARCH_ALT=aarch64
+#TODO OPS-592: Leverage uname -m to determine ARCH instead of passing it as an arg
+ARG ARCH=amd64
+ARG ARCH_ALT=x86_64
+
+# SCCACHE configuration
+ARG USE_SCCACHE
+ARG SCCACHE_BUCKET=""
+ARG SCCACHE_REGION=""
+
+# NIXL configuration
+ARG NIXL_UCX_REF
+ARG NIXL_REF
+ARG NIXL_GDRCOPY_REF
 
 ARG RUNTIME_IMAGE="nvcr.io/nvidia/cuda"
 ARG RUNTIME_IMAGE_TAG="12.9.0-runtime-ubuntu24.04"
@@ -38,31 +80,267 @@ ARG FLASHINF_REF="v0.5.3"
 ARG DEEPGEMM_REF=""
 ARG LMCACHE_REF="0.3.10"
 
-# sccache configuration - inherit from base build
-ARG USE_SCCACHE
-ARG SCCACHE_BUCKET=""
-ARG SCCACHE_REGION=""
+##################################
+########## Base Image ############
+##################################
 
-# Define general architecture ARGs for supporting both x86 and aarch64 builds.
-#   ARCH: Used for package suffixes (e.g., amd64, arm64)
-#   ARCH_ALT: Used for Rust targets, manylinux suffix (e.g., x86_64, aarch64)
-#
-# Default values are for x86/amd64:
-#   --build-arg ARCH=amd64 --build-arg ARCH_ALT=x86_64
-#
-# For arm64/aarch64, build with:
-#   --build-arg ARCH=arm64 --build-arg ARCH_ALT=aarch64
-#
-# NOTE: There isn't an easy way to define one of these values based on the other value
-# without adding if statements everywhere, so just define both as ARGs for now.
-ARG ARCH=amd64
-ARG ARCH_ALT=x86_64
+FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG} AS dynamo_base
 
-ARG DYNAMO_BASE_IMAGE="dynamo:latest-none"
-FROM ${DYNAMO_BASE_IMAGE} AS dynamo_base
+ARG ARCH
+ARG ARCH_ALT
+
+USER root
+WORKDIR /opt/dynamo
+
+# Install uv package manager
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
+
+# Install NATS server
+ENV NATS_VERSION="v2.10.28"
+RUN --mount=type=cache,target=/var/cache/apt \
+    wget --tries=3 --waitretry=5 https://github.com/nats-io/nats-server/releases/download/${NATS_VERSION}/nats-server-${NATS_VERSION}-${ARCH}.deb && \
+    dpkg -i nats-server-${NATS_VERSION}-${ARCH}.deb && rm nats-server-${NATS_VERSION}-${ARCH}.deb
+
+# Install etcd
+ENV ETCD_VERSION="v3.5.21"
+RUN wget --tries=3 --waitretry=5 https://github.com/etcd-io/etcd/releases/download/$ETCD_VERSION/etcd-$ETCD_VERSION-linux-${ARCH}.tar.gz -O /tmp/etcd.tar.gz && \
+    mkdir -p /usr/local/bin/etcd && \
+    tar -xvf /tmp/etcd.tar.gz -C /usr/local/bin/etcd --strip-components=1 && \
+    rm /tmp/etcd.tar.gz
+ENV PATH=/usr/local/bin/etcd/:$PATH
+
+# Rust Setup
+# Rust environment setup
+ENV RUSTUP_HOME=/usr/local/rustup \
+    CARGO_HOME=/usr/local/cargo \
+    PATH=/usr/local/cargo/bin:$PATH \
+    RUST_VERSION=1.90.0
 
-# Copy cuda tools and libs from base image
-FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG} AS base
+# Define Rust target based on ARCH_ALT ARG
+ARG RUSTARCH=${ARCH_ALT}-unknown-linux-gnu
+
+# Install Rust
+RUN wget --tries=3 --waitretry=5 "https://static.rust-lang.org/rustup/archive/1.28.1/${RUSTARCH}/rustup-init" && \
+    chmod +x rustup-init && \
+    ./rustup-init -y --no-modify-path --profile minimal --default-toolchain $RUST_VERSION --default-host ${RUSTARCH} && \
+    rm rustup-init && \
+    chmod -R a+w $RUSTUP_HOME $CARGO_HOME
+
+
+##################################
+##### Wheel Build Image ##########
+##################################
+
+# Redeclare ARCH_ALT ARG so it's available for interpolation in the FROM instruction
+ARG ARCH_ALT
+
+FROM quay.io/pypa/manylinux_2_28_${ARCH_ALT} AS wheel_builder
+
+# Redeclare ARGs for this stage
+ARG ARCH
+ARG ARCH_ALT
+ARG CARGO_BUILD_JOBS
+
+WORKDIR /workspace
+
+# Copy CUDA from base stage
+COPY --from=dynamo_base /usr/local/cuda /usr/local/cuda
+COPY --from=dynamo_base /etc/ld.so.conf.d/hpcx.conf /etc/ld.so.conf.d/hpcx.conf
+
+# Set environment variables first so they can be used in COPY commands
+ENV CARGO_BUILD_JOBS=${CARGO_BUILD_JOBS:-16} \
+    RUSTUP_HOME=/usr/local/rustup \
+    CARGO_HOME=/usr/local/cargo \
+    CARGO_TARGET_DIR=/opt/dynamo/target \
+    PATH=/usr/local/cargo/bin:$PATH
+
+# Copy artifacts from base stage
+COPY --from=dynamo_base $RUSTUP_HOME $RUSTUP_HOME
+COPY --from=dynamo_base $CARGO_HOME $CARGO_HOME
+# Install system dependencies
+RUN yum groupinstall -y 'Development Tools' &&  \
+    dnf install -y almalinux-release-synergy && \
+    dnf config-manager --set-enabled powertools && \
+    dnf install -y \
+        # Build tools
+        cmake \
+        ninja-build \
+        clang-devel \
+        gcc-c++ \
+        flex \
+        wget \
+        # Kernel module build dependencies
+        dkms \
+        # Protobuf support
+        protobuf-compiler \
+        # RDMA/InfiniBand support (required for UCX build with --with-verbs)
+        libibverbs \
+        libibverbs-devel \
+        rdma-core \
+        rdma-core-devel \
+        libibumad \
+        libibumad-devel \
+        librdmacm-devel \
+        numactl-devel
+
+# Ensure a modern protoc is available (required for --experimental_allow_proto3_optional)
+RUN set -eux; \
+    PROTOC_VERSION=25.3; \
+    case "${ARCH_ALT}" in \
+      x86_64) PROTOC_ZIP="protoc-${PROTOC_VERSION}-linux-x86_64.zip" ;; \
+      aarch64) PROTOC_ZIP="protoc-${PROTOC_VERSION}-linux-aarch_64.zip" ;; \
+      *) echo "Unsupported architecture: ${ARCH_ALT}" >&2; exit 1 ;; \
+    esac; \
+    wget --tries=3 --waitretry=5 -O /tmp/protoc.zip "https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOC_VERSION}/${PROTOC_ZIP}"; \
+    rm -f /usr/local/bin/protoc /usr/bin/protoc; \
+    unzip -o /tmp/protoc.zip -d /usr/local bin/protoc include/*; \
+    chmod +x /usr/local/bin/protoc; \
+    ln -s /usr/local/bin/protoc /usr/bin/protoc; \
+    protoc --version
+
+# Point build tools explicitly at the modern protoc
+ENV PROTOC=/usr/local/bin/protoc
+
+ENV CUDA_PATH=/usr/local/cuda \
+    PATH=/usr/local/cuda/bin:$PATH \
+    LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/lib:/usr/local/lib64:$LD_LIBRARY_PATH \
+    NVIDIA_DRIVER_CAPABILITIES=video,compute,utility
+
+# Create virtual environment for building wheels
+ARG PYTHON_VERSION
+ENV VIRTUAL_ENV=/workspace/.venv
+RUN uv venv ${VIRTUAL_ENV} --python $PYTHON_VERSION && \
+    uv pip install --upgrade meson pybind11 patchelf maturin[patchelf]
+
+ARG NIXL_UCX_REF
+ARG NIXL_REF
+ARG NIXL_GDRCOPY_REF
+
+# Build and install gdrcopy
+RUN git clone --depth 1 --branch ${NIXL_GDRCOPY_REF} https://github.com/NVIDIA/gdrcopy.git && \
+    cd gdrcopy/packages && \
+    CUDA=/usr/local/cuda ./build-rpm-packages.sh && \
+    rpm -Uvh gdrcopy-kmod-*.el8.noarch.rpm && \
+    rpm -Uvh gdrcopy-*.el8.${ARCH_ALT}.rpm && \
+    rpm -Uvh gdrcopy-devel-*.el8.noarch.rpm
+
+# Install SCCACHE if requested
+ARG USE_SCCACHE
+ARG SCCACHE_BUCKET
+ARG SCCACHE_REGION
+COPY container/use-sccache.sh /tmp/use-sccache.sh
+RUN if [ "$USE_SCCACHE" = "true" ]; then \
+        /tmp/use-sccache.sh install; \
+    fi
+
+# Set SCCACHE environment variables
+ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET}} \
+    SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION}}
+
+# Build and install UCX
+RUN --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
+    --mount=type=secret,id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY \
+    export SCCACHE_S3_KEY_PREFIX="${SCCACHE_S3_KEY_PREFIX:-${ARCH}}" && \
+    export CMAKE_C_COMPILER_LAUNCHER="sccache" && \
+    export CMAKE_CXX_COMPILER_LAUNCHER="sccache" && \
+    export CMAKE_CUDA_COMPILER_LAUNCHER="sccache" && \
+    cd /usr/local/src && \
+     git clone https://github.com/openucx/ucx.git && \
+     cd ucx && 			     \
+     git checkout $NIXL_UCX_REF &&	 \
+     ./autogen.sh &&      \
+     ./contrib/configure-release    \
+        --prefix=/usr/local/ucx     \
+        --enable-shared             \
+        --disable-static            \
+        --disable-doxygen-doc       \
+        --enable-optimizations      \
+        --enable-cma                \
+        --enable-devel-headers      \
+        --with-cuda=/usr/local/cuda \
+        --with-verbs                \
+        --with-dm                   \
+        --with-gdrcopy=/usr/local   \
+        --with-efa                  \
+        --enable-mt &&              \
+     make -j &&                      \
+     make -j install-strip &&        \
+     /tmp/use-sccache.sh show-stats "UCX" && \
+     echo "/usr/local/ucx/lib" > /etc/ld.so.conf.d/ucx.conf && \
+     echo "/usr/local/ucx/lib/ucx" >> /etc/ld.so.conf.d/ucx.conf && \
+     ldconfig
+
+# build and install nixl
+RUN --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
+    --mount=type=secret,id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY \
+    export SCCACHE_S3_KEY_PREFIX="${SCCACHE_S3_KEY_PREFIX:-${ARCH}}" && \
+    export CC=$(which gcc) && \
+    export CXX=$(which g++) && \
+    source ${VIRTUAL_ENV}/bin/activate && \
+    git clone --depth 1 --branch ${NIXL_REF} "https://github.com/ai-dynamo/nixl.git" && \
+    cd nixl && \
+    mkdir build && \
+    meson setup build/ --prefix=/opt/nvidia/nvda_nixl --buildtype=release \
+    -Dcudapath_lib="/usr/local/cuda/lib64" \
+    -Dcudapath_inc="/usr/local/cuda/include" \
+    -Ducx_path="/usr/local/ucx" && \
+    cd build && \
+    ninja && \
+    ninja install && \
+    /tmp/use-sccache.sh show-stats "NIXL"
+
+ENV NIXL_LIB_DIR=/opt/nvidia/nvda_nixl/lib64  \
+    NIXL_PLUGIN_DIR=/opt/nvidia/nvda_nixl/lib64/plugins \
+    NIXL_PREFIX=/opt/nvidia/nvda_nixl
+ENV LD_LIBRARY_PATH=${NIXL_LIB_DIR}:${NIXL_PLUGIN_DIR}:/usr/local/ucx/lib:/usr/local/ucx/lib/ucx:${LD_LIBRARY_PATH}
+
+RUN echo "$NIXL_LIB_DIR" > /etc/ld.so.conf.d/nixl.conf && \
+    echo "$NIXL_PLUGIN_DIR" >> /etc/ld.so.conf.d/nixl.conf && \
+    ldconfig
+
+RUN --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
+    --mount=type=secret,id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY \
+    export CC=$(which gcc) && \
+    export CXX=$(which g++) && \
+    export SCCACHE_S3_KEY_PREFIX="${SCCACHE_S3_KEY_PREFIX:-${ARCH}}" && \
+    cd /workspace/nixl && \
+    uv build . --out-dir /opt/dynamo/dist/nixl --python $PYTHON_VERSION
+
+# Copy source code (order matters for layer caching)
+COPY pyproject.toml README.md LICENSE Cargo.toml Cargo.lock rust-toolchain.toml hatch_build.py /opt/dynamo/
+COPY launch/ /opt/dynamo/launch/
+COPY lib/ /opt/dynamo/lib/
+COPY components/ /opt/dynamo/components/
+
+# Build dynamo wheels
+ARG ENABLE_KVBM
+RUN --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
+    --mount=type=secret,id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY \
+    export SCCACHE_S3_KEY_PREFIX=${SCCACHE_S3_KEY_PREFIX:-${ARCH}} && \
+    export CMAKE_C_COMPILER_LAUNCHER="sccache" && \
+    export CMAKE_CXX_COMPILER_LAUNCHER="sccache" && \
+    export RUSTC_WRAPPER="sccache" && \
+    source ${VIRTUAL_ENV}/bin/activate && \
+    cd /opt/dynamo && \
+    uv build --wheel --out-dir /opt/dynamo/dist && \
+    cd /opt/dynamo/lib/bindings/python && \
+    if [ "$ENABLE_MEDIA_NIXL" == "true" ]; then \
+        maturin build --release --features dynamo-llm/media-nixl --out /opt/dynamo/dist; \
+    else \
+        maturin build --release --out /opt/dynamo/dist; \
+    fi && \
+    if [ "$ENABLE_KVBM" == "true" ]; then \
+        cd /opt/dynamo/lib/bindings/kvbm && \
+        maturin build --release --out target/wheels && \
+        auditwheel repair \
+            --exclude libnixl.so \
+            --exclude libnixl_build.so \
+            --exclude libnixl_common.so \
+            --plat manylinux_2_28_${ARCH_ALT} \
+            --wheel-dir /opt/dynamo/dist \
+            target/wheels/*.whl; \
+    fi && \
+    /tmp/use-sccache.sh show-stats "Dynamo"
 
 ########################################################
 ########## Framework Development Image ################
@@ -147,16 +425,17 @@ RUN if [ "$USE_SCCACHE" = "true" ]; then \
 
 # Set environment variables - they'll be empty strings if USE_SCCACHE=false
 ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET}} \
-    SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION}} \
-    CMAKE_C_COMPILER_LAUNCHER=${USE_SCCACHE:+sccache} \
-    CMAKE_CXX_COMPILER_LAUNCHER=${USE_SCCACHE:+sccache} \
-    CMAKE_CUDA_COMPILER_LAUNCHER=${USE_SCCACHE:+sccache}
+    SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION}}
+
 # Install VLLM and related dependencies
 RUN --mount=type=bind,source=./container/deps/,target=/tmp/deps \
     --mount=type=cache,target=/root/.cache/uv \
     --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
     --mount=type=secret,id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY \
     export SCCACHE_S3_KEY_PREFIX=${SCCACHE_S3_KEY_PREFIX:-${ARCH}} && \
+    export CMAKE_C_COMPILER_LAUNCHER="sccache" && \
+    export CMAKE_CXX_COMPILER_LAUNCHER="sccache" && \
+    export CMAKE_CUDA_COMPILER_LAUNCHER="sccache" && \
         cp /tmp/deps/vllm/install_vllm.sh /tmp/install_vllm.sh && \
         chmod +x /tmp/install_vllm.sh && \
         /tmp/install_vllm.sh --vllm-ref $VLLM_REF --max-jobs $MAX_JOBS --arch $ARCH --installation-dir /opt ${DEEPGEMM_REF:+--deepgemm-ref "$DEEPGEMM_REF"} ${FLASHINF_REF:+--flashinf-ref "$FLASHINF_REF"} ${LMCACHE_REF:+--lmcache-ref "$LMCACHE_REF"} --cuda-version $CUDA_VERSION && \
@@ -196,13 +475,13 @@ ENV PATH="${VIRTUAL_ENV}/bin:${PATH}"
 ENV CUDA_DEVICE_ORDER=PCI_BUS_ID
 
 # Copy CUDA development tools (nvcc, headers, dependencies, etc.) from base devel image
-COPY --from=base /usr/local/cuda/bin/nvcc /usr/local/cuda/bin/nvcc
-COPY --from=base /usr/local/cuda/bin/cudafe++ /usr/local/cuda/bin/cudafe++
-COPY --from=base /usr/local/cuda/bin/ptxas /usr/local/cuda/bin/ptxas
-COPY --from=base /usr/local/cuda/bin/fatbinary /usr/local/cuda/bin/fatbinary
-COPY --from=base /usr/local/cuda/include/ /usr/local/cuda/include/
-COPY --from=base /usr/local/cuda/nvvm /usr/local/cuda/nvvm
-COPY --from=base /usr/local/cuda/lib64/libcudart.so* /usr/local/cuda/lib64/
+COPY --from=dynamo_base /usr/local/cuda/bin/nvcc /usr/local/cuda/bin/nvcc
+COPY --from=dynamo_base /usr/local/cuda/bin/cudafe++ /usr/local/cuda/bin/cudafe++
+COPY --from=dynamo_base /usr/local/cuda/bin/ptxas /usr/local/cuda/bin/ptxas
+COPY --from=dynamo_base /usr/local/cuda/bin/fatbinary /usr/local/cuda/bin/fatbinary
+COPY --from=dynamo_base /usr/local/cuda/include/ /usr/local/cuda/include/
+COPY --from=dynamo_base /usr/local/cuda/nvvm /usr/local/cuda/nvvm
+COPY --from=dynamo_base /usr/local/cuda/lib64/libcudart.so* /usr/local/cuda/lib64/
 RUN ln -s /usr/local/cuda/lib64/libcublas.so.12 /usr/local/cuda/lib64/libcublas.so
 RUN ln -s /usr/local/cuda/lib64/libcublasLt.so.12 /usr/local/cuda/lib64/libcublasLt.so
 
@@ -276,8 +555,11 @@ COPY --chmod=775 --chown=dynamo:0 --from=framework ${VIRTUAL_ENV} ${VIRTUAL_ENV}
 COPY --chown=dynamo:0 --from=framework /opt/vllm /opt/vllm
 
 # Copy UCX and NIXL to system directories (read-only, no group-write needed)
-COPY --chown=dynamo:0 --from=dynamo_base /usr/local/ucx /usr/local/ucx
-COPY --chown=dynamo:0 --from=dynamo_base $NIXL_PREFIX $NIXL_PREFIX
+COPY --from=wheel_builder /usr/local/ucx /usr/local/ucx
+COPY --chown=dynamo: --from=wheel_builder $NIXL_PREFIX $NIXL_PREFIX
+COPY --chown=dynamo: --from=wheel_builder /opt/nvidia/nvda_nixl/lib64/. ${NIXL_LIB_DIR}/
+COPY --chown=dynamo: --from=wheel_builder /opt/dynamo/dist/nixl/ /opt/dynamo/wheelhouse/nixl/
+COPY --chown=dynamo: --from=wheel_builder /workspace/nixl/build/src/bindings/python/nixl-meta/nixl-*.whl /opt/dynamo/wheelhouse/nixl/
 ENV PATH=/usr/local/ucx/bin:$PATH
 
 ENV LD_LIBRARY_PATH=\
@@ -296,7 +578,7 @@ COPY --chmod=775 --chown=dynamo:0 benchmarks/ /workspace/benchmarks/
 # Install dynamo, NIXL, and dynamo-specific dependencies
 # Pattern: COPY --chmod=775 <path>; chmod g+w <path> done later as root because COPY --chmod only affects <path>/*, not <path>
 ARG ENABLE_KVBM
-COPY --chmod=775 --chown=dynamo:0 --from=dynamo_base /opt/dynamo/wheelhouse/ /opt/dynamo/wheelhouse/
+COPY --chmod=775 --chown=dynamo:0 --from=wheel_builder /opt/dynamo/dist/*.whl /opt/dynamo/wheelhouse/
 RUN uv pip install \
       /opt/dynamo/wheelhouse/ai_dynamo_runtime*.whl \
       /opt/dynamo/wheelhouse/ai_dynamo*any.whl \

container/README.md

@@ -18,6 +18,48 @@ The NVIDIA Dynamo project uses containerized development and deployment to maint
   - `Dockerfile.frontend` - For Kubernetes Gateway API Inference Extension integration with EPP
   - `Dockerfile.epp` - For building the Endpoint Picker (EPP) image
 
+### Stage Summary for Frameworks
+
+<details>
+<summary>Show Stage Summary Table</summary>
+Dockerfile.${FRAMEWORK} General Structure
+
+Below is a summary of the general file structure for the framework Dockerfile stages. Some exceptions exist.
+
+| Stage/Filepath | Target |
+| --- | --- |
+| **STAGE dynamo_base** | **FROM ${BASE_IMAGE}** |
+| /bin/uv, /bin/uvx | COPY from ghcr.io/astral-sh/uv:latest (→ framework, runtime) |
+|  /usr/bin/nats-server | Downloaded from GitHub (→ runtime) |
+|  /usr/local/bin/etcd/ | Downloaded from GitHub (→ runtime) |
+|  /usr/local/rustup/ | Installed via rustup-init (→ wheel_builder, dev) |
+|  /usr/local/cargo/ | Installed via rustup-init (→ wheel_builder, dev) |
+|  /usr/local/cuda/ | Inherited from BASE_IMAGE (→ wheel_builder, runtime) |
+| **STAGE: wheel_builder** | **FROM quay.io/pypa/manylinux_2_28_${ARCH_ALT}** |
+|  /usr/local/ucx/ | Built from source (→ runtime)
+|  /opt/nvidia/nvda_nixl/ | Built from source (→ runtime)
+|  /opt/nvidia/nvda_nixl/lib64/ | Built from source (→ runtime)
+|  /opt/dynamo/target/ | Cargo build output (→ runtime)
+|  /opt/dynamo/dist/*.whl | Built wheels (→ runtime)
+|  /opt/dynamo/dist/nixl/ | Built nixl wheels (→ runtime)
+| **STAGE: framework** | **FROM ${BASE_IMAGE}** |
+|  /opt/dynamo/venv/ | Created with uv venv (→ runtime)
+|  /${FRAMEWORK_INSTALL} | Built framework (→ runtime)
+| **STAGE: runtime** | **FROM ${RUNTIME_IMAGE}** |
+|  /usr/local/cuda/{bin,include,nvvm}/ | COPY from dynamo_base |
+|  /usr/bin/nats-server | COPY from dynamo_runtime |
+|  /usr/local/bin/etcd/ | COPY from dynamo_runtime |
+|  /usr/local/ucx/ | COPY from dynamo_runtime |
+|  /opt/nvidia/nvda_nixl/ | COPY from wheel_builder |
+|  /opt/dynamo/wheelhouse/ | COPY from wheel_builder |
+|  /opt/dynamo/venv/ | COPY from framework |
+|  /opt/vllm/ | COPY from framework |
+|  /workspace/{tests,examples,deploy}/ |COPY from build context |
+| **STAGE: dev** | **FROM runtime** |
+|  /usr/local/rustup/ | COPY from dynamo_runtime |
+|  /usr/local/cargo/ | COPY from dynamo_runtime |
+</details>
+
 ### Why Containerization?
 
 Each inference framework (vLLM, TensorRT-LLM, SGLang) has specific CUDA versions, Python dependencies, and system libraries. Containers provide consistent environments, framework isolation, and proper GPU configurations across development and production.

container/build.sh

@@ -122,6 +122,7 @@ SGLANG_FRAMEWORK_IMAGE_TAG="${SGLANG_CUDA_VERSION}-cudnn-devel-ubuntu24.04"
 NIXL_REF=0.7.1
 NIXL_UCX_REF=v1.19.0
 NIXL_UCX_EFA_REF=9d2b88a1f67faf9876f267658bd077b379b8bb76
+NIXL_GDRCOPY_REF=v2.5.1
 
 NO_CACHE=""
 
@@ -276,7 +277,7 @@ get_options() {
             ;;
         --cache-from)
             if [ "$2" ]; then
-                CACHE_FROM="--cache-from $2"
+                CACHE_FROM+="--cache-from $2 "
                 shift
             else
                 missing_requirement "$1"
@@ -284,7 +285,7 @@ get_options() {
             ;;
         --cache-to)
             if [ "$2" ]; then
-                CACHE_TO="--cache-to $2"
+                CACHE_TO+="--cache-to $2 "
                 shift
             else
                 missing_requirement "$1"
@@ -830,12 +831,17 @@ if [ -z "${ENABLE_MEDIA_NIXL}" ]; then
 fi
 BUILD_ARGS+=" --build-arg ENABLE_MEDIA_NIXL=${ENABLE_MEDIA_NIXL} "
 
-# NIXL_UCX_REF: Used in base Dockerfile only.
-#               Passed to framework Dockerfile.{vllm,sglang,...} where it's NOT used.
+# NIXL_UCX_REF: Used in dynamo base stages.
 if [ -n "${NIXL_UCX_REF}" ]; then
     BUILD_ARGS+=" --build-arg NIXL_UCX_REF=${NIXL_UCX_REF} "
 fi
 
+# NIXL_GDRCOPY_REF: Used in dynamo base stages.
+if [ -n "${NIXL_GDRCOPY_REF}" ]; then
+    BUILD_ARGS+=" --build-arg NIXL_GDRCOPY_REF=${NIXL_GDRCOPY_REF} "
+
+fi
+
 # MAX_JOBS is only used by Dockerfile.vllm
 if [ -n "${MAX_JOBS}" ]; then
     BUILD_ARGS+=" --build-arg MAX_JOBS=${MAX_JOBS} "
@@ -880,89 +886,26 @@ fi
 
 # Skip Build 1 and Build 2 if DEV_IMAGE_INPUT is set (we'll handle it at the bottom)
 if [[ -z "${DEV_IMAGE_INPUT:-}" ]]; then
-    # Follow 2-step build process for all frameworks
-    if [[ $FRAMEWORK != "NONE" ]]; then
-        # Define base image tag with framework suffix to prevent clobbering
-        # Different frameworks require different base configurations:
-        # - VLLM: Python 3.12, ENABLE_KVBM=true, BASE_IMAGE=cuda-dl-base
-        # - SGLANG: Python 3.10, BASE_IMAGE=cuda-dl-base
-        # - TRTLLM: Python 3.12, ENABLE_KVBM=true, BASE_IMAGE=pytorch
-        # Without unique tags, building different frameworks would overwrite each other's names
-        DYNAMO_BASE_IMAGE="dynamo-base:${VERSION}-${FRAMEWORK,,}"
-        # Start base image build
-        echo "======================================"
-        echo "Starting Build 1: Base Image"
-        echo "======================================"
-
-        # Create build log directory for BuildKit reports
-        BUILD_LOG_DIR="${BUILD_CONTEXT}/build-logs"
-        mkdir -p "${BUILD_LOG_DIR}"
-        BASE_BUILD_LOG="${BUILD_LOG_DIR}/base-image-build.log"
-
-        # Use BuildKit for enhanced metadata
-        if [ -z "$RUN_PREFIX" ]; then
-            if docker buildx version &>/dev/null; then
-                docker buildx build --builder default --progress=plain --load -f "${SOURCE_DIR}/Dockerfile" --target runtime $PLATFORM $BUILD_ARGS $CACHE_FROM $CACHE_TO --tag $DYNAMO_BASE_IMAGE $BUILD_CONTEXT_ARG $BUILD_CONTEXT $NO_CACHE 2>&1 | tee "${BASE_BUILD_LOG}"
-                BUILD_EXIT_CODE=${PIPESTATUS[0]}
-            else
-                DOCKER_BUILDKIT=1 docker build --progress=plain -f "${SOURCE_DIR}/Dockerfile" --target runtime $PLATFORM $BUILD_ARGS $CACHE_FROM $CACHE_TO --tag $DYNAMO_BASE_IMAGE $BUILD_CONTEXT_ARG $BUILD_CONTEXT $NO_CACHE 2>&1 | tee "${BASE_BUILD_LOG}"
-                BUILD_EXIT_CODE=${PIPESTATUS[0]}
-            fi
+    # Create build log directory for BuildKit reports
+    BUILD_LOG_DIR="${BUILD_CONTEXT}/build-logs"
+    mkdir -p "${BUILD_LOG_DIR}"
+    SINGLE_BUILD_LOG="${BUILD_LOG_DIR}/single-stage-build.log"
 
-            if [ ${BUILD_EXIT_CODE} -ne 0 ]; then
-                exit ${BUILD_EXIT_CODE}
-            fi
+    # Use BuildKit for enhanced metadata
+    if [ -z "$RUN_PREFIX" ]; then
+        if docker buildx version &>/dev/null; then
+            docker buildx build --progress=plain --load -f $DOCKERFILE $TARGET_STR $PLATFORM $BUILD_ARGS $CACHE_FROM $CACHE_TO $TAG $LATEST_TAG $BUILD_CONTEXT_ARG $BUILD_CONTEXT $NO_CACHE 2>&1 | tee "${SINGLE_BUILD_LOG}"
+            BUILD_EXIT_CODE=${PIPESTATUS[0]}
         else
-            $RUN_PREFIX docker build -f "${SOURCE_DIR}/Dockerfile" --target runtime $PLATFORM $BUILD_ARGS $CACHE_FROM $CACHE_TO --tag $DYNAMO_BASE_IMAGE $BUILD_CONTEXT_ARG $BUILD_CONTEXT $NO_CACHE
+            DOCKER_BUILDKIT=1 docker build --progress=plain -f $DOCKERFILE $TARGET_STR $PLATFORM $BUILD_ARGS $CACHE_FROM $CACHE_TO $TAG $LATEST_TAG $BUILD_CONTEXT_ARG $BUILD_CONTEXT $NO_CACHE 2>&1 | tee "${SINGLE_BUILD_LOG}"
+            BUILD_EXIT_CODE=${PIPESTATUS[0]}
         fi
 
-        # Start framework build
-        echo "======================================"
-        echo "Starting Build 2: Framework Image"
-        echo "======================================"
-
-        FRAMEWORK_BUILD_LOG="${BUILD_LOG_DIR}/framework-${FRAMEWORK,,}-build.log"
-
-        BUILD_ARGS+=" --build-arg DYNAMO_BASE_IMAGE=${DYNAMO_BASE_IMAGE}"
-
-        # Use BuildKit for enhanced metadata
-        if [ -z "$RUN_PREFIX" ]; then
-            if docker buildx version &>/dev/null; then
-                docker buildx build --builder default --progress=plain --load -f $DOCKERFILE $TARGET_STR $PLATFORM $BUILD_ARGS $CACHE_FROM $CACHE_TO $TAG $LATEST_TAG $BUILD_CONTEXT_ARG $BUILD_CONTEXT $NO_CACHE 2>&1 | tee "${FRAMEWORK_BUILD_LOG}"
-                BUILD_EXIT_CODE=${PIPESTATUS[0]}
-            else
-                DOCKER_BUILDKIT=1 docker build --progress=plain -f $DOCKERFILE $TARGET_STR $PLATFORM $BUILD_ARGS $CACHE_FROM $CACHE_TO $TAG $LATEST_TAG $BUILD_CONTEXT_ARG $BUILD_CONTEXT $NO_CACHE 2>&1 | tee "${FRAMEWORK_BUILD_LOG}"
-                BUILD_EXIT_CODE=${PIPESTATUS[0]}
-            fi
-
-            if [ ${BUILD_EXIT_CODE} -ne 0 ]; then
-                exit ${BUILD_EXIT_CODE}
-            fi
-        else
-            $RUN_PREFIX docker build -f $DOCKERFILE $TARGET_STR $PLATFORM $BUILD_ARGS $CACHE_FROM $CACHE_TO $TAG $LATEST_TAG $BUILD_CONTEXT_ARG $BUILD_CONTEXT $NO_CACHE
+        if [ ${BUILD_EXIT_CODE} -ne 0 ]; then
+            exit ${BUILD_EXIT_CODE}
         fi
     else
-        # Create build log directory for BuildKit reports
-        BUILD_LOG_DIR="${BUILD_CONTEXT}/build-logs"
-        mkdir -p "${BUILD_LOG_DIR}"
-        SINGLE_BUILD_LOG="${BUILD_LOG_DIR}/single-stage-build.log"
-
-        # Use BuildKit for enhanced metadata
-        if [ -z "$RUN_PREFIX" ]; then
-            if docker buildx version &>/dev/null; then
-                docker buildx build --builder default --progress=plain --load -f $DOCKERFILE $TARGET_STR $PLATFORM $BUILD_ARGS $CACHE_FROM $CACHE_TO $TAG $LATEST_TAG $BUILD_CONTEXT_ARG $BUILD_CONTEXT $NO_CACHE 2>&1 | tee "${SINGLE_BUILD_LOG}"
-                BUILD_EXIT_CODE=${PIPESTATUS[0]}
-            else
-                DOCKER_BUILDKIT=1 docker build --progress=plain -f $DOCKERFILE $TARGET_STR $PLATFORM $BUILD_ARGS $CACHE_FROM $CACHE_TO $TAG $LATEST_TAG $BUILD_CONTEXT_ARG $BUILD_CONTEXT $NO_CACHE 2>&1 | tee "${SINGLE_BUILD_LOG}"
-                BUILD_EXIT_CODE=${PIPESTATUS[0]}
-            fi
-
-            if [ ${BUILD_EXIT_CODE} -ne 0 ]; then
-                exit ${BUILD_EXIT_CODE}
-            fi
-        else
-            $RUN_PREFIX docker build -f $DOCKERFILE $TARGET_STR $PLATFORM $BUILD_ARGS $CACHE_FROM $CACHE_TO $TAG $LATEST_TAG $BUILD_CONTEXT_ARG $BUILD_CONTEXT $NO_CACHE
-        fi
+        $RUN_PREFIX docker build -f $DOCKERFILE $TARGET_STR $PLATFORM $BUILD_ARGS $CACHE_FROM $CACHE_TO $TAG $LATEST_TAG $BUILD_CONTEXT_ARG $BUILD_CONTEXT $NO_CACHE
     fi
 fi