multiflex start

kddubey · kddubey · commit 73fc4003a42c · 2026-05-20T14:13:24.000-07:00
diff --git a/bin/_startup.sh b/bin/_startup.sh
@@ -16,6 +16,35 @@ GROUPING_TRAINER_BUCKET=$(curl -fsS -H "Metadata-Flavor: Google" \
 export GROUPING_TRAINER_BUCKET
 
 
+# ----------------------------------------------------------------------------------------------------------------------
+# Multi-flex-start lock race
+# ----------------------------------------------------------------------------------------------------------------------
+# When gt.launch fans out flex-start submits across multiple zones (--multi_flex_start), every sibling VM gets the same
+# `launch-id` metadata. The first to reach this point claims the GCS object atomically via --if-generation-match=0;
+# losers see a 412 and self-delete.
+#
+# Any non-zero `gcloud storage cp` (412 race-loss or a transient gcloud error) self-deletes. Prefer over-deleting (user
+# needs to retry the launch) to under-deleting (bunch of colliding work).
+LAUNCH_ID=$(curl -fsS -H "Metadata-Flavor: Google" \
+    http://metadata.google.internal/computeMetadata/v1/instance/attributes/launch-id 2>/dev/null || true)
+if [ -n "$LAUNCH_ID" ]; then
+    LOCK_PATH="gs://${GROUPING_TRAINER_BUCKET}/launches/${LAUNCH_ID}/winner"
+    if hostname | gcloud storage cp - "$LOCK_PATH" --if-generation-match=0; then
+        echo "Congratulations! You won the race to $LOCK_PATH"
+    else
+        echo "You lost the race to $LOCK_PATH. Self-deleting. Better luck next time"
+        # || true so a metadata-server hiccup doesn't trip `set -e` before we get to `gcloud compute instances delete`.
+        NAME=$(curl -fsS -H 'Metadata-Flavor: Google' \
+            http://metadata.google.internal/computeMetadata/v1/instance/name) || NAME=$(hostname)
+        ZONE=$(curl -fsS -H 'Metadata-Flavor: Google' \
+            http://metadata.google.internal/computeMetadata/v1/instance/zone 2>/dev/null \
+            | awk -F/ '{print $NF}') || ZONE=""
+        gcloud compute instances delete "$NAME" --zone="$ZONE" --quiet
+        exit 0
+    fi
+fi
+
+
 # ----------------------------------------------------------------------------------------------------------------------
 # Set up python env
 # ----------------------------------------------------------------------------------------------------------------------
diff --git a/src/grouping_trainer/launch.py b/src/grouping_trainer/launch.py
@@ -20,13 +20,16 @@
 import sys
 import tempfile
 import time
+import uuid
 from datetime import datetime
 from enum import StrEnum
-from typing import Literal
+from typing import Literal, NoReturn
 
 from pydantic import BaseModel, ConfigDict
 from tap import tapify
 
+from grouping_trainer.logging import configure_logging
+
 logger = logging.getLogger(__name__)
 
 _REMOTE_ENV_VAR = "GROUPING_TRAINER_REMOTE"
@@ -87,9 +90,11 @@ class GpuConfig(BaseModel):
     reservation_affinity: Literal["none", "any"]
     is_for_training: bool
     n_gpu: int
+    boot_disk_type: str = "pd-balanced"
+    """Boot disk type. A3 Ultra (a3-ultragpu-8g) rejects pd-balanced and needs hyperdisk-balanced."""
 
 
-# gcloud compute accelerator-types list --filter="name=nvidia-l4" --format='value(zone)'
+# gcloud compute accelerator-types list --filter="name=nvidia-l4" --format='value(zone)' | sort -u
 L4_ZONES = (
     "us-central1-a",
     "us-central1-b",
@@ -168,6 +173,18 @@ class GpuConfig(BaseModel):
     "us-east4-c",
     "europe-west4-a",
 )
+# gcloud compute accelerator-types list --filter="name=nvidia-h200-141gb" --format='value(zone)'
+H200_ZONES = (
+    "us-central1-b",
+    "us-east4-b",
+    "us-east5-a",
+    "us-south1-b",
+    "us-west1-c",
+    "asia-south1-b",
+    "asia-south2-c",
+    "europe-west1-b",
+    "europe-west4-a",
+)
 
 GpuType = Literal[
     "l4",
@@ -177,6 +194,7 @@ class GpuConfig(BaseModel):
     "a100",
     "a100-ddp-2",
     "a100-ddp-4",
+    "h200-ddp-8",
 ]
 
 gpu_type_to_config: dict[GpuType, GpuConfig] = {
@@ -257,6 +275,18 @@ class GpuConfig(BaseModel):
         is_for_training=True,
         n_gpu=4,
     ),
+    "h200-ddp-8": GpuConfig(
+        flex_start_zone="us-central1-a",
+        standard_zones=H200_ZONES,
+        machine_type="a3-ultragpu-8g",
+        accelerator=None,
+        max_run_duration="172800s",
+        install_nvidia_driver=True,
+        reservation_affinity="none",
+        is_for_training=True,
+        n_gpu=8,
+        boot_disk_type="hyperdisk-balanced",
+    ),
 }
 
 
@@ -314,6 +344,17 @@ def _is_stockout(stderr: str) -> bool:
     return "ZONE_RESOURCE_POOL_EXHAUSTED" in stderr
 
 
+def _raise_gce_create_failure(result: subprocess.CompletedProcess[str], args: list[str]) -> NoReturn:
+    """Log the gcloud stderr (CalledProcessError's repr drops it) and raise."""
+    logger.error(f"gcloud failed (exit {result.returncode}):\n{result.stderr}")
+    raise subprocess.CalledProcessError(
+        result.returncode,
+        args,
+        output=result.stdout,
+        stderr=result.stderr,
+    )
+
+
 def _gce_create_cmd(
     config: GpuConfig,
     instance_name: str,
@@ -322,11 +363,14 @@ def _gce_create_cmd(
     provisioning_model: Literal["FLEX_START", "STANDARD"],
     wait_for_instance_creation: bool,
     path_to_metadata_script: dict[str, str],
+    launch_id: str | None = None,
 ) -> list[str]:
     metadata = {"gcs-bucket": os.environ["GROUPING_TRAINER_BUCKET"]}
     if config.install_nvidia_driver:
         metadata["enable-osconfig"] = "TRUE"
         metadata["install-nvidia-driver"] = "True"
+    if launch_id is not None:
+        metadata["launch-id"] = launch_id
 
     args = [
         "gcloud",
@@ -346,7 +390,7 @@ def _gce_create_cmd(
         "--scopes=https://www.googleapis.com/auth/cloud-platform",
         (
             f"--create-disk=auto-delete=yes,boot=yes,device-name={instance_name},"
-            f"image={_IMAGE},mode=rw,size=200,type=pd-balanced"
+            f"image={_IMAGE},mode=rw,size=200,type={config.boot_disk_type}"
         ),
         "--no-shielded-secure-boot",
         "--shielded-vtpm",
@@ -362,13 +406,69 @@ def _gce_create_cmd(
     return args
 
 
+def _gce_multi_flex_start(
+    *,
+    config: GpuConfig,
+    base_instance_name: str,
+    num_zones: int,
+    path_to_metadata_script: dict[str, str],
+) -> None:
+    """
+    Fan out async FLEX_START submits across the first `num_zones` of `config.standard_zones`, all sharing the same
+    `launch-id` metadata. First VM to boot claims a GCS atomic-create lock (see `bin/_startup.sh`); losers self-delete.
+    """
+    launch_id = uuid.uuid4().hex[:12]
+    zones = config.standard_zones[:num_zones]
+    logger.info(f"Multi-flex-start launch-id={launch_id}, fanning out to {len(zones)} zones: {zones}")
+    n_submitted = 0
+    last_stockout_stderr = ""
+    for zone in zones:
+        gce_create_args = _gce_create_cmd(
+            config,
+            base_instance_name,  # instance names only need to be unique within (project, zone).
+            zone,
+            provisioning_model="FLEX_START",
+            wait_for_instance_creation=False,
+            path_to_metadata_script=path_to_metadata_script,
+            launch_id=launch_id,
+        )
+        result = subprocess.run(gce_create_args, capture_output=True, text=True)
+        if result.returncode == 0:
+            logger.info(f"Flex-started {base_instance_name} in zone {zone}")
+            n_submitted += 1
+            continue
+
+        if _is_stockout(result.stderr):  # I've gotten a stockout on flex-starts before
+            logger.warning(f"Stockout in {zone}. Continuing with remaining zones")
+            last_stockout_stderr = result.stderr
+            continue
+
+        # Non-stockout error: fail loudly. Any earlier siblings already submitted will still race for the lock.
+        if n_submitted > 0:
+            logger.warning(
+                f"{n_submitted} sibling(s) already submitted will still race for the lock; "
+                "you may still end up with a VM despite this error."
+            )
+        _raise_gce_create_failure(result, gce_create_args)
+
+    if n_submitted == 0:
+        suffix = f" Last stderr:\n{last_stockout_stderr}" if last_stockout_stderr else ""
+        raise RuntimeError(f"Multi-flex-start: all {len(zones)} zones stocked out.{suffix}")
+    logger.info(
+        f"Multi-flex-start launched {n_submitted}/{len(zones)} VMs (launch-id={launch_id}). "
+        f"First to boot claims the winner lock, others self-delete."
+    )
+
+
 def gce_vm(
     *,
     gpu: GpuType,
     # TODO: support StrEnum-typed params in tap
     job_type: Literal[tuple(member.value for member in JobType)] = JobType.SSH,  # type: ignore[valid-type]
     name_suffix: str = "",
     sync_start: bool = False,
+    multi_flex_start: bool = False,
+    multi_flex_start_num_zones: int = 10,
     command: str | None = None,
     zone: str | None = None,
     num_cycles_through_zones: int = 5,
@@ -379,7 +479,9 @@ def gce_vm(
     `_startup.sh` `eval`s it after env setup. If not, the instance just sets up the env and stops for when you want to
     SSH in and iterate.
 
-    The instance's name is `grouping-trainer-{gpu}-{job_type}[-{name_suffix}]` (the suffix is dropped when empty).
+    The instance's name is `grouping-trainer-{gpu}-{job_type}[-{name_suffix}]` (the suffix is dropped when empty). In
+    `--multi_flex_start` mode, siblings share the same name — GCE instance names are unique per (project, zone), and
+    the zone column of `gcloud compute instances list` already distinguishes them.
 
     Parameters
     ----------
@@ -393,24 +495,35 @@ def gce_vm(
     sync_start
         If False (default), flex-starts the instance—GCP waits up to 1h to find one. `--sync_start` uses on-demand
         pricing and finds an instance in any zone, as flex-starting often can't find instances in time.
+    multi_flex_start
+        Fan out async FLEX_START submits across the first `multi_flex_start_num_zones` zones of `standard_zones`. First
+        VM to boot wins via a GCS atomic-create lock; losers self-delete. Mutually exclusive with `sync_start`. Not
+        applicable to GPU types without a `flex_start_zone` (e.g. L4).
+    multi_flex_start_num_zones
+        How many zones to fan out to in multi-flex mode. Capped at the length of the GPU config's `standard_zones`.
     command
         The command to run on the remote instance. If not given, the instance just sets up the env and stops for when
         you want to SSH in and iterate.
     zone
         Override the GCP zone. In FLEX_START mode pins the single submit zone; in STANDARD mode pins the loop to that
-        zone (still retried `num_cycles_through_zones` times).
+        zone (still retried `num_cycles_through_zones` times). Ignored in `--multi_flex_start` mode.
     num_cycles_through_zones
         No-op for sync_start. Otherwise, loop through zones this many times before giving up.
     seconds_between_gce_create_attempts
         No-op for sync_start. Otherwise, sleep this many seconds b/t consecutive zone attempts.
     """
     config = gpu_type_to_config[gpu]
+    if multi_flex_start and sync_start:
+        raise ValueError("--multi_flex_start and --sync_start are mutually exclusive")
+    if multi_flex_start and config.flex_start_zone is None:
+        raise ValueError(f"GPU type {gpu!r} does not support FLEX_START, so --multi_flex_start is not applicable")
+
     instance_name = f"{_INSTANCE_NAME_PREFIX}-{gpu}-{job_type}"
     if name_suffix:
         # GCE instance names must match `[a-z]([-a-z0-9]*[a-z0-9])?`, so swap underscores for hyphens.
         instance_name += f"-{name_suffix.replace('_', '-')}"
 
-    if (not sync_start) and (config.flex_start_zone is not None):
+    if (not sync_start) and (not multi_flex_start) and (config.flex_start_zone is not None):
         provisioning_model: Literal["FLEX_START", "STANDARD"] = "FLEX_START"
         wait_for_instance_creation = False
         zones_to_try: tuple[str, ...] = (zone or config.flex_start_zone,)
@@ -431,6 +544,15 @@ def gce_vm(
             cmd_file.flush()
             path_to_metadata_script["command"] = cmd_file.name
 
+        if multi_flex_start:
+            _gce_multi_flex_start(
+                config=config,
+                base_instance_name=instance_name,
+                num_zones=multi_flex_start_num_zones,
+                path_to_metadata_script=path_to_metadata_script,
+            )
+            return
+
         for zone_idx, zone in enumerate(zones_to_try):
             # Attempt creation in this zone
             gce_create_args = _gce_create_cmd(
@@ -444,7 +566,8 @@ def gce_vm(
             logger.info(f"Attempting to create {instance_name} in zone {zone}")
             gce_create_cmd_result = subprocess.run(gce_create_args, capture_output=True, text=True)
             if gce_create_cmd_result.returncode == 0:
-                logger.info(f"Created {instance_name} in zone {zone}")
+                creation_type = "Flex-started" if provisioning_model == "FLEX_START" else "Created"
+                logger.info(f"{creation_type} {instance_name} in zone {zone}")
                 return
 
             # Retry next zone if stockout
@@ -455,16 +578,9 @@ def gce_vm(
                 continue
 
             # Fail
-            logger.error(f"gcloud failed (exit {gce_create_cmd_result.returncode}):\n{gce_create_cmd_result.stderr}")
             if not sync_start:
-                logger.warning("You may have success with --sync_start and no GCP zone pinning.")
-            # CalledProcessError's repr drops stderr, so log it
-            raise subprocess.CalledProcessError(
-                gce_create_cmd_result.returncode,
-                gce_create_args,
-                output=gce_create_cmd_result.stdout,
-                stderr=gce_create_cmd_result.stderr,
-            )
+                logger.warning("You may have success with --multi_flex_start if you're fine waiting, else --sync_start")
+            _raise_gce_create_failure(gce_create_cmd_result, gce_create_args)
 
 
 def run_argv_remotely(
@@ -473,6 +589,7 @@ def run_argv_remotely(
     job_type: JobType,
     name_suffix: str,
     sync_start: bool = False,
+    multi_flex_start: bool = False,
     zone: str | None = None,
     env_var_to_value: dict[str, str] | None = None,
 ) -> None:
@@ -513,7 +630,7 @@ def run_argv_remotely(
     # Script path and args
     script_path = os.path.relpath(os.path.abspath(sys.argv[0]), _repo_root())
     argv_remote = _strip_flags_and_their_values(argv=sys.argv[1:], flags=("--gpu", "--zone"))
-    argv_remote = _strip_bool_flags(argv=argv_remote, flags=("--sync_start",))
+    argv_remote = _strip_bool_flags(argv=argv_remote, flags=("--sync_start", "--multi_flex_start"))
     args_for_remote = shlex.join(argv_remote)
 
     command_parts.append(script_path)
@@ -527,10 +644,12 @@ def run_argv_remotely(
         job_type=job_type,
         name_suffix=name_suffix,
         sync_start=sync_start,
+        multi_flex_start=multi_flex_start,
         command=command,
         zone=zone,
     )
 
 
 if __name__ == "__main__":
+    configure_logging(process_type="launch")
     tapify(gce_vm)
