Support for managed job submission to Google Cloud

launch/README.md

@@ -72,6 +72,49 @@ pimm submit \
 Use `--dry-run` to print the submitit manifest and `--output PATH` to write it.
 `--submit.host iana` can be used when submission should happen from a remote login host.
 
+## Managed Google Cloud Batch Submission
+
+Submit to Google Cloud Batch with `--site gcloud`. Batch queues the job,
+provisions a single A100 VM, runs the published pimm container image, writes
+artifacts to a gcsfuse-mounted `gs://` bucket, and tears the VM down.
+
+Prerequisites: an authenticated `gcloud` CLI on the submit host, `gsutil` (for
+code staging), a GCS bucket for output, A100 quota in the chosen region, and the
+pimm image published to a registry the VM can pull (e.g.
+`ghcr.io/deeplearnphysics/pimm:main`, built by `.github/workflows/docker.yml`).
+
+Edit `launch/sites/gcloud.yaml` for your project (project, location, machine
+type, `gs://` exp_root, and data paths), then submit:
+
+```bash
+# WANDB_API_KEY is read automatically from a `WANDB_API_KEY=...` line in the
+# repo `.env`, so with that in place the flag below can be omitted.
+pimm submit \
+  --site gcloud \
+  --resources.time 04:00:00 \
+  --train.config panda/pretrain/pretrain-sonata-v1m1-pilarnet-smallmask \
+  --run.wandb-api-key "$WANDB_API_KEY"   # optional; overrides the `.env` value
+```
+
+Notes:
+
+- `paths.exp_root` must be a `gs://` URI. The backend derives the bucket, mounts
+  it at `resources.scheduler_options.gcs_mount_path`, and rewrites `EXP_ROOT` to
+  that local mount, so training and checkpoint code need no cloud-storage
+  awareness.
+- `resources.nproc_per_node` must be an explicit GPU count (`auto` is local-only).
+- `resources.time` becomes the Batch `maxRunDuration`.
+- `resources.scheduler_options.stage_code: true` rsyncs the local checkout to the
+  bucket at submit time so code edits take effect without rebuilding the image;
+  the image then supplies only the environment.
+- `WANDB_API_KEY` is required; submission fails fast if it is unset. Provide it
+  EITHER via a `WANDB_API_KEY=...` line in the repo `.env` (read automatically at
+  submit time) OR via `--run.wandb-api-key` (which overrides `.env`). The submit
+  host injects the key into the rendered Batch job — the `.env` file itself is
+  never staged to the VM.
+- Use `--dry-run` to print the rendered Batch job JSON and `--output PATH` to
+  write it.
+
 ## Container Repo Mounts
 
 The Docker images ship only the locked environment - no pimm source is baked
@@ -165,6 +208,9 @@ pimm submit \
   `scripts/nersc_env.sh`.
 - `launch/sites/nersc-container.yaml`: containerized NERSC alternative
   (Shifter, frozen image environment) for large-scale or pinned runs.
+- `launch/sites/gcloud.yaml`: Google Cloud Batch profile (A100 VM, published
+  pimm image, `gs://` output mounted via gcsfuse); submitted with `pimm submit
+  --site gcloud`.
 - `container.repo_mount`: in-container path where `paths.repo_root` is mounted
   so `pimm` imports resolve to the checkout; defaults to `/opt/pimm/src`.
 - `launch/sites/local.yaml`: no scheduler/container wrapper; runs directly on
@@ -173,5 +219,5 @@ pimm submit \
   choices, not model architecture.
 
 All topology and scheduler settings live under `resources`; site profiles set
-`resources.scheduler` to `local` or `slurm`. The legacy `slurm:` YAML group and
+`resources.scheduler` to `local`, `slurm`, or `gcloud`. The legacy `slurm:` YAML group and
 `--slurm.*` flags warn and are removed in pimm 0.6.0.

launch/sites/gcloud.yaml

@@ -0,0 +1,113 @@
+site: gcloud
+
+# Submit training to Google Cloud Batch (a managed, queue-based batch service --
+# the closest analog to a Slurm queue). A submitted job queues, provisions a
+# single A100 VM, runs the pimm container image, and tears the VM down.
+#
+# Output artifacts (checkpoints, logs, config) are written to a GCS bucket that
+# Batch mounts on the VM via gcsfuse. Set `paths.exp_root` to a gs:// URI; the
+# gcloud backend derives the bucket, mounts it at `scheduler_options.gcs_mount_path`,
+# and rewrites EXP_ROOT to that local mount -- so the training/checkpoint code
+# needs no cloud-storage awareness.
+#
+# Prerequisites (not automated):
+#   - An authenticated `gcloud` CLI on the machine you run `pimm submit` from.
+#   - The pimm container image published to a registry the VM can pull, e.g.
+#     ghcr.io/deeplearnphysics/pimm:main (built by .github/workflows/docker.yml).
+#   - A GCS bucket for output, and (for A100) quota in the chosen region.
+#
+# Fill in every REPLACE_ME_* value below before submitting.
+
+paths:
+  # The source is baked into the dev image at /opt/pimm/src; the rendered job
+  # runs `cd /opt/pimm/src && sh scripts/train.sh ...` inside the container.
+  repo_root: /opt/pimm/src
+  # gs:// output. Mounted read-write via gcsfuse; EXP_ROOT becomes
+  # <gcs_mount_path>/<prefix> on the VM.
+  exp_root: gs://lartpc-artifacts/pimm_exp
+
+resources:
+  scheduler: gcloud
+  nnodes: 1
+  nproc_per_node: 1 # a2-highgpu-1g bundles 1x A100 40GB
+  cpus_per_proc: 12
+  time: "04:00:00" # -> Cloud Batch maxRunDuration
+  scheduler_options:
+    project: gcp-physics
+    location: us-west1 # Oregon (has A100 40GB quota in zone us-west1-b)
+    machine_type: a2-highgpu-1g # 1x A100 40GB (a2-ultragpu-1g would be 80GB)
+    provisioning_model: STANDARD # or SPOT for preemptible/cheaper
+    boot_disk_gb: 350 # room for the gcsfuse file cache (~157 GB dataset) + OS/image
+    gcs_mount_path: /mnt/disks/gcs # where the exp_root bucket is fuse-mounted
+    # NOTE: Cloud Batch only auto-creates GCS mount dirs under /mnt/disks/; a
+    # path elsewhere (e.g. /mnt/gcs) fails gcsfuse with "stat: no such file or
+    # directory".
+    # Code staging: rsync the local checkout to gs://<bucket>/_pimm_code/<run>
+    # at submit time; the job copies it off the mount and runs from local disk,
+    # so edits take effect without rebuilding the image (the image then supplies
+    # only the environment/deps). Requires gsutil on the submit host.
+    stage_code: true
+    # code_prefix: _pimm_code   # bucket sub-path for staged source
+    # stage_dir: /tmp/pimm_src  # where the VM copies the source before running
+    # gcsfuse file cache: the dataset is read off the mount as many small
+    # random reads; the cache pulls each object to local disk on first touch so
+    # later reads (and later epochs) skip GCS. On by default (max-size-mb=-1 =
+    # fill free disk, LRU-evicted). Raise boot_disk_gb to cache more of the
+    # ~141 GB train set; set gcs_file_cache: false to disable.
+    # gcs_file_cache: true
+    # gcs_cache_dir: /mnt/disks/gcsfuse-cache
+    # gcs_cache_max_size_mb: -1
+    # gcsfuse HTTP connections per host (bounds the host-side gcsfuse process's
+    # fd usage so a checkpoint write burst can't exhaust it -> EMFILE on the
+    # mount). Default 100; set 0/false to use the gcsfuse default (unbounded).
+    # gcs_max_conns_per_host: 100
+    # DataLoader workers pass tensors via /dev/shm, which defaults to 64 MB in a
+    # container and overflows ("unable to allocate shared memory (shm) ... (11)").
+    # The container runs with --ipc=host by default (shm bounded by host RAM);
+    # set shm_size to size an isolated /dev/shm instead (docker --shm-size value).
+    # shm_size: 16g
+    # Only needed for non-a2 machine types (a2-* bundle their A100s):
+    # accelerator_type: nvidia-tesla-a100
+    # accelerator_count: 1
+    # service_account: my-batch-sa@REPLACE_ME_PROJECT.iam.gserviceaccount.com
+    # network: projects/REPLACE_ME_PROJECT/global/networks/default
+    # subnetwork: projects/REPLACE_ME_PROJECT/regions/us-central1/subnetworks/default
+
+container:
+  # Cloud Batch itself runs the image, so there is no nested `docker run`; keep
+  # runtime `none`. `image` is the published pimm container image.
+  runtime: none
+  image: ghcr.io/deeplearnphysics/pimm:main
+
+env:
+  PYTHONFAULTHANDLER: "1"
+  # The local .env is NOT staged to the VM (it holds s3df paths), so set the
+  # gcloud environment here. Paths point into the gcsfuse-mounted bucket
+  # (gcs_mount_path=/mnt/disks/gcs == gs://lartpc-artifacts).
+  #
+  # Training data: upload the PILArNet v3 set to gs://lartpc-artifacts/pilarnet
+  # (with train/ val/ test/ subdirs of *.h5) so the loader finds it here.
+  PILARNET_DATA_ROOT_V3: /mnt/disks/gcs/pilarnet
+  # W&B scratch/cache/artifacts on local VM disk (fast). gcsfuse is poor for
+  # W&B's many small writes; the run uploads to wandb.ai regardless, and these
+  # dirs are just ephemeral working space, so local /tmp is the right home.
+  WANDB_DIR: /tmp/wandb
+  WANDB_CACHE_DIR: /tmp/wandb/cache
+  WANDB_ARTIFACT_DIR: /tmp/wandb/artifacts
+  # WANDB_API_KEY is a secret -- do NOT commit it here. It is REQUIRED: the
+  # gcloud VM does not see the local .env, so submission fails fast if the key
+  # is empty/unset (a run cannot authenticate to W&B without it). Provide it in
+  # EITHER of two ways -- the submit host reads the key and injects it into the
+  # rendered Batch job (the `.env` file itself is never staged to the VM):
+  #   1. Keep a `WANDB_API_KEY=...` line in the repo `.env` (read automatically), or
+  #   2. Pass it explicitly at submit time (wins over `.env`):
+  #        pimm submit --site gcloud ... --run.wandb-api-key "$WANDB_API_KEY"
+
+# Per-run defaults are valid here, but reusable run choices usually belong in
+# launch/runs/*.yaml so the site profile remains portable.
+# run:
+#   wandb_project: my-project
+# train:
+#   config: panda/panseg/detector-v5-pt-v3m2-ft-joint-pxpypz-fft
+#   weight: hf://DeepLearnPhysics/panda-particle
+#   options: {}

pimm/launch/config.py

@@ -237,6 +237,29 @@ def load_config(
     return cfg
 
 
+def wandb_api_key_from_dotenv(root: Path = ROOT) -> str | None:
+    """Read only WANDB_API_KEY from the repo `.env`, ignoring all other vars.
+
+    Batch launches kept failing when the shell's $WANDB_API_KEY was empty and
+    got dropped/misparsed on the command line. Reading the single key straight
+    from `.env` here makes a real key the default without exposing any other
+    secrets from that file to the remote job.
+    """
+    env_path = root / ".env"
+    if not env_path.exists():
+        return None
+    for line in env_path.read_text(encoding="utf-8").splitlines():
+        line = line.strip()
+        if not line or line.startswith("#"):
+            continue
+        key, sep, value = line.partition("=")
+        if not sep or key.strip() != "WANDB_API_KEY":
+            continue
+        value = value.strip().strip("'").strip('"').strip()
+        return value or None
+    return None
+
+
 def finalize_config(
     cfg: dict[str, Any],
     *,
@@ -252,8 +275,17 @@ def finalize_config(
     run_cfg = cfg.setdefault("run", {})
     if run_cfg.get("wandb_project"):
         train_cfg.setdefault("options", {})["wandb_project"] = run_cfg["wandb_project"]
-    if run_cfg.get("wandb_api_key"):
-        cfg.setdefault("env", {})["WANDB_API_KEY"] = run_cfg["wandb_api_key"]
+    # Precedence: explicit --run.wandb-api-key, then an already-set env value,
+    # then the single WANDB_API_KEY from the repo `.env`. Only this one var is
+    # ever taken from `.env` — no other secrets leak into the job env. This is
+    # what lets a gcloud submit authenticate to W&B either by passing
+    # --run.wandb-api-key or by keeping WANDB_API_KEY=... in the repo `.env`
+    # (the `.env` file itself is never staged to the VM).
+    wandb_key = run_cfg.get("wandb_api_key") or cfg.get("env", {}).get("WANDB_API_KEY")
+    if not str(wandb_key or "").strip():
+        wandb_key = wandb_api_key_from_dotenv()
+    if str(wandb_key or "").strip():
+        cfg.setdefault("env", {})["WANDB_API_KEY"] = wandb_key
 
     rdzv_cfg = cfg.get("rdzv") or {}
     env = cfg.setdefault("env", {})
@@ -308,22 +340,47 @@ def validate_launch_config(cfg: dict[str, Any]) -> None:
         "resources.cpus_per_proc",
         "container.runtime",
     ]
-    if scheduler(cfg) == "slurm":
+    active_scheduler = scheduler(cfg)
+    if active_scheduler in {"slurm", "gcloud"} and (
+        cfg.get("resources", {}).get("nproc_per_node") == "auto"
+    ):
+        raise SystemExit(
+            "resources.nproc_per_node='auto' is only valid for the local "
+            "executor; set an explicit GPU count for Slurm/gcloud (batch)."
+        )
+    if active_scheduler == "slurm":
         required.extend(["resources.time", "resources.gpu_directive"])
-        if cfg.get("resources", {}).get("nproc_per_node") == "auto":
-            raise SystemExit(
-                "resources.nproc_per_node='auto' is only valid for the local "
-                "executor; set an explicit GPU count for Slurm (batch/interactive)."
-            )
+    elif active_scheduler == "gcloud":
+        # Cloud Batch runs the image directly (no host checkout to bind), so the
+        # image is mandatory; the GCS bucket is derived from a gs:// exp_root.
+        required.extend(
+            [
+                "container.image",
+                "resources.scheduler_options.project",
+                "resources.scheduler_options.location",
+                "resources.scheduler_options.machine_type",
+                "resources.time",
+            ]
+        )
     for dotted_path in required:
         require_path(cfg, dotted_path)
 
-    gpu_directive = cfg.get("resources", {}).get("gpu_directive")
-    if gpu_directive not in {"gres", "gpus-per-node"}:
-        raise SystemExit(
-            "resources.gpu_directive must be 'gres' or 'gpus-per-node', "
-            f"got {gpu_directive!r}"
-        )
+    # gpu_directive is a Slurm concept only; skip the check for other schedulers.
+    if active_scheduler == "slurm":
+        gpu_directive = cfg.get("resources", {}).get("gpu_directive")
+        if gpu_directive not in {"gres", "gpus-per-node"}:
+            raise SystemExit(
+                "resources.gpu_directive must be 'gres' or 'gpus-per-node', "
+                f"got {gpu_directive!r}"
+            )
+
+    if active_scheduler == "gcloud":
+        exp_root = str(cfg.get("paths", {}).get("exp_root", ""))
+        if not exp_root.startswith("gs://"):
+            raise SystemExit(
+                "gcloud site requires paths.exp_root to be a gs:// URI "
+                f"(mounted via gcsfuse on the Batch VM), got {exp_root!r}"
+            )
 
     runtime = cfg.get("container", {}).get("runtime")
     if runtime in {"apptainer", "singularity", "shifter", "docker"}: