fix(macrobenchmarks): retry GKE control plane in other zones on stockout

cloudbuild/macrobenchmarks/README.md

@@ -116,7 +116,7 @@ Before creating the triggers, set up the following in your GCP project.
 | `init-variables` | Validates all substitutions and the operator-supplied buckets, derives the region and run identifiers. **Fails fast before anything is provisioned or billed.** |
 | `cleanup-leaked-resources` | Best-effort reap of orphaned infra from prior failed builds with the same `_INFRA_PREFIX`. |
 | `create-buckets` | Creates the per-run checkpoint bucket and per-run dataset bucket (copying `_DATASET_PATH` in), and ensures the shared results bucket exists. |
-| `create-cluster` | Creates the dedicated VPC/subnet, the GKE cluster + a `_MACHINE_TYPE` node pool (with TIER_1 networking when enabled), and installs the JobSet controller. |
+| `create-cluster` | Creates the dedicated VPC/subnet, the GKE cluster + a `_MACHINE_TYPE` node pool in `_ZONE` (with TIER_1 networking when enabled), and installs the JobSet controller. If the control plane can't be created in `_ZONE`, it retries in the region's other zones. |
 | `seed-checkpoint` | When `_SEED_CHECKPOINT=true` and no external checkpoint is given, generates a per-run checkpoint so the measured run exercises the **restore** path. |
 | `run-workload` | `helm install`s the chart and polls the JobSet to completion, recording start/end timestamps for the scrape. |
 | `scrape-metrics` | Parses Cloud Logging into raw CSVs, aggregates them into one summary row (retrying for ingestion lag), folds in system metrics, and uploads the summary to the results bucket. |
@@ -130,7 +130,7 @@ Before creating the triggers, set up the following in your GCP project.
 | Substitution | Description |
 | :----------- | :---------- |
 | `_INFRA_PREFIX` | Prefix for all created resources (cluster, network, buckets, BigQuery dataset). Also selects the shared results bucket `gs://<_INFRA_PREFIX>-macrobench-results`. |
-| `_ZONE` | GCP zone for the cluster and zonal buckets (e.g. `us-central1-a`). The region is derived from it. |
+| `_ZONE` | GCP zone for the workload node pool and zonal buckets (e.g. `us-central1-a`), and the preferred zone for the cluster's control plane (which may fall back to another zone in the region). The region is derived from it. |
 | `_GKE_SERVICE_ACCOUNT` | Email of the node service account (see prerequisites). |
 | `_DATASET_PATH` | `gs://` directory of `*.parquet` shards (with a `text` column) to train on. Must match the run's region and bucket type. |
 | `_REQUIREMENTS` | Whitespace-separated immutable pip specs/URLs installed last. Every workload requires the gcsfs build under test; the Ray workload needs nothing else. |

cloudbuild/macrobenchmarks/metrics/tests/test_create_cluster.py

@@ -0,0 +1,177 @@
+"""Tests for the control-plane zone fallback in create_cluster.sh and the
+matching lookup in cleanup_cluster.sh.
+
+The scripts run for real under bash, with stub ``gcloud``/``kubectl``/``sleep``
+executables on PATH that log every call and fail on demand.
+"""
+
+import os
+import shutil
+import subprocess
+from pathlib import Path
+
+import pytest
+
+SCRIPTS = Path(__file__).resolve().parents[2] / "scripts"
+
+pytestmark = pytest.mark.skipif(shutil.which("bash") is None, reason="needs bash")
+
+GCLOUD_STUB = r"""#!/usr/bin/env bash
+echo "gcloud $*" >> "$STUB_LOG"
+case "$*" in
+  "compute zones list"*)
+    [ -n "$STUB_ZONES_LIST_FAIL" ] && exit 1
+    printf '%s\n' $STUB_ZONES ;;
+  "container clusters create"*)
+    for z in $STUB_FAIL_ZONES; do
+      [[ "$*" == *"--zone=$z "* ]] && exit 1
+    done ;;
+  "container clusters list"*)
+    [ -n "$STUB_LIST_FAIL" ] && exit 1
+    [ -n "$STUB_CLUSTER_LOCATIONS" ] && printf '%s\n' $STUB_CLUSTER_LOCATIONS ;;
+esac
+exit 0
+"""
+
+
+@pytest.fixture
+def env(tmp_path):
+    bin_dir = tmp_path / "bin"
+    bin_dir.mkdir()
+    for name, body in {
+        "gcloud": GCLOUD_STUB,
+        "kubectl": '#!/usr/bin/env bash\necho "kubectl $*" >> "$STUB_LOG"\n',
+        "sleep": "#!/usr/bin/env bash\n",
+    }.items():
+        stub = bin_dir / name
+        stub.write_text(body)
+        stub.chmod(0o755)
+
+    build_vars = tmp_path / "build_vars.env"
+    build_vars.write_text(
+        "export CLUSTER_NAME=test-cluster\n"
+        "export NETWORK_NAME=test-net\n"
+        "export SUBNET_NAME=test-subnet\n"
+        "export REGION=region1\n"
+    )
+    return {
+        **os.environ,
+        "PATH": f"{bin_dir}:{os.environ['PATH']}",
+        "STUB_LOG": str(tmp_path / "calls.log"),
+        "FAILED_FILE": str(tmp_path / "FAILED"),
+        "BUILD_VARS_FILE": str(build_vars),
+        "PROJECT_ID": "test-project",
+        "_ZONE": "region1-a",
+        "_MACHINE_TYPE": "test-machine",
+        "_NODES": "1",
+        "_GKE_SERVICE_ACCOUNT": "sa@example.com",
+        "_JOBSET_VERSION": "v0.0.0",
+        "STUB_ZONES": "region1-a region1-b region1-c",
+    }
+
+
+def _run(script, env, **overrides):
+    env = {**env, **overrides}
+    proc = subprocess.run(
+        ["bash", str(SCRIPTS / script)], env=env, capture_output=True, text=True
+    )
+    log = Path(env["STUB_LOG"])
+    calls = log.read_text().splitlines() if log.exists() else []
+    return proc, calls
+
+
+def _calls(calls, prefix):
+    return [c for c in calls if c.startswith(prefix)]
+
+
+def _failures(env):
+    failed = Path(env["FAILED_FILE"])
+    return failed.read_text().split() if failed.exists() else []
+
+
+def _cluster_location(env):
+    for line in Path(env["BUILD_VARS_FILE"]).read_text().splitlines():
+        if line.startswith("export CLUSTER_LOCATION="):
+            return line.split("=", 1)[1]
+    return None
+
+
+def test_create_succeeds_in_preferred_zone(env):
+    proc, calls = _run("create_cluster.sh", env)
+
+    assert proc.returncode == 0, proc.stderr
+    creates = _calls(calls, "gcloud container clusters create")
+    assert len(creates) == 1 and "--zone=region1-a " in creates[0]
+    assert not _calls(calls, "gcloud container clusters delete")
+    assert _cluster_location(env) == "region1-a"
+    assert _failures(env) == []
+
+
+def test_create_falls_back_and_keeps_node_pool_in_preferred_zone(env):
+    proc, calls = _run("create_cluster.sh", env, STUB_FAIL_ZONES="region1-a")
+
+    assert proc.returncode == 0, proc.stderr
+    creates = _calls(calls, "gcloud container clusters create")
+    assert ["--zone=region1-a " in c for c in creates] == [True, False]
+    assert "--zone=region1-b " in creates[1]
+    deletes = _calls(calls, "gcloud container clusters delete")
+    assert len(deletes) == 1 and "--zone=region1-a" in deletes[0]
+    (pool,) = _calls(calls, "gcloud container node-pools create")
+    assert "--zone=region1-b" in pool and "--node-locations=region1-a" in pool
+    (creds,) = _calls(calls, "gcloud container clusters get-credentials")
+    assert "--zone=region1-b" in creds
+    assert _cluster_location(env) == "region1-b"
+    assert _failures(env) == []
+
+
+def test_create_records_failure_once_when_all_zones_fail(env):
+    proc, calls = _run(
+        "create_cluster.sh", env, STUB_FAIL_ZONES="region1-a region1-b region1-c"
+    )
+
+    assert proc.returncode == 1
+    assert len(_calls(calls, "gcloud container clusters create")) == 3
+    assert not _calls(calls, "gcloud container node-pools create")
+    assert _failures(env) == ["create-cluster"]
+    assert _cluster_location(env) is None
+    assert "could not create test-cluster" in proc.stderr
+
+
+def test_create_tries_only_preferred_zone_when_zone_list_fails(env):
+    proc, calls = _run(
+        "create_cluster.sh", env, STUB_ZONES_LIST_FAIL="1", STUB_FAIL_ZONES="region1-a"
+    )
+
+    assert proc.returncode == 1
+    creates = _calls(calls, "gcloud container clusters create")
+    assert len(creates) == 1 and "--zone=region1-a " in creates[0]
+    assert _failures(env) == ["create-cluster"]
+
+
+def test_cleanup_deletes_cluster_in_every_location_found(env):
+    proc, calls = _run(
+        "cleanup_cluster.sh", env, STUB_CLUSTER_LOCATIONS="region1-a region1-b"
+    )
+
+    assert proc.returncode == 0, proc.stderr
+    deletes = _calls(calls, "gcloud container clusters delete")
+    assert len(deletes) == 2
+    assert "--zone=region1-a" in deletes[0] and "--zone=region1-b" in deletes[1]
+
+
+def test_cleanup_skips_delete_when_no_cluster_found(env):
+    proc, calls = _run("cleanup_cluster.sh", env)
+
+    assert proc.returncode == 0, proc.stderr
+    assert not _calls(calls, "gcloud container clusters delete")
+    assert "No cluster test-cluster found" in proc.stdout
+
+
+def test_cleanup_falls_back_to_recorded_location_when_lookup_fails(env):
+    proc, calls = _run(
+        "cleanup_cluster.sh", env, STUB_LIST_FAIL="1", CLUSTER_LOCATION="region1-c"
+    )
+
+    assert proc.returncode == 0, proc.stderr
+    (delete,) = _calls(calls, "gcloud container clusters delete")
+    assert "--zone=region1-c" in delete

cloudbuild/macrobenchmarks/scripts/cleanup_cluster.sh

@@ -7,19 +7,30 @@ fi
 source "$(dirname "$0")/lib.sh"
 source "${BUILD_VARS_FILE}"
 echo "--- Deleting GKE cluster: ${CLUSTER_NAME} ---"
+# create-cluster may have tried several zones, and a failed attempt's delete is
+# best-effort, so look the cluster up by name and delete it wherever it exists.
+# If the lookup itself fails, fall back to the recorded location. An empty
+# lookup means there is nothing to delete.
+if LOCATIONS=$(gcloud container clusters list --project="${PROJECT_ID}" \
+    --filter="name=${CLUSTER_NAME}" --format="value(location)"); then
+  [ -z "$LOCATIONS" ] && echo "No cluster ${CLUSTER_NAME} found; skipping."
+else
+  LOCATIONS="${CLUSTER_LOCATION:-${_ZONE}}"
+fi
 # Best-effort: just retry the delete until it succeeds. A delete can transiently
 # fail (e.g. FAILED_PRECONDITION while a cluster operation is still settling);
 # retrying after a pause handles that without inspecting cluster/operation state.
-# A not-found cluster is already the desired end state. Retry generously -- this
-# is the last cleanup step and blocks nothing, so it's worth being patient to
-# release the cluster's resources rather than leaking them.
-for i in {1..20}; do
-  if gcloud container clusters delete "$CLUSTER_NAME" --zone="${_ZONE}" --project="${PROJECT_ID}" --quiet; then
-    echo "Cluster ${CLUSTER_NAME} deleted successfully."
-    break
-  fi
-  echo "Attempt $i to delete cluster ${CLUSTER_NAME} failed. Retrying in 30 seconds..."
-  sleep 30
+# Retry generously -- this is the last cleanup step and blocks nothing, so it's
+# worth being patient to release the cluster's resources rather than leaking them.
+for loc in $LOCATIONS; do
+  for i in {1..20}; do
+    if gcloud container clusters delete "$CLUSTER_NAME" --zone="${loc}" --project="${PROJECT_ID}" --quiet; then
+      echo "Cluster ${CLUSTER_NAME} in ${loc} deleted successfully."
+      break
+    fi
+    echo "Attempt $i to delete cluster ${CLUSTER_NAME} in ${loc} failed. Retrying in 30 seconds..."
+    sleep 30
+  done
 done
 
 echo "--- Deleting dedicated subnetwork: ${SUBNET_NAME} ---"

cloudbuild/macrobenchmarks/scripts/cleanup_helm.sh

@@ -7,7 +7,7 @@ fi
 source "$(dirname "$0")/lib.sh"
 source "${BUILD_VARS_FILE}"
 curl -fsSL https://raw.githubusercontent.com/helm/helm/main/scripts/get-helm-3 | bash || true
-gcloud container clusters get-credentials "$CLUSTER_NAME" --zone="${_ZONE}" --project="${PROJECT_ID}" || true
+gcloud container clusters get-credentials "$CLUSTER_NAME" --zone="${CLUSTER_LOCATION:-${_ZONE}}" --project="${PROJECT_ID}" || true
 helm uninstall "$RUN_ID" || true
 # Best-effort: the seed release is normally uninstalled by the seed-checkpoint
 # step; clean up here in case that step failed mid-way. Whole-cluster deletion

cloudbuild/macrobenchmarks/scripts/create_cluster.sh

@@ -21,25 +21,49 @@ gcloud compute networks subnets create "${SUBNET_NAME}" \
   --range="10.0.0.0/20" \
   --enable-private-ip-google-access --quiet
 
+# The GKE control plane of a zonal cluster runs in the cluster's zone and can hit
+# GCE_STOCKOUT there independently of the node pools. If that happens, retry the
+# control plane in the region's other zones; the ${_MACHINE_TYPE} pool stays in
+# ${_ZONE}.
+ZONES="${_ZONE} $(gcloud compute zones list --project="${PROJECT_ID}" --filter="region:${REGION}" --format="value(name)" | grep -vx "${_ZONE}" || true)"
+CLUSTER_LOCATION=""
+for zone in $ZONES; do
+  echo "--- Creating GKE cluster ${CLUSTER_NAME} in ${zone} ---"
+  if gcloud container clusters create "$CLUSTER_NAME" \
+    --project="${PROJECT_ID}" --zone="${zone}" \
+    --machine-type="e2-standard-4" --num-nodes="1" \
+    --service-account="${_GKE_SERVICE_ACCOUNT}" \
+    --scopes="https://www.googleapis.com/auth/cloud-platform" \
+    --private-ipv6-google-access-type=outbound-only \
+    --network="${NETWORK_NAME}" --subnetwork="${SUBNET_NAME}" \
+    --enable-autoupgrade --quiet; then
+    CLUSTER_LOCATION="$zone"
+    break
+  fi
+  echo "Cluster creation in ${zone} failed; deleting the failed cluster and trying the next zone."
+  gcloud container clusters delete "$CLUSTER_NAME" --zone="${zone}" --project="${PROJECT_ID}" --quiet || true
+done
+if [ -z "$CLUSTER_LOCATION" ]; then
+  echo "ERROR: could not create ${CLUSTER_NAME} in any of: ${ZONES}" >&2
+  record_failure create-cluster
+  exit 1
+fi
+echo "Cluster ${CLUSTER_NAME} created in ${CLUSTER_LOCATION}."
+echo "export CLUSTER_LOCATION=${CLUSTER_LOCATION}" >> "${BUILD_VARS_FILE}"
+
 # `gcloud container clusters create` has no --node-pool flag and always names
 # its initial pool "default-pool", which would not match the workload's
 # nodeSelector (cloud.google.com/gke-nodepool=c4-standard-192). Create a small
 # system pool for the cluster, then a dedicated ${_MACHINE_TYPE} pool whose name
 # the GKE node label (and the chart's nodeSelector) keys off of. --enable-gvnic
 # puts the nodes on gVNIC, which C4 requires and which TIER_1 egress (the direct
-# high-bandwidth path) is gated on.
-gcloud container clusters create "$CLUSTER_NAME" \
-  --project="${PROJECT_ID}" --zone="${_ZONE}" \
-  --machine-type="e2-standard-4" --num-nodes="1" \
-  --service-account="${_GKE_SERVICE_ACCOUNT}" \
-  --scopes="https://www.googleapis.com/auth/cloud-platform" \
-  --private-ipv6-google-access-type=outbound-only \
-  --network="${NETWORK_NAME}" --subnetwork="${SUBNET_NAME}" \
-  --enable-autoupgrade --quiet
+# high-bandwidth path) is gated on. --zone selects the cluster (wherever its
+# control plane landed); --node-locations keeps the workers in ${_ZONE}.
 NODE_POOL_ARGS=(
   --cluster="$CLUSTER_NAME"
   --project="${PROJECT_ID}"
-  --zone="${_ZONE}"
+  --zone="${CLUSTER_LOCATION}"
+  --node-locations="${_ZONE}"
   --machine-type="${_MACHINE_TYPE}"
   --num-nodes="${_NODES}"
   --disk-size="200"
@@ -54,7 +78,7 @@ if [ "${_ENABLE_TIER1_NETWORKING:-true}" = "true" ]; then
   NODE_POOL_ARGS+=(--network-performance-configs="total-egress-bandwidth-tier=TIER_1")
 fi
 gcloud container node-pools create "${_MACHINE_TYPE}" "${NODE_POOL_ARGS[@]}"
-gcloud container clusters get-credentials "$CLUSTER_NAME" --zone="${_ZONE}" --project="${PROJECT_ID}"
+gcloud container clusters get-credentials "$CLUSTER_NAME" --zone="${CLUSTER_LOCATION}" --project="${PROJECT_ID}"
 kubectl apply --server-side -f "https://github.com/kubernetes-sigs/jobset/releases/download/${_JOBSET_VERSION}/manifests.yaml"
 kubectl rollout status deployment/jobset-controller-manager -n jobset-system --timeout=300s
 

cloudbuild/macrobenchmarks/scripts/run_workload.sh

@@ -8,7 +8,7 @@ trap 'record_failure run-workload' ERR
 skip_if_failed
 source "${BUILD_VARS_FILE}"
 curl -fsSL https://raw.githubusercontent.com/helm/helm/main/scripts/get-helm-3 | bash
-gcloud container clusters get-credentials "$CLUSTER_NAME" --zone="${_ZONE}" --project="${PROJECT_ID}"
+gcloud container clusters get-credentials "$CLUSTER_NAME" --zone="${CLUSTER_LOCATION:-${_ZONE}}" --project="${PROJECT_ID}"
 CHART="gcsfs/tests/perf/macrobenchmarks/workloads/${_WORKLOAD}/helm_chart"
 # Restore precedence: an external checkpoint wins; otherwise, when seeding is on,
 # restore the per-run seed produced by the seed-checkpoint step. Empty => fresh

cloudbuild/macrobenchmarks/scripts/seed_checkpoint.sh

@@ -20,7 +20,7 @@ fi
 SEED_RUN_ID="${RUN_ID}-seed"
 SEED_CKPT_DIR="gs://$CHECKPOINT_BUCKET/seed"
 curl -fsSL https://raw.githubusercontent.com/helm/helm/main/scripts/get-helm-3 | bash
-gcloud container clusters get-credentials "$CLUSTER_NAME" --zone="${_ZONE}" --project="${PROJECT_ID}"
+gcloud container clusters get-credentials "$CLUSTER_NAME" --zone="${CLUSTER_LOCATION:-${_ZONE}}" --project="${PROJECT_ID}"
 CHART="gcsfs/tests/perf/macrobenchmarks/workloads/${_WORKLOAD}/helm_chart"
 
 # One optimizer step (gradient accumulation is fixed at 1 in launcher.sh) with