From fa422a3f449662bf5d4e4374e014ac14cd90b783 Mon Sep 17 00:00:00 2001 From: andyluo7 <43718156+andyluo7@users.noreply.github.com> Date: Sat, 19 Sep 2026 18:51:14 -0700 Subject: [PATCH] [CI/Build][ROCm] Cache AMD test images in registry (#7712) Signed-off-by: andyluo7 --- .buildkite/amd/scripts/build-ci-image.sh | 69 ++++++++++++++++++++++++ .buildkite/amd/test-template-amd-omni.j2 | 8 +-- docker/Dockerfile.rocm | 47 ++++++++++------ docker/Dockerfile.rocm.dockerignore | 35 ++++++++++++ tests/buildkite/test_rocm_dockerfile.py | 63 +++++++++++++++++++--- 5 files changed, 195 insertions(+), 27 deletions(-) create mode 100755 .buildkite/amd/scripts/build-ci-image.sh create mode 100644 docker/Dockerfile.rocm.dockerignore diff --git a/.buildkite/amd/scripts/build-ci-image.sh b/.buildkite/amd/scripts/build-ci-image.sh new file mode 100755 index 000000000..3954cd567 --- /dev/null +++ b/.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" . diff --git a/.buildkite/amd/test-template-amd-omni.j2 b/.buildkite/amd/test-template-amd-omni.j2 index 11a4086b6..cca99f509 100644 --- a/.buildkite/amd/test-template-amd-omni.j2 +++ b/.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 diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index adc767691..48631994f 100644 --- a/docker/Dockerfile.rocm +++ b/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 diff --git a/docker/Dockerfile.rocm.dockerignore b/docker/Dockerfile.rocm.dockerignore new file mode 100644 index 000000000..f68737749 --- /dev/null +++ b/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/ diff --git a/tests/buildkite/test_rocm_dockerfile.py b/tests/buildkite/test_rocm_dockerfile.py index 9171b3d53..8101b2157 100644 --- a/tests/buildkite/test_rocm_dockerfile.py +++ b/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: