[ROCm][CI] Optimize ROCm Docker build: registry cache, DeepEP, and ci-bake script by AndreasKaratzas · Pull Request #36949 · vllm-project/vllm

AndreasKaratzas · 2026-03-13T03:19:02Z

Summary

Implements the three-tier Docker build for ROCm CI. Every PR currently rebuilds RIXL, DeepEP, rocshmem, torchcodec, and RDMA libraries from scratch, costing a total of 26 minutes on average per build. This PR introduces a pre-built Tier-1 ci_base image that absorbs those stable layers. Per-PR builds then only rebuild the thin vLLM wheel + workspace layer.

Image registry layout after this PR:

Tag	Built by	Frequency
`rocm/vllm-dev:base`	`Dockerfile.rocm_base`	Monthly
`rocm/vllm-dev:ci_base`	`ci_base` stage (this PR)	Weekly
`rocm/vllm-ci:$COMMIT`	`test` stage (this PR)	Per PR

This PR is connected to: vllm-project/ci-infra#307
These two PRs should likely be merged simultaneously.

cc @kenroche @okakarpa @tjtanaa @gshtras @khluu

gemini-code-assist

Code Review

This pull request introduces significant optimizations to the ROCm Docker build process by leveraging docker bake, multi-stage builds, and caching mechanisms like ccache. The new ci-bake.sh script centralizes and improves the CI build logic, enhancing build times and reliability. The changes are well-structured and thoughtful. I've identified a couple of critical issues related to missing runtime dependencies in the Dockerfile and a high-severity issue regarding configuration consistency in the new bake script.

gemini-code-assist · 2026-03-13T03:27:53Z

