Merge branch 'main' into async-launch-pipelining

Author: kevmo314

.github/workflows/gpu-integration-gcloud.yml

@@ -16,16 +16,17 @@ permissions:
 
 jobs:
   gpu-integration:
-    name: CUDA samples and PyTorch on GCE T4
+    name: CUDA samples and PyTorch on GCE L4
     if: github.event.pull_request.head.repo.full_name == github.repository
     runs-on: ubuntu-latest
     timeout-minutes: 300
 
     env:
       GCP_PROJECT_ID: ${{ vars.GCP_PROJECT_ID }}
       GCP_ZONE: us-central1-a
+      GCP_FALLBACK_ZONES: us-central1-c us-central1-b
       GCP_NETWORK: default
-      MACHINE_TYPE: n1-standard-4
+      MACHINE_TYPE: g2-standard-4
       CUDA_VERSION: 12.9.1
       UBUNTU_VERSION: "24.04"
       PYTORCH_INDEX_URL: https://download.pytorch.org/whl/cu128
@@ -76,7 +77,8 @@ jobs:
           test -n "$GCP_PROJECT_ID"
           test -n "${{ vars.GCP_WORKLOAD_IDENTITY_PROVIDER }}"
           test -n "${{ vars.GCP_SERVICE_ACCOUNT }}"
-          test "$MACHINE_TYPE" = n1-standard-4
+          test "$GCP_FALLBACK_ZONES" = "us-central1-c us-central1-b"
+          test "$MACHINE_TYPE" = g2-standard-4
           test "$VM_IMAGE_PROJECT" = ml-images
           test "$VM_IMAGE_FAMILY" = common-cu129-ubuntu-2404-nvidia-580
           test "$VM_MAX_RUN_DURATION" = 260m
@@ -171,9 +173,9 @@ jobs:
           ln -sf libcuda.so.1 "$PWD/build/libcuda.so"
           ln -sf libnvidia-ml.so.1 "$PWD/build/libnvidia-ml.so"
 
-          # CUDA_SAMPLES_ARCH=75: only compile device code for the T4 (sm_75)
+          # CUDA_SAMPLES_ARCH=89: only compile device code for the L4 (sm_89)
           # instead of the 15 architectures the samples build by default.
-          BUILD_ONLY=1 BUILD_SAMPLES=1 SAMPLE_SUITE=extended CUDA_SAMPLES_ARCH=75 \
+          BUILD_ONLY=1 BUILD_SAMPLES=1 SAMPLE_SUITE=extended CUDA_SAMPLES_ARCH=89 \
             CUDA_HOME="$CUDA_HOME" CUDA_LIB_DIR="$CUDA_LIB_DIR" \
             "$PWD/test/run_cuda_samples.sh"
 
@@ -213,7 +215,7 @@ jobs:
             tcp:1-65535,udp:1-65535,icmp,esp,ah,sctp \
             0.0.0.0/0
 
-      - name: Create T4 server VM
+      - name: Create L4 server VM
         run: |
           set -euo pipefail
           startup_script="$RUNNER_TEMP/gce-startup.sh"
