diff --git a/.github/workflows/build-gpu-image.yml b/.github/workflows/build-gpu-image.yml index 7d952c6a7..8044002a6 100644 --- a/.github/workflows/build-gpu-image.yml +++ b/.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 }} diff --git a/scripts/build-gpu-image.sh b/scripts/build-gpu-image.sh index 4e8d78949..23fbad0c1 100755 --- a/scripts/build-gpu-image.sh +++ b/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 diff --git a/tests/unit/test_gpu_image_build_context.py b/tests/unit/test_gpu_image_build_context.py index a3279ada6..da94b3ccd 100644 --- a/tests/unit/test_gpu_image_build_context.py +++ b/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