+RUN --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
    uv pip install --system /deep_install/*.whl


The rocshmem library appears to be a runtime dependency for deepep. This test stage installs the deepep wheel but no longer copies the rocshmem installation from the build stage. This could lead to runtime errors if the deepep wheel does not bundle the rocshmem shared libraries. Please restore the copy of the rocshmem directory from the build_rocshmem stage to ensure deepep can function correctly.

RUN --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \ uv pip install --system /deep_install/*.whl # Copy rocshmem runtime libraries COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem

gemini-code-assist · 2026-03-13T03:27:53Z

+RUN --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
+    uv pip install --system /deep_install/*.whl


Similar to the test stage, the final stage now installs the deepep wheel but is missing the rocshmem runtime libraries which are likely a runtime dependency. This is likely to cause runtime failures. Please add a COPY instruction to include the rocshmem installation from the build_rocshmem stage.

RUN --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \ uv pip install --system /deep_install/*.whl # Copy rocshmem runtime libraries COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem

gemini-code-assist · 2026-03-13T03:27:53Z

+    # Check if baked-vllm-builder already exists and is using the socket
+    if docker buildx inspect baked-vllm-builder >/dev/null 2>&1; then
+        echo "Using existing baked-vllm-builder"
+        docker buildx use baked-vllm-builder
+    else
+        echo "Creating baked-vllm-builder with remote driver"
+        docker buildx create \
+            --name baked-vllm-builder \
+            --driver remote \
+            --use \
+            "unix://${BUILDKIT_SOCKET}"
+    fi


There's an inconsistency in the buildx builder naming. The script accepts a BUILDER_NAME environment variable (defaulting to vllm-builder), but when a local buildkitd socket is detected, it hardcodes the builder name to baked-vllm-builder. This could lead to confusion and incorrect builder usage if BUILDER_NAME is customized. For consistency, please use the ${BUILDER_NAME} variable throughout the script.

Suggested change

# Check if baked-vllm-builder already exists and is using the socket

if docker buildx inspect baked-vllm-builder >/dev/null 2>&1; then

echo "Using existing baked-vllm-builder"

docker buildx use baked-vllm-builder

else

echo "Creating baked-vllm-builder with remote driver"

docker buildx create \

--name baked-vllm-builder \

--driver remote \

--use \

"unix://${BUILDKIT_SOCKET}"

fi

# Check if ${BUILDER_NAME} already exists and is using the socket

if docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then

echo "Using existing builder: ${BUILDER_NAME}"

docker buildx use "${BUILDER_NAME}"

else

echo "Creating builder '${BUILDER_NAME}' with remote driver"

docker buildx create \

--name "${BUILDER_NAME}" \

--driver remote \

--use \

"unix://${BUILDKIT_SOCKET}"

fi

gshtras · 2026-03-13T17:58:32Z

-    apt-transport-https ca-certificates wget curl
+    apt-transport-https ca-certificates wget curl \
+    ccache mold \
+    && update-alternatives --install /usr/bin/ld ld /usr/bin/mold 100


Changing the system loader hardly falls under "Install some basic utilities"
Could you at least provide the motivation for this in the PR description?

Correct, mb, I updated the comment there as well.

gshtras · 2026-03-13T18:03:07Z

+RUN --mount=type=cache,target=/root/.cache/ccache \
+    --mount=type=cache,target=/root/.cache/uv \
+    cd vllm \
+    && uv pip install --system -r requirements/rocm-build.txt \


Why is rocm-build.txt being used in the docker build?

That's an oversight on my part, thought it was just rocm.txt. I updated that as well.

gshtras · 2026-03-13T18:08:24Z

+COPY requirements/rocm-build.txt requirements/rocm-build.txt
+COPY pyproject.toml setup.py CMakeLists.txt ./
+COPY cmake cmake/
+COPY csrc csrc/


Are you copying host files here? The point of REMOTE_VLLM is exactly to not do this

I refactored based on offline conversation to bring back the old way of doing this and avoid any trouble. I also integrated another recommended point which is a per-arch build so that we then use a specific docker dependency and not an all-arch dependency. Hope it looks better now :)

AndreasKaratzas · 2026-03-15T03:26:07Z

First successful build with vllm-project/ci-infra#307: https://buildkite.com/vllm/amd-ci/builds/6536/steps/canvas

mergify · 2026-03-19T20:02:39Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @AndreasKaratzas.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

AndreasKaratzas · 2026-03-23T03:02:09Z

@mawong-amd Let's check if Kernels Core Operation Test passes as well. We may need to bring back the compilation of triton_kernels. Not sure yet.

mergify · 2026-03-30T12:54:39Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @AndreasKaratzas.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

…er changes with caching (vllm-project#36949) Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

…ker_build Signed-off-by: Andreas Karatzas <akaratza@amd.com>

mawong-amd

Two comments to start with:

Does it make sense to also add the CI base content SHA to the image tag, rather than using rocm/vllm-dev:ci_base generically? The situation I have in mind is a PR where we're updating say rocSHMEM.
Each commit pushed to this PR branch would trigger a fresh CI run and the "ensure CI base" stage would rebuild and push rocm/vllm-dev:ci_base to a version consistent with this PR branch, but inconsistent with main. And vice-versa. So we might see ping-pong invalidation of rocm/vllm-dev:ci_base between main and the PR branch.
This may not be a major issue if caching somehow means that subsequent rebuilds of CI base images are fast, but I have not yet reasoned as to whether this should be the case. Have we tested it?
Do we need all the Docker builder related additions?

mawong-amd · 2026-05-20T20:42:49Z

+setup_builder() {
+    echo "--- :buildkite: Setting up buildx builder"
+
+    if [[ -S "${BUILDKIT_SOCKET}" ]]; then
+        echo "Found local buildkitd socket at ${BUILDKIT_SOCKET}"
+        echo "Using remote driver to connect to buildkitd"
+
+        if docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then
+            use_existing_builder
+        else
+            create_and_bootstrap_builder remote "unix://${BUILDKIT_SOCKET}"
+        fi
+    elif docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then
+        use_existing_builder
+    else
+        echo "No local buildkitd found, using docker-container driver"
+        create_and_bootstrap_builder docker-container
+    fi
+
+    echo "Active builder:"
+    docker buildx ls | grep -E '^\*|^NAME' || docker buildx ls
+}


What do we need the Docker builder changes here and elsewhere for?

…ker_build Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

mergify · 2026-05-22T04:30:35Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @AndreasKaratzas.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

…ker_build Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

mergify · 2026-05-28T06:04:43Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @AndreasKaratzas.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

…ker_build Signed-off-by: Andreas Karatzas <akaratza@amd.com>

…ker_build

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

…argets. Signed-off-by: Andreas Karatzas <akaratza@amd.com>

AndreasKaratzas · 2026-06-02T06:57:15Z

Latest CI run with vllm-project/ci-infra#307 integrated: https://buildkite.com/vllm/ci/builds/69415/canvas

AndreasKaratzas requested review from gshtras and tjtanaa as code owners March 13, 2026 03:19

AndreasKaratzas mentioned this pull request Mar 13, 2026

[ROCm][CI] Add ROCm Docker Hub registry cache and weekly cleanup pipeline vllm-project/ci-infra#307

Open

mergify Bot added ci/build rocm Related to AMD ROCm labels Mar 13, 2026

github-project-automation Bot added this to AMD Mar 13, 2026

github-project-automation Bot moved this to Todo in AMD Mar 13, 2026

gemini-code-assist Bot reviewed Mar 13, 2026

View reviewed changes

gshtras reviewed Mar 13, 2026

View reviewed changes

AndreasKaratzas marked this pull request as draft March 13, 2026 18:27

AndreasKaratzas mentioned this pull request Mar 15, 2026

[ROCm][CI] Revamping AMD mirrors #35897

Closed

AndreasKaratzas marked this pull request as ready for review March 18, 2026 05:38

AndreasKaratzas requested review from LucasWilkinson and tlrmchlsmth as code owners March 18, 2026 05:38

mergify Bot added the needs-rebase label Mar 19, 2026

mergify Bot removed the needs-rebase label Mar 26, 2026

AndreasKaratzas requested a review from hmellor as a code owner March 26, 2026 09:32

AndreasKaratzas mentioned this pull request Mar 28, 2026

[ROCm][CI] Add K8s-hardened minimal Python CI runner with JUnit exit-code fix and resource watchdog #38444

Open

mergify Bot added the needs-rebase label Mar 30, 2026

AndreasKaratzas added a commit to ROCm/vllm that referenced this pull request Mar 31, 2026

Merging new CI executioner (vllm-project#38444) and fine-grained dock…

4663a78

…er changes with caching (vllm-project#36949) Signed-off-by: Andreas Karatzas <akaratza@amd.com>

mergify Bot removed the needs-rebase label Mar 31, 2026

AndreasKaratzas force-pushed the akaratza_optimize_docker_build branch from 4b4f357 to 012f864 Compare April 3, 2026 08:36

mergify Bot added the nvidia label Apr 3, 2026

github-project-automation Bot added this to NVIDIA Apr 3, 2026

AndreasKaratzas added 2 commits May 6, 2026 21:37

Docker cache best effort patch

422ac91

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Merge remote-tracking branch 'origin/main' into akaratza_optimize_doc…

a3c95e2

…ker_build Signed-off-by: Andreas Karatzas <akaratza@amd.com>

AndreasKaratzas requested a review from Harry-Chen as a code owner May 19, 2026 18:22

mawong-amd reviewed May 20, 2026

View reviewed changes

AndreasKaratzas added 2 commits May 21, 2026 17:39

Merge remote-tracking branch 'origin/main' into akaratza_optimize_doc…

f207534

…ker_build Signed-off-by: Andreas Karatzas <akaratza@amd.com>

[ROCm][CI] Optimize ROCm Docker CI caching and artifact tests

c7772d1

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

AndreasKaratzas requested a review from dllehr-amd as a code owner May 22, 2026 04:29

mergify Bot added the needs-rebase label May 22, 2026

Merge remote-tracking branch 'origin/main' into akaratza_optimize_doc…

5831245

…ker_build Signed-off-by: Andreas Karatzas <akaratza@amd.com>

mergify Bot removed the needs-rebase label May 22, 2026

[ROCm][CI] Unsetting arch completely

f1cc80e

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

depthfirst-app Bot reviewed May 22, 2026

View reviewed changes

Comment thread .buildkite/scripts/ci-bake-rocm.sh Outdated

AndreasKaratzas added 4 commits May 21, 2026 23:55

[ROCm][CI] Use cache-capable Buildx builder

0b1177f

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

[ROCm][CI] Remove remote ROCm CI HCL loading

7e0ae4f

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Fix wheel artifact archive issue

4e0fd9b

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Keep ROCm test images and improve cache reuse

2b86a2d

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

depthfirst-app Bot reviewed May 22, 2026

View reviewed changes

Comment thread docker/Dockerfile.rocm

Unsetting sscache

44cb1d3

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

mergify Bot added the needs-rebase label May 28, 2026

Merge remote-tracking branch 'origin/main' into akaratza_optimize_doc…

9ea1ffc

…ker_build Signed-off-by: Andreas Karatzas <akaratza@amd.com>

AndreasKaratzas requested a review from khluu as a code owner June 1, 2026 18:04

mergify Bot removed the needs-rebase label Jun 1, 2026

AndreasKaratzas added 3 commits June 1, 2026 22:33

Merge remote-tracking branch 'origin/main' into akaratza_optimize_doc…

d61a8b9

…ker_build

Adding retry logic for docker build

23cdacd

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Avoid concurrent Docker Hub cache exports for ROCm dependency cache t…

f10c35d

…argets. Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Harry-Chen approved these changes Jun 2, 2026

View reviewed changes

github-project-automation Bot moved this to Ready in NVIDIA Jun 2, 2026

		RUN --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
		uv pip install --system /deep_install/*.whl

Uh oh!

Conversation

AndreasKaratzas commented Mar 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Mar 13, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Mar 13, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Mar 13, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

AndreasKaratzas commented Mar 15, 2026

Uh oh!

mergify Bot commented Mar 19, 2026

Uh oh!

AndreasKaratzas commented Mar 23, 2026

Uh oh!

mergify Bot commented Mar 30, 2026

Uh oh!

mawong-amd left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mergify Bot commented May 22, 2026

Uh oh!

Uh oh!

Uh oh!

mergify Bot commented May 28, 2026

Uh oh!

AndreasKaratzas commented Jun 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

AndreasKaratzas commented Mar 13, 2026 •

edited

Loading

AndreasKaratzas commented Jun 2, 2026 •

edited

Loading