@@ -226,38 +228,60 @@ jobs:
           rm -rf /var/lib/apt/lists/*
           EOF
 
-          create_args=(
-            "$VM_NAME"
-            "--project=$GCP_PROJECT_ID"
-            "--zone=$GCP_ZONE"
-            "--machine-type=$MACHINE_TYPE"
-            "--network=$GCP_NETWORK"
-            "--tags=$VM_TAG"
-            "--image-family=$VM_IMAGE_FAMILY"
-            "--image-project=$VM_IMAGE_PROJECT"
-            "--boot-disk-size=100GB"
-            "--boot-disk-type=pd-balanced"
-            "--accelerator=type=nvidia-tesla-t4,count=1"
-            "--maintenance-policy=TERMINATE"
-            "--max-run-duration=$VM_MAX_RUN_DURATION"
-            "--instance-termination-action=DELETE"
-            "--no-service-account"
-            "--no-scopes"
-            "--metadata=block-project-ssh-keys=TRUE,enable-oslogin=FALSE,ssh-keys=gha:$(cat "$SSH_DIR/id_ed25519.pub")"
-            "--metadata-from-file=startup-script=$startup_script"
-            "--shielded-vtpm"
-            "--shielded-integrity-monitoring"
-            "--no-shielded-secure-boot"
-          )
+          # shellcheck disable=SC2206
+          fallback_zones=($GCP_FALLBACK_ZONES)
+          candidate_zones=("$GCP_ZONE" "${fallback_zones[@]}")
+
+          vm_ip=""
+          selected_zone=""
+          for zone in "${candidate_zones[@]}"; do
+            echo "Creating $MACHINE_TYPE in $zone"
+            create_args=(
+              "$VM_NAME"
+              "--project=$GCP_PROJECT_ID"
+              "--zone=$zone"
+              "--machine-type=$MACHINE_TYPE"
+              "--network=$GCP_NETWORK"
+              "--tags=$VM_TAG"
+              "--image-family=$VM_IMAGE_FAMILY"
+              "--image-project=$VM_IMAGE_PROJECT"
+              "--boot-disk-size=100GB"
+              "--boot-disk-type=pd-balanced"
+              "--maintenance-policy=TERMINATE"
+              "--max-run-duration=$VM_MAX_RUN_DURATION"
+              "--instance-termination-action=DELETE"
+              "--no-service-account"
+              "--no-scopes"
+              "--metadata=block-project-ssh-keys=TRUE,enable-oslogin=FALSE,ssh-keys=gha:$(cat "$SSH_DIR/id_ed25519.pub")"
+              "--metadata-from-file=startup-script=$startup_script"
+              "--shielded-vtpm"
+              "--shielded-integrity-monitoring"
+              "--no-shielded-secure-boot"
+            )
+
+            if [[ "$USE_SPOT" == "true" ]]; then
+              create_args+=("--provisioning-model=SPOT")
+            fi
 
-          if [[ "$USE_SPOT" == "true" ]]; then
-            create_args+=("--provisioning-model=SPOT")
-          fi
+            if vm_ip="$(gcloud compute instances create "${create_args[@]}" \
+              --format='value(networkInterfaces[0].accessConfigs[0].natIP)')"; then
+              selected_zone="$zone"
+              break
+            fi
+
+            echo "Could not create $MACHINE_TYPE in $zone; trying next candidate zone" >&2
+            gcloud compute instances delete "$VM_NAME" \
+              --project="$GCP_PROJECT_ID" \
+              --zone="$zone" \
+              --quiet >/dev/null 2>&1 || true
+          done
 
-          vm_ip="$(gcloud compute instances create "${create_args[@]}" \
-            --format='value(networkInterfaces[0].accessConfigs[0].natIP)')"
           test -n "$vm_ip"
-          echo "VM_IP=$vm_ip" >> "$GITHUB_ENV"
+          test -n "$selected_zone"
+          {
+            echo "VM_IP=$vm_ip"
+            echo "GCP_ZONE=$selected_zone"
+          } >> "$GITHUB_ENV"
 
       - name: Wait for SSH and NVIDIA driver
         run: |
@@ -290,7 +314,7 @@ jobs:
             done
           '
 
-      - name: Run GitHub runner client against T4 server
+      - name: Run GitHub runner client against L4 server
         run: |
           set -euo pipefail
 
@@ -371,10 +395,16 @@ jobs:
         run: |
           set -euo pipefail
 
-          gcloud compute instances delete "$VM_NAME" \
-            --project="$GCP_PROJECT_ID" \
-            --zone="$GCP_ZONE" \
-            --quiet || true
+          # shellcheck disable=SC2206
+          fallback_zones=($GCP_FALLBACK_ZONES)
+          candidate_zones=("$GCP_ZONE" "${fallback_zones[@]}")
+
+          for zone in "${candidate_zones[@]}"; do
+            gcloud compute instances delete "$VM_NAME" \
+              --project="$GCP_PROJECT_ID" \
+              --zone="$zone" \
+              --quiet || true
+          done
 
           for rule in "$FIREWALL_ALLOW_RULE" "$FIREWALL_DENY_RULE"; do
             gcloud compute firewall-rules delete "$rule" \
@@ -383,12 +413,14 @@ jobs:
           done
 
           leftovers=0
-          if gcloud compute instances describe "$VM_NAME" \
-            --project="$GCP_PROJECT_ID" \
-            --zone="$GCP_ZONE" >/dev/null 2>&1; then
-            echo "leftover VM still exists: $VM_NAME" >&2
-            leftovers=1
-          fi
+          for zone in "${candidate_zones[@]}"; do
+            if gcloud compute instances describe "$VM_NAME" \
+              --project="$GCP_PROJECT_ID" \
+              --zone="$zone" >/dev/null 2>&1; then
+              echo "leftover VM still exists in $zone: $VM_NAME" >&2
+              leftovers=1
+            fi
+          done
 
           for rule in "$FIREWALL_ALLOW_RULE" "$FIREWALL_DENY_RULE"; do
             if gcloud compute firewall-rules describe "$rule" \