Clean up cancelled GPU image prewarm pods

.github/workflows/build-gpu-image.yml

@@ -72,6 +72,7 @@ jobs:
     env:
       SKY_INFRA: k8s/cks-wb3
       BUILD_CLUSTER_NAME: art-gpu-build-${{ github.run_id }}-${{ github.run_attempt }}
+      PREWARM_RUN_UID: ${{ github.run_id }}-${{ github.run_attempt }}
       PREWARM_INFRAS: ${{ inputs.prewarm_infras || 'k8s/cks-wb3 k8s/ext-collab2 k8s/serverless-training-central' }}
 
     steps:
@@ -219,11 +220,23 @@ jobs:
 
           bash scripts/build-gpu-image.sh "${args[@]}"
 
-      - name: Remove temporary image builder
+      - name: Remove temporary image builder and prewarm pods
         if: ${{ always() }}
         run: |
+          prewarm_selector="art.openpipe/prewarm-name=art-gpu-image-prewarm,art.openpipe/prewarm-run=${PREWARM_RUN_UID}"
           kubectl --context cks-wb3 delete pod -n default "${BUILD_CLUSTER_NAME}" \
-            --ignore-not-found --wait=false
+            --ignore-not-found --wait=false --request-timeout=30s || true
+          while IFS= read -r infra; do
+            case "${infra}" in
+              k8s/*) context="${infra#k8s/}" ;;
+              kubernetes/*) context="${infra#kubernetes/}" ;;
+              *) continue ;;
+            esac
+            kubectl --context "${context}" delete pod -n default \
+              -l "${prewarm_selector}" \
+              --ignore-not-found --wait=false --request-timeout=30s || true
+          done < <(printf '%s\n' "${PREWARM_INFRAS}" \
+            | awk 'BEGIN { RS = "[[:space:],]+" } NF { print }')
 
       - name: Smoke launch prewarmed image
         if: ${{ github.event_name != 'workflow_dispatch' || inputs.smoke_launch }}

scripts/build-gpu-image.sh

@@ -43,6 +43,7 @@ if [[ -n "${PREWARM_INFRAS:-}" ]]; then
 fi
 prewarm_namespace="${PREWARM_NAMESPACE:-default}"
 prewarm_name="${PREWARM_NAME:-art-gpu-image-prewarm}"
+prewarm_run_uid="${PREWARM_RUN_UID:-}"
 prewarm_image_pull_secret="${PREWARM_IMAGE_PULL_SECRET:-art-gpu-registry-auth}"
 prewarm_node_selector="${PREWARM_NODE_SELECTOR:-node.coreweave.cloud/class=gpu}"
 prewarm_hypervisor_label="node.coreweave.cloud/hypervisor"
@@ -162,6 +163,14 @@ prewarm_node_selector_value="${prewarm_node_selector#*=}"
 art_sha="$(git -C "${repo_root}" rev-parse HEAD)"
 art_short_sha="$(git -C "${repo_root}" rev-parse --short=12 HEAD)"
 timestamp="$(date +%m%d-%H%M%S)"
+if [[ -z "${prewarm_run_uid}" ]]; then
+  prewarm_run_uid="${timestamp}-${art_short_sha}-$$"
+fi
+if (( ${#prewarm_run_uid} > 63 )) ||
+  [[ ! "${prewarm_run_uid}" =~ ^[a-z0-9]([-a-z0-9_.]*[a-z0-9])?$ ]]; then
+  echo "PREWARM_RUN_UID must be a valid Kubernetes label value, got: ${prewarm_run_uid}" >&2
+  exit 1
+fi
 
 if [[ -z "${cluster_name}" ]]; then
   cluster_name="art-gpu-build-${timestamp}"
@@ -287,7 +296,20 @@ registry_auth_json_path="$(mktemp "${TMPDIR:-/tmp}/art-gpu-auth.XXXXXX")"
 build_command_path="$(mktemp "${TMPDIR:-/tmp}/art-gpu-build-command.XXXXXX")"
 build_log_snapshot_path="$(mktemp "${TMPDIR:-/tmp}/art-gpu-build-log.XXXXXX")"
 build_log_offset_path="$(mktemp "${TMPDIR:-/tmp}/art-gpu-build-log-offset.XXXXXX")"
+cleanup_prewarm_pods() {
+  local context
+  local selector
+
+  selector="art.openpipe/prewarm-name=${prewarm_name},art.openpipe/prewarm-run=${prewarm_run_uid}"
+  for context in "${prewarm_contexts[@]}"; do
+    kubectl --context "${context}" delete pod -n "${prewarm_namespace}" \
+      -l "${selector}" \
+      --ignore-not-found --wait=false --request-timeout=30s \
+      >/dev/null 2>&1 || true
+  done
+}
 cleanup() {
+  cleanup_prewarm_pods
   rm -rf "${context_dir}"
   rm -f "${buildkit_manifest_path}" "${registry_auth_json_path}" \
     "${build_command_path}" "${build_log_snapshot_path}" "${build_log_offset_path}"
@@ -808,6 +830,7 @@ metadata:
   labels:
     app: ${prewarm_name}-oneshot
     art.openpipe/prewarm-name: ${prewarm_name}
+    art.openpipe/prewarm-run: "${prewarm_run_uid}"
     art.openpipe/prewarm-token: "${timestamp}-${art_short_sha}"
 spec:
   restartPolicy: Never
@@ -874,6 +897,7 @@ metadata:
   labels:
     app: ${prewarm_name}-oneshot
     art.openpipe/prewarm-name: ${prewarm_name}
+    art.openpipe/prewarm-run: "${prewarm_run_uid}"
     art.openpipe/prewarm-token: "${timestamp}-${art_short_sha}"
 spec:
   restartPolicy: Never

tests/unit/test_gpu_image_build_context.py

@@ -13,3 +13,22 @@ def test_gpu_image_build_context_copies_every_local_docker_source() -> None:
             continue
         for source in shlex.split(line)[1:-1]:
             assert f"${{repo_root}}/{source}" in build_script
+
+
+def test_gpu_image_build_cleans_only_its_oneshot_prewarm_pods() -> None:
+    build_script = (ROOT / "scripts/build-gpu-image.sh").read_text()
+    workflow = (ROOT / ".github/workflows/build-gpu-image.yml").read_text()
+
+    run_label = 'art.openpipe/prewarm-run: "${prewarm_run_uid}"'
+    run_selector = (
+        "art.openpipe/prewarm-name=${prewarm_name},"
+        "art.openpipe/prewarm-run=${prewarm_run_uid}"
+    )
+    assert build_script.count(run_label) == 2
+    assert build_script.count(run_selector) == 1
+    assert "PREWARM_RUN_UID must be a valid Kubernetes label value" in build_script
+    assert "PREWARM_RUN_UID: ${{ github.run_id }}-${{ github.run_attempt }}" in workflow
+    assert (
+        "art.openpipe/prewarm-name=art-gpu-image-prewarm,"
+        "art.openpipe/prewarm-run=${PREWARM_RUN_UID}"
+    ) in workflow