@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 87954eb50e
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 2, 02:15 UTC

[ROCm][CI] Optimize ROCm Docker build: registry cache, DeepEP, and ci-bake script (#36949)

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

Andreas Karatzas committed Jun 3, 2026 at 01:43 UTC 87954eb50ec90caaae8fb75124138dd8819cc49b
10 files changed +2745 -157
.buildkite/ci_config_rocm.yaml new
+23
@@ -0,0 +1,23 @@
1 +name: vllm_rocm_ci
2 +job_dirs:
3 + - ".buildkite/hardware_tests"
4 +run_all_patterns:
5 + - "docker/Dockerfile.rocm"
6 + - "docker/Dockerfile.rocm_base"
7 + - "docker/ci-rocm.hcl"
8 + - "docker/docker-bake-rocm.hcl"
9 + - ".buildkite/hardware_tests/amd.yaml"
10 + - ".buildkite/scripts/ci-bake-rocm.sh"
11 + - ".buildkite/scripts/hardware_ci/run-amd-test.py"
12 + - ".buildkite/scripts/hardware_ci/run-amd-test.sh"
13 + - "CMakeLists.txt"
14 + - "requirements/common.txt"
15 + - "requirements/rocm.txt"
16 + - "requirements/build/rocm.txt"
17 + - "requirements/test/rocm.txt"
18 + - "setup.py"
19 + - "csrc/"
20 + - "cmake/"
21 +run_all_exclude_patterns:
22 + - "csrc/cpu/"
23 + - "cmake/cpu_extension.cmake"
.buildkite/hardware_tests/amd.yaml
+65 -34
@@ -1,42 +1,73 @@
1 -group: Hardware - AMD Build
1 +group: Hardware - AMD Build
2 steps:
3 - - label: "AMD: :docker: build image"
4 - key: image-build-amd
3 + # Ensure ci_base is up-to-date before building the test image.
4 + # Compares a content hash of ci_base-affecting files against the remote
5 + # image label. If hashes match the build is skipped (< 30 s); if they
6 + # differ ci_base is rebuilt and pushed automatically.
7 + - label: "AMD: :docker: ensure ci_base"
8 + key: ensure-ci-base-amd
9 depends_on: []
10 device: amd_cpu
11 no_plugin: true
12 commands:
9 - - >
10 - docker build
11 - --build-arg max_jobs=16
12 - --build-arg REMOTE_VLLM=1
13 - --build-arg ARG_PYTORCH_ROCM_ARCH='gfx90a;gfx942;gfx950'
14 - --build-arg VLLM_BRANCH=$BUILDKITE_COMMIT
15 - --tag "rocm/vllm-ci:${BUILDKITE_COMMIT}"
16 - -f docker/Dockerfile.rocm
17 - --target test
18 - --no-cache
19 - --progress plain .
20 - - |
21 - docker run --rm --network=none --entrypoint /bin/bash "rocm/vllm-ci:${BUILDKITE_COMMIT}" -ec '
22 - if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi
23 - if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi
24 - if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi
25 - if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi
26 - command -v python3
27 - command -v uv
28 - command -v pytest
29 - if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then
30 - echo No ROCm CLI found in image >&2
31 - exit 1
13 + - bash .buildkite/scripts/ci-bake-rocm.sh ci-base-rocm-ci-with-deps
14 + env:
15 + DOCKER_BUILDKIT: "1"
16 + VLLM_BAKE_FILE: "docker/docker-bake-rocm.hcl"
17 + PYTORCH_ROCM_ARCH: "gfx90a;gfx942;gfx950"
18 + REMOTE_VLLM: "1"
19 + VLLM_BRANCH: "$BUILDKITE_COMMIT"
20 + retry:
21 + automatic:
22 + - exit_status: -1 # Agent was lost
23 + limit: 1
24 + - exit_status: -10 # Agent was lost
25 + limit: 1
26 +
27 + - label: "AMD: :docker: build test image and artifacts"
28 + key: image-build-amd
29 + depends_on:
30 + - ensure-ci-base-amd
31 + device: amd_cpu
32 + no_plugin: true
33 + commands:
34 + - |
35 + if [[ "${ROCM_CI_ARTIFACT_ONLY:-0}" == "1" ]]; then
36 + echo "ROCM_CI_ARTIFACT_ONLY=1; building ROCm wheel artifact only"
37 + IMAGE_TAG="" bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-artifacts
38 + else
39 + bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-wheel
40 fi
33 - python3 - <<PY
34 - import torch, vllm
35 - print(torch.__version__)
36 - print(vllm.__version__)
37 - PY
38 - echo AMD image smoke OK
39 - '
40 - - docker push "rocm/vllm-ci:${BUILDKITE_COMMIT}"
41 + - |
42 + docker run --rm --network=none --entrypoint /bin/bash "rocm/vllm-ci:${BUILDKITE_COMMIT}" -ec '
43 + if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi
44 + if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi
45 + if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi
46 + if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi
47 + command -v python3
48 + command -v uv
49 + command -v pytest
50 + if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then
51 + echo No ROCm CLI found in image >&2
52 + exit 1
53 + fi
54 + python3 - <<PY
55 + import torch, vllm
56 + print(torch.__version__)
57 + print(vllm.__version__)
58 + PY
59 + echo AMD image smoke OK
60 + '
61 env:
62 DOCKER_BUILDKIT: "1"
63 + VLLM_BAKE_FILE: "docker/docker-bake-rocm.hcl"
64 + PYTORCH_ROCM_ARCH: "gfx90a;gfx942;gfx950"
65 + IMAGE_TAG: "rocm/vllm-ci:$BUILDKITE_COMMIT"
66 + REMOTE_VLLM: "1"
67 + VLLM_BRANCH: "$BUILDKITE_COMMIT"
68 + retry:
69 + automatic:
70 + - exit_status: -1 # Agent was lost
71 + limit: 1
72 + - exit_status: -10 # Agent was lost
73 + limit: 1
.buildkite/scripts/ci-bake-rocm.sh new
+1749
@@ -0,0 +1,1749 @@
1 +#!/bin/bash
2 +# ci-bake-rocm.sh - Docker buildx bake wrapper for ROCm CI builds.
3 +#
4 +# The wrapper keeps three build concerns separate:
5 +# * ci_base builds are content-addressed by vllm.ci_base.content_hash.
6 +# * test image builds are commit-addressed by org.opencontainers.image.revision.
7 +# * ROCm install artifacts are uploaded only for targets that export wheels.
8 +#
9 +# Usage:
10 +# ci-bake-rocm.sh [TARGET]
11 +#
12 +# Set BAKE_PRINT_ONLY=1 to stop after docker buildx bake --print.
13 +
14 +set -euo pipefail
15 +
16 +DEFAULT_REPO_SLUG="vllm-project/vllm"
17 +DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl"
18 +DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base tools/install_torchcodec_rocm.sh tests/vllm_test_utils"
19 +DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm"
20 +DEFAULT_CI_BASE_DOCKERFILE_STAGES="base build_rixl build_rocshmem build_deepep mori_base ci_base"
21 +IMAGE_EXISTED_BEFORE_BUILD=0
22 +
23 +TARGET=""
24 +CI_HCL_SOURCE="${CI_HCL_SOURCE:-}"
25 +CI_HCL_PATH=""
26 +CI_BASE_LABEL_OVERRIDE_PATH=""
27 +CSRC_CACHE_OVERRIDE_PATH=""
28 +ROCM_ARG_OVERRIDE_PATH=""
29 +SCRIPT_TMP_DIR=""
30 +BAKE_CONFIG_FILE=""
31 +BAKE_FILES=()
32 +BAKE_TARGETS=()
33 +DEPENDENCY_CACHE_TARGETS=()
34 +
35 +cleanup() {
36 + if [[ -n "${SCRIPT_TMP_DIR}" && -d "${SCRIPT_TMP_DIR}" ]]; then
37 + rm -rf "${SCRIPT_TMP_DIR}"
38 + fi
39 +}
40 +trap cleanup EXIT
41 +
42 +clean_docker_tag() {
43 + local input="$1"
44 + echo "${input}" | sed 's/[^a-zA-Z0-9._-]/_/g' | cut -c1-128
45 +}
46 +
47 +is_url_like() {
48 + local value="${1:-}"
49 + [[ "${value}" =~ ^[a-zA-Z][a-zA-Z0-9+.-]*:// || "${value}" == git@*:* ]]
50 +}
51 +
52 +is_full_git_sha() {
53 + local value="${1:-}"
54 + [[ "${value}" =~ ^[0-9a-fA-F]{40}$ ]]
55 +}
56 +
57 +select_cache_branch_name() {
58 + local candidate=""
59 + local var=""
60 +
61 + for var in \
62 + ROCM_CACHE_BRANCH_NAME \
63 + BUILDKITE_PULL_REQUEST_HEAD_BRANCH \
64 + BUILDKITE_HEAD_BRANCH \
65 + BUILDKITE_BRANCH \
66 + VLLM_BRANCH; do
67 + candidate="${!var:-}"
68 + [[ -n "${candidate}" ]] || continue
69 + is_url_like "${candidate}" && continue
70 + is_full_git_sha "${candidate}" && continue
71 + printf '%s\n' "${candidate}"
72 + return 0
73 + done
74 +}
75 +
76 +cache_scope_suffix() {
77 + local arch_hash=""
78 + arch_hash=$(printf '%s' "${PYTORCH_ROCM_ARCH:-default}" | sha256sum | cut -c1-12)
79 + printf 'arch-%s\n' "${arch_hash}"
80 +}
81 +
82 +compose_cache_branch_tag() {
83 + local repo_slug="$1"
84 + local branch="$2"
85 + local suffix=""
86 + local prefix=""
87 + local max_prefix_len=0
88 +
89 + suffix="$(cache_scope_suffix)"
90 + prefix="$(clean_docker_tag "${repo_slug}")-$(clean_docker_tag "${branch}")"
91 + max_prefix_len=$((128 - ${#suffix} - 1))
92 + if (( max_prefix_len < 1 )); then
93 + max_prefix_len=1
94 + fi
95 + printf '%s-%s\n' "${prefix:0:${max_prefix_len}}" "${suffix}"
96 +}
97 +
98 +parse_repo_slug() {
99 + local repo_url="${1:-}"
100 + local repo_slug=""
101 +
102 + if [[ -z "${repo_url}" ]]; then
103 + printf '%s\n' "${DEFAULT_REPO_SLUG}"
104 + return 0
105 + fi
106 +
107 + repo_slug=$(echo "${repo_url}" | sed -E 's#(git@|https?://)([^/:]+)[:/]([^/]+/[^/.]+)(\.git)?$#\3#')
108 + if [[ "${repo_slug}" != */* ]]; then
109 + repo_slug="${DEFAULT_REPO_SLUG}"
110 + fi
111 + printf '%s\n' "${repo_slug}"
112 +}
113 +
114 +get_buildkite_repo_slug() {
115 + parse_repo_slug "${BUILDKITE_PULL_REQUEST_REPO:-${BUILDKITE_REPO:-}}"
116 +}
117 +
118 +get_buildkite_target_repo_slug() {
119 + parse_repo_slug "${BUILDKITE_REPO:-}"
120 +}
121 +
122 +get_buildkite_target_repo_url() {
123 + local repo_url="${BUILDKITE_REPO:-}"
124 +
125 + if [[ -n "${repo_url}" ]] && is_url_like "${repo_url}"; then
126 + printf '%s\n' "${repo_url}"
127 + return 0
128 + fi
129 +
130 + printf 'https://github.com/%s.git\n' "${DEFAULT_REPO_SLUG}"
131 +}
132 +
133 +git_fetch_for_cache() {
134 + local timeout_secs="${ROCM_CACHE_GIT_FETCH_TIMEOUT:-60}"
135 +
136 + if command -v timeout >/dev/null 2>&1; then
137 + timeout "${timeout_secs}s" git fetch "$@" 2>/dev/null
138 + else
139 + git fetch "$@" 2>/dev/null
140 + fi
141 +}
142 +
143 +hash_string_short() {
144 + printf '%s' "$1" | sha256sum | cut -c1-16
145 +}
146 +
147 +compute_content_hash() {
148 + local path
149 + local file
150 +
151 + for path in "$@"; do
152 + if [[ -d "${path}" ]]; then
153 + while IFS= read -r -d '' file; do
154 + printf 'file:%s\n' "${file}"
155 + sha256sum "${file}"
156 + done < <(find "${path}" -type f -print0 | sort -z)
157 + elif [[ -f "${path}" ]]; then
158 + printf 'file:%s\n' "${path}"
159 + sha256sum "${path}"
160 + else
161 + printf 'missing:%s\n' "${path}"
162 + fi
163 + done | sha256sum | cut -d' ' -f1
164 +}
165 +
166 +compose_dependency_cache_key() {
167 + local prefix="$1"
168 + local material="$2"
169 + local cleaned_prefix=""
170 +
171 + cleaned_prefix=$(clean_docker_tag "${prefix}" | cut -c1-96)
172 + printf '%s-%s\n' "${cleaned_prefix}" "$(hash_string_short "${material}")"
173 +}
174 +
175 +hash_dockerfile_stages() {
176 + local dockerfile="$1"
177 + local stages="$2"
178 +
179 + awk -v wanted_stages="${stages}" '
180 + BEGIN {
181 + split(wanted_stages, stage_list, /[[:space:]]+/)
182 + for (idx in stage_list) {
183 + if (stage_list[idx] != "") {
184 + wanted[stage_list[idx]] = 1
185 + }
186 + }
187 + emit = 1
188 + }
189 + $1 == "FROM" {
190 + stage = ""
191 + for (idx = 1; idx <= NF; idx++) {
192 + if (tolower($idx) == "as" && idx < NF) {
193 + stage = $(idx + 1)
194 + }
195 + }
196 + emit = (stage in wanted)
197 + }
198 + emit {
199 + print
200 + }
201 + ' "${dockerfile}"
202 +}
203 +
204 +discover_dockerfile_stage_args() {
205 + local dockerfile="$1"
206 + local stages="$2"
207 +
208 + [[ -f "${dockerfile}" ]] || return 0
209 +
210 + awk -v wanted_stages="${stages}" '
211 + function add_arg(name) {
212 + if (name != "" && !(name in seen)) {
213 + seen[name] = 1
214 + args[++arg_count] = name
215 + }
216 + }
217 + BEGIN {
218 + split(wanted_stages, stage_list, /[[:space:]]+/)
219 + for (idx in stage_list) {
220 + if (stage_list[idx] != "") {
221 + wanted[stage_list[idx]] = 1
222 + }
223 + }
224 + emit = 1
225 + }
226 + {
227 + line = $0
228 + if ($1 == "FROM") {
229 + stage = ""
230 + for (idx = 1; idx <= NF; idx++) {
231 + if (tolower($idx) == "as" && idx < NF) {
232 + stage = $(idx + 1)
233 + }
234 + }
235 + emit = (stage in wanted)
236 + }
237 + if (emit) {
238 + lines[++line_count] = line
239 + }
240 + }
241 + END {
242 + for (idx = 1; idx <= line_count; idx++) {
243 + line = lines[idx]
244 + arg_name = line
245 + sub(/^[[:space:]]*ARG[[:space:]]+/, "", arg_name)
246 + if (arg_name != line) {
247 + sub(/[=[:space:]].*/, "", arg_name)
248 + if (arg_name ~ /^[A-Za-z_][A-Za-z0-9_]*$/) {
249 + add_arg(arg_name)
250 + }
251 + }
252 + }
253 +
254 + for (idx = 1; idx <= line_count; idx++) {
255 + line = lines[idx]
256 + for (arg_idx = 1; arg_idx <= arg_count; arg_idx++) {
257 + name = args[arg_idx]
258 + if (line ~ "\\$\\{" name "([}:][^}]*)?\\}" \
259 + || line ~ "\\$" name "([^A-Za-z0-9_]|$)") {
260 + used[name] = 1
261 + }
262 + }
263 + }
264 +
265 + for (arg_idx = 1; arg_idx <= arg_count; arg_idx++) {
266 + name = args[arg_idx]
267 + if (used[name]) {
268 + print name
269 + }
270 + }
271 + }
272 + ' "${dockerfile}"
273 +}
274 +
275 +get_content_arg_names() {
276 + local dockerfile="$1"
277 + local stages="$2"
278 + local explicit_args="${3:-}"
279 +
280 + if [[ -n "${explicit_args}" ]]; then
281 + tr ' ' '\n' <<< "${explicit_args}"
282 + else
283 + discover_dockerfile_stage_args "${dockerfile}" "${stages}"
284 + fi | awk 'NF && !seen[$0]++'
285 +}
286 +
287 +compute_ci_base_content_hash() {
288 + local -a content_paths=()
289 + local -a content_args=()
290 + local dockerfile="${CI_BASE_DOCKERFILE:-}"
291 + local stages="${CI_BASE_DOCKERFILE_STAGES:-}"
292 +
293 + read -r -a content_paths <<< "${CI_BASE_CONTENT_FILES}"
294 + mapfile -t content_args < <(
295 + get_content_arg_names "${dockerfile}" "${stages}" "${CI_BASE_CONTENT_ARGS:-}"
296 + )
297 +
298 + {
299 + printf 'content-files-hash:%s\n' "$(compute_content_hash "${content_paths[@]}")"
300 + if [[ -n "${dockerfile}" ]]; then
301 + printf 'dockerfile:%s\n' "${dockerfile}"
302 + printf 'resolved-build-args:\n'
303 + hash_dockerfile_arg_values "${dockerfile}" "${content_args[@]}"
304 + if [[ -n "${stages}" ]]; then
305 + printf 'dockerfile-stages:%s\n' "${stages}"
306 + if [[ -f "${dockerfile}" ]]; then
307 + hash_dockerfile_stages "${dockerfile}" "${stages}"
308 + else
309 + printf 'missing:%s\n' "${dockerfile}"
310 + fi
311 + fi
312 + fi
313 + } | sha256sum | cut -d' ' -f1
314 +}
315 +
316 +extract_dockerfile_arg_default() {
317 + local dockerfile="$1"
318 + local arg_name="$2"
319 +
320 + sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \
321 + "${dockerfile}" | head -1
322 +}
323 +
324 +resolve_image_digest() {
325 + local image_ref="$1"
326 +
327 + docker buildx imagetools inspect "${image_ref}" 2>/dev/null \
328 + | sed -n -E 's/^Digest:[[:space:]]+//p' \
329 + | head -1
330 +}
331 +
332 +resolve_dockerfile_arg_value() {
333 + local dockerfile="$1"
334 + local arg_name="$2"
335 + local env_name="${arg_name}"
336 + local value=""
337 +
338 + case "${arg_name}" in
339 + ARG_PYTORCH_ROCM_ARCH)
340 + env_name="PYTORCH_ROCM_ARCH"
341 + ;;
342 + esac
343 +
344 + value="${!env_name:-}"
345 + if [[ -z "${value}" && "${env_name}" != "${arg_name}" ]]; then
346 + value="${!arg_name:-}"
347 + fi
348 + if [[ -z "${value}" && -f "${dockerfile}" ]]; then
349 + value=$(extract_dockerfile_arg_default "${dockerfile}" "${arg_name}")
350 + fi
351 +
352 + printf '%s\n' "${value}"
353 +}
354 +
355 +hash_dockerfile_arg_values() {
356 + local dockerfile="$1"
357 + local arg_name=""
358 + local arg_value=""
359 + local digest=""
360 + shift || true
361 +
362 + for arg_name in "$@"; do
363 + [[ -n "${arg_name}" ]] || continue
364 + arg_value=$(resolve_dockerfile_arg_value "${dockerfile}" "${arg_name}")
365 + printf 'arg:%s=%s\n' "${arg_name}" "${arg_value:-<empty>}"
366 + if [[ "${arg_name}" == "BASE_IMAGE" && -n "${arg_value}" ]]; then
367 + digest=$(resolve_image_digest "${arg_value}")
368 + printf 'arg:%s.digest=%s\n' "${arg_name}" "${digest:-unknown}"
369 + fi
370 + done
371 +}
372 +
373 +is_ci_base_target() {
374 + [[ "${TARGET}" == *"ci-base-rocm"* ]]
375 +}
376 +
377 +is_commit_image_target() {
378 + [[ -n "${IMAGE_TAG:-}" && -n "${BUILDKITE_COMMIT:-}" ]] || return 1
379 + is_ci_base_target && return 1
380 + return 0
381 +}
382 +
383 +image_tag_is_commit_scoped() {
384 + [[ -n "${IMAGE_TAG:-}" && -n "${BUILDKITE_COMMIT:-}" ]] || return 1
385 + [[ "${IMAGE_TAG}" == *"${BUILDKITE_COMMIT}"* ]]
386 +}
387 +
388 +should_upload_wheel_artifacts() {
389 + [[ "${UPLOAD_ROCM_WHEEL_ARTIFACTS:-0}" == "1" ]] && return 0
390 + [[ "${TARGET}" == *"with-wheel"* \
391 + || "${TARGET}" == *"export-wheel"* \
392 + || "${TARGET}" == *"artifact"* ]]
393 +}
394 +
395 +get_remote_image_label() {
396 + local image_ref="$1"
397 + local label_key="$2"
398 +
399 + docker buildx imagetools inspect "${image_ref}" --raw 2>/dev/null \
400 + | python3 -c '
401 +import json
402 +import subprocess
403 +import sys
404 +import urllib.parse
405 +import urllib.request
406 +
407 +image_ref = sys.argv[1]
408 +label_key = sys.argv[2]
409 +
410 +
411 +def docker_hub_repo(image_name):
412 + image_name = image_name.split("@", 1)[0]
413 + last_component = image_name.rsplit("/", 1)[-1]
414 + if ":" in last_component:
415 + image_name = image_name.rsplit(":", 1)[0]
416 +
417 + parts = image_name.split("/")
418 + if len(parts) > 1 and (
419 + "." in parts[0] or ":" in parts[0] or parts[0] == "localhost"
420 + ):
421 + registry = parts[0]
422 + if registry not in {
423 + "docker.io",
424 + "index.docker.io",
425 + "registry-1.docker.io",
426 + }:
427 + return None
428 + image_name = "/".join(parts[1:])
429 + elif len(parts) == 1:
430 + image_name = f"library/{image_name}"
431 +
432 + return image_name
433 +
434 +
435 +try:
436 + data = json.load(sys.stdin)
437 + if data.get("manifests"):
438 + manifest = next(
439 + (
440 + entry
441 + for entry in data["manifests"]
442 + if entry.get("platform", {}).get("os") != "unknown"
443 + and entry.get("platform", {}).get("architecture") != "unknown"
444 + ),
445 + data["manifests"][0],
446 + )
447 + digest = manifest["digest"]
448 + result = subprocess.run(
449 + [
450 + "docker",
451 + "buildx",
452 + "imagetools",
453 + "inspect",
454 + image_ref + "@" + digest,
455 + "--raw",
456 + ],
457 + capture_output=True,
458 + text=True,
459 + check=False,
460 + )
461 + if result.returncode != 0 or not result.stdout:
462 + raise RuntimeError("digest inspect failed")
463 + data = json.loads(result.stdout)
464 +
465 + annotations = data.get("annotations", {})
466 + if label_key in annotations:
467 + print(annotations[label_key])
468 + raise SystemExit(0)
469 +
470 + config_digest = data.get("config", {}).get("digest")
471 + if not config_digest:
472 + print("")
473 + raise SystemExit(0)
474 +
475 + image_name = docker_hub_repo(image_ref)
476 + if not image_name:
477 + print("")
478 + raise SystemExit(0)
479 +
480 + token_url = (
481 + "https://auth.docker.io/token?"
482 + + urllib.parse.urlencode(
483 + {
484 + "service": "registry.docker.io",
485 + "scope": f"repository:{image_name}:pull",
486 + }
487 + )
488 + )
489 + with urllib.request.urlopen(token_url, timeout=30) as response:
490 + token = json.load(response)["token"]
491 +
492 + request = urllib.request.Request(
493 + f"https://registry-1.docker.io/v2/{image_name}/blobs/{config_digest}",
494 + headers={"Authorization": f"Bearer {token}"},
495 + )
496 + with urllib.request.urlopen(request, timeout=30) as response:
497 + config_blob = json.load(response)
498 +
499 + labels = config_blob.get("config", {}).get("Labels", {})
500 + print(labels.get(label_key, ""))
501 +except Exception:
502 + print("")
503 +' "${image_ref}" "${label_key}" 2>/dev/null || echo ""
504 +}
505 +
506 +get_remote_image_label_with_retry() {
507 + local image_ref="$1"
508 + local label_key="$2"
509 + local attempts="${3:-6}"
510 + local delay_secs="${4:-5}"
511 + local label_value=""
512 + local attempt
513 +
514 + for ((attempt = 1; attempt <= attempts; attempt++)); do
515 + label_value=$(get_remote_image_label "${image_ref}" "${label_key}")
516 + if [[ -n "${label_value}" ]]; then
517 + printf '%s\n' "${label_value}"
518 + return 0
519 + fi
520 + if [[ ${attempt} -lt ${attempts} ]]; then
521 + sleep "${delay_secs}"
522 + fi
523 + done
524 +
525 + return 0
526 +}
527 +
528 +remote_image_exists() {
529 + local image_ref="$1"
530 + docker manifest inspect "${image_ref}" >/dev/null 2>&1
531 +}
532 +
533 +use_existing_builder() {
534 + echo "Using existing builder: ${BUILDER_NAME}"
535 + docker buildx use "${BUILDER_NAME}"
536 + docker buildx inspect --bootstrap
537 +}
538 +
539 +buildx_driver() {
540 + local builder="${1:-}"
541 +
542 + if [[ -n "${builder}" ]]; then
543 + docker buildx inspect "${builder}" 2>/dev/null
544 + else
545 + docker buildx inspect 2>/dev/null
546 + fi | awk -F': *' '$1 == "Driver" { print $2; exit }'
547 +}
548 +
549 +builder_supports_registry_cache() {
550 + local driver="$1"
551 +
552 + [[ -n "${driver}" && "${driver}" != "docker" ]]
553 +}
554 +
555 +create_and_bootstrap_builder() {
556 + local driver="$1"
557 + local endpoint="${2:-}"
558 +
559 + echo "Creating builder '${BUILDER_NAME}' with ${driver} driver"
560 + if [[ -n "${endpoint}" ]]; then
561 + docker buildx create \
562 + --name "${BUILDER_NAME}" \
563 + --driver "${driver}" \
564 + --use \
565 + "${endpoint}"
566 + else
567 + docker buildx create --name "${BUILDER_NAME}" --driver "${driver}" --use
568 + fi
569 + docker buildx inspect --bootstrap
570 +}
571 +
572 +init_config() {
573 + TARGET="${1:-test-ci}"
574 + BAKE_TARGETS=("${TARGET}")
575 + DEPENDENCY_CACHE_TARGETS=()
576 + CI_HCL_SOURCE="${CI_HCL_SOURCE:-${CI_HCL_FILE:-${DEFAULT_CI_HCL_SOURCE}}}"
577 + VLLM_BAKE_FILE="${VLLM_BAKE_FILE:-docker/docker-bake-rocm.hcl}"
578 + BUILDER_NAME="${BUILDER_NAME:-vllm-builder}"
579 + BUILDKIT_SOCKET="${BUILDKIT_SOCKET:-/run/buildkit/buildkitd.sock}"
580 + PYTORCH_ROCM_ARCH="${PYTORCH_ROCM_ARCH:-gfx90a;gfx942;gfx950}"
581 + CI_BASE_CONTENT_FILES="${CI_BASE_CONTENT_FILES:-${DEFAULT_CI_BASE_CONTENT_FILES}}"
582 + CI_BASE_DOCKERFILE="${CI_BASE_DOCKERFILE:-${DEFAULT_CI_BASE_DOCKERFILE}}"
583 + CI_BASE_DOCKERFILE_STAGES="${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}"
584 + CI_BASE_IMAGE_TAG="${CI_BASE_IMAGE_TAG:-rocm/vllm-dev:ci_base}"
585 + export PYTORCH_ROCM_ARCH
586 +
587 + SCRIPT_TMP_DIR=$(mktemp -d -t ci-bake-rocm.XXXXXX)
588 + CI_HCL_PATH="${SCRIPT_TMP_DIR}/ci.hcl"
589 + CI_BASE_LABEL_OVERRIDE_PATH="${SCRIPT_TMP_DIR}/ci-base-label-override.hcl"
590 + CSRC_CACHE_OVERRIDE_PATH="${SCRIPT_TMP_DIR}/rocm-csrc-cache-override.hcl"
591 + ROCM_ARG_OVERRIDE_PATH="${SCRIPT_TMP_DIR}/rocm-arg-override.hcl"
592 + BAKE_CONFIG_FILE="bake-config-build-${BUILDKITE_BUILD_NUMBER:-local}.json"
593 +}
594 +
595 +print_header() {
596 + echo "--- :docker: Setting up Docker buildx bake"
597 + echo "Target: ${TARGET}"
598 + echo "CI HCL source: ${CI_HCL_SOURCE}"
599 + echo "vLLM bake file: ${VLLM_BAKE_FILE}"
600 + if is_ci_base_target; then
601 + echo "Build mode: ci_base"
602 + elif is_commit_image_target; then
603 + echo "Build mode: commit image"
604 + else
605 + echo "Build mode: generic"
606 + fi
607 + if [[ "${USE_SCCACHE:-0}" == "1" ]]; then
608 + echo "Compiler cache: sccache enabled"
609 + fi
610 +}
611 +
612 +validate_inputs() {
613 + if [[ ! -f "${VLLM_BAKE_FILE}" ]]; then
614 + echo "Error: vLLM bake file not found at ${VLLM_BAKE_FILE}"
615 + echo "Make sure you're running from the vLLM repository root"
616 + exit 1
617 + fi
618 +
619 + if [[ -n "${CI_HCL_SOURCE:-}" ]] && is_url_like "${CI_HCL_SOURCE}"; then
620 + echo "Error: remote CI HCL sources are not supported: ${CI_HCL_SOURCE}"
621 + echo "Use the vLLM-owned docker/ci-rocm.hcl or set CI_HCL_SOURCE to a local file."
622 + exit 1
623 + fi
624 +
625 + if [[ -n "${CI_HCL_SOURCE:-}" && ! -f "${CI_HCL_SOURCE}" ]]; then
626 + echo "Error: CI HCL file not found at ${CI_HCL_SOURCE}"
627 + echo "Set CI_HCL_SOURCE to a local file if you need an override."
628 + exit 1
629 + fi
630 +}
631 +
632 +load_ci_hcl() {
633 + echo "--- :page_facing_up: Loading ci.hcl"
634 + cp "${CI_HCL_SOURCE}" "${CI_HCL_PATH}"
635 + echo "Copied ${CI_HCL_SOURCE} to ${CI_HCL_PATH}"
636 +}
637 +
638 +compute_ci_base_hash_if_needed() {
639 + if [[ -z "${CI_BASE_CONTENT_FILES:-}" ]]; then
640 + return 0
641 + fi
642 +
643 + CI_BASE_CONTENT_HASH=$(compute_ci_base_content_hash)
644 + export CI_BASE_CONTENT_HASH
645 + echo "ci_base content hash: ${CI_BASE_CONTENT_HASH:0:16}..."
646 +}
647 +
648 +should_push_stable_ci_base_tag() {
649 + if [[ "${CI_BASE_PUSH_STABLE_TAG:-}" == "1" ]]; then
650 + return 0
651 + fi
652 + if [[ "${CI_BASE_PUSH_STABLE_TAG:-}" == "0" ]]; then
653 + return 1
654 + fi
655 +
656 + [[ "${NIGHTLY:-0}" == "1" && "${BUILDKITE_BRANCH:-}" == "${CI_BASE_STABLE_BRANCH:-main}" ]]
657 +}
658 +
659 +ci_base_tag_with_suffix() {
660 + local base_tag="$1"
661 + local suffix="$2"
662 +
663 + printf '%s-%s\n' "${base_tag}" "$(clean_docker_tag "${suffix}")"
664 +}
665 +
666 +configure_ci_base_image_refs() {
667 + local stable_tag="${CI_BASE_IMAGE_TAG:-rocm/vllm-dev:ci_base}"
668 + local content_tag=""
669 + local commit_tag=""
670 + local primary_tag=""
671 +
672 + if [[ -z "${CI_BASE_CONTENT_HASH:-}" ]]; then
673 + CI_BASE_IMAGE="${CI_BASE_IMAGE:-${stable_tag}}"
674 + export CI_BASE_IMAGE
675 + return 0
676 + fi
677 +
678 + content_tag=$(ci_base_tag_with_suffix "${stable_tag}" "${CI_BASE_CONTENT_HASH}")
679 + if [[ -n "${BUILDKITE_COMMIT:-}" ]]; then
680 + commit_tag=$(ci_base_tag_with_suffix "${stable_tag}" "${BUILDKITE_COMMIT}")
681 + CI_BASE_IMAGE_TAG_COMMIT="${commit_tag}"
682 + export CI_BASE_IMAGE_TAG_COMMIT
683 + fi
684 +
685 + if should_push_stable_ci_base_tag; then
686 + primary_tag="${content_tag}"
687 + CI_BASE_IMAGE_TAG_STABLE="${stable_tag}"
688 + else
689 + primary_tag="${commit_tag:-${content_tag}}"
690 + CI_BASE_IMAGE_TAG_STABLE=""
691 + fi
692 + CI_BASE_IMAGE_TAG="${primary_tag}"
693 + if [[ "${primary_tag}" == "${content_tag}" ]]; then
694 + CI_BASE_IMAGE_TAG_CONTENT=""
695 + else
696 + CI_BASE_IMAGE_TAG_CONTENT="${content_tag}"
697 + fi
698 + export CI_BASE_IMAGE_TAG CI_BASE_IMAGE_TAG_CONTENT CI_BASE_IMAGE_TAG_STABLE
699 +
700 + if is_ci_base_target; then
701 + IMAGE_TAG="${primary_tag}"
702 + export IMAGE_TAG
703 +
704 + echo "ci_base primary image tag: ${CI_BASE_IMAGE_TAG}"
705 + if [[ -n "${CI_BASE_IMAGE_TAG_COMMIT:-}" ]]; then
706 + echo "ci_base commit image tag: ${CI_BASE_IMAGE_TAG_COMMIT}"
707 + fi
708 + echo "ci_base content image tag: ${content_tag}"
709 + if [[ -n "${CI_BASE_IMAGE_TAG_STABLE}" ]]; then
710 + echo "ci_base stable alias will also be pushed: ${CI_BASE_IMAGE_TAG_STABLE}"
711 + else
712 + echo "ci_base stable alias will not be pushed for this build"
713 + echo "Set NIGHTLY=1 on ${CI_BASE_STABLE_BRANCH:-main} to refresh ${stable_tag}"
714 + fi
715 + return 0
716 + fi
717 +
718 + if [[ -z "${CI_BASE_IMAGE:-}" || "${CI_BASE_IMAGE}" == "${stable_tag}" ]]; then
719 + CI_BASE_IMAGE="${primary_tag}"
720 + export CI_BASE_IMAGE
721 + echo "Using ci_base image: ${CI_BASE_IMAGE}"
722 + else
723 + echo "Using provided CI_BASE_IMAGE override: ${CI_BASE_IMAGE}"
724 + fi
725 +}
726 +
727 +ci_base_candidate_refs() {
728 + printf '%s\n' \
729 + "${IMAGE_TAG:-}" \
730 + "${CI_BASE_IMAGE_TAG:-}" \
731 + "${CI_BASE_IMAGE_TAG_COMMIT:-}" \
732 + "${CI_BASE_IMAGE_TAG_CONTENT:-}" \
733 + "${CI_BASE_IMAGE_TAG_STABLE:-}" \
734 + | awk 'NF && !seen[$0]++'
735 +}
736 +
737 +find_matching_ci_base_ref() {
738 + local candidate=""
739 + local candidate_hash=""
740 +
741 + while IFS= read -r candidate; do
742 + [[ -n "${candidate}" ]] || continue
743 + remote_image_exists "${candidate}" || continue
744 + candidate_hash=$(get_remote_image_label "${candidate}" "vllm.ci_base.content_hash")
745 + if [[ "${candidate_hash}" == "${CI_BASE_CONTENT_HASH}" ]]; then
746 + printf '%s\n' "${candidate}"
747 + return 0
748 + fi
749 + done < <(ci_base_candidate_refs)
750 +
751 + return 1
752 +}
753 +
754 +refresh_ci_base_tags_from_ref() {
755 + local source_ref="$1"
756 + local tag=""
757 + local tag_hash=""
758 +
759 + while IFS= read -r tag; do
760 + [[ -n "${tag}" ]] || continue
761 + [[ "${tag}" != "${source_ref}" ]] || continue
762 + tag_hash=$(get_remote_image_label "${tag}" "vllm.ci_base.content_hash")
763 + if [[ "${tag_hash}" == "${CI_BASE_CONTENT_HASH}" ]]; then
764 + echo "ci_base tag is already current: ${tag}"
765 + continue
766 + fi
767 + echo "Updating ci_base tag ${tag} -> ${source_ref}"
768 + docker buildx imagetools create -t "${tag}" "${source_ref}"
769 + done < <(ci_base_candidate_refs)
770 +}
771 +
772 +maybe_skip_existing_image() {
773 + local remote_hash=""
774 + local remote_revision=""
775 + local matching_ref=""
776 +
777 + if [[ -z "${IMAGE_TAG:-}" ]]; then
778 + return 0
779 + fi
780 +
781 + if [[ "${FORCE_BUILD:-0}" == "1" ]]; then
782 + echo "FORCE_BUILD=1 set; skipping existing-image check"
783 + return 0
784 + fi
785 +
786 + echo "--- :mag: Checking image tag"
787 + echo "Image tag: ${IMAGE_TAG}"
788 +
789 + if ! remote_image_exists "${IMAGE_TAG}"; then
790 + if is_ci_base_target && [[ -n "${CI_BASE_CONTENT_HASH:-}" ]]; then
791 + matching_ref=$(find_matching_ci_base_ref || true)
792 + if [[ -n "${matching_ref}" ]]; then
793 + echo "Found existing ci_base image with matching content hash: ${matching_ref}"
794 + if ! refresh_ci_base_tags_from_ref "${matching_ref}"; then
795 + echo "ci_base tag refresh failed; rebuilding to push expected tags"
796 + return 0
797 + fi
798 + echo "Content hashes match -- ci_base is current"
799 + echo "Skipping build"
800 + exit 0
801 + fi
802 + fi
803 + echo "Image not found, proceeding with build"
804 + return 0
805 + fi
806 +
807 + IMAGE_EXISTED_BEFORE_BUILD=1
808 +
809 + if is_ci_base_target; then
810 + if [[ -z "${CI_BASE_CONTENT_HASH:-}" ]]; then
811 + echo "ci_base image already exists and no content hash was configured"
812 + echo "Skipping build"
813 + exit 0
814 + fi
815 +
816 + remote_hash=$(get_remote_image_label "${IMAGE_TAG}" "vllm.ci_base.content_hash")
817 + if [[ -n "${remote_hash}" ]]; then
818 + echo "Remote ci_base content hash: ${remote_hash:0:16}..."
819 + if [[ "${remote_hash}" == "${CI_BASE_CONTENT_HASH}" ]]; then
820 + if ! refresh_ci_base_tags_from_ref "${IMAGE_TAG}"; then
821 + echo "ci_base tag refresh failed; rebuilding to push expected tags"
822 + return 0
823 + fi
824 + echo "Content hashes match -- ci_base is current"
825 + echo "Skipping build"
826 + exit 0
827 + fi
828 +
829 + echo "Content hashes differ -- ci_base is stale, rebuilding"
830 + return 0
831 + fi
832 +
833 + echo "Remote ci_base has no content-hash label; rebuilding to add one"
834 + return 0
835 + fi
836 +
837 + if is_commit_image_target; then
838 + remote_revision=$(get_remote_image_label "${IMAGE_TAG}" "org.opencontainers.image.revision")
839 + if [[ -n "${remote_revision}" && "${remote_revision}" != "${BUILDKITE_COMMIT}" ]]; then
840 + echo "Existing image revision does not match ${BUILDKITE_COMMIT}"
841 + echo " found revision: ${remote_revision}"
842 + echo "Rebuilding image"
843 + return 0
844 + fi
845 +
846 + if should_upload_wheel_artifacts; then
847 + echo "Commit image already exists: ${IMAGE_TAG}"
848 + echo "Continuing build because this target uploads per-build ROCm artifacts"
849 + return 0
850 + fi
851 +
852 + echo "Commit image already exists: ${IMAGE_TAG}"
853 + echo "Skipping build"
854 + exit 0
855 + fi
856 +
857 + echo "Image already exists: ${IMAGE_TAG}"
858 + echo "Skipping build"
859 + exit 0
860 +}
861 +
862 +setup_builder() {
863 + echo "--- :buildkite: Setting up buildx builder"
864 +
865 + local setup_mode="${ROCM_SETUP_BUILDX_BUILDER:-auto}"
866 + local current_driver=""
867 + local named_driver=""
868 +
869 + if [[ "${setup_mode}" == "0" || "${setup_mode}" == "false" ]]; then
870 + echo "Using current Docker buildx builder"
871 + echo "ROCM_SETUP_BUILDX_BUILDER=${setup_mode}; cache exporters may fail if the driver is docker"
872 + docker buildx inspect --bootstrap
873 + echo "Active builder:"
874 + docker buildx ls | grep -E '^\*|^NAME' || docker buildx ls
875 + return 0
876 + fi
877 +
878 + current_driver=$(buildx_driver || true)
879 + if [[ "${setup_mode}" != "1" ]] && builder_supports_registry_cache "${current_driver}"; then
880 + echo "Using current Docker buildx builder with ${current_driver} driver"
881 + docker buildx inspect --bootstrap
882 + echo "Active builder:"
883 + docker buildx ls | grep -E '^\*|^NAME' || docker buildx ls
884 + return 0
885 + fi
886 +
887 + if [[ "${setup_mode}" != "1" ]]; then
888 + echo "Current buildx driver '${current_driver:-unknown}' cannot export registry caches"
889 + echo "Creating or using a cache-capable builder: ${BUILDER_NAME}"
890 + fi
891 +
892 + if docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then
893 + named_driver=$(buildx_driver "${BUILDER_NAME}" || true)
894 + if ! builder_supports_registry_cache "${named_driver}"; then
895 + echo "Builder '${BUILDER_NAME}' uses ${named_driver:-unknown} driver; using ${BUILDER_NAME}-cache instead"
896 + BUILDER_NAME="${BUILDER_NAME}-cache"
897 + fi
898 + fi
899 +
900 + if [[ -S "${BUILDKIT_SOCKET}" ]]; then
901 + echo "Found local buildkitd socket at ${BUILDKIT_SOCKET}"
902 + echo "Using remote driver to connect to buildkitd"
903 +
904 + if docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then
905 + use_existing_builder
906 + else
907 + create_and_bootstrap_builder remote "unix://${BUILDKIT_SOCKET}"
908 + fi
909 + elif docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then
910 + use_existing_builder
911 + else
912 + echo "No local buildkitd found, using docker-container driver"
913 + create_and_bootstrap_builder docker-container
914 + fi
915 +
916 + echo "Active builder:"
917 + docker buildx ls | grep -E '^\*|^NAME' || docker buildx ls
918 +}
919 +
920 +prepare_git_cache_metadata() {
921 + local cache_branch_name=""
922 + local cache_base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}"
923 + local target_repo_slug=""
924 + local target_repo_url=""
925 + local merge_base_ref=""
926 +
927 + if [[ -z "${PARENT_COMMIT:-}" || -z "${VLLM_MERGE_BASE_COMMIT:-}" ]] \
928 + && git rev-parse --is-shallow-repository 2>/dev/null | grep -q "true"; then
929 + echo "Shallow clone detected - deepening for cache key computation"
930 + git_fetch_for_cache --deepen=1 origin || true
931 + fi
932 +
933 + if [[ -z "${PARENT_COMMIT:-}" ]]; then
934 + PARENT_COMMIT=$(git rev-parse HEAD~1 2>/dev/null || echo "")
935 + if [[ -n "${PARENT_COMMIT}" ]]; then
936 + export PARENT_COMMIT
937 + echo "Computed parent commit for cache fallback: ${PARENT_COMMIT}"
938 + else
939 + echo "Could not determine parent commit"
940 + fi
941 + else
942 + echo "Using provided PARENT_COMMIT: ${PARENT_COMMIT}"
943 + fi
944 +
945 + if [[ -z "${ROCM_CACHE_BRANCH_TAG:-}" ]]; then
946 + cache_branch_name=$(select_cache_branch_name)
947 + if [[ -z "${cache_branch_name}" && "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then
948 + cache_branch_name="pr-${BUILDKITE_PULL_REQUEST}"
949 + echo "Using pull request number for ROCm branch cache tag: ${cache_branch_name}"
950 + fi
951 + fi
952 +
953 + if [[ -z "${ROCM_CACHE_BRANCH_TAG:-}" && -n "${cache_branch_name}" ]]; then
954 + ROCM_CACHE_BRANCH_TAG=$(
955 + compose_cache_branch_tag "$(get_buildkite_repo_slug)" "${cache_branch_name}"
956 + )
957 + export ROCM_CACHE_BRANCH_TAG
958 + echo "Computed ROCm branch cache tag: ${ROCM_CACHE_BRANCH_TAG} (from ${cache_branch_name})"
959 + elif [[ -n "${ROCM_CACHE_BRANCH_TAG:-}" ]]; then
960 + echo "Using provided ROCM_CACHE_BRANCH_TAG: ${ROCM_CACHE_BRANCH_TAG}"
961 + elif [[ -n "${BUILDKITE_BRANCH:-}" ]]; then
962 + echo "Skipping ROCm branch cache tag: no usable branch name found"
963 + echo " BUILDKITE_BRANCH=${BUILDKITE_BRANCH}"
964 + fi
965 +
966 + if [[ -z "${ROCM_CACHE_UPSTREAM_BRANCH_TAG:-}" \
967 + && -n "${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-}" \
968 + && "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then
969 + target_repo_slug=$(get_buildkite_target_repo_slug)
970 + ROCM_CACHE_UPSTREAM_BRANCH_TAG=$(
971 + compose_cache_branch_tag "${target_repo_slug}" "${BUILDKITE_PULL_REQUEST_BASE_BRANCH}"
972 + )
973 + export ROCM_CACHE_UPSTREAM_BRANCH_TAG
974 + echo "Computed ROCm upstream branch cache tag: ${ROCM_CACHE_UPSTREAM_BRANCH_TAG}"
975 + elif [[ -n "${ROCM_CACHE_UPSTREAM_BRANCH_TAG:-}" ]]; then
976 + echo "Using provided ROCM_CACHE_UPSTREAM_BRANCH_TAG: ${ROCM_CACHE_UPSTREAM_BRANCH_TAG}"
977 + fi
978 +
979 + if [[ -z "${VLLM_MERGE_BASE_COMMIT:-}" ]]; then
980 + target_repo_url=$(get_buildkite_target_repo_url)
981 + merge_base_ref="refs/remotes/vllm-cache-upstream/${cache_base_branch}"
982 + git_fetch_for_cache --no-tags --depth=200 "${target_repo_url}" \
983 + "+refs/heads/${cache_base_branch}:${merge_base_ref}" 2>/dev/null || true
984 + VLLM_MERGE_BASE_COMMIT=$(git merge-base HEAD "${merge_base_ref}" 2>/dev/null || echo "")
985 + if [[ -z "${VLLM_MERGE_BASE_COMMIT}" ]]; then
986 + git_fetch_for_cache --no-tags --deepen=1000 "${target_repo_url}" \
987 + "+refs/heads/${cache_base_branch}:${merge_base_ref}" 2>/dev/null || true
988 + VLLM_MERGE_BASE_COMMIT=$(git merge-base HEAD "${merge_base_ref}" 2>/dev/null || echo "")
989 + fi
990 + if [[ -n "${VLLM_MERGE_BASE_COMMIT}" ]]; then
991 + export VLLM_MERGE_BASE_COMMIT
992 + echo "Computed merge base commit for cache fallback: ${VLLM_MERGE_BASE_COMMIT}"
993 + else
994 + echo "Could not determine merge base with ${cache_base_branch}"
995 + fi
996 + else
997 + echo "Using provided VLLM_MERGE_BASE_COMMIT: ${VLLM_MERGE_BASE_COMMIT}"
998 + fi
999 +}
1000 +
1001 +write_ci_base_label_override() {
1002 + local target_name=""
1003 + local -a ci_base_targets=()
1004 +
1005 + BAKE_FILES=(-f "${VLLM_BAKE_FILE}" -f "${CI_HCL_PATH}")
1006 +
1007 + if [[ -z "${CI_BASE_CONTENT_HASH:-}" ]]; then
1008 + return 0
1009 + fi
1010 +
1011 + mapfile -t ci_base_targets < <(
1012 + {
1013 + printf '%s\n' "ci-base-rocm"
1014 + sed -n -E 's/^target "(ci-base-rocm[^"]+)".*/\1/p' "${CI_HCL_PATH}" 2>/dev/null || true
1015 + } | awk '!seen[$0]++'
1016 + )
1017 +
1018 + if [[ ${#ci_base_targets[@]} -eq 0 ]]; then
1019 + return 0
1020 + fi
1021 +
1022 + : > "${CI_BASE_LABEL_OVERRIDE_PATH}"
1023 + for target_name in "${ci_base_targets[@]}"; do
1024 + cat >> "${CI_BASE_LABEL_OVERRIDE_PATH}" <<EOF
1025 +target "${target_name}" {
1026 + annotations = [
1027 + "manifest:org.opencontainers.image.revision=",
1028 + ]
1029 + labels = {
1030 + "org.opencontainers.image.revision" = ""
1031 + "vllm.ci_base.content_hash" = "${CI_BASE_CONTENT_HASH}"
1032 + }
1033 +}
1034 +
1035 +EOF
1036 + done
1037 +
1038 + BAKE_FILES+=(-f "${CI_BASE_LABEL_OVERRIDE_PATH}")
1039 + echo "Appended ci_base content-hash label override for targets: ${ci_base_targets[*]}"
1040 +}
1041 +
1042 +uses_rocm_csrc_cache() {
1043 + case "${TARGET}" in
1044 + csrc-rocm-ci|test-rocm-ci|test-rocm-ci-with-wheel|test-rocm-ci-with-artifacts|export-wheel-rocm)
1045 + return 0
1046 + ;;
1047 + *)
1048 + return 1
1049 + ;;
1050 + esac
1051 +}
1052 +
1053 +compute_rocm_csrc_content_hash() {
1054 + local bake_dir=""
1055 + local dockerfile_rocm=""
1056 + local -a content_paths=(
1057 + "requirements/common.txt"
1058 + "requirements/rocm.txt"
1059 + "setup.py"
1060 + "CMakeLists.txt"
1061 + "cmake"
1062 + "csrc"
1063 + "vllm/envs.py"
1064 + "vllm/__init__.py"
1065 + )
1066 + local -a content_args=()
1067 +
1068 + bake_dir=$(dirname "${VLLM_BAKE_FILE}")
1069 + dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
1070 + mapfile -t content_args < <(
1071 + get_content_arg_names "${dockerfile_rocm}" "base csrc-build" "${ROCM_CSRC_CONTENT_ARGS:-}"
1072 + )
1073 +
1074 + {
1075 + printf 'csrc-input-files-hash:%s\n' "$(compute_content_hash "${content_paths[@]}")"
1076 + printf 'dockerfile:%s\n' "${dockerfile_rocm}"
1077 + printf 'resolved-build-args:\n'
1078 + hash_dockerfile_arg_values "${dockerfile_rocm}" "${content_args[@]}"
1079 + printf 'dockerfile-stages:base csrc-build\n'
1080 + if [[ -f "${dockerfile_rocm}" ]]; then
1081 + hash_dockerfile_stages "${dockerfile_rocm}" "base csrc-build"
1082 + else
1083 + printf 'missing:%s\n' "${dockerfile_rocm}"
1084 + fi
1085 + } | sha256sum | cut -d' ' -f1
1086 +}
1087 +
1088 +compute_rocm_csrc_content_hash_if_needed() {
1089 + local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"
1090 +
1091 + if [[ "${ROCM_CSRC_CONTENT_CACHE:-1}" == "0" ]] || ! uses_rocm_csrc_cache; then
1092 + return 0
1093 + fi
1094 +
1095 + ROCM_CSRC_CONTENT_HASH=$(compute_rocm_csrc_content_hash)
1096 + ROCM_CSRC_CONTENT_CACHE_REF="${cache_repo}:csrc-rocm-input-${ROCM_CSRC_CONTENT_HASH}"
1097 + export ROCM_CSRC_CONTENT_HASH
1098 + export ROCM_CSRC_CONTENT_CACHE_REF
1099 + echo "ROCm csrc content cache ref: ${ROCM_CSRC_CONTENT_CACHE_REF}"
1100 +}
1101 +
1102 +write_hcl_string_list_entries() {
1103 + local indent="$1"
1104 + local value=""
1105 + shift
1106 +
1107 + for value in "$@"; do
1108 + value="${value//\\/\\\\}"
1109 + value="${value//\"/\\\"}"
1110 + printf '%s"%s",\n' "${indent}" "${value}"
1111 + done
1112 +}
1113 +
1114 +hcl_escape_string() {
1115 + local value="$1"
1116 +
1117 + value="${value//\\/\\\\}"
1118 + value="${value//\"/\\\"}"
1119 + printf '%s' "${value}"
1120 +}
1121 +
1122 +write_hcl_string_list() {
1123 + local indent="$1"
1124 + shift
1125 +
1126 + printf '%s[\n' "${indent}"
1127 + write_hcl_string_list_entries "${indent} " "$@"
1128 + printf '%s]\n' "${indent}"
1129 +}
1130 +
1131 +write_rocm_build_arg_override() {
1132 + local bake_dir=""
1133 + local dockerfile_rocm=""
1134 + local -a arg_names=()
1135 + local arg_name=""
1136 + local arg_value=""
1137 +
1138 + bake_dir=$(dirname "${VLLM_BAKE_FILE}")
1139 + dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
1140 + mapfile -t arg_names < <(
1141 + {
1142 + get_content_arg_names \
1143 + "${dockerfile_rocm}" \
1144 + "${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}" \
1145 + "${CI_BASE_CONTENT_ARGS:-}"
1146 + get_content_arg_names "${dockerfile_rocm}" "base csrc-build" "${ROCM_CSRC_CONTENT_ARGS:-}"
1147 + } | awk 'NF && !seen[$0]++'
1148 + )
1149 +
1150 + {
1151 + cat <<EOF
1152 +target "_common-rocm" {
1153 + args = {
1154 +EOF
1155 + for arg_name in "${arg_names[@]}"; do
1156 + [[ -n "${arg_name}" ]] || continue
1157 + arg_value=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "${arg_name}")
1158 + [[ -n "${arg_value}" ]] || continue
1159 + printf ' %s = "%s"\n' "${arg_name}" "$(hcl_escape_string "${arg_value}")"
1160 + done
1161 + cat <<EOF
1162 + }
1163 +}
1164 +EOF
1165 + } > "${ROCM_ARG_OVERRIDE_PATH}"
1166 +
1167 + BAKE_FILES+=(-f "${ROCM_ARG_OVERRIDE_PATH}")
1168 + echo "Appended resolved ROCm Docker ARG override"
1169 +}
1170 +
1171 +write_hcl_string_list_attr() {
1172 + local indent="$1"
1173 + local attr="$2"
1174 + shift 2
1175 +
1176 + printf '%s%s = [\n' "${indent}" "${attr}"
1177 + write_hcl_string_list_entries "${indent} " "$@"
1178 + printf '%s]\n' "${indent}"
1179 +}
1180 +
1181 +validate_cache_export_mode() {
1182 + local mode="$1"
1183 + local env_name="$2"
1184 +
1185 + case "${mode}" in
1186 + min|max)
1187 + ;;
1188 + *)
1189 + echo "Error: ${env_name} must be one of: min, max"
1190 + exit 1
1191 + ;;
1192 + esac
1193 +}
1194 +
1195 +write_rocm_cache_override() {
1196 + local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"
1197 + local csrc_cache_to_mode="${ROCM_CSRC_CACHE_TO_MODE:-max}"
1198 + local rocm_cache_to_mode="${ROCM_FINAL_CACHE_TO_MODE:-min}"
1199 + local -a content_cache_from=()
1200 + local -a csrc_cache_to=()
1201 + local -a rocm_cache_to=()
1202 + local -a export_wheel_cache_to=()
1203 +
1204 + if ! uses_rocm_csrc_cache; then
1205 + return 0
1206 + fi
1207 +
1208 + validate_cache_export_mode "${csrc_cache_to_mode}" "ROCM_CSRC_CACHE_TO_MODE"
1209 + validate_cache_export_mode "${rocm_cache_to_mode}" "ROCM_FINAL_CACHE_TO_MODE"
1210 + echo "ROCm csrc cache export mode: ${csrc_cache_to_mode}"
1211 + echo "ROCm final image cache export mode: ${rocm_cache_to_mode}"
1212 +
1213 + if [[ -n "${ROCM_CSRC_CONTENT_CACHE_REF:-}" ]]; then
1214 + content_cache_from+=("type=registry,ref=${ROCM_CSRC_CONTENT_CACHE_REF}")
1215 + csrc_cache_to+=(
1216 + "type=registry,ref=${ROCM_CSRC_CONTENT_CACHE_REF},mode=${csrc_cache_to_mode},ignore-error=true"
1217 + )
1218 + fi
1219 +
1220 + # Docker Hub cache exports are best-effort. A cache-only target failure can
1221 + # otherwise cancel the sibling image target before its manifest is pushed.
1222 + if [[ -n "${BUILDKITE_COMMIT:-}" ]]; then
1223 + csrc_cache_to+=(
1224 + "type=registry,ref=${cache_repo}:csrc-rocm-${BUILDKITE_COMMIT},mode=${csrc_cache_to_mode},ignore-error=true"
1225 + )
1226 + rocm_cache_to+=(
1227 + "type=registry,ref=${cache_repo}:rocm-${BUILDKITE_COMMIT},mode=${rocm_cache_to_mode},ignore-error=true"
1228 + )
1229 + fi
1230 +
1231 + if [[ -n "${ROCM_CACHE_BRANCH_TAG:-}" ]]; then
1232 + csrc_cache_to+=(
1233 + "type=registry,ref=${cache_repo}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${csrc_cache_to_mode},ignore-error=true"
1234 + )
1235 + rocm_cache_to+=(
1236 + "type=registry,ref=${cache_repo}:rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${rocm_cache_to_mode},ignore-error=true"
1237 + )
1238 + fi
1239 +
1240 + if [[ "${TARGET}" == "test-rocm-ci-with-wheel" ]]; then
1241 + export_wheel_cache_to=()
1242 + else
1243 + export_wheel_cache_to=("${rocm_cache_to[@]}")
1244 + fi
1245 +
1246 + {
1247 + cat <<EOF
1248 +target "csrc-rocm-ci" {
1249 + cache-from = concat(
1250 + get_cache_from_rocm_csrc(),
1251 +EOF
1252 + write_hcl_string_list " " "${content_cache_from[@]}"
1253 + cat <<EOF
1254 + )
1255 +EOF
1256 + write_hcl_string_list_attr " " "cache-to" "${csrc_cache_to[@]}"
1257 + cat <<EOF
1258 +}
1259 +
1260 +target "test-rocm-ci" {
1261 + cache-from = concat(
1262 + get_cache_from_rocm(),
1263 +EOF
1264 + write_hcl_string_list " " "${content_cache_from[@]}"
1265 + cat <<EOF
1266 + )
1267 +EOF
1268 + write_hcl_string_list_attr " " "cache-to" "${rocm_cache_to[@]}"
1269 + cat <<EOF
1270 +}
1271 +
1272 +target "export-wheel-rocm" {
1273 + cache-from = concat(
1274 + get_cache_from_rocm(),
1275 +EOF
1276 + write_hcl_string_list " " "${content_cache_from[@]}"
1277 + cat <<EOF
1278 + )
1279 +EOF
1280 + write_hcl_string_list_attr " " "cache-to" "${export_wheel_cache_to[@]}"
1281 + cat <<EOF
1282 +}
1283 +EOF
1284 + } > "${CSRC_CACHE_OVERRIDE_PATH}"
1285 +
1286 + BAKE_FILES+=(-f "${CSRC_CACHE_OVERRIDE_PATH}")
1287 + echo "Appended ROCm cache override with non-fatal registry exports"
1288 +}
1289 +
1290 +extract_dependency_pins() {
1291 + local bake_dir=""
1292 + local dockerfile_rocm=""
1293 + local var=""
1294 + local val=""
1295 +
1296 + bake_dir=$(dirname "${VLLM_BAKE_FILE}")
1297 + dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
1298 + if [[ ! -f "${dockerfile_rocm}" ]]; then
1299 + return 0
1300 + fi
1301 +
1302 + for var in RIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
1303 + if [[ -n "${!var:-}" ]]; then
1304 + echo "Using provided ${var}: ${!var}"
1305 + continue
1306 + fi
1307 +
1308 + val=$(
1309 + sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${var}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \
1310 + "${dockerfile_rocm}" | head -1
1311 + )
1312 + if [[ -n "${val}" ]]; then
1313 + export "${var}=${val}"
1314 + echo "Extracted ${var}=${val} from Dockerfile.rocm"
1315 + fi
1316 + done
1317 +}
1318 +
1319 +compute_dependency_cache_keys() {
1320 + local bake_dir=""
1321 + local dockerfile_rocm=""
1322 + local rixl_branch=""
1323 + local ucx_branch=""
1324 + local rocshmem_branch=""
1325 + local deepep_branch=""
1326 + local rixl_material=""
1327 + local rocshmem_material=""
1328 + local deepep_material=""
1329 +
1330 + bake_dir=$(dirname "${VLLM_BAKE_FILE}")
1331 + dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
1332 + rixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "RIXL_BRANCH")
1333 + ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH")
1334 + rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH")
1335 + deepep_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "DEEPEP_BRANCH")
1336 +
1337 + if [[ -n "${rixl_branch}" && -n "${ucx_branch}" ]]; then
1338 + rixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rixl")
1339 + RIXL_CACHE_KEY=$(
1340 + compose_dependency_cache_key \
1341 + "${rixl_branch}-ucx-${ucx_branch}" \
1342 + "${rixl_material}"
1343 + )
1344 + export RIXL_CACHE_KEY
1345 + echo "RIXL dependency cache key: ${RIXL_CACHE_KEY}"
1346 + fi
1347 +
1348 + if [[ -n "${rocshmem_branch}" ]]; then
1349 + rocshmem_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rocshmem")
1350 + ROCSHMEM_CACHE_KEY=$(
1351 + compose_dependency_cache_key \
1352 + "${rocshmem_branch}" \
1353 + "${rocshmem_material}"
1354 + )
1355 + export ROCSHMEM_CACHE_KEY
1356 + echo "ROCShmem dependency cache key: ${ROCSHMEM_CACHE_KEY}"
1357 + fi
1358 +
1359 + if [[ -n "${deepep_branch}" && -n "${rocshmem_branch}" ]]; then
1360 + deepep_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rocshmem build_deepep")
1361 + DEEPEP_CACHE_KEY=$(
1362 + compose_dependency_cache_key \
1363 + "${deepep_branch}-rocshmem-${rocshmem_branch}" \
1364 + "${deepep_material}"
1365 + )
1366 + export DEEPEP_CACHE_KEY
1367 + echo "DeepEP dependency cache key: ${DEEPEP_CACHE_KEY}"
1368 + fi
1369 +}
1370 +
1371 +compose_stage_cache_material() {
1372 + local dockerfile="$1"
1373 + local stages="$2"
1374 + local -a content_args=()
1375 +
1376 + mapfile -t content_args < <(get_content_arg_names "${dockerfile}" "${stages}" "")
1377 + {
1378 + printf 'dockerfile:%s\n' "${dockerfile}"
1379 + printf 'dockerfile-stages:%s\n' "${stages}"
1380 + hash_dockerfile_stages "${dockerfile}" "${stages}"
1381 + printf 'resolved-build-args:\n'
1382 + hash_dockerfile_arg_values "${dockerfile}" "${content_args[@]}"
1383 + }
1384 +}
1385 +
1386 +dependency_cache_ref_exists() {
1387 + local cache_ref="$1"
1388 + docker buildx imagetools inspect "${cache_ref}" >/dev/null 2>&1
1389 +}
1390 +
1391 +dependency_cache_ref_for_target() {
1392 + local target="$1"
1393 + local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"
1394 +
1395 + case "${target}" in
1396 + rixl-rocm-ci)
1397 + if [[ -n "${RIXL_CACHE_KEY:-}" ]]; then
1398 + printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_CACHE_KEY}"
1399 + elif [[ -n "${RIXL_BRANCH:-}" ]]; then
1400 + printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
1401 + fi
1402 + ;;
1403 + rocshmem-rocm-ci)
1404 + if [[ -n "${ROCSHMEM_CACHE_KEY:-}" ]]; then
1405 + printf '%s\n' "${cache_repo}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY}"
1406 + elif [[ -n "${ROCSHMEM_BRANCH:-}" ]]; then
1407 + printf '%s\n' "${cache_repo}:rocshmem-rocm-${ROCSHMEM_BRANCH}"
1408 + fi
1409 + ;;
1410 + deepep-rocm-ci)
1411 + if [[ -n "${DEEPEP_CACHE_KEY:-}" ]]; then
1412 + printf '%s\n' "${cache_repo}:deepep-rocm-${DEEPEP_CACHE_KEY}"
1413 + elif [[ -n "${DEEPEP_BRANCH:-}" ]]; then
1414 + printf '%s\n' "${cache_repo}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH:-}"
1415 + fi
1416 + ;;
1417 + esac
1418 +}
1419 +
1420 +add_dependency_cache_target() {
1421 + local target="$1"
1422 +
1423 + if printf '%s\n' "${DEPENDENCY_CACHE_TARGETS[@]}" | grep -qx "${target}"; then
1424 + return 0
1425 + fi
1426 + DEPENDENCY_CACHE_TARGETS+=("${target}")
1427 +}
1428 +
1429 +resolve_ci_base_dependency_targets() {
1430 + local mode="${ROCM_DEP_CACHE_EXPORT_MODE:-missing}"
1431 + local rixl_ref=""
1432 + local rocshmem_ref=""
1433 + local deepep_ref=""
1434 +
1435 + [[ "${TARGET}" == "ci-base-rocm-ci-with-deps" ]] || return 0
1436 +
1437 + case "${mode}" in
1438 + always)
1439 + echo "ROCM_DEP_CACHE_EXPORT_MODE=always; exporting all dependency caches serially"
1440 + for target in rixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
1441 + if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then
1442 + add_dependency_cache_target "${target}"
1443 + fi
1444 + done
1445 + ;;
1446 + never)
1447 + BAKE_TARGETS=("ci-base-rocm-ci")
1448 + DEPENDENCY_CACHE_TARGETS=()
1449 + echo "ROCM_DEP_CACHE_EXPORT_MODE=never; building ci_base without dependency cache exports"
1450 + return 0
1451 + ;;
1452 + missing|"")
1453 + ;;
1454 + *)
1455 + echo "Error: ROCM_DEP_CACHE_EXPORT_MODE must be one of: missing, always, never"
1456 + exit 1
1457 + ;;
1458 + esac
1459 +
1460 + if [[ "${mode}" != "always" && -n "${RIXL_CACHE_KEY:-}" ]]; then
1461 + rixl_ref=$(dependency_cache_ref_for_target "rixl-rocm-ci")
1462 + if dependency_cache_ref_exists "${rixl_ref}"; then
1463 + echo "RIXL dependency cache exists: ${rixl_ref}"
1464 + else
1465 + echo "RIXL dependency cache missing; will seed: ${rixl_ref}"
1466 + add_dependency_cache_target "rixl-rocm-ci"
1467 + fi
1468 + fi
1469 +
1470 + if [[ "${mode}" != "always" && -n "${ROCSHMEM_CACHE_KEY:-}" ]]; then
1471 + rocshmem_ref=$(dependency_cache_ref_for_target "rocshmem-rocm-ci")
1472 + if dependency_cache_ref_exists "${rocshmem_ref}"; then
1473 + echo "ROCShmem dependency cache exists: ${rocshmem_ref}"
1474 + else
1475 + echo "ROCShmem dependency cache missing; will seed: ${rocshmem_ref}"
1476 + add_dependency_cache_target "rocshmem-rocm-ci"
1477 + fi
1478 + fi
1479 +
1480 + if [[ "${mode}" != "always" && -n "${DEEPEP_CACHE_KEY:-}" ]]; then
1481 + deepep_ref=$(dependency_cache_ref_for_target "deepep-rocm-ci")
1482 + if dependency_cache_ref_exists "${deepep_ref}"; then
1483 + echo "DeepEP dependency cache exists: ${deepep_ref}"
1484 + else
1485 + echo "DeepEP dependency cache missing; will seed: ${deepep_ref}"
1486 + add_dependency_cache_target "deepep-rocm-ci"
1487 + fi
1488 + fi
1489 +
1490 + # DeepEP inherits from ROCShmem. If ROCShmem is being seeded, seed DeepEP too
1491 + # so the pair stays consistent for future ci_base rebuilds.
1492 + if printf '%s\n' "${DEPENDENCY_CACHE_TARGETS[@]}" | grep -qx "rocshmem-rocm-ci" \
1493 + && ! printf '%s\n' "${DEPENDENCY_CACHE_TARGETS[@]}" | grep -qx "deepep-rocm-ci" \
1494 + && [[ -n "${DEEPEP_BRANCH:-}" ]]; then
1495 + echo "ROCShmem cache is missing; also seeding DeepEP cache"
1496 + add_dependency_cache_target "deepep-rocm-ci"
1497 + fi
1498 +
1499 + BAKE_TARGETS=("ci-base-rocm-ci")
1500 + if [[ ${#DEPENDENCY_CACHE_TARGETS[@]} -eq 0 ]]; then
1501 + echo "All dependency caches exist; building ci_base without dependency cache exports"
1502 + else
1503 + echo "Resolved dependency cache seed targets: ${DEPENDENCY_CACHE_TARGETS[*]}"
1504 + echo "Resolved ci_base bake targets: ${BAKE_TARGETS[*]}"
1505 + fi
1506 +}
1507 +
1508 +bake_config_targets() {
1509 + printf '%s\n' "${DEPENDENCY_CACHE_TARGETS[@]}" "${BAKE_TARGETS[@]}" \
1510 + | awk 'NF && !seen[$0]++'
1511 +}
1512 +
1513 +print_bake_config() {
1514 + local -a print_targets=()
1515 +
1516 + echo "--- :page_facing_up: Resolved bake configuration"
1517 + mapfile -t print_targets < <(bake_config_targets)
1518 + docker buildx bake "${BAKE_FILES[@]}" --print "${print_targets[@]}" | tee "${BAKE_CONFIG_FILE}"
1519 +
1520 + if command -v buildkite-agent >/dev/null 2>&1 && [[ -n "${BUILDKITE_BUILD_NUMBER:-}" ]]; then
1521 + buildkite-agent artifact upload "${BAKE_CONFIG_FILE}" || true
1522 + echo "Uploaded ${BAKE_CONFIG_FILE} as Buildkite artifact"
1523 + else
1524 + echo "Saved bake config to ${BAKE_CONFIG_FILE} (not in Buildkite, skipping upload)"
1525 + fi
1526 +}
1527 +
1528 +confirm_remote_image_push() {
1529 + local image_ref="$1"
1530 + local remote_hash=""
1531 + local remote_revision=""
1532 +
1533 + if ! remote_image_exists "${image_ref}"; then
1534 + return 1
1535 + fi
1536 +
1537 + if is_ci_base_target; then
1538 + if [[ -z "${CI_BASE_CONTENT_HASH:-}" ]]; then
1539 + return 0
1540 + fi
1541 +
1542 + remote_hash=$(get_remote_image_label_with_retry "${image_ref}" "vllm.ci_base.content_hash")
1543 + if [[ -n "${remote_hash}" && "${remote_hash}" == "${CI_BASE_CONTENT_HASH}" ]]; then
1544 + return 0
1545 + fi
1546 +
1547 + echo "Remote image exists but does not have the expected ci_base content hash."
1548 + echo " expected: ${CI_BASE_CONTENT_HASH:0:16}..."
1549 + echo " found: ${remote_hash:0:16}..."
1550 + return 1
1551 + fi
1552 +
1553 + if is_commit_image_target; then
1554 + remote_revision=$(get_remote_image_label_with_retry "${image_ref}" "org.opencontainers.image.revision")
1555 + if [[ -n "${remote_revision}" && "${remote_revision}" == "${BUILDKITE_COMMIT}" ]]; then
1556 + return 0
1557 + fi
1558 +
1559 + if [[ -z "${remote_revision}" \
1560 + && ${IMAGE_EXISTED_BEFORE_BUILD} -eq 0 \
1561 + && image_tag_is_commit_scoped ]]; then
1562 + echo "Remote image exists under a commit-scoped tag; accepting push despite missing revision label."
1563 + return 0
1564 + fi
1565 +
1566 + echo "Remote image exists but revision label does not match ${BUILDKITE_COMMIT}."
1567 + echo " found revision: ${remote_revision:-<missing>}"
1568 + return 1
1569 + fi
1570 +
1571 + return 0
1572 +}
1573 +
1574 +verify_dependency_cache_ref() {
1575 + local cache_ref="$1"
1576 + local attempts="${ROCM_DEP_CACHE_VERIFY_ATTEMPTS:-6}"
1577 + local delay_secs="${ROCM_DEP_CACHE_VERIFY_DELAY:-5}"
1578 + local attempt
1579 +
1580 + for ((attempt = 1; attempt <= attempts; attempt++)); do
1581 + if dependency_cache_ref_exists "${cache_ref}"; then
1582 + echo "Dependency cache confirmed: ${cache_ref}"
1583 + return 0
1584 + fi
1585 + if [[ ${attempt} -lt ${attempts} ]]; then
1586 + echo "Dependency cache not visible yet (${attempt}/${attempts}): ${cache_ref}"
1587 + sleep "${delay_secs}"
1588 + fi
1589 + done
1590 +
1591 + echo "ERROR: dependency cache was not confirmed after upload: ${cache_ref}"
1592 + return 1
1593 +}
1594 +
1595 +seed_dependency_caches_if_needed() {
1596 + local target=""
1597 + local cache_ref=""
1598 +
1599 + if [[ "${TARGET}" != "ci-base-rocm-ci-with-deps" ]]; then
1600 + return 0
1601 + fi
1602 + if [[ ${#DEPENDENCY_CACHE_TARGETS[@]} -eq 0 ]]; then
1603 + return 0
1604 + fi
1605 +
1606 + echo "--- :docker: Seeding ROCm dependency caches"
1607 + echo "Dependency cache uploads are required for this build."
1608 + echo "Seeding serially to avoid concurrent Docker Hub cache exporters."
1609 +
1610 + for target in "${DEPENDENCY_CACHE_TARGETS[@]}"; do
1611 + cache_ref=$(dependency_cache_ref_for_target "${target}")
1612 + if [[ -z "${cache_ref}" ]]; then
1613 + echo "ERROR: could not resolve dependency cache ref for ${target}"
1614 + return 1
1615 + fi
1616 +
1617 + echo "--- :docker: Seeding ${target}"
1618 + echo "Expected cache ref: ${cache_ref}"
1619 + docker buildx bake "${BAKE_FILES[@]}" --progress plain "${target}"
1620 + verify_dependency_cache_ref "${cache_ref}"
1621 + done
1622 +}
1623 +
1624 +annotate_cache_export_warning() {
1625 + local build_rc="$1"
1626 +
1627 + if ! command -v buildkite-agent >/dev/null 2>&1; then
1628 + return 0
1629 + fi
1630 +
1631 + buildkite-agent annotate \
1632 + --style warning \
1633 + --context "cache-export-warning" \
1634 + "### :warning: Docker cache export failed (non-fatal)
1635 +
1636 +Image was pushed successfully: \`${IMAGE_TAG}\`
1637 +
1638 +The BuildKit build returned exit code ${build_rc}, but the expected image
1639 +is present in the registry. Treating this as a registry cache export failure
1640 +so tests can continue with the pushed image." 2>/dev/null || true
1641 +}
1642 +
1643 +run_bake() {
1644 + local build_rc=0
1645 +
1646 + echo "--- :docker: Building ${TARGET}"
1647 + docker buildx bake "${BAKE_FILES[@]}" --progress plain "${BAKE_TARGETS[@]}" || build_rc=$?
1648 +
1649 + if [[ ${build_rc} -eq 0 ]]; then
1650 + echo "--- :white_check_mark: Build complete"
1651 + return 0
1652 + fi
1653 +
1654 + echo ""
1655 + echo "WARNING: docker buildx bake exited with code ${build_rc}"
1656 +
1657 + if [[ -n "${IMAGE_TAG:-}" ]]; then
1658 + echo "Checking if image was pushed successfully..."
1659 + if confirm_remote_image_push "${IMAGE_TAG}"; then
1660 + echo ""
1661 + echo "WARNING: Build reported failure (rc=${build_rc}) but the"
1662 + echo " image was pushed successfully: ${IMAGE_TAG}"
1663 + echo ""
1664 + echo " Treating this as a non-fatal registry cache export failure."
1665 + echo " The image is usable, but registry cache may be cold on the next build."
1666 + echo ""
1667 + annotate_cache_export_warning "${build_rc}"
1668 + echo "--- :white_check_mark: Build complete"
1669 + return 0
1670 + fi
1671 +
1672 + echo ""
1673 + echo "ERROR: Build failed and image was NOT confirmed: ${IMAGE_TAG}"
1674 + echo " This is a real build failure, not a cache export warning."
1675 + echo ""
1676 + fi
1677 +
1678 + return "${build_rc}"
1679 +}
1680 +
1681 +upload_wheel_artifacts_if_present() {
1682 + local wheel_dir="./wheel-export"
1683 + local artifact_dir="artifacts/vllm-rocm-install"
1684 + local archive_name="vllm-rocm-install.tar.gz"
1685 + local whl=""
1686 + local whl_name=""
1687 +
1688 + if ! should_upload_wheel_artifacts; then
1689 + return 0
1690 + fi
1691 +
1692 + if [[ ! -d "${wheel_dir}" ]] || ! ls "${wheel_dir}"/*.whl >/dev/null 2>&1; then
1693 + echo "No ROCm wheel artifacts found in ${wheel_dir}"
1694 + return 0
1695 + fi
1696 +
1697 + echo "--- :package: Uploading ROCm vLLM install artifact"
1698 + mkdir -p "${artifact_dir}"
1699 +
1700 + tar -C "${wheel_dir}" -czf "${artifact_dir}/${archive_name}" .
1701 + echo "Created ${archive_name}: $(du -sh "${artifact_dir}/${archive_name}" | cut -f1)"
1702 + printf '%s\n' "${CI_BASE_IMAGE:-}" > "${artifact_dir}/ci-base-image.txt"
1703 + printf '%s\n' "${IMAGE_TAG:-}" > "${artifact_dir}/fallback-image.txt"
1704 +
1705 + for whl in "${wheel_dir}"/*.whl; do
1706 + [[ -f "${whl}" ]] || continue
1707 + whl_name=$(basename "${whl}")
1708 + cp "${whl}" "${artifact_dir}/${whl_name}"
1709 + echo "Copied ${whl_name}: $(du -sh "${artifact_dir}/${whl_name}" | cut -f1)"
1710 + done
1711 +
1712 + if command -v buildkite-agent >/dev/null 2>&1; then
1713 + buildkite-agent artifact upload "${artifact_dir}/*"
1714 + echo "ROCm vLLM install artifacts uploaded to ${artifact_dir}/"
1715 + else
1716 + echo "Not in Buildkite, skipping artifact upload"
1717 + fi
1718 +
1719 + rm -rf "${wheel_dir}"
1720 +}
1721 +
1722 +main() {
1723 + init_config "$@"
1724 + print_header
1725 + validate_inputs
1726 + load_ci_hcl
1727 + compute_ci_base_hash_if_needed
1728 + configure_ci_base_image_refs
1729 + maybe_skip_existing_image
1730 + setup_builder
1731 + prepare_git_cache_metadata
1732 + write_ci_base_label_override
1733 + extract_dependency_pins
1734 + write_rocm_build_arg_override
1735 + compute_dependency_cache_keys
1736 + compute_rocm_csrc_content_hash_if_needed
1737 + write_rocm_cache_override
1738 + resolve_ci_base_dependency_targets
1739 + print_bake_config
1740 + if [[ "${BAKE_PRINT_ONLY:-0}" == "1" ]]; then
1741 + echo "BAKE_PRINT_ONLY=1 set; skipping build"
1742 + return 0
1743 + fi
1744 + seed_dependency_caches_if_needed
1745 + run_bake
1746 + upload_wheel_artifacts_if_present
1747 +}
1748 +
1749 +main "$@"
.buildkite/scripts/hardware_ci/run-amd-test.sh
+121 -11
@@ -52,6 +52,108 @@ cleanup_network() {
52 fi
53 }
54
55 +prepare_artifact_image() {
56 + if [[ "${VLLM_CI_USE_ARTIFACTS:-0}" != "1" ]]; then
57 + return 1
58 + fi
59 + if ! command -v buildkite-agent >/dev/null 2>&1; then
60 + echo "buildkite-agent not found; cannot download ROCm wheel artifact"
61 + return 1
62 + fi
63 +
64 + local artifact_glob="${VLLM_CI_ARTIFACT_GLOB:-artifacts/vllm-rocm-install/vllm-rocm-install.tar.gz}"
65 + local archive=""
66 + local metadata_file=""
67 + local base_image="${VLLM_CI_BASE_IMAGE:-rocm/vllm-dev:ci_base}"
68 + local artifact_image=""
69 + local artifact_key=""
70 + local base_digest=""
71 + local wheel_dir=""
72 + local context_dir=""
73 + local workspace_dir=""
74 +
75 + artifact_work_dir=$(mktemp -d -t vllm-rocm-artifact.XXXXXX)
76 + wheel_dir="${artifact_work_dir}/wheels"
77 + context_dir="${artifact_work_dir}/context"
78 + workspace_dir="${context_dir}/workspace"
79 + mkdir -p "${wheel_dir}" "${context_dir}/wheels" "${workspace_dir}"
80 +
81 + echo "--- Downloading ROCm wheel artifact"
82 + if ! buildkite-agent artifact download "${artifact_glob}" "${artifact_work_dir}"; then
83 + echo "Failed to download ${artifact_glob}"
84 + return 1
85 + fi
86 + buildkite-agent artifact download \
87 + "artifacts/vllm-rocm-install/ci-base-image.txt" \
88 + "${artifact_work_dir}" >/dev/null 2>&1 || true
89 +
90 + archive=$(find "${artifact_work_dir}" -name "vllm-rocm-install.tar.gz" -type f | head -1)
91 + if [[ -z "${archive}" || ! -f "${archive}" ]]; then
92 + echo "ROCm wheel artifact archive was not found"
93 + return 1
94 + fi
95 +
96 + metadata_file=$(find "${artifact_work_dir}" -name "ci-base-image.txt" -type f | head -1)
97 + if [[ -n "${metadata_file}" && -s "${metadata_file}" ]]; then
98 + base_image=$(tr -d '[:space:]' < "${metadata_file}")
99 + fi
100 +
101 + echo "--- Preparing local ROCm test image"
102 + echo "Base image: ${base_image}"
103 + docker pull "${base_image}" || return 1
104 + base_digest=$(
105 + docker image inspect \
106 + --format='{{if .RepoDigests}}{{index .RepoDigests 0}}{{else}}{{.Id}}{{end}}' \
107 + "${base_image}" 2>/dev/null || printf '%s' "${base_image}"
108 + )
109 +
110 + artifact_key=$(
111 + {
112 + printf 'base-image:%s\n' "${base_digest}"
113 + sha256sum "${archive}"
114 + } | sha256sum | cut -c1-24
115 + )
116 + artifact_image="rocm/vllm-ci-artifact:${artifact_key}"
117 +
118 + if docker image inspect "${artifact_image}" >/dev/null 2>&1; then
119 + echo "Using existing local ROCm artifact image: ${artifact_image}"
120 + image_name="${artifact_image}"
121 + return 0
122 + fi
123 +
124 + tar -xzf "${archive}" -C "${wheel_dir}" || return 1
125 + if ! ls "${wheel_dir}"/*.whl >/dev/null 2>&1; then
126 + echo "ROCm wheel artifact did not contain a wheel"
127 + return 1
128 + fi
129 + if [[ ! -d "${wheel_dir}/tests" ]]; then
130 + echo "ROCm wheel artifact did not contain the test workspace"
131 + return 1
132 + fi
133 +
134 + cp "${wheel_dir}"/*.whl "${context_dir}/wheels/" || return 1
135 + tar -C "${wheel_dir}" --exclude='*.whl' -cf - . \
136 + | tar -C "${workspace_dir}" -xf - || return 1
137 + cat > "${context_dir}/Dockerfile" <<'EOF'
138 +ARG BASE_IMAGE
139 +FROM ${BASE_IMAGE}
140 +COPY wheels/ /tmp/vllm-wheels/
141 +COPY workspace/ /vllm-workspace/
142 +RUN python3 -m pip install --no-deps --force-reinstall /tmp/vllm-wheels/*.whl \
143 + && rm -rf /tmp/vllm-wheels
144 +WORKDIR /vllm-workspace
145 +EOF
146 +
147 + echo "--- Building local ROCm test image"
148 + docker build \
149 + --pull=false \
150 + --build-arg "BASE_IMAGE=${base_image}" \
151 + -t "${artifact_image}" \
152 + "${context_dir}" || return 1
153 + image_name="${artifact_image}"
154 + return 0
155 +}
156 +
157 is_multi_node() {
158 local cmds="$1"
159 # Primary signal: NUM_NODES environment variable set by the pipeline
@@ -243,22 +345,30 @@ report_docker_usage
345
346 # --- Pull test image ---
347 echo "--- Pulling container"
246 -image_name="rocm/vllm-ci:${BUILDKITE_COMMIT}"
348 +image_name="${VLLM_CI_FALLBACK_IMAGE:-rocm/vllm-ci:${BUILDKITE_COMMIT:-local}}"
349 +artifact_work_dir=""
350 container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
248 -docker pull "${image_name}"
351
352 remove_docker_container() {
251 - # docker run uses --rm, so the container is normally already gone when the
252 - # EXIT trap runs. Cleanup is best-effort and must not affect the test result.
253 - docker rm -f "${container_name}" >/dev/null 2>&1 || true
353 + if docker container inspect "${container_name}" >/dev/null 2>&1; then
354 + docker rm -f "${container_name}" || true
355 + fi
356 + if [[ "${VLLM_CI_REMOVE_TEST_IMAGE:-0}" == "1" ]]; then
357 + docker image rm -f "${image_name}" || true
358 + else
359 + # Keep images by default so later jobs on the same AMD node can reuse layers.
360 + echo "Keeping ROCm test image locally: ${image_name}"
361 + fi
362 + if [[ -n "${artifact_work_dir}" ]]; then
363 + rm -rf "${artifact_work_dir}"
364 + fi
365 }
366 +trap remove_docker_container EXIT
367
256 -on_exit() {
257 - local exit_code=$?
258 - remove_docker_container
259 - exit "$exit_code"
260 -}
261 -trap on_exit EXIT
368 +if ! prepare_artifact_image; then
369 + echo "Using full ROCm CI image: ${image_name}"
370 + docker pull "${image_name}" || exit 1
371 +fi
372
373 # --- Prepare commands ---
374 echo "--- Running container"
.dockerignore
+7
@@ -33,3 +33,10 @@ share/python-wheels/
33 *.egg
34 MANIFEST
35 rust/target/
36 +# Not needed in Docker builds
37 +docs/
38 +.github/
39 +.pre-commit-config.yaml
40 +.clang-format
41 +.gitattributes
42 +format.sh
cmake/utils.cmake
+8
@@ -81,6 +81,14 @@ function (hipify_sources_target OUT_SRCS NAME ORIG_SRCS)
81 set_property(GLOBAL APPEND PROPERTY VLLM_HIPIFY_ALL_SRCS ${SRCS})
82 set_property(GLOBAL APPEND PROPERTY VLLM_HIPIFY_ALL_BYPRODUCTS ${HIP_SRCS})
83
84 + # Chain hipify targets so they run sequentially. Parallel hipify
85 + # invocations race on shutil.copytree, overwriting .hip files
86 + # produced by another target back to .cu originals.
87 + if (DEFINED _VLLM_LAST_HIPIFY_TARGET)
88 + add_dependencies(hipify${NAME} ${_VLLM_LAST_HIPIFY_TARGET})
89 + endif()
90 + set(_VLLM_LAST_HIPIFY_TARGET "hipify${NAME}" PARENT_SCOPE)
91 +
92 # Swap out original extension sources with hipified sources.
93 list(APPEND HIP_SRCS ${CXX_SRCS})
94 set(${OUT_SRCS} ${HIP_SRCS} PARENT_SCOPE)
docker/Dockerfile.rocm
+200 -108
@@ -2,6 +2,7 @@
2 ARG REMOTE_VLLM="0"
3 ARG COMMON_WORKDIR=/app
4 ARG BASE_IMAGE=rocm/vllm-dev:base
5 +ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base
6 # NIC backend for MoRI RDMA support.
7 # By default (all), drivers and userspace libraries for all supported NIC types
8 # (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime.
@@ -16,7 +17,8 @@ ARG NIC_BACKEND=all
17 ARG AINIC_VERSION=1.117.3-hydra
18 ARG UBUNTU_CODENAME=jammy
19
19 -# Sccache configuration (only used in release pipeline)
20 +# Sccache configuration. Release builds use this today; CI can opt in when a
21 +# shared S3-compatible cache backend is available.
22 ARG USE_SCCACHE
23 ARG SCCACHE_DOWNLOAD_URL
24 ARG SCCACHE_ENDPOINT
@@ -29,12 +31,16 @@ FROM ${BASE_IMAGE} AS base
31 ARG ARG_PYTORCH_ROCM_ARCH
32 ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
33
32 -# Install some basic utilities
34 +# Install build dependencies and utilities
35 RUN apt-get update -q -y && apt-get install -q -y \
36 sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \
37 apt-transport-https ca-certificates wget curl \
36 - libnuma-dev
37 -RUN python3 -m pip install --upgrade pip
38 + libnuma-dev ccache mold
39 +RUN --mount=type=cache,target=/root/.cache/pip \
40 + python3 -m pip install --upgrade pip
41 +# Note: mold is installed but not set as the system default linker because
42 +# some packages use JIT compilation at runtime with flags mold does not support.
43 +# Build stages opt in via LDFLAGS="-fuse-ld=mold".
44 # Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_base)
45 ARG USE_SCCACHE
46 RUN if [ "$USE_SCCACHE" != "1" ]; then \
@@ -55,6 +61,12 @@ ENV UV_HTTP_TIMEOUT=500
61 ENV UV_INDEX_STRATEGY="unsafe-best-match"
62 # Use copy mode to avoid hardlink failures with Docker cache mounts
63 ENV UV_LINK_MODE=copy
64 +# ccache directory - persisted across layer rebuilds via cache mounts.
65 +ENV CCACHE_DIR=/root/.cache/ccache
66 +ENV CCACHE_COMPILERCHECK=content
67 +# Empty by default so build steps fall back to $(nproc); CI can override.
68 +ARG max_jobs
69 +ENV MAX_JOBS=${max_jobs}
70
71 # Install sccache if USE_SCCACHE is enabled (for release builds)
72 ARG USE_SCCACHE
@@ -86,6 +98,7 @@ RUN if [ "$USE_SCCACHE" = "1" ]; then \
98 ARG USE_SCCACHE
99 ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}}
100 ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
101 +ENV SCCACHE_ENDPOINT=${USE_SCCACHE:+${SCCACHE_ENDPOINT}}
102 ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
103 ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
104
@@ -114,8 +127,7 @@ FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm
127 # -----------------------
128 # Rust build stage
129 # Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages
117 -# don't need the rust toolchain or protoc. Runs in parallel with the main wheel
118 -# build for faster end-to-end builds.
130 +# don't need the rust toolchain or protoc.
131 FROM fetch_vllm AS rust-build
132 ARG COMMON_WORKDIR
133
@@ -144,24 +156,74 @@ ENV RUSTUP_MAX_RETRIES=10
156 # layer for later COPY --from=rust-build.
157 RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \
158 --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \
159 + --mount=type=cache,id=vllm-rocm-cargo-target,target=${COMMON_WORKDIR}/vllm/rust/target,sharing=locked \
160 cd ${COMMON_WORKDIR}/vllm \
161 && VLLM_RS_TARGET_PATH=/tmp/vllm-rs bash build_rust.sh \
162 && test -x /tmp/vllm-rs
163
164 # -----------------------
152 -# vLLM build stages
165 +# vLLM native build stages
166 +#
167 +# csrc-build intentionally copies only files that affect ROCm native extension
168 +# compilation. That keeps unrelated CI/test/docs edits from invalidating the
169 +# expensive HIP/C++ build layer.
170 +FROM base AS csrc-build
171 +ARG COMMON_WORKDIR
172 +WORKDIR ${COMMON_WORKDIR}/vllm
173 +
174 +COPY requirements/rocm.txt requirements/rocm.txt
175 +COPY requirements/common.txt requirements/common.txt
176 +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
177 + uv pip install --system -r requirements/rocm.txt
178 +
179 +# pyproject.toml is bind-mounted in the RUN step so metadata-only changes do
180 +# not invalidate the expensive native build layer.
181 +COPY setup.py CMakeLists.txt ./
182 +COPY cmake cmake/
183 +COPY csrc csrc/
184 +COPY vllm/envs.py vllm/envs.py
185 +COPY vllm/__init__.py vllm/__init__.py
186 +
187 +ENV VLLM_TARGET_DEVICE=rocm
188 +ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+rocm.csrc.build"
189 +
190 +RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproject.toml \
191 + --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
192 + export CCACHE_BASEDIR="$PWD" \
193 + && echo "=== ccache stats before ROCm native build ===" \
194 + && (ccache --show-stats || true) \
195 + && (ccache --zero-stats || true) \
196 + && EFFECTIVE_MAX_JOBS="${MAX_JOBS:-$(nproc)}" \
197 + && echo "Building ROCm native extension wheel with MAX_JOBS=${EFFECTIVE_MAX_JOBS}" \
198 + && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${EFFECTIVE_MAX_JOBS}" python3 setup.py bdist_wheel --dist-dir=dist \
199 + && test -d dist \
200 + && ls dist/*.whl >/dev/null \
201 + && echo "=== ccache stats after ROCm native build ===" \
202 + && (ccache --show-stats || true)
203 +
204 +# Build the full vLLM ROCm wheel by reusing the native extension wheel from
205 +# csrc-build. This stage still rebuilds for Python/package changes, but skips
206 +# the expensive HIP/C++ compile when native inputs are unchanged.
207 FROM fetch_vllm AS build_vllm
208 ARG COMMON_WORKDIR
209 +ENV VLLM_TARGET_DEVICE=rocm
210 +
211 +COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels
212
213 # Drop the pre-built rust frontend binary into the source tree. setup.py
214 # detects it and ships it as-is, skipping the local cargo build.
215 COPY --from=rust-build /tmp/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
216
160 -# Build vLLM (setup.py auto-detects sccache in PATH)
161 -RUN cd vllm \
162 - && python3 -m pip install -r requirements/rocm.txt \
163 - && python3 setup.py clean --all \
164 - && python3 setup.py bdist_wheel --dist-dir=dist
217 +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
218 + cd vllm \
219 + && uv pip install --system -r requirements/rocm.txt \
220 + && export VLLM_USE_PRECOMPILED=1 \
221 + && export VLLM_PRECOMPILED_WHEEL_LOCATION="$(ls /precompiled-wheels/*.whl)" \
222 + && export VLLM_DOCKER_BUILD_CONTEXT=1 \
223 + && echo "Packaging vLLM ROCm wheel using precompiled extensions from ${VLLM_PRECOMPILED_WHEEL_LOCATION}" \
224 + && python3 setup.py bdist_wheel --dist-dir=dist \
225 + && test -d dist \
226 + && ls dist/*.whl >/dev/null
227 FROM scratch AS export_vllm
228 ARG COMMON_WORKDIR
229 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/dist/*.whl /
@@ -171,6 +233,7 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests
233 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples
234 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
235 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
236 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
237 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
238
239 # RIXL/UCX build stages
@@ -201,14 +264,17 @@ RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
264 ibverbs-providers \
265 && rm -rf /var/lib/apt/lists/*
266
204 -RUN uv pip install --system meson auditwheel patchelf tomlkit
267 +RUN --mount=type=cache,target=/root/.cache/uv \
268 + uv pip install --system meson auditwheel patchelf tomlkit
269
206 -RUN cd /usr/local/src && \
270 +RUN --mount=type=cache,target=/root/.cache/ccache \
271 + cd /usr/local/src && \
272 git clone ${UCX_REPO} && \
273 cd ucx && \
274 git checkout ${UCX_BRANCH} && \
275 ./autogen.sh && \
276 mkdir build && cd build && \
277 + CC="ccache gcc" CXX="ccache g++" \
278 ../configure \
279 --prefix=/usr/local/ucx \
280 --enable-shared \
@@ -220,20 +286,22 @@ RUN cd /usr/local/src && \
286 --with-verbs \
287 --with-dm \
288 --enable-mt && \
223 - make -j && \
289 + make -j$(nproc) && \
290 make install
291
292 ENV PATH=/usr/local/ucx/bin:$PATH
293 ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
294
229 -RUN git clone ${RIXL_REPO} /opt/rixl && \
295 +RUN --mount=type=cache,target=/root/.cache/ccache \
296 + git clone ${RIXL_REPO} /opt/rixl && \
297 cd /opt/rixl && \
298 git checkout ${RIXL_BRANCH} && \
299 + CC="ccache gcc" CXX="ccache g++" \
300 meson setup build --prefix=${RIXL_HOME} \
301 -Ducx_path=${UCX_HOME} \
302 -Drocm_path=${ROCM_PATH} && \
303 cd build && \
236 - ninja && \
304 + ninja -j$(nproc) && \
305 ninja install
306
307 # Generate RIXL wheel
@@ -250,30 +318,44 @@ RUN cd /opt/rixl && \
318 --ucx-plugins-dir ${UCX_HOME}/lib/ucx \
319 --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
320
253 -# DeepEP build stage
254 -FROM base AS build_deep
321 +# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not
322 +# invalidate the slow ROCShmem build.
323 +FROM base AS build_rocshmem
324 ARG ROCSHMEM_BRANCH="f0acb0c6"
325 ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
257 -ARG DEEPEP_BRANCH="a9ea9774"
258 -ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
259 -ARG DEEPEP_NIC="cx7"
326 +# DeepEP only supports gfx942 and gfx950; build ROCShmem for the same set so
327 +# it can be linked against DeepEP without arch mismatches.
328 ARG DEEPEP_ROCM_ARCH="gfx942;gfx950"
329 +ENV ROCM_PATH=/opt/rocm
330 ENV ROCSHMEM_DIR=/opt/rocshmem
331
263 -RUN git clone ${ROCSHMEM_REPO} \
332 +RUN --mount=type=cache,target=/root/.cache/ccache \
333 + git clone --no-checkout --filter=blob:none ${ROCSHMEM_REPO} \
334 && cd rocm-systems \
335 + && git sparse-checkout set --cone projects/rocshmem \
336 && git checkout ${ROCSHMEM_BRANCH} \
337 && mkdir -p projects/rocshmem/build \
338 && cd projects/rocshmem/build \
268 - && INSTALL_PREFIX=${ROCSHMEM_DIR} \
269 - ../scripts/build_configs/all_backends -DUSE_EXTERNAL_MPI=OFF
339 + && CC="ccache gcc" CXX="ccache g++" INSTALL_PREFIX=${ROCSHMEM_DIR} \
340 + bash ../scripts/build_configs/all_backends \
341 + -DROCM_PATH=${ROCM_PATH} \
342 + -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \
343 + -DUSE_EXTERNAL_MPI=OFF
344 +
345 +# DeepEP build stage - depends on ROCShmem, builds the HIP kernel wheel.
346 +FROM build_rocshmem AS build_deepep
347 +ARG DEEPEP_BRANCH="a9ea9774"
348 +ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
349 +ARG DEEPEP_NIC="cx7"
350
271 -# Build DeepEP wheel.
272 -# DeepEP looks for rocshmem at ROCSHMEM_DIR.
273 -RUN git clone ${DEEPEP_REPO} \
351 +# Build DeepEP wheel. DeepEP looks for rocshmem at ROCSHMEM_DIR.
352 +# DeepEP only supports gfx942 and gfx950, so avoid gfx90a in the default list.
353 +RUN --mount=type=cache,target=/root/.cache/ccache \
354 + export PYTORCH_ROCM_ARCH="gfx942;gfx950" \
355 + && git clone ${DEEPEP_REPO} \
356 && cd DeepEP \
357 && git checkout ${DEEPEP_BRANCH} \
276 - && python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install
358 + && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install
359
360 # MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do
361 # not force users to rebuild the long-lived Dockerfile.rocm_base image.
@@ -372,8 +454,9 @@ RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \
454 # Extract version from git BEFORE any modifications (pin_rocm_dependencies.py modifies requirements/rocm.txt)
455 # This ensures setuptools_scm sees clean repo state for version detection
456 RUN --mount=type=bind,source=.git,target=vllm/.git \
457 + --mount=type=cache,target=/root/.cache/uv \
458 cd vllm \
376 - && pip install setuptools_scm regex \
459 + && uv pip install --system setuptools_scm regex \
460 && VLLM_VERSION=$(python3 -c "import setuptools_scm; print(setuptools_scm.get_version())") \
461 && echo "Detected vLLM version: ${VLLM_VERSION}" \
462 && echo "${VLLM_VERSION}" > /tmp/vllm_version.txt
@@ -409,18 +492,20 @@ RUN echo "Pinning vLLM dependencies to custom wheel versions..." \
492 && python3 /tmp/pin_rocm_dependencies.py /install ${COMMON_WORKDIR}/vllm/requirements/rocm.txt
493
494 # Install dependencies using custom wheels from /install
412 -RUN cd vllm \
495 +RUN --mount=type=cache,target=/root/.cache/uv \
496 + cd vllm \
497 && echo "Building vLLM with custom wheels from /install" \
414 - && python3 -m pip install --find-links /install -r requirements/rocm.txt \
415 - && python3 setup.py clean --all
498 + && uv pip install --system --find-links /install -r requirements/rocm.txt
499
500 # Build wheel using pre-extracted version to avoid dirty state from modified requirements/rocm.txt
418 -# (setup.py auto-detects sccache in PATH)
501 +# (setup.py auto-detects ccache/sccache in PATH)
502 RUN --mount=type=bind,source=.git,target=vllm/.git \
503 + --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
504 cd vllm \
505 + && export CCACHE_BASEDIR="$PWD" \
506 && export SETUPTOOLS_SCM_PRETEND_VERSION=$(cat /tmp/vllm_version.txt) \
507 && echo "Building wheel with version: ${SETUPTOOLS_SCM_PRETEND_VERSION}" \
423 - && python3 setup.py bdist_wheel --dist-dir=dist
508 + && MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist
509
510 FROM scratch AS export_vllm_wheel_release
511 ARG COMMON_WORKDIR
@@ -431,112 +516,118 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests
516 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples
517 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
518 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
519 +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
520 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
521
522 # -----------------------
437 -# Test vLLM image
438 -FROM mori_base AS test
439 -
440 -RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
441 -
442 -# Install vLLM using uv (inherited from base stage)
443 -# Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version
444 -RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
445 - --mount=type=cache,target=/root/.cache/uv \
446 - cd /install \
447 - && uv pip install --system -r requirements/rocm.txt \
448 - && uv pip install --system -r requirements/test/rocm.txt \
449 - && pip uninstall -y vllm \
450 - && uv pip install --system *.whl
451 -
452 -# Persist the built wheel in the image so python_only_compile_rocm.sh can
453 -# reinstall it after removing compilers. The bind-mounted /install contents
454 -# above are not available once that RUN step completes.
455 -COPY --from=export_vllm /*.whl /opt/vllm-wheels/
523 +# CI base image (Tier 1) - stable, rarely changing CI dependencies.
524 +# Per-PR test builds pull this as CI_BASE_IMAGE so the test stage only layers
525 +# in the vLLM artifacts for the current commit.
526 +FROM mori_base AS ci_base
527 +ARG COMMON_WORKDIR
528
457 -# Update rdma-core to support latest rocshmem
529 +# Update rdma-core to support latest rocshmem.
530 ARG DEEPEP_NIC
531 RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
532 git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \
533 cd /tmp/rdma-core && \
534 mkdir -p build && cd build && \
535 cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \
464 - ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
536 + ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
537 fi
538
467 -# Install RIXL wheel
539 +# Install RIXL + DeepEP wheels.
540 RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
469 - uv pip install --system /rixl_install/*.whl
541 + --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
542 + uv pip install --system /rixl_install/*.whl /deep_install/*.whl
543
471 -# Install DeepEP wheel
472 -RUN --mount=type=bind,from=build_deep,src=/app/deep_install,target=/deep_install \
473 - uv pip install --system /deep_install/*.whl
474 -COPY --from=build_deep /opt/rocshmem /opt/rocshmem
544 +# Copy ROCShmem runtime libraries.
545 +COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem
546
476 -# RIXL/MoRIIO runtime dependencies (RDMA userspace libraries)
477 -RUN apt-get update -q -y && apt-get install -q -y \
547 +# RDMA userspace libraries plus FFmpeg dev libs needed by torchcodec.
548 +RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
549 librdmacm1 \
550 libibverbs1 \
551 ibverbs-providers \
552 ibverbs-utils \
553 + pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \
554 + libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \
555 && rm -rf /var/lib/apt/lists/*
556
484 -WORKDIR /vllm-workspace
485 -ARG COMMON_WORKDIR
486 -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace
487 -
488 -# install development dependencies (for testing)
489 -RUN cd /vllm-workspace \
490 - && python3 -m pip install -e tests/vllm_test_utils \
491 - && python3 -m pip install pytest-shard
492 -
493 -# enable fast downloads from hf (for testing)
494 -ENV HF_XET_HIGH_PERFORMANCE=1
495 -
496 -# increase timeout for hf downloads (for testing)
497 -ENV HF_HUB_DOWNLOAD_TIMEOUT 60
498 -
499 -# install audio decode package `torchcodec` from source (required due to
500 -# ROCm and torch version mismatch) for tests with datasets package
557 +# Install torchcodec from source for ROCm/torch ABI compatibility.
558 COPY tools/install_torchcodec_rocm.sh /tmp/install_torchcodec.sh
502 -RUN bash /tmp/install_torchcodec.sh \
559 +RUN --mount=type=cache,target=/root/.cache/uv \
560 + --mount=type=cache,target=/root/.cache/pip \
561 + --mount=type=cache,target=/root/.cache/torchcodec-wheels \
562 + bash /tmp/install_torchcodec.sh \
563 && rm /tmp/install_torchcodec.sh \
504 - && apt-get clean \
505 - && rm -rf /var/lib/apt/lists/*
564 + && apt-get clean && rm -rf /var/lib/apt/lists/*
565
507 -# Copy in the v1 package (for python-only install test group)
508 -COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1
566 +# Pre-install shared ROCm runtime dependencies.
567 +COPY requirements/common.txt requirements/rocm.txt /tmp/ci-base-requirements/
568 +RUN --mount=type=cache,target=/root/.cache/uv \
569 + uv pip install --system -r /tmp/ci-base-requirements/rocm.txt \
570 + && rm -rf /tmp/ci-base-requirements
571 +
572 +# Enable fast and less brittle model downloads in tests.
573 +ENV HF_XET_HIGH_PERFORMANCE=1
574 +ENV HF_HUB_DOWNLOAD_TIMEOUT=60
575
510 -# Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel
576 +# Pre-install vLLM test dependencies.
577 +COPY requirements/test/rocm.txt /tmp/rocm-test-reqs.txt
578 +RUN --mount=type=cache,target=/root/.cache/uv \
579 + uv pip install --system -r /tmp/rocm-test-reqs.txt
580 +
581 +# Rebuild fastsafetensors from source so its C++ extension is compiled with
582 +# USE_ROCM and can detect libamdhip64.so at runtime.
583 +RUN --mount=type=cache,target=/root/.cache/pip \
584 + FASTSAFETENSORS_REQ="$(grep -E '^fastsafetensors(==| @ )' /tmp/rocm-test-reqs.txt | head -1)" \
585 + && test -n "${FASTSAFETENSORS_REQ}" \
586 + && python3 -m pip install --force-reinstall --no-deps \
587 + --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}" \
588 + && rm /tmp/rocm-test-reqs.txt
589 +
590 +# Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel.
591 # See: https://github.com/pytorch/pytorch/issues/169857
592 ENV MIOPEN_DEBUG_CONV_DIRECT=0
593 ENV MIOPEN_DEBUG_CONV_GEMM=0
594
515 -# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
595 +# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc.
596 # See: https://github.com/ROCm/rocm-libraries/issues/6266
597 ENV HSA_ENABLE_IPC_MODE_LEGACY=1
598
519 -# Source code is used in the `python_only_compile.sh` test
520 -# We hide it inside `src/` so that this source code
521 -# will not be imported by other tests
522 -RUN mkdir src && mv vllm src/vllm
599 +# ROCm profiler limits workaround.
600 +RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
601 +ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
602
524 -# This is a workaround to ensure pytest exits with the correct status code in CI tests.
525 -RUN printf '%s\n' \
526 - 'import os' \
527 - '' \
528 - '_exit_code = 1' \
529 - '' \
530 - 'def pytest_sessionfinish(session, exitstatus):' \
531 - ' global _exit_code' \
532 - ' _exit_code = int(exitstatus)' \
533 - '' \
534 - 'def pytest_unconfigure(config):' \
535 - ' import sys' \
536 - ' sys.stdout.flush()' \
537 - ' sys.stderr.flush()' \
538 - ' os._exit(_exit_code)' \
539 - > /vllm-workspace/conftest.py
603 +# Install vllm_test_utils in ci_base for ci_base + wheel parity.
604 +COPY tests/vllm_test_utils /tmp/vllm_test_utils
605 +RUN --mount=type=cache,target=/root/.cache/uv \
606 + uv pip install --system /tmp/vllm_test_utils \
607 + && rm -rf /tmp/vllm_test_utils
608 +
609 +# -----------------------
610 +# Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base.
611 +FROM ${CI_BASE_IMAGE} AS test
612 +ARG COMMON_WORKDIR
613 +
614 +# Install the vLLM wheel (--no-deps: all deps already in ci_base).
615 +RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
616 + --mount=type=cache,target=/root/.cache/uv \
617 + cd /install \
618 + && uv pip install --system --no-deps *.whl
619 +
620 +# Store the vLLM wheel in the image for python-only install tests.
621 +COPY --from=export_vllm /*.whl /opt/vllm-wheels/
622 +
623 +WORKDIR /vllm-workspace
624 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace
625 +
626 +# Copy in the v1 package (for python-only install test group).
627 +COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1
628 +
629 +# Hide source under src/ so it won't shadow the installed package in tests.
630 +RUN mkdir src && mv vllm src/vllm
631
632 # -----------------------
633 # Final vLLM image
@@ -553,6 +644,7 @@ RUN rm -f /usr/bin/sccache || true \
644 # This prevents S3 bucket config from leaking into production images
645 ENV SCCACHE_BUCKET=
646 ENV SCCACHE_REGION=
647 +ENV SCCACHE_ENDPOINT=
648 ENV SCCACHE_S3_NO_CREDENTIALS=
649 ENV SCCACHE_IDLE_TIMEOUT=
650
docker/ci-rocm.hcl new
+376
@@ -0,0 +1,376 @@
1 +# ci-rocm.hcl - CI-specific configuration for vLLM ROCm Docker builds
2 +#
3 +# This file lives in the vLLM repo at docker/ci-rocm.hcl so ROCm Docker
4 +# build mechanics can evolve with Dockerfile.rocm and docker-bake-rocm.hcl.
5 +# Used with: docker buildx bake -f docker/docker-bake-rocm.hcl -f docker/ci-rocm.hcl test-rocm-ci
6 +#
7 +# Registry cache: Docker Hub (rocm/vllm-ci-cache) is used exclusively.
8 +# AMD build agents already have Docker Hub credentials (they push the test
9 +# image to rocm/vllm-ci), so no additional credential setup is required.
10 +# ROCm CI uses Docker Hub for BuildKit layer cache by default. A separate
11 +# compiler cache can be enabled with USE_SCCACHE=1 when AMD provides a shared
12 +# S3-compatible cache endpoint.
13 +
14 +# CI metadata
15 +
16 +variable "BUILDKITE_COMMIT" {
17 + default = ""
18 +}
19 +
20 +variable "BUILDKITE_BUILD_NUMBER" {
21 + default = ""
22 +}
23 +
24 +variable "BUILDKITE_BUILD_ID" {
25 + default = ""
26 +}
27 +
28 +variable "PARENT_COMMIT" {
29 + default = ""
30 +}
31 +
32 +# Merge-base of HEAD with main - provides a more stable cache fallback than
33 +# parent commit for long-lived PRs. Mirrors the VLLM_MERGE_BASE_COMMIT
34 +# pattern used in the shared ci.hcl file. Auto-computed by ci-bake-rocm.sh
35 +# when unset.
36 +variable "VLLM_MERGE_BASE_COMMIT" {
37 + default = ""
38 +}
39 +
40 +# Bridge to vLLM's COMMIT variable for OCI labels
41 +variable "COMMIT" {
42 + default = BUILDKITE_COMMIT
43 +}
44 +
45 +# Image tags (set by CI)
46 +
47 +variable "IMAGE_TAG" {
48 + default = ""
49 +}
50 +
51 +variable "IMAGE_TAG_LATEST" {
52 + default = ""
53 +}
54 +
55 +# ROCm-specific GPU architecture targets
56 +
57 +variable "PYTORCH_ROCM_ARCH" {
58 + default = "gfx90a;gfx942;gfx950"
59 +}
60 +
61 +# Pre-built CI base image (Tier 1). Per-PR builds pull this instead of
62 +# rebuilding RIXL/DeepEP/torchcodec from scratch. The ci_base stage in
63 +# Dockerfile.rocm inherits from base, so CI_BASE_IMAGE only affects the test
64 +# stage and is irrelevant when building --target ci_base itself.
65 +variable "CI_BASE_IMAGE" {
66 + default = "rocm/vllm-dev:ci_base"
67 +}
68 +
69 +# Leave CI_MAX_JOBS empty so the Dockerfile falls back to $(nproc) and uses
70 +# the full builder parallelism. Operators can still override this per build.
71 +variable "CI_MAX_JOBS" {
72 + default = ""
73 +}
74 +
75 +# Upstream dependency commit pins -- extracted from Dockerfile.rocm by
76 +# ci-bake-rocm.sh at build time. Empty defaults are safe: the cache
77 +# functions produce no entries when the variable is empty.
78 +variable "RIXL_BRANCH" {
79 + default = ""
80 +}
81 +
82 +variable "UCX_BRANCH" {
83 + default = ""
84 +}
85 +
86 +variable "ROCSHMEM_BRANCH" {
87 + default = ""
88 +}
89 +
90 +variable "DEEPEP_BRANCH" {
91 + default = ""
92 +}
93 +
94 +variable "RIXL_CACHE_KEY" {
95 + default = ""
96 +}
97 +
98 +variable "ROCSHMEM_CACHE_KEY" {
99 + default = ""
100 +}
101 +
102 +variable "DEEPEP_CACHE_KEY" {
103 + default = ""
104 +}
105 +
106 +# Docker Hub registry cache for AMD builds.
107 +#
108 +# A separate repo (rocm/vllm-ci-cache) is used for BuildKit layer cache.
109 +# Final-image cache exports use mode=min to reduce the volume of data pushed.
110 +# Source-scoped csrc cache exports default to mode=max so fresh workers can
111 +# recover more of the native build graph when ROCm extension inputs change.
112 +# NOTE: mode=min still includes all layers referenced by the final image
113 +# manifest, including inherited base layers (~7.25GB ROCm runtime).
114 +# Docker Hub auto-creates the repo on first push.
115 +#
116 +# Final-image cache stays commit-scoped. Branch-to-branch reuse for the test
117 +# image comes from importing the parent and merge-base commit cache refs.
118 +#
119 +# The source-scoped native cache is exported both per-commit and per-branch so
120 +# ROCm extension rebuilds are shareable within the same commit reruns and across
121 +# consecutive commits on the same branch without depending on a single global
122 +# latest tag.
123 +
124 +variable "DOCKERHUB_CACHE_REPO" {
125 + default = "rocm/vllm-ci-cache"
126 +}
127 +
128 +variable "DOCKERHUB_CACHE_TO" {
129 + default = ""
130 +}
131 +
132 +variable "ROCM_CACHE_BRANCH_TAG" {
133 + default = ""
134 +}
135 +
136 +variable "ROCM_CACHE_UPSTREAM_BRANCH_TAG" {
137 + default = ""
138 +}
139 +
140 +variable "ROCM_CSRC_CACHE_TO_MODE" {
141 + default = "max"
142 +}
143 +
144 +variable "ROCM_FINAL_CACHE_TO_MODE" {
145 + default = "min"
146 +}
147 +
148 +# Functions
149 +
150 +function "get_cache_from_rocm" {
151 + params = []
152 + result = compact([
153 + # Exact commit hit - fastest cache on re-runs of the same commit
154 + BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-${BUILDKITE_COMMIT}" : "",
155 + # Parent commit - useful cache for incremental changes
156 + PARENT_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-${PARENT_COMMIT}" : "",
157 + # Merge-base with main - stable fallback for long-lived or rebased PRs;
158 + # maps to a real main-branch commit whose cache layers are likely warm
159 + VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-${VLLM_MERGE_BASE_COMMIT}" : "",
160 + # Import the source-scoped native build cache as well so builds whose
161 + # Python/package layers changed can still reuse compiled ROCm objects.
162 + BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${BUILDKITE_COMMIT}" : "",
163 + PARENT_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${PARENT_COMMIT}" : "",
164 + VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${VLLM_MERGE_BASE_COMMIT}" : "",
165 + ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "",
166 + ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "",
167 + # Branch-scoped full image cache - fallback when parent-commit cache is evicted
168 + ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "",
169 + ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "",
170 + ])
171 +}
172 +
173 +function "get_cache_to_rocm" {
174 + params = []
175 + result = compact([
176 + # Commit-scoped cache for exact re-runs.
177 + BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-${BUILDKITE_COMMIT},mode=${ROCM_FINAL_CACHE_TO_MODE}" : "",
178 + # Branch-scoped cache so later commits on the same branch can reuse the full
179 + # image layers when the parent-commit cache is evicted. Unlike the old
180 + # rocm-latest tag (which caused duplicate exporter 400s), this is per-branch.
181 + ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${ROCM_FINAL_CACHE_TO_MODE}" : "",
182 + ])
183 +}
184 +
185 +function "get_cache_from_rocm_csrc" {
186 + params = []
187 + result = compact([
188 + BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${BUILDKITE_COMMIT}" : "",
189 + PARENT_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${PARENT_COMMIT}" : "",
190 + VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${VLLM_MERGE_BASE_COMMIT}" : "",
191 + ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "",
192 + ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "",
193 + ])
194 +}
195 +
196 +function "get_cache_to_rocm_csrc" {
197 + params = []
198 + result = compact([
199 + # Export the exact-commit native cache for same-commit reruns.
200 + BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${BUILDKITE_COMMIT},mode=${ROCM_CSRC_CACHE_TO_MODE}" : "",
201 + # Export the branch-scoped native cache so later commits on the same branch
202 + # can reuse compiled ROCm objects even when the exact parent cache is absent.
203 + ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${ROCM_CSRC_CACHE_TO_MODE}" : "",
204 + ])
205 +}
206 +
207 +# Cache functions for upstream dependency stages (RIXL/UCX, ROCShmem, DeepEP).
208 +# These stages are pinned to specific upstream commit hashes, so cache keys use
209 +# those hashes rather than the Buildkite commit. This means the cache persists
210 +# across all vLLM commits as long as the upstream dependency pins don't change.
211 +
212 +function "get_cache_from_rocm_deps" {
213 + params = []
214 + result = compact([
215 + RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY}" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""),
216 + ROCSHMEM_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY}" : (ROCSHMEM_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_BRANCH}" : ""),
217 + DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY}" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH}" : ""),
218 + ])
219 +}
220 +
221 +function "get_cache_to_rocm_rixl" {
222 + params = []
223 + result = compact([
224 + RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY},mode=min" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""),
225 + ])
226 +}
227 +
228 +function "get_cache_to_rocm_rocshmem" {
229 + params = []
230 + result = compact([
231 + ROCSHMEM_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY},mode=min" : (ROCSHMEM_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_BRANCH},mode=min" : ""),
232 + ])
233 +}
234 +
235 +function "get_cache_to_rocm_deepep" {
236 + params = []
237 + result = compact([
238 + DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY},mode=min" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH},mode=min" : ""),
239 + ])
240 +}
241 +
242 +# CI targets
243 +
244 +target "_ci-rocm" {
245 + annotations = [
246 + "manifest:vllm.buildkite.build_number=${BUILDKITE_BUILD_NUMBER}",
247 + "manifest:vllm.buildkite.build_id=${BUILDKITE_BUILD_ID}",
248 + ]
249 + args = {
250 + ARG_PYTORCH_ROCM_ARCH = PYTORCH_ROCM_ARCH
251 + CI_BASE_IMAGE = CI_BASE_IMAGE
252 + max_jobs = CI_MAX_JOBS
253 + }
254 +}
255 +
256 +target "test-rocm-ci" {
257 + inherits = ["_common-rocm", "_ci-rocm", "_labels"]
258 + target = "test"
259 + cache-from = get_cache_from_rocm()
260 + cache-to = get_cache_to_rocm()
261 + tags = compact([
262 + IMAGE_TAG,
263 + IMAGE_TAG_LATEST,
264 + ])
265 + output = ["type=registry"]
266 +}
267 +
268 +# Cache-only target for the source-scoped ROCm native build stage.
269 +# This persists the csrc-build stage in the registry cache even though the
270 +# final test image only consumes it indirectly while packaging the wheel.
271 +target "csrc-rocm-ci" {
272 + inherits = ["_common-rocm", "_ci-rocm"]
273 + target = "csrc-build"
274 + cache-from = get_cache_from_rocm_csrc()
275 + cache-to = get_cache_to_rocm_csrc()
276 + output = ["type=cacheonly"]
277 +}
278 +
279 +# Keep wheel export on the same CI graph as the test image build so the
280 +# shared build_vllm/export_vllm stages resolve identically within one bake
281 +# invocation. Without this, export-wheel-rocm uses the plain local target
282 +# args while test-rocm-ci uses CI-only args, which can lead to separate
283 +# cache lineages and inconsistent export_vllm results.
284 +target "export-wheel-rocm" {
285 + inherits = ["_common-rocm", "_ci-rocm"]
286 + target = "export_vllm"
287 + cache-from = get_cache_from_rocm()
288 + cache-to = get_cache_to_rocm()
289 + output = ["type=local,dest=./wheel-export"]
290 +}
291 +
292 +# Artifact-only vLLM build. GPU test jobs consume this artifact on top of
293 +# ci_base, avoiding a per-commit multi-GB image push/pull.
294 +group "test-rocm-ci-with-artifacts" {
295 + targets = ["csrc-rocm-ci", "export-wheel-rocm"]
296 +}
297 +
298 +# Full test image + wheel export. Kept for fallback/debugging when a pushed
299 +# per-commit image is useful.
300 +group "test-rocm-ci-with-wheel" {
301 + targets = ["csrc-rocm-ci", "test-rocm-ci", "export-wheel-rocm"]
302 +}
303 +
304 +# Image tags for the ci_base build. ci-bake-rocm.sh rewrites CI_BASE_IMAGE_TAG
305 +# to the primary tag for this build. Non-nightly builds use a commit-scoped tag
306 +# and also publish a content tag for reuse. NIGHTLY=1 builds on the stable branch
307 +# can additionally set CI_BASE_IMAGE_TAG_STABLE to refresh rocm/vllm-dev:ci_base.
308 +variable "CI_BASE_IMAGE_TAG" {
309 + default = "rocm/vllm-dev:ci_base"
310 +}
311 +
312 +variable "CI_BASE_IMAGE_TAG_CONTENT" {
313 + default = ""
314 +}
315 +
316 +variable "CI_BASE_IMAGE_TAG_STABLE" {
317 + default = ""
318 +}
319 +
320 +# Cache-only targets for upstream dependency stages. These persist each stage
321 +# in the registry cache keyed by its upstream commit hash. When ci_base rebuilds
322 +# (e.g., requirements change), these stages are cache hits if their upstream
323 +# pins haven't changed -- saving ~35min of compilation.
324 +target "rixl-rocm-ci" {
325 + inherits = ["_common-rocm", "_ci-rocm"]
326 + target = "build_rixl"
327 + cache-from = get_cache_from_rocm_deps()
328 + cache-to = get_cache_to_rocm_rixl()
329 + output = ["type=cacheonly"]
330 +}
331 +
332 +target "rocshmem-rocm-ci" {
333 + inherits = ["_common-rocm", "_ci-rocm"]
334 + target = "build_rocshmem"
335 + cache-from = get_cache_from_rocm_deps()
336 + cache-to = get_cache_to_rocm_rocshmem()
337 + output = ["type=cacheonly"]
338 +}
339 +
340 +target "deepep-rocm-ci" {
341 + inherits = ["_common-rocm", "_ci-rocm"]
342 + target = "build_deepep"
343 + cache-from = get_cache_from_rocm_deps()
344 + cache-to = get_cache_to_rocm_deepep()
345 + output = ["type=cacheonly"]
346 +}
347 +
348 +# Builds only the ci_base stage (RIXL, DeepEP, torchcodec, etc.)
349 +# Invoked by the ensure-ci-base step when the content hash of ci_base-affecting
350 +# files drifts from the remote image label. Per-PR builds then pull the result
351 +# as CI_BASE_IMAGE instead of rebuilding those slow layers on every commit.
352 +# Uses inline cache metadata on the ci_base image itself instead of exporting a
353 +# separate registry cache artifact.
354 +target "ci-base-rocm-ci" {
355 + inherits = ["_common-rocm", "_ci-rocm", "_labels"]
356 + target = "ci_base"
357 + cache-from = concat(
358 + compact([
359 + CI_BASE_IMAGE_TAG != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG}" : "",
360 + CI_BASE_IMAGE_TAG_CONTENT != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_CONTENT}" : "",
361 + CI_BASE_IMAGE_TAG_STABLE != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_STABLE}" : "",
362 + ]),
363 + # Import upstream dependency caches so RIXL/ROCShmem/DeepEP stages
364 + # are cache hits even when ci_base itself needs rebuilding.
365 + get_cache_from_rocm_deps(),
366 + )
367 + cache-to = ["type=inline"]
368 + tags = compact([CI_BASE_IMAGE_TAG, CI_BASE_IMAGE_TAG_CONTENT, CI_BASE_IMAGE_TAG_STABLE])
369 + output = ["type=registry"]
370 +}
371 +
372 +# Group for ci_base builds -- exports dependency stage caches alongside the
373 +# ci_base image so future rebuilds can reuse them independently.
374 +group "ci-base-rocm-ci-with-deps" {
375 + targets = ["rixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
376 +}
docker/docker-bake-rocm.hcl new
+143
@@ -0,0 +1,143 @@
1 +# docker-bake-rocm.hcl - vLLM ROCm Docker build configuration
2 +#
3 +# This file lives in the vLLM repo at docker/docker-bake-rocm.hcl
4 +# Equivalent of docker-bake.hcl for ROCm builds.
5 +#
6 +# Usage:
7 +# docker buildx bake -f docker/docker-bake-rocm.hcl # Build test (default)
8 +# docker buildx bake -f docker/docker-bake-rocm.hcl final-rocm # Build final image
9 +# docker buildx bake -f docker/docker-bake-rocm.hcl --print # Show resolved config
10 +#
11 +# CI usage (with the vLLM-owned CI overlay):
12 +# docker buildx bake -f docker/docker-bake-rocm.hcl -f docker/ci-rocm.hcl test-rocm-ci
13 +
14 +variable "MAX_JOBS" {
15 + # Empty string lets the Dockerfile fall back to $(nproc) via
16 + # MAX_JOBS="${MAX_JOBS:-$(nproc)}" in each RUN step, which uses all
17 + # available cores on whatever machine the build runs on.
18 + # Override with --set '*.args.max_jobs=8' for local builds on small machines.
19 + default = ""
20 +}
21 +
22 +variable "PYTORCH_ROCM_ARCH" {
23 + default = "gfx90a;gfx942;gfx950"
24 +}
25 +
26 +variable "COMMIT" {
27 + default = ""
28 +}
29 +
30 +# Content hash of ci_base-affecting files. Computed by ci-bake-rocm.sh and
31 +# embedded as a label so future builds can compare without rebuilding.
32 +variable "CI_BASE_CONTENT_HASH" {
33 + default = ""
34 +}
35 +
36 +# REMOTE_VLLM=0: use local source via Docker build context (ONBUILD COPY ./ vllm/)
37 +# REMOTE_VLLM=1: clone from GitHub at VLLM_BRANCH (standalone builds without local source)
38 +variable "REMOTE_VLLM" {
39 + default = "0"
40 +}
41 +
42 +variable "VLLM_BRANCH" {
43 + default = "main"
44 +}
45 +
46 +# CI_BASE_IMAGE: pre-built ci_base image for per-PR test builds.
47 +# Defaults to the local "ci_base" stage for standalone/local builds.
48 +# CI overrides this to "rocm/vllm-dev:ci_base" via environment variable.
49 +variable "CI_BASE_IMAGE" {
50 + default = "rocm/vllm-dev:ci_base"
51 +}
52 +
53 +# Upstream dependency commit pins. Plain local bake builds use the Dockerfile
54 +# ARG defaults. ci-bake-rocm.sh resolves those defaults (plus any env
55 +# overrides) and writes a small HCL override before invoking CI targets.
56 +variable "RIXL_BRANCH" {
57 + default = ""
58 +}
59 +
60 +variable "UCX_BRANCH" {
61 + default = ""
62 +}
63 +
64 +variable "ROCSHMEM_BRANCH" {
65 + default = ""
66 +}
67 +
68 +variable "DEEPEP_BRANCH" {
69 + default = ""
70 +}
71 +
72 +group "default" {
73 + targets = ["test-rocm"]
74 +}
75 +
76 +target "_common-rocm" {
77 + dockerfile = "docker/Dockerfile.rocm"
78 + context = "."
79 + args = {
80 + max_jobs = MAX_JOBS
81 + ARG_PYTORCH_ROCM_ARCH = PYTORCH_ROCM_ARCH
82 + REMOTE_VLLM = REMOTE_VLLM
83 + VLLM_BRANCH = VLLM_BRANCH
84 + CI_BASE_IMAGE = CI_BASE_IMAGE
85 + }
86 +}
87 +
88 +target "_labels" {
89 + labels = {
90 + "org.opencontainers.image.source" = "https://github.com/vllm-project/vllm"
91 + "org.opencontainers.image.vendor" = "vLLM"
92 + "org.opencontainers.image.title" = "vLLM ROCm"
93 + "org.opencontainers.image.description" = "vLLM: A high-throughput and memory-efficient inference and serving engine for LLMs (ROCm)"
94 + "org.opencontainers.image.licenses" = "Apache-2.0"
95 + "org.opencontainers.image.revision" = COMMIT
96 + }
97 + annotations = [
98 + "manifest:org.opencontainers.image.revision=${COMMIT}",
99 + ]
100 +}
101 +
102 +target "test-rocm" {
103 + inherits = ["_common-rocm", "_labels"]
104 + target = "test"
105 + tags = ["rocm/vllm:test"]
106 + output = ["type=docker"]
107 +}
108 +
109 +# CI base image target - builds only the ci_base stage (RIXL, DeepEP,
110 +# torchcodec, requirements, etc.). Used by the weekly scheduled build and
111 +# the auto-rebuild trigger when requirements change in a PR.
112 +target "ci-base-rocm" {
113 + inherits = ["_common-rocm", "_labels"]
114 + target = "ci_base"
115 + labels = {
116 + "vllm.ci_base.content_hash" = CI_BASE_CONTENT_HASH
117 + }
118 + tags = ["rocm/vllm-dev:ci_base"]
119 + output = ["type=docker"]
120 +}
121 +
122 +# Wheel export target - extracts the built vLLM wheel + test workspace
123 +# to local disk. Used by CI to upload the wheel as a Buildkite artifact
124 +# so test jobs can assemble images locally from ci_base + wheel instead
125 +# of pulling the full large image from Docker Hub.
126 +#
127 +# Usage:
128 +# docker buildx bake -f docker/docker-bake-rocm.hcl export-wheel-rocm
129 +# # Creates ./wheel-export/*.whl, ./wheel-export/requirements/, etc.
130 +#
131 +# After a full bake build, BuildKit cache makes this nearly instant.
132 +target "export-wheel-rocm" {
133 + inherits = ["_common-rocm"]
134 + target = "export_vllm"
135 + output = ["type=local,dest=./wheel-export"]
136 +}
137 +
138 +target "final-rocm" {
139 + inherits = ["_common-rocm", "_labels"]
140 + target = "final"
141 + tags = ["rocm/vllm:latest"]
142 + output = ["type=docker"]
143 +}
tools/install_torchcodec_rocm.sh
+53 -4
@@ -3,12 +3,16 @@
3 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
5 # Script to install TorchCodec from source (required for ROCm compatibility)
6 +# The PyPI wheel is built against upstream PyTorch and has ABI mismatches with
7 +# ROCm's custom torch build, so we must compile from source.
8
9 set -e
10
11 TORCHCODEC_REPO="${TORCHCODEC_REPO:-https://github.com/pytorch/torchcodec.git}"
12 # Pin to a specific release for reproducibility; update as needed.
13 TORCHCODEC_BRANCH="${TORCHCODEC_BRANCH:-v0.10.0}"
14 +# Cache directory for pre-built wheels to avoid redundant recompilation.
15 +TORCHCODEC_WHEEL_CACHE="${TORCHCODEC_WHEEL_CACHE:-/root/.cache/torchcodec-wheels}"
16
17 echo "=== TorchCodec Installation Script ==="
18
@@ -18,9 +22,26 @@ if python3 -c "from torchcodec.decoders import VideoDecoder" 2>/dev/null; then
22 exit 0
23 fi
24
25 +# Try to install from cached wheel first
26 +ARCH_TAG="${PYTORCH_ROCM_ARCH:-all}"
27 +# Normalize arch tag (replace ; with _) for use in filename
28 +ARCH_TAG="${ARCH_TAG//;/_}"
29 +CACHED_WHEEL="${TORCHCODEC_WHEEL_CACHE}/torchcodec-${TORCHCODEC_BRANCH}-${ARCH_TAG}.whl"
30 +
31 +if [ -f "$CACHED_WHEEL" ]; then
32 + echo "Found cached wheel: $CACHED_WHEEL"
33 + pip install "$CACHED_WHEEL" && {
34 + echo "Installed from cached wheel."
35 + echo "=== TorchCodec installation complete ==="
36 + exit 0
37 + }
38 + echo "Cached wheel installation failed, rebuilding from source..."
39 +fi
40 +
41 echo "TorchCodec not found. Installing from source..."
42
23 -# Install system dependencies (FFmpeg + pkg-config)
43 +# Install system dependencies (FFmpeg + pkg-config) if not already present.
44 +# The Docker test image pre-installs these, so this is a fallback for other envs.
45 install_system_deps() {
46 if command -v apt-get &> /dev/null; then
47 echo "Installing system dependencies..."
@@ -56,6 +77,12 @@ export pybind11_DIR=$(python3 -c "import pybind11; print(pybind11.get_cmake_dir(
77 export CMAKE_PREFIX_PATH="${pybind11_DIR}:${CMAKE_PREFIX_PATH}"
78 echo "pybind11_DIR set to: $pybind11_DIR"
79
80 +# Limit GPU architectures to only what this image targets.
81 +# The default builds for all supported archs which is very slow.
82 +if [ -n "$PYTORCH_ROCM_ARCH" ]; then
83 + echo "Building for PYTORCH_ROCM_ARCH=$PYTORCH_ROCM_ARCH"
84 +fi
85 +
86 # Create temp directory for build
87 BUILD_DIR=$(mktemp -d -t torchcodec-XXXXXX)
88 echo "Building in temporary directory: $BUILD_DIR"
@@ -77,9 +104,31 @@ cd torchcodec
104 export TORCHCODEC_CMAKE_BUILD_DIR="${PWD}/build"
105 export TORCHCODEC_DISABLE_COMPILE_WARNING_AS_ERROR=1
106 export I_CONFIRM_THIS_IS_NOT_A_LICENSE_VIOLATION=1
107 +# Use ninja for faster builds and parallelize compilation
108 +export CMAKE_GENERATOR=Ninja
109 +export MAX_JOBS="${MAX_JOBS:-$(nproc)}"
110 +# Use ccache if available to speed up recompilation
111 +if command -v ccache &> /dev/null; then
112 + export CMAKE_C_COMPILER_LAUNCHER=ccache
113 + export CMAKE_CXX_COMPILER_LAUNCHER=ccache
114 +fi
115
81 -echo "Building TorchCodec..."
82 -pip install . --no-build-isolation
116 +echo "Building TorchCodec (MAX_JOBS=$MAX_JOBS)..."
117 +pip wheel . --no-build-isolation --no-deps -w "$BUILD_DIR/dist"
118 +
119 +# Install the built wheel
120 +BUILT_WHEEL=$(ls "$BUILD_DIR/dist"/torchcodec-*.whl 2>/dev/null | head -1)
121 +if [ -z "$BUILT_WHEEL" ]; then
122 + echo "Error: No wheel produced"
123 + exit 1
124 +fi
125 +
126 +pip install "$BUILT_WHEEL"
127 +
128 +# Cache the wheel for future runs
129 +mkdir -p "$TORCHCODEC_WHEEL_CACHE"
130 +cp "$BUILT_WHEEL" "$CACHED_WHEEL"
131 +echo "Cached wheel to: $CACHED_WHEEL"
132
133 # Verify installation
134 echo "Verifying installation..."
@@ -88,4 +137,4 @@ if python3 -c "from torchcodec.decoders import VideoDecoder; print('TorchCodec i
137 else
138 echo "Error: TorchCodec installation failed verification"
139 exit 1
91 -fi
\ No newline at end of file
140 +fi