[ROCm][CI] Add ci_base metadata for external cache orchestration (#46886)
Signed-off-by: Andreas Karatzas <akaratza@amd.com> Signed-off-by: Codex <codex@example.invalid> Co-authored-by: Codex <codex@example.invalid>
Andreas Karatzas committed
Jun 27, 2026 at 23:51 UTC
a65f93fb2e295e501b929df3c291ec89c27d39e8
2 files changed
+194
-21
.buildkite/scripts/ci-bake-rocm.sh
+179
-15
@@ -18,6 +18,7 @@ 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
+DEFAULT_CI_BASE_METADATA_VERSION="1"
22
IMAGE_EXISTED_BEFORE_BUILD=0
23
24
TARGET=""
@@ -525,6 +526,22 @@ get_remote_image_label_with_retry() {
526
return 0
527
}
528
529
+remote_ci_base_metadata_is_current() {
530
+ local image_ref="$1"
531
+ local metadata_version=""
532
+
533
+ metadata_version=$(get_remote_image_label "${image_ref}" "vllm.ci_base.metadata_version")
534
+ [[ "${metadata_version}" == "${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}" ]]
535
+}
536
+
537
+remote_ci_base_metadata_is_current_with_retry() {
538
+ local image_ref="$1"
539
+ local metadata_version=""
540
+
541
+ metadata_version=$(get_remote_image_label_with_retry "${image_ref}" "vllm.ci_base.metadata_version")
542
+ [[ "${metadata_version}" == "${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}" ]]
543
+}
544
+
545
remote_image_exists() {
546
local image_ref="$1"
547
docker manifest inspect "${image_ref}" >/dev/null 2>&1
@@ -581,6 +598,7 @@ init_config() {
598
CI_BASE_CONTENT_FILES="${CI_BASE_CONTENT_FILES:-${DEFAULT_CI_BASE_CONTENT_FILES}}"
599
CI_BASE_DOCKERFILE="${CI_BASE_DOCKERFILE:-${DEFAULT_CI_BASE_DOCKERFILE}}"
600
CI_BASE_DOCKERFILE_STAGES="${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}"
601
+ CI_BASE_METADATA_VERSION="${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}"
602
CI_BASE_IMAGE_TAG="${CI_BASE_IMAGE_TAG:-rocm/vllm-dev:ci_base}"
603
export PYTORCH_ROCM_ARCH
604
@@ -635,6 +653,10 @@ load_ci_hcl() {
653
echo "Copied ${CI_HCL_SOURCE} to ${CI_HCL_PATH}"
654
}
655
656
+init_bake_files() {
657
+ BAKE_FILES=(-f "${VLLM_BAKE_FILE}" -f "${CI_HCL_PATH}")
658
+}
659
+
660
compute_ci_base_hash_if_needed() {
661
if [[ -z "${CI_BASE_CONTENT_FILES:-}" ]]; then
662
return 0
@@ -676,12 +698,14 @@ configure_ci_base_image_refs() {
698
fi
699
700
content_tag=$(ci_base_tag_with_suffix "${stable_tag}" "${CI_BASE_CONTENT_HASH}")
701
+ CI_BASE_IMAGE_TAG_CONTENT_REF="${content_tag}"
702
if [[ -n "${BUILDKITE_COMMIT:-}" ]]; then
703
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
704
fi
705
+ CI_BASE_IMAGE_TAG_COMMIT_REF="${commit_tag}"
706
707
+ # *_REF is the logical tag recorded in metadata. *_EXTRA is only passed to
708
+ # bake when that tag is not already the primary tag, avoiding duplicates.
709
if should_push_stable_ci_base_tag; then
710
primary_tag="${content_tag}"
711
CI_BASE_IMAGE_TAG_STABLE="${stable_tag}"
@@ -691,19 +715,33 @@ configure_ci_base_image_refs() {
715
fi
716
CI_BASE_IMAGE_TAG="${primary_tag}"
717
if [[ "${primary_tag}" == "${content_tag}" ]]; then
694
- CI_BASE_IMAGE_TAG_CONTENT=""
718
+ CI_BASE_IMAGE_TAG_CONTENT_EXTRA=""
719
+ else
720
+ CI_BASE_IMAGE_TAG_CONTENT_EXTRA="${content_tag}"
721
+ fi
722
+ if [[ -n "${commit_tag}" && "${commit_tag}" != "${primary_tag}" ]]; then
723
+ CI_BASE_IMAGE_TAG_COMMIT_EXTRA="${commit_tag}"
724
else
696
- CI_BASE_IMAGE_TAG_CONTENT="${content_tag}"
725
+ CI_BASE_IMAGE_TAG_COMMIT_EXTRA=""
726
fi
698
- export CI_BASE_IMAGE_TAG CI_BASE_IMAGE_TAG_CONTENT CI_BASE_IMAGE_TAG_STABLE
727
+ export CI_BASE_IMAGE_TAG
728
+ export CI_BASE_IMAGE_TAG_COMMIT_EXTRA
729
+ export CI_BASE_IMAGE_TAG_CONTENT_EXTRA
730
+ export CI_BASE_IMAGE_TAG_CONTENT_REF
731
+ export CI_BASE_IMAGE_TAG_COMMIT_REF
732
+ export CI_BASE_IMAGE_TAG_STABLE
733
734
if is_ci_base_target; then
735
IMAGE_TAG="${primary_tag}"
736
export IMAGE_TAG
737
738
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}"
739
+ if [[ -n "${commit_tag}" ]]; then
740
+ if [[ "${commit_tag}" == "${primary_tag}" ]]; then
741
+ echo "ci_base commit image tag: ${commit_tag} (primary)"
742
+ else
743
+ echo "ci_base commit image tag: ${commit_tag}"
744
+ fi
745
fi
746
echo "ci_base content image tag: ${content_tag}"
747
if [[ -n "${CI_BASE_IMAGE_TAG_STABLE}" ]]; then
@@ -728,8 +766,8 @@ ci_base_candidate_refs() {
766
printf '%s\n' \
767
"${IMAGE_TAG:-}" \
768
"${CI_BASE_IMAGE_TAG:-}" \
731
- "${CI_BASE_IMAGE_TAG_COMMIT:-}" \
732
- "${CI_BASE_IMAGE_TAG_CONTENT:-}" \
769
+ "${CI_BASE_IMAGE_TAG_COMMIT_EXTRA:-}" \
770
+ "${CI_BASE_IMAGE_TAG_CONTENT_EXTRA:-}" \
771
"${CI_BASE_IMAGE_TAG_STABLE:-}" \
772
| awk 'NF && !seen[$0]++'
773
}
@@ -743,6 +781,10 @@ find_matching_ci_base_ref() {
781
remote_image_exists "${candidate}" || continue
782
candidate_hash=$(get_remote_image_label "${candidate}" "vllm.ci_base.content_hash")
783
if [[ "${candidate_hash}" == "${CI_BASE_CONTENT_HASH}" ]]; then
784
+ if ! remote_ci_base_metadata_is_current "${candidate}"; then
785
+ echo "Found matching ci_base content hash but stale metadata: ${candidate}" >&2
786
+ continue
787
+ fi
788
printf '%s\n' "${candidate}"
789
return 0
790
fi
@@ -817,6 +859,10 @@ maybe_skip_existing_image() {
859
if [[ -n "${remote_hash}" ]]; then
860
echo "Remote ci_base content hash: ${remote_hash:0:16}..."
861
if [[ "${remote_hash}" == "${CI_BASE_CONTENT_HASH}" ]]; then
862
+ if ! remote_ci_base_metadata_is_current "${IMAGE_TAG}"; then
863
+ echo "Content hashes match but ci_base metadata is stale; rebuilding to refresh metadata"
864
+ return 0
865
+ fi
866
if ! refresh_ci_base_tags_from_ref "${IMAGE_TAG}"; then
867
echo "ci_base tag refresh failed; rebuilding to push expected tags"
868
return 0
@@ -998,12 +1044,104 @@ prepare_git_cache_metadata() {
1044
fi
1045
}
1046
1047
+ci_base_metadata_pairs() {
1048
+ local dockerfile="${CI_BASE_DOCKERFILE:-${DEFAULT_CI_BASE_DOCKERFILE}}"
1049
+ local stages="${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}"
1050
+ local content_files="${CI_BASE_CONTENT_FILES:-${DEFAULT_CI_BASE_CONTENT_FILES}}"
1051
+ local content_files_hash=""
1052
+ local base_image=""
1053
+ local base_image_digest=""
1054
+ local git_branch=""
1055
+ local -a content_paths=()
1056
+ local -a content_args=()
1057
+
1058
+ read -r -a content_paths <<< "${content_files}"
1059
+ if [[ ${#content_paths[@]} -gt 0 ]]; then
1060
+ content_files_hash=$(compute_content_hash "${content_paths[@]}")
1061
+ fi
1062
+ mapfile -t content_args < <(
1063
+ get_content_arg_names "${dockerfile}" "${stages}" "${CI_BASE_CONTENT_ARGS:-}"
1064
+ )
1065
+
1066
+ base_image=$(resolve_dockerfile_arg_value "${dockerfile}" "BASE_IMAGE")
1067
+ if [[ -n "${base_image}" ]]; then
1068
+ base_image_digest=$(resolve_image_digest "${base_image}")
1069
+ fi
1070
+ git_branch="${BUILDKITE_BRANCH:-${VLLM_BRANCH:-}}"
1071
+
1072
+ metadata_pair "vllm.ci_base.metadata_version" "${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}"
1073
+ metadata_pair "vllm.ci_base.content_hash" "${CI_BASE_CONTENT_HASH:-}"
1074
+ metadata_pair "vllm.ci_base.content_files_hash" "${content_files_hash}"
1075
+ metadata_pair "vllm.ci_base.content_files" "${content_files}"
1076
+ metadata_pair "vllm.ci_base.content_args" "$(join_words "${content_args[@]}")"
1077
+ metadata_pair "vllm.ci_base.dockerfile" "${dockerfile}"
1078
+ metadata_pair "vllm.ci_base.dockerfile_stages" "${stages}"
1079
+ metadata_pair "vllm.ci_base.image.primary" "${CI_BASE_IMAGE_TAG:-}"
1080
+ metadata_pair "vllm.ci_base.image.content" "${CI_BASE_IMAGE_TAG_CONTENT_REF:-${CI_BASE_IMAGE_TAG_CONTENT_EXTRA:-}}"
1081
+ metadata_pair "vllm.ci_base.image.commit" "${CI_BASE_IMAGE_TAG_COMMIT_REF:-${CI_BASE_IMAGE_TAG_COMMIT_EXTRA:-}}"
1082
+ metadata_pair "vllm.ci_base.image.stable" "${CI_BASE_IMAGE_TAG_STABLE:-}"
1083
+ metadata_pair "vllm.ci_base.git_commit" "${BUILDKITE_COMMIT:-}"
1084
+ metadata_pair "vllm.ci_base.git_branch" "${git_branch}"
1085
+ metadata_pair "vllm.ci_base.vllm_branch" "${VLLM_BRANCH:-}"
1086
+ metadata_pair "vllm.ci_base.stable_branch" "${CI_BASE_STABLE_BRANCH:-main}"
1087
+
1088
+ metadata_pair "vllm.rocm.base_image" "${base_image}"
1089
+ metadata_pair "vllm.rocm.base_image_digest" "${base_image_digest}"
1090
+ metadata_pair "vllm.rocm.pytorch_rocm_arch" "${PYTORCH_ROCM_ARCH:-}"
1091
+ metadata_pair "vllm.rocm.nic_backend" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIC_BACKEND")"
1092
+ metadata_pair "vllm.rocm.ainic_version" "$(resolve_dockerfile_arg_value "${dockerfile}" "AINIC_VERSION")"
1093
+ metadata_pair "vllm.rocm.ubuntu_codename" "$(resolve_dockerfile_arg_value "${dockerfile}" "UBUNTU_CODENAME")"
1094
+ metadata_pair "vllm.rocm.rixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_REPO")"
1095
+ metadata_pair "vllm.rocm.rixl_commit" "${RIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_BRANCH")}"
1096
+ metadata_pair "vllm.rocm.ucx_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_REPO")"
1097
+ metadata_pair "vllm.rocm.ucx_commit" "${UCX_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_BRANCH")}"
1098
+ metadata_pair "vllm.rocm.rocshmem_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_REPO")"
1099
+ metadata_pair "vllm.rocm.rocshmem_commit" "${ROCSHMEM_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_BRANCH")}"
1100
+ metadata_pair "vllm.rocm.deepep_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_REPO")"
1101
+ metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}"
1102
+ metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")"
1103
+ metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")"
1104
+ metadata_pair "vllm.rocm.rixl_cache_key" "${RIXL_CACHE_KEY:-}"
1105
+ metadata_pair "vllm.rocm.rocshmem_cache_key" "${ROCSHMEM_CACHE_KEY:-}"
1106
+ metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}"
1107
+
1108
+ metadata_pair "vllm.buildkite.build_number" "${BUILDKITE_BUILD_NUMBER:-}"
1109
+ metadata_pair "vllm.buildkite.build_id" "${BUILDKITE_BUILD_ID:-}"
1110
+}
1111
+
1112
+write_ci_base_metadata_annotations() {
1113
+ local metadata="$1"
1114
+ local key=""
1115
+ local value=""
1116
+ local annotation=""
1117
+
1118
+ [[ -n "${metadata}" ]] || return 0
1119
+ while IFS=$'\t' read -r key value; do
1120
+ [[ -n "${key}" && -n "${value}" ]] || continue
1121
+ annotation="manifest:${key}=${value}"
1122
+ printf ' "%s",\n' "$(hcl_escape_string "${annotation}")"
1123
+ done <<< "${metadata}"
1124
+}
1125
+
1126
+write_ci_base_metadata_labels() {
1127
+ local metadata="$1"
1128
+ local key=""
1129
+ local value=""
1130
+
1131
+ [[ -n "${metadata}" ]] || return 0
1132
+ while IFS=$'\t' read -r key value; do
1133
+ [[ -n "${key}" && -n "${value}" ]] || continue
1134
+ printf ' "%s" = "%s"\n' \
1135
+ "$(hcl_escape_string "${key}")" \
1136
+ "$(hcl_escape_string "${value}")"
1137
+ done <<< "${metadata}"
1138
+}
1139
+
1140
write_ci_base_label_override() {
1141
local target_name=""
1142
+ local metadata=""
1143
local -a ci_base_targets=()
1144
1005
- BAKE_FILES=(-f "${VLLM_BAKE_FILE}" -f "${CI_HCL_PATH}")
1006
-
1145
if [[ -z "${CI_BASE_CONTENT_HASH:-}" ]]; then
1146
return 0
1147
fi
@@ -1019,16 +1157,23 @@ write_ci_base_label_override() {
1157
return 0
1158
fi
1159
1160
+ metadata=$(ci_base_metadata_pairs)
1161
+
1162
: > "${CI_BASE_LABEL_OVERRIDE_PATH}"
1163
for target_name in "${ci_base_targets[@]}"; do
1164
cat >> "${CI_BASE_LABEL_OVERRIDE_PATH}" <<EOF
1165
target "${target_name}" {
1166
annotations = [
1167
"manifest:org.opencontainers.image.revision=",
1168
+EOF
1169
+ write_ci_base_metadata_annotations "${metadata}" >> "${CI_BASE_LABEL_OVERRIDE_PATH}"
1170
+ cat >> "${CI_BASE_LABEL_OVERRIDE_PATH}" <<EOF
1171
]
1172
labels = {
1173
"org.opencontainers.image.revision" = ""
1031
- "vllm.ci_base.content_hash" = "${CI_BASE_CONTENT_HASH}"
1174
+EOF
1175
+ write_ci_base_metadata_labels "${metadata}" >> "${CI_BASE_LABEL_OVERRIDE_PATH}"
1176
+ cat >> "${CI_BASE_LABEL_OVERRIDE_PATH}" <<EOF
1177
}
1178
}
1179
@@ -1036,7 +1181,7 @@ EOF
1181
done
1182
1183
BAKE_FILES+=(-f "${CI_BASE_LABEL_OVERRIDE_PATH}")
1039
- echo "Appended ci_base content-hash label override for targets: ${ci_base_targets[*]}"
1184
+ echo "Appended ci_base metadata label override for targets: ${ci_base_targets[*]}"
1185
}
1186
1187
uses_rocm_csrc_cache() {
@@ -1119,6 +1264,18 @@ hcl_escape_string() {
1264
printf '%s' "${value}"
1265
}
1266
1267
+join_words() {
1268
+ local IFS=" "
1269
+ printf '%s' "$*"
1270
+}
1271
+
1272
+metadata_pair() {
1273
+ local key="$1"
1274
+ local value="${2:-}"
1275
+
1276
+ printf '%s\t%s\n' "${key}" "${value}"
1277
+}
1278
+
1279
write_hcl_string_list() {
1280
local indent="$1"
1281
shift
@@ -1541,7 +1698,13 @@ confirm_remote_image_push() {
1698
1699
remote_hash=$(get_remote_image_label_with_retry "${image_ref}" "vllm.ci_base.content_hash")
1700
if [[ -n "${remote_hash}" && "${remote_hash}" == "${CI_BASE_CONTENT_HASH}" ]]; then
1544
- return 0
1701
+ if remote_ci_base_metadata_is_current_with_retry "${image_ref}"; then
1702
+ return 0
1703
+ fi
1704
+
1705
+ echo "Remote image exists with the expected ci_base content hash but stale metadata."
1706
+ echo " expected metadata version: ${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}"
1707
+ return 1
1708
fi
1709
1710
echo "Remote image exists but does not have the expected ci_base content hash."
@@ -1724,15 +1887,16 @@ main() {
1887
print_header
1888
validate_inputs
1889
load_ci_hcl
1890
+ init_bake_files
1891
compute_ci_base_hash_if_needed
1892
configure_ci_base_image_refs
1893
maybe_skip_existing_image
1894
setup_builder
1895
prepare_git_cache_metadata
1732
- write_ci_base_label_override
1896
extract_dependency_pins
1897
write_rocm_build_arg_override
1898
compute_dependency_cache_keys
1899
+ write_ci_base_label_override
1900
compute_rocm_csrc_content_hash_if_needed
1901
write_rocm_cache_override
1902
resolve_ci_base_dependency_targets
docker/ci-rocm.hcl
+15
-6
@@ -302,14 +302,22 @@ group "test-rocm-ci-with-wheel" {
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.
305
+# to the primary tag for this build. Builds always publish a content-scoped tag
306
+# when the ci_base content hash is available. Builds with BUILDKITE_COMMIT also
307
+# publish a commit-scoped tag, either as the primary tag or an additional alias.
308
+# NIGHTLY=1 builds on the stable branch can additionally set
309
+# CI_BASE_IMAGE_TAG_STABLE to refresh rocm/vllm-dev:ci_base.
310
variable "CI_BASE_IMAGE_TAG" {
311
default = "rocm/vllm-dev:ci_base"
312
}
313
312
-variable "CI_BASE_IMAGE_TAG_CONTENT" {
314
+# Supplemental tags only. ci-bake-rocm.sh leaves these empty when the same ref
315
+# is already the primary CI_BASE_IMAGE_TAG.
316
+variable "CI_BASE_IMAGE_TAG_COMMIT_EXTRA" {
317
+ default = ""
318
+}
319
+
320
+variable "CI_BASE_IMAGE_TAG_CONTENT_EXTRA" {
321
default = ""
322
}
323
@@ -357,7 +365,8 @@ target "ci-base-rocm-ci" {
365
cache-from = concat(
366
compact([
367
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}" : "",
368
+ CI_BASE_IMAGE_TAG_COMMIT_EXTRA != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_COMMIT_EXTRA}" : "",
369
+ CI_BASE_IMAGE_TAG_CONTENT_EXTRA != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_CONTENT_EXTRA}" : "",
370
CI_BASE_IMAGE_TAG_STABLE != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_STABLE}" : "",
371
]),
372
# Import upstream dependency caches so RIXL/ROCShmem/DeepEP stages
@@ -365,7 +374,7 @@ target "ci-base-rocm-ci" {
374
get_cache_from_rocm_deps(),
375
)
376
cache-to = ["type=inline"]
368
- tags = compact([CI_BASE_IMAGE_TAG, CI_BASE_IMAGE_TAG_CONTENT, CI_BASE_IMAGE_TAG_STABLE])
377
+ tags = compact([CI_BASE_IMAGE_TAG, CI_BASE_IMAGE_TAG_COMMIT_EXTRA, CI_BASE_IMAGE_TAG_CONTENT_EXTRA, CI_BASE_IMAGE_TAG_STABLE])
378
output = ["type=registry"]
379
}
380