[CI/Build][ROCm] Cache AMD test images in registry (#7712)

Signed-off-by: andyluo7 <andy.luo@amd.com>

.buildkite/amd/scripts/build-ci-image.sh

@@ -0,0 +1,69 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+set -euo pipefail
+
+# Build and push the ROCm CI image with a registry-backed BuildKit cache. The
+# AMD builders are ephemeral, so a local Docker cache cannot survive between
+# jobs. Keeping the cache beside the runtime image also reuses the registry
+# credentials already configured on the amd-cpu agents.
+
+REGISTRY="${ROCM_CI_IMAGE_REGISTRY:-rocm/vllm-omni}"
+DOCKERFILE="docker/Dockerfile.rocm"
+BUILDER_NAME="vllm-omni-rocm-builder"
+ROCM_ARCH="${PYTORCH_ROCM_ARCH:-gfx942}"
+
+if [[ -z "${BUILDKITE_COMMIT:-}" ]]; then
+    echo "ERROR: BUILDKITE_COMMIT is not set" >&2
+    exit 1
+fi
+
+dependency_files=(
+    "$DOCKERFILE"
+    "${DOCKERFILE}.dockerignore"
+    pyproject.toml
+    setup.py
+    tools/install_torchcodec_rocm.sh
+)
+while IFS= read -r dependency_file; do
+    dependency_files+=("$dependency_file")
+done < <(find requirements -maxdepth 1 -type f -name '*.txt' -print | LC_ALL=C sort)
+
+for dependency_file in "${dependency_files[@]}"; do
+    if [[ ! -f "$dependency_file" ]]; then
+        echo "ERROR: cache-key input does not exist: $dependency_file" >&2
+        exit 1
+    fi
+done
+
+CACHE_KEY=$(
+    {
+        printf 'rocm-arch=%s\n' "$ROCM_ARCH"
+        for dependency_file in "${dependency_files[@]}"; do
+            printf 'path=%s\n' "$dependency_file"
+            sha256sum "$dependency_file"
+        done
+    } | sha256sum | cut -c1-16
+)
+CACHE_REF="${REGISTRY}:rocm-deps-cache-${CACHE_KEY}"
+IMAGE_REF="${REGISTRY}:${BUILDKITE_COMMIT}"
+
+echo "ROCm dependency cache: ${CACHE_REF}"
+echo "ROCm CI image: ${IMAGE_REF}"
+
+# The docker-container driver supports exporting the full cache to a registry;
+# the default docker driver on CI does not.
+if ! docker buildx inspect "$BUILDER_NAME" >/dev/null 2>&1; then
+    docker buildx create --name "$BUILDER_NAME" --driver docker-container
+fi
+docker buildx use "$BUILDER_NAME"
+docker buildx inspect --bootstrap
+
+docker buildx build --push --pull --progress=plain \
+    --target test \
+    --build-arg "ARG_PYTORCH_ROCM_ARCH=${ROCM_ARCH}" \
+    --cache-from "type=registry,ref=${CACHE_REF}" \
+    --cache-to "type=registry,ref=${CACHE_REF},mode=max,compression=zstd,ignore-error=true" \
+    --file "$DOCKERFILE" \
+    --tag "$IMAGE_REF" .

.buildkite/amd/test-template-amd-omni.j2

@@ -3,7 +3,8 @@
    Last synced: 2026-08-04
    Modifications: Use the shared MI300 queues and Omni command-passing contract
 #}
-{% set docker_image_amd = "rocm/vllm-omni:$BUILDKITE_COMMIT" %}
+{% set docker_registry_amd = "rocm/vllm-omni" %}
+{% set docker_image_amd = docker_registry_amd ~ ":$BUILDKITE_COMMIT" %}
 {% set default_working_dir = "/app/vllm-omni" %}
 
 {# MI300 runs tests natively in Kubernetes; Docker-in-Docker is unavailable. #}
@@ -74,11 +75,12 @@ steps:
         depends_on: ~
         soft_fail: false
         commands:
-          - "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target test --build-arg ARG_PYTORCH_ROCM_ARCH=gfx942 --progress plain ."
-          - "docker push {{ docker_image_amd }}"
+          - "bash .buildkite/amd/scripts/build-ci-image.sh"
         key: "amd-build"
         env:
           DOCKER_BUILDKIT: "1"
+          PYTORCH_ROCM_ARCH: "gfx942"
+          ROCM_CI_IMAGE_REGISTRY: "{{ docker_registry_amd }}"
         retry:
           automatic:
             - exit_status: -1  # Agent was lost

docker/Dockerfile.rocm

@@ -2,7 +2,11 @@
 # docker/Dockerfile.ci. A stale base image can build successfully while failing
 # at runtime when vLLM-Omni imports APIs added by the newer vLLM release.
 ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.29.0
-FROM ${BASE_IMAGE} AS base
+
+# Keep the dependency-heavy portion independent from the source tree. The
+# registry cache is keyed on these inputs, while the final source copy can be
+# linked onto the cached image without materializing its large ROCm layers.
+FROM ${BASE_IMAGE} AS deps
 
 # NOTE: REMINDER to vLLM-Omni rebase maintainer
 # BASE_IMAGE pins the default prebuilt runtime when USE_NIGHTLY_BUILD=0.
@@ -64,9 +68,24 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
-# Step 3: Copy vllm-omni code and install without uv
-RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
-COPY . ${COMMON_WORKDIR}/vllm-omni
+# Step 3: Install project dependencies against a placeholder package tree.
+# The editable install keeps resolving imports from this directory after the
+# final COPY --link replaces the placeholder with the real source tree.
+WORKDIR ${COMMON_WORKDIR}/vllm-omni
+RUN mkdir -p vllm_omni && touch vllm_omni/__init__.py
+COPY pyproject.toml setup.py ./
+COPY requirements/ requirements/
+ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0
+
+RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" \
+        --no-cache-dir -e ".[dev]" --no-build-isolation && \
+    python3 -m pip uninstall -y onnxruntime && \
+    python3 -m pip install --no-cache-dir --force-reinstall --no-deps \
+        "onnxruntime-rocm>=1.22.2" && \
+    python3 -c "import importlib.metadata as m; import onnxruntime as ort; assert 'ROCMExecutionProvider' in ort.get_available_providers(), ort.get_available_providers(); assert m.version('onnxruntime-rocm'); assert not any(d.metadata['Name'] == 'onnxruntime' for d in m.distributions())" && \
+    python3 -c "from torchcodec.decoders import VideoDecoder"
+
+RUN ln -sf /usr/bin/python3 /usr/bin/python
 
 # This is a workaround to ensure pytest exits with the correct status code in CI tests.
 RUN printf '%s\n' \
@@ -85,19 +104,13 @@ RUN printf '%s\n' \
     '    os._exit(_exit_code)' \
     > ${COMMON_WORKDIR}/vllm-omni/conftest.py
 
-# Some optional dependencies pull in the CPU ONNX Runtime distribution. The
-# CPU and ROCm wheels own the same import path and cannot coexist, so remove
-# the CPU wheel and reinstall the ROCm wheel after dependency resolution.
-RUN cd ${COMMON_WORKDIR}/vllm-omni && \
-    uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" \
-        --no-cache-dir ".[dev]" --no-build-isolation && \
-    python3 -m pip uninstall -y onnxruntime && \
-    python3 -m pip install --no-cache-dir --force-reinstall --no-deps \
-        "onnxruntime-rocm>=1.22.2" && \
-    python3 -c "import importlib.metadata as m; import onnxruntime as ort; assert 'ROCMExecutionProvider' in ort.get_available_providers(), ort.get_available_providers(); assert m.version('onnxruntime-rocm'); assert not any(d.metadata['Name'] == 'onnxruntime' for d in m.distributions())" && \
-    python3 -c "from torchcodec.decoders import VideoDecoder"
-
-RUN ln -sf /usr/bin/python3 /usr/bin/python
+# Keep source changes in one lightweight linked layer. Do not add RUN, COPY, or
+# ADD instructions that depend on the resulting filesystem after this point;
+# doing so would force BuildKit to download and extract the cached ROCm layers.
+FROM deps AS base
+ARG COMMON_WORKDIR=/app
+WORKDIR ${COMMON_WORKDIR}/vllm-omni
+COPY --link . .
 
 FROM base AS test
 

docker/Dockerfile.rocm.dockerignore

@@ -0,0 +1,35 @@
+# Keep the linked source layer small and deterministic.
+.git
+.venv
+/build
+dist
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+.mypy_cache
+
+# Distribution / packaging
+.Python
+/build/
+cmake-build-*/
+CMakeUserPresets.json
+develop-eggs/
+/dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+rust/target/

tests/buildkite/test_rocm_dockerfile.py

@@ -9,8 +9,10 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 REPO_ROOT = Path(__file__).resolve().parents[2]
 AMD_TEMPLATE = REPO_ROOT / ".buildkite/amd/test-template-amd-omni.j2"
+AMD_BUILD_SCRIPT = REPO_ROOT / ".buildkite/amd/scripts/build-ci-image.sh"
 CI_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.ci"
 ROCM_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.rocm"
+ROCM_DOCKERIGNORE = REPO_ROOT / "docker/Dockerfile.rocm.dockerignore"
 
 
 def _docker_arg(path: Path, name: str) -> str:
@@ -42,16 +44,63 @@ def test_rocm_defaults_to_prebuilt_base_image() -> None:
     assert _docker_arg(ROCM_DOCKERFILE, "USE_NIGHTLY_BUILD") == "0"
 
 
-def test_amd_build_uses_rocm_dockerfile_defaults() -> None:
+def test_amd_build_uses_registry_cache() -> None:
     template = AMD_TEMPLATE.read_text(encoding="utf-8")
-    build_commands = [line.strip() for line in template.splitlines() if '"docker build ' in line]
+    script_command = '"bash .buildkite/amd/scripts/build-ci-image.sh"'
+    script = AMD_BUILD_SCRIPT.read_text(encoding="utf-8")
 
-    assert len(build_commands) == 1, f"expected one AMD image build command, found {len(build_commands)}"
-    build_command = build_commands[0]
-    assert "-f docker/Dockerfile.rocm" in build_command
+    assert template.count(script_command) == 1
+    assert "docker buildx build --push" in script
+    assert '--cache-from "type=registry,ref=${CACHE_REF}"' in script
+    assert '--cache-to "type=registry,ref=${CACHE_REF}' in script
+    assert "mode=max" in script
+
+
+def test_amd_cache_key_tracks_dependency_inputs_and_architecture() -> None:
+    script = AMD_BUILD_SCRIPT.read_text(encoding="utf-8")
+
+    for cache_input in (
+        "docker/Dockerfile.rocm",
+        "${DOCKERFILE}.dockerignore",
+        "pyproject.toml",
+        "setup.py",
+        "requirements",
+        "tools/install_torchcodec_rocm.sh",
+        "ROCM_ARCH",
+    ):
+        assert cache_input in script
+
+
+def test_amd_build_uses_rocm_dockerfile_defaults() -> None:
+    build_script = AMD_BUILD_SCRIPT.read_text(encoding="utf-8")
+
+    assert 'DOCKERFILE="docker/Dockerfile.rocm"' in build_script
     for arg_name in ("BASE_IMAGE", "USE_NIGHTLY_BUILD"):
-        assert f"--build-arg {arg_name}" not in build_command
-        assert f"--build-arg={arg_name}" not in build_command
+        assert f"--build-arg {arg_name}" not in build_script
+        assert f"--build-arg={arg_name}" not in build_script
+
+
+def test_rocm_source_copy_preserves_dependency_cache() -> None:
+    dockerfile = ROCM_DOCKERFILE.read_text(encoding="utf-8")
+    lines = dockerfile.splitlines()
+    source_copy_index = _line_index(lines, "COPY --link . .")
+    test_stage_index = _line_index(lines, "FROM base AS test")
+
+    assert "FROM deps AS base" in lines[:source_copy_index]
+    assert any('-e ".[dev]"' in line for line in lines[:source_copy_index])
+    assert not any(
+        line.startswith(("RUN ", "COPY ", "ADD ")) for line in lines[source_copy_index + 1 : test_stage_index]
+    )
+
+
+def test_rocm_build_context_excludes_git_history() -> None:
+    patterns = {
+        line.strip()
+        for line in ROCM_DOCKERIGNORE.read_text(encoding="utf-8").splitlines()
+        if line.strip() and not line.lstrip().startswith("#")
+    }
+
+    assert ".git" in patterns
 
 
 def test_rocm_source_ref_tracks_ci_vllm_release() -> None: