From c9b44a5628c3ba5bf0064832ee326957c976032f Mon Sep 17 00:00:00 2001 From: Gregor Krzmanc Date: Fri, 24 Jul 2026 21:52:44 -0700 Subject: [PATCH 1/3] launch: add Google Cloud Batch job submission support Add a `gcloud` launch scheduler so `pimm submit` can queue training on Google Cloud Batch: it provisions a single A100 VM, runs the published pimm container image, writes artifacts to a gcsfuse-mounted gs:// bucket (EXP_ROOT is rewritten to the local mount so training code stays storage-agnostic), and tears the VM down. - pimm/launch/gcloud.py: Batch job builder + submit driver, gs:// URI parsing, optional GCS code staging, job-name sanitization. - launch/sites/gcloud.yaml: documented gcloud site config. - config/schema/utils/submit: wire "gcloud" through validation, the scheduler enum, and the submit entrypoint. - tests/unit/test_launch_gcloud.py: unit coverage for the builder. Co-Authored-By: Claude Opus 4.8 --- launch/sites/gcloud.yaml | 113 ++++++++++ pimm/launch/config.py | 49 +++- pimm/launch/gcloud.py | 371 +++++++++++++++++++++++++++++++ pimm/launch/schema.py | 2 +- pimm/launch/submit.py | 19 +- pimm/launch/utils.py | 4 +- tests/unit/test_launch_gcloud.py | 245 ++++++++++++++++++++ 7 files changed, 786 insertions(+), 17 deletions(-) create mode 100644 launch/sites/gcloud.yaml create mode 100644 pimm/launch/gcloud.py create mode 100644 tests/unit/test_launch_gcloud.py diff --git a/launch/sites/gcloud.yaml b/launch/sites/gcloud.yaml new file mode 100644 index 0000000..02e3327 --- /dev/null +++ b/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 + # / 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:///_pimm_code/ + # 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. Pass it at submit time: + # uv run pimm submit --site gcloud ... --run.wandb-api-key "$WANDB_API_KEY" + # (the submit scripts' --gcloud branch already forwards $WANDB_API_KEY). + # + # It is REQUIRED: the gcloud VM does not see the local .env, so submission + # fails fast if WANDB_API_KEY is empty/unset (a run cannot authenticate to + # W&B without it). Export it before submitting: + # export WANDB_API_KEY=$(grep '^WANDB_API_KEY=' .env | cut -d= -f2-) + +# 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: {} diff --git a/pimm/launch/config.py b/pimm/launch/config.py index 70f7a67..a8f173b 100644 --- a/pimm/launch/config.py +++ b/pimm/launch/config.py @@ -308,22 +308,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"}: diff --git a/pimm/launch/gcloud.py b/pimm/launch/gcloud.py new file mode 100644 index 0000000..c1503ec --- /dev/null +++ b/pimm/launch/gcloud.py @@ -0,0 +1,371 @@ +"""Google Cloud Batch executor for ``pimm submit --site gcloud``. + +Cloud Batch is GCP's managed, queue-based batch service -- the closest analog to +a Slurm queue. A submitted job queues, provisions an A100 VM, runs the pimm dev +container image, and tears the VM down. Output artifacts are written to a GCS +bucket that Batch mounts on the VM via gcsfuse, so the training/checkpoint code +sees an ordinary local directory and needs no cloud-storage awareness. + +The user writes ``paths.exp_root: gs:///`` in the site config; +this module derives the bucket, mounts it at ``gcs_mount_path`` on the VM, and +rewrites ``EXP_ROOT`` to the mounted local path before rendering the training +script. Everything else reuses the shared renderers in ``local.py``. +""" + +from __future__ import annotations + +import copy +import json +import re +import shlex +import shutil +import subprocess +import tempfile +from typing import Any + +from .local import build_train_sh_command, redact_script, render_script +from .utils import ROOT, as_bool, scheduler, shell_join, slurm_time_to_minutes, write_text + +# Cloud Batch only auto-creates gcsfuse mount dirs under /mnt/disks/; mounting +# elsewhere fails with "mount: stat : no such file or directory". +DEFAULT_MOUNT_PATH = "/mnt/disks/gcs" +DEFAULT_PROVISIONING_MODEL = "STANDARD" +DEFAULT_BOOT_DISK_GB = 200 +# PyTorch DataLoader workers pass tensors to the main process through POSIX +# shared memory (/dev/shm). A container's /dev/shm defaults to 64 MB, which a +# multi-worker loader exhausts almost immediately -- the failure surfaces as +# `RuntimeError: unable to allocate shared memory (shm) ... (11)`. Share the +# host IPC namespace by default so shm is bounded by host RAM, not 64 MB; set +# `scheduler_options.shm_size` (e.g. "16g") to size an isolated /dev/shm instead. +DEFAULT_CONTAINER_OPTIONS = "--ipc=host" +# Open-file limit for the Batch container. Across one training epoch the rank-0 +# process's open-fd count climbs well past the old 65536 cap (worker-shared +# tensors, gcsfuse read/cache handles, CUDA/NCCL) and the first checkpoint +# publish on the gcsfuse mount then trips `OSError: [Errno 24] Too many open +# files`. The same code runs fine on s3df, whose login shells give a 262144 +# limit -- i.e. the working set fits under 262144 but not 65536. Set the +# container hard+soft limit high (1048576, the usual Linux fs.nr_open ceiling +# and a value docker `--ulimit` reliably accepts) so the VM matches/exceeds the +# s3df headroom; train.sh's `ulimit -n` and train.py's setrlimit then raise the +# soft limit to this hard cap. +DEFAULT_NOFILE_LIMIT = 1048576 +# gcsfuse file cache. Training data is read straight off this mount as many +# small random-offset reads (one HDF5 event per __getitem__), which is slow +# over the network. Enabling the file cache pulls each object to local disk on +# first touch, so the rest of that shard's reads -- and every later epoch -- +# hit local disk instead of GCS. `--cache-dir` is what turns the cache on (it +# is a host path: gcsfuse runs on the VM, not in the container, and must live +# under /mnt/disks on Batch VMs; gcsfuse creates it if missing). +# `--file-cache-max-size-mb=-1` lets gcsfuse fill the free space on that disk +# with LRU eviction, so it can't overflow the boot disk. To cache the full +# dataset raise `resources.scheduler_options.boot_disk_gb` accordingly (train +# alone is ~141 GB); disable with `scheduler_options.gcs_file_cache=false`. +DEFAULT_GCS_CACHE_DIR = "/mnt/disks/gcsfuse-cache" +DEFAULT_GCS_CACHE_MAX_SIZE_MB = -1 +# Cap gcsfuse's per-host HTTP connections so its host-side fd budget (which the +# container ulimit cannot raise) is not exhausted by unbounded connections plus +# the checkpoint write burst. 100 matches gcsfuse's historical default and is +# ample for this workload; override via scheduler_options.gcs_max_conns_per_host +# (falsy leaves the gcsfuse default). +DEFAULT_GCS_MAX_CONNS_PER_HOST = 100 +DEFAULT_CODE_PREFIX = "_pimm_code" +DEFAULT_STAGE_DIR = "/tmp/pimm_src" +# Files never worth shipping to GCS with the source (a Python regex for +# `gsutil rsync -x`): VCS/venv/caches, large data/checkpoint artifacts, and +# `.env` -- the local .env is site-specific (e.g. s3df paths) and train.sh +# would source it and clobber the gcloud site's env; set gcloud env in the site +# config instead. +# +# NOTE: `gsutil rsync -x` matches the pattern against the *relative* path with +# `re.match`, i.e. anchored at the start (see `gsutil help rsync`). So patterns +# for anything that can appear nested must allow a leading path prefix +# (`(.*/)?` for dirs, `.*` for suffixes) -- a bare `__pycache__/` or `\.pyc$` +# would only ever match at the repo root and silently ship nested copies. +STAGE_EXCLUDE = ( + r"(.*/)?(\.git|\.venv|__pycache__|\.pytest_cache|\.mypy_cache)/|" + r"(exp|slurm_logs)/|" + r"(.*/)?\.env$|" + r".*\.(pyc|h5|pth)$" +) + + +def parse_gs_uri(uri: str) -> tuple[str, str]: + """Split a ``gs://bucket/prefix`` URI into ``(bucket, prefix)``. + + ``prefix`` has no leading/trailing slashes and may be empty (bucket root). + """ + if not uri.startswith("gs://"): + raise SystemExit(f"Expected a gs:// URI, got {uri!r}") + rest = uri[len("gs://"):] + bucket, _, prefix = rest.partition("/") + if not bucket: + raise SystemExit(f"gs:// URI is missing a bucket: {uri!r}") + return bucket, prefix.strip("/") + + +def sanitize_job_name(run_name: str) -> str: + """Coerce a run name into a valid Cloud Batch job id. + + Batch job names must match ``^[a-z]([a-z0-9-]{0,61})?$`` (<=63 chars, + lowercase, start with a letter). Run names contain uppercase, underscores + (timestamp), and dots, so normalize them. + """ + name = re.sub(r"[^a-z0-9-]", "-", run_name.lower()) + name = re.sub(r"-+", "-", name).strip("-") + if not name or not name[0].isalpha(): + name = f"pimm-{name}".strip("-") + return name[:63].rstrip("-") + + +def build_batch_job( + cfg: dict[str, Any], run_name: str +) -> tuple[dict[str, Any], str, tuple[str, str] | None]: + """Build a single-node A100 Batch job. + + Returns ``(job_spec, inner_script, stage_plan)`` where ``stage_plan`` is + ``(local_source_root, gs_dest)`` when code staging is enabled (else None) -- + the caller performs the actual upload. + """ + opts = cfg.get("resources", {}).get("scheduler_options", {}) or {} + mount_path = str(opts.get("gcs_mount_path") or DEFAULT_MOUNT_PATH).rstrip("/") + + exp_root = str(cfg.get("paths", {}).get("exp_root", "")) + bucket, prefix = parse_gs_uri(exp_root) + # EXP_ROOT on the VM points into the gcsfuse mount, not the gs:// URI. + exp_root_local = "/".join([mount_path, prefix]).rstrip("/") if prefix else mount_path + + # Cloud Batch is the container runner: render the *inner* training script + # with no nested `docker run` (runtime="none") and the mounted EXP_ROOT. + render_cfg = copy.deepcopy(cfg) + render_cfg.setdefault("container", {})["runtime"] = "none" + render_cfg.setdefault("paths", {})["exp_root"] = exp_root_local + + # Code staging: rsync the local checkout to the bucket at submit time, then + # have the job copy it off the gcsfuse mount to a local dir and run from + # there -- so edits take effect without rebuilding the image (the image then + # supplies only the environment/deps). On by default; opt out per site. + stage = as_bool(opts.get("stage_code", True)) + stage_plan: tuple[str, str] | None = None + stage_dir = str(opts.get("stage_dir") or DEFAULT_STAGE_DIR).rstrip("/") + code_mount = "" + if stage: + code_prefix = str(opts.get("code_prefix") or DEFAULT_CODE_PREFIX).strip("/") + code_rel = f"{code_prefix}/{run_name}" + code_mount = f"{mount_path}/{code_rel}" + # Run from the staged copy and make it win over the baked-in install. + render_cfg["paths"]["repo_root"] = stage_dir + render_cfg.setdefault("env", {})["PYTHONPATH"] = stage_dir + stage_plan = (str(ROOT), f"gs://{bucket}/{code_rel}") + + train_cmd = build_train_sh_command(render_cfg, run_name) + script = render_script(render_cfg, train_cmd, run_name) + + if stage: + # Copy the staged source off the (high-latency) gcsfuse mount to local + # disk before cd-ing into it. Injected right before render_script's + # `cd ` line. + cd_line = f"cd {shlex.quote(stage_dir)}" + copy_block = "\n".join( + [ + "echo '# staging pimm source from gcsfuse mount'", + f"mkdir -p {shlex.quote(stage_dir)}", + f"cp -a {shlex.quote(code_mount)}/. {shlex.quote(stage_dir)}/", + ] + ) + if cd_line not in script: + raise SystemExit("could not inject code-staging step into rendered script") + script = script.replace(cd_line, f"{copy_block}\n{cd_line}", 1) + + image = cfg.get("container", {}).get("image") + provisioning = str(opts.get("provisioning_model") or DEFAULT_PROVISIONING_MODEL) + boot_disk_gb = int(opts.get("boot_disk_gb") or DEFAULT_BOOT_DISK_GB) + minutes = slurm_time_to_minutes(cfg.get("resources", {}).get("time", "24:00:00")) + max_run_seconds = f"{minutes * 60}s" + + instance_policy: dict[str, Any] = { + "machineType": opts["machine_type"], + "provisioningModel": provisioning, + "bootDisk": {"sizeGb": boot_disk_gb}, + } + # a2-* machine types bundle their A100s automatically. For other machine + # types (e.g. n1-*), require an explicit accelerator spec. + accel_type = opts.get("accelerator_type") + accel_count = opts.get("accelerator_count") + if accel_type and accel_count: + instance_policy["accelerators"] = [ + {"type": accel_type, "count": int(accel_count)} + ] + + # DataLoader workers need more shared memory than a container's 64 MB + # default; pass docker-run flags to the Batch container to raise it. An + # explicit shm_size sizes an isolated /dev/shm; otherwise share host IPC. + shm_size = opts.get("shm_size") + shm_option = f"--shm-size={shm_size}" if shm_size else DEFAULT_CONTAINER_OPTIONS + nofile = int(opts.get("nofile_limit") or DEFAULT_NOFILE_LIMIT) + container_options = f"{shm_option} --ulimit nofile={nofile}:{nofile}" + + # Mount the bucket, optionally with the gcsfuse file cache turned on so the + # dataset's random-offset reads are served from local disk after first + # touch (see DEFAULT_GCS_CACHE_DIR). Batch passes mountOptions to gcsfuse as + # CLI flags ("--flag value" strings). + gcs_volume: dict[str, Any] = { + "gcs": {"remotePath": bucket}, + "mountPath": mount_path, + } + mount_options: list[str] = [] + if as_bool(opts.get("gcs_file_cache", True)): + cache_dir = str(opts.get("gcs_cache_dir") or DEFAULT_GCS_CACHE_DIR).rstrip("/") + cache_max_mb = int( + opts.get("gcs_cache_max_size_mb", DEFAULT_GCS_CACHE_MAX_SIZE_MB) + ) + mount_options += [ + f"--cache-dir {cache_dir}", + f"--file-cache-max-size-mb {cache_max_mb}", + ] + # Bound the gcsfuse process's HTTP connections per host. gcsfuse runs on the + # VM host with its own (~1024) fd budget that the container `--ulimit` cannot + # raise; unbounded connections plus the checkpoint write burst can exhaust it + # (EMFILE on the mount). A finite cap keeps gcsfuse's socket fds bounded. Set + # falsy to leave gcsfuse's default. + max_conns = opts.get("gcs_max_conns_per_host", DEFAULT_GCS_MAX_CONNS_PER_HOST) + if max_conns: + mount_options.append(f"--max-conns-per-host {int(max_conns)}") + if mount_options: + gcs_volume["mountOptions"] = mount_options + + task_spec: dict[str, Any] = { + "runnables": [ + { + "container": { + "imageUri": image, + "entrypoint": "/bin/bash", + "commands": ["-lc", script], + # gcsfuse mount lives on the host; expose it to the container. + "volumes": [f"{mount_path}:{mount_path}"], + "options": container_options, + } + } + ], + "volumes": [gcs_volume], + "maxRunDuration": max_run_seconds, + "maxRetryCount": 0, + } + + # W&B key: the key (--run.wandb-api-key, moved into env by finalize_config) + # is rendered as an env export in the script. It must be a non-empty value: + # the gcloud VM does not see the local .env, so a missing key means the run + # cannot authenticate to W&B. Fail loudly at build time rather than letting + # the job start and die. build_batch_job is the single chokepoint every + # submit (and --dry-run render) goes through, so this guards them all. + if not str(cfg.get("env", {}).get("WANDB_API_KEY", "")).strip(): + raise SystemExit( + "WANDB_API_KEY is empty; a gcloud run needs a real W&B key.\n" + "It is read automatically from a `WANDB_API_KEY=...` line in the " + "repo `.env`; add that line, or pass --run.wandb-api-key explicitly." + ) + + allocation_instance: dict[str, Any] = { + "policy": instance_policy, + "installGpuDrivers": True, + } + network = opts.get("network") + subnetwork = opts.get("subnetwork") + allocation_policy: dict[str, Any] = {"instances": [allocation_instance]} + if opts.get("service_account"): + allocation_policy["serviceAccount"] = {"email": opts["service_account"]} + if network or subnetwork: + interface: dict[str, Any] = {} + if network: + interface["network"] = network + if subnetwork: + interface["subnetwork"] = subnetwork + allocation_policy["network"] = {"networkInterfaces": [interface]} + + job_spec: dict[str, Any] = { + "taskGroups": [{"taskCount": 1, "taskSpec": task_spec}], + "allocationPolicy": allocation_policy, + "logsPolicy": {"destination": "CLOUD_LOGGING"}, + } + return job_spec, script, stage_plan + + +def stage_code(local_root: str, gs_dest: str, *, dry_run: bool) -> None: + """Mirror the local checkout to ``gs_dest`` with ``gsutil rsync``.""" + argv = [ + "gsutil", "-m", "rsync", "-r", "-x", STAGE_EXCLUDE, local_root, gs_dest, + ] + if dry_run: + print(f"# would stage code: {shell_join(argv)}") + return + if shutil.which("gsutil") is None: + raise SystemExit( + "gsutil not found; it is required to stage code to GCS " + "(install the Google Cloud SDK, or set scheduler_options.stage_code=false)" + ) + print(f"# staging code -> {gs_dest}") + rc = subprocess.run(argv).returncode + if rc != 0: + raise SystemExit(f"gsutil rsync failed with exit code {rc}") + + +def run_gcloud( + cfg: dict[str, Any], + run_name: str, + *, + dry_run: bool, + output: str | None, +) -> int: + """Render and submit a Google Cloud Batch job for a training run.""" + if scheduler(cfg) != "gcloud": + raise SystemExit("run_gcloud requires resources.scheduler='gcloud'") + + opts = cfg.get("resources", {}).get("scheduler_options", {}) or {} + project = opts["project"] + location = opts["location"] + + job_spec, script, stage_plan = build_batch_job(cfg, run_name) + job_name = sanitize_job_name(run_name) + job_json = json.dumps(job_spec, indent=2) + + # Upload the local checkout before submitting (a no-op print under --dry-run). + if stage_plan is not None: + stage_code(stage_plan[0], stage_plan[1], dry_run=dry_run) + + # The inner training script is embedded as a JSON string, so its `export` + # lines can't be redacted after serialization. Redact the script itself and + # splice it into a display-only copy of the spec. + display_spec = copy.deepcopy(job_spec) + display_spec["taskGroups"][0]["taskSpec"]["runnables"][0]["container"][ + "commands" + ][1] = redact_script(script) + display_json = json.dumps(display_spec, indent=2) + if output: + path = write_text(output, display_json) + print(f"# wrote Cloud Batch job spec: {path}") + if dry_run: + print(f"# gcloud batch jobs submit {job_name} " + f"--project {project} --location {location} --config ") + print(display_json) + return 0 + + with tempfile.NamedTemporaryFile( + "w", suffix=".json", prefix=f"{job_name}-", delete=False + ) as handle: + handle.write(job_json) + config_path = handle.name + + argv = [ + "gcloud", + "batch", + "jobs", + "submit", + job_name, + "--project", + project, + "--location", + location, + "--config", + config_path, + ] + print(f"# submitting Cloud Batch job {job_name} to {project}/{location}") + return subprocess.run(argv).returncode diff --git a/pimm/launch/schema.py b/pimm/launch/schema.py index ddbb353..e42b48a 100644 --- a/pimm/launch/schema.py +++ b/pimm/launch/schema.py @@ -23,7 +23,7 @@ class Resources: ``pimm submit`` requires ``scheduler='slurm'`` and consumes them. """ - scheduler: Literal["local", "slurm"] = "local" + scheduler: Literal["local", "slurm", "gcloud"] = "local" nnodes: int = 1 nproc_per_node: Annotated[ int | Literal["auto"], diff --git a/pimm/launch/submit.py b/pimm/launch/submit.py index dada5a0..b260a80 100644 --- a/pimm/launch/submit.py +++ b/pimm/launch/submit.py @@ -543,12 +543,27 @@ def run_submit( """Validate, dry-run, or submit a managed Slurm launch.""" validate_launch_config(cfg) validate_training_config(cfg) - if scheduler(cfg) != "slurm": - raise SystemExit("pimm submit requires resources.scheduler='slurm'") + active_scheduler = scheduler(cfg) + if active_scheduler not in {"slurm", "gcloud"}: + raise SystemExit( + "pimm submit requires resources.scheduler='slurm' or 'gcloud'" + ) run_name = build_run_name(cfg, launch_timestamp) if not run_name: raise SystemExit("Could not determine run name") + if active_scheduler == "gcloud": + # Google Cloud Batch: managed queue that provisions an A100 VM, runs the + # pimm image, and writes artifacts to a gcsfuse-mounted gs:// bucket. + from .gcloud import run_gcloud + + return run_gcloud( + cfg, + run_name, + dry_run=dry_run, + output=output, + ) + if as_bool(cfg.get("interactive", False)): # A chained interactive run needs its foreground driver to survive the # login session/node; host it under a scron watchdog instead (unless we diff --git a/pimm/launch/utils.py b/pimm/launch/utils.py index 01f246b..a5d220d 100644 --- a/pimm/launch/utils.py +++ b/pimm/launch/utils.py @@ -65,9 +65,9 @@ def shell_join(parts: list[Any]) -> str: def scheduler(cfg: dict[str, Any]) -> str: """Return the scheduler used by the current launcher command.""" configured = cfg.get("resources", {}).get("scheduler") - if configured not in {"local", "slurm"}: + if configured not in {"local", "slurm", "gcloud"}: raise SystemExit( - "resources.scheduler must be 'local' or 'slurm', " + "resources.scheduler must be 'local', 'slurm', or 'gcloud', " f"got {configured!r}" ) executor = cfg.get("executor") diff --git a/tests/unit/test_launch_gcloud.py b/tests/unit/test_launch_gcloud.py new file mode 100644 index 0000000..8854202 --- /dev/null +++ b/tests/unit/test_launch_gcloud.py @@ -0,0 +1,245 @@ +import json +import re + +import pytest + +from pimm.launch.config import finalize_config, load_config, validate_launch_config +from pimm.launch.gcloud import ( + STAGE_EXCLUDE, + build_batch_job, + parse_gs_uri, + sanitize_job_name, +) + + +TIMESTAMP = "2026-01-02_03-04-05" + + +def _config(**overrides): + cfg = load_config(site="gcloud", recipe=None, launch_timestamp=TIMESTAMP) + cfg["executor"] = "batch" + cfg["resources"]["scheduler_options"].update( + project="my-proj", location="us-central1", machine_type="a2-highgpu-1g" + ) + cfg["paths"]["exp_root"] = "gs://my-bucket/pimm_exp" + cfg["resources"]["time"] = "24:00:00" # pin: independent of the site yaml + cfg["run"] = {"name": "gcloud-render", "timestamp": False} + cfg["train"]["config"] = "tests/tiny_semseg" + # A gcloud run requires a real W&B key (the VM does not see the local .env). + cfg.setdefault("env", {})["WANDB_API_KEY"] = "test-key" + for key, value in overrides.items(): + cfg[key] = value + return finalize_config(cfg, launch_timestamp=TIMESTAMP, require_config=True) + + +def test_stage_exclude_matches_nested_paths(): + # `gsutil rsync -x` anchors the pattern at the start of the relative path + # (re.match), so the exclude must catch caches/artifacts at any depth, not + # just the repo root. + pat = re.compile(STAGE_EXCLUDE) + + excluded = [ + "pimm/__pycache__/utils.cpython-311.pyc", + "__pycache__/foo.pyc", + "pimm/models/foo.pyc", + ".git/config", + "pimm/sub/.git/config", + ".venv/lib/x.py", + ".env", + "conf/.env", + "exp/run1/ckpt.pth", + "data/x.h5", + "pimm/weights/model.pth", + ] + kept = [ + "pimm/launch/gcloud.py", + "pimm/__init__.py", + "README.md", + "configs/base.py", + ] + for p in excluded: + assert pat.match(p), f"expected {p!r} to be excluded" + for p in kept: + assert not pat.match(p), f"expected {p!r} to be kept" + + +def test_parse_gs_uri(): + assert parse_gs_uri("gs://b/p/q") == ("b", "p/q") + assert parse_gs_uri("gs://b") == ("b", "") + assert parse_gs_uri("gs://b/") == ("b", "") + with pytest.raises(SystemExit): + parse_gs_uri("/local/path") + + +def test_sanitize_job_name(): + raw = "detector-v5_FT.joint-2026-07-20_15-16-57" + name = sanitize_job_name(raw) + assert name == "detector-v5-ft-joint-2026-07-20-15-16-57" + assert len(name) <= 63 + assert name[0].isalpha() + # a leading digit gets a letter prefix + assert sanitize_job_name("2026-run").startswith("pimm-") + + +def test_build_batch_job_structure(): + cfg = _config() + cfg["resources"]["scheduler_options"]["stage_code"] = False # baked-image path + job, script, stage_plan = build_batch_job(cfg, cfg["run"]["name"]) + + task_spec = job["taskGroups"][0]["taskSpec"] + container = task_spec["runnables"][0]["container"] + instance = job["allocationPolicy"]["instances"][0] + + # gs:// bucket is mounted; EXP_ROOT points at the mount, not the URI. + assert task_spec["volumes"][0]["gcs"]["remotePath"] == "my-bucket" + assert task_spec["volumes"][0]["mountPath"] == "/mnt/disks/gcs" + assert "export EXP_ROOT=/mnt/disks/gcs/pimm_exp" in script + assert "gs://" not in script + + # Cloud Batch runs the image directly: no nested docker, train.sh invoked. + assert container["imageUri"] == "ghcr.io/deeplearnphysics/pimm:main" + assert container["entrypoint"] == "/bin/bash" + assert "docker run" not in script + # Staging off: run from the baked-in source, no copy step. + assert stage_plan is None + assert "sh /opt/pimm/src/scripts/train.sh -m 1 -g 1" in script + assert "gcsfuse mount" not in script + + # A100 VM with drivers installed, and the time budget becomes a duration. + assert instance["policy"]["machineType"] == "a2-highgpu-1g" + assert instance["installGpuDrivers"] is True + assert task_spec["maxRunDuration"] == "86400s" + + # spec is JSON-serializable + json.dumps(job) + + +def test_code_staging_default(): + cfg = _config() # stage_code defaults to True + job, script, stage_plan = build_batch_job(cfg, "gcloud-render") + + # Staged to a sibling prefix of exp_root in the SAME bucket. + assert stage_plan is not None + local_root, gs_dest = stage_plan + assert gs_dest == "gs://my-bucket/_pimm_code/gcloud-render" + assert local_root # the submit-host checkout root + + # The job copies the source off the mount and runs from local disk, with + # PYTHONPATH shadowing the baked-in install. + assert "cp -a /mnt/disks/gcs/_pimm_code/gcloud-render/. /tmp/pimm_src/" in script + assert "export PYTHONPATH=/tmp/pimm_src" in script + assert "cd /tmp/pimm_src" in script + assert "sh /tmp/pimm_src/scripts/train.sh -m 1 -g 1" in script + + +def test_container_shm_options(): + # Default: share host IPC so DataLoader workers aren't capped at 64 MB shm, + # and raise the open-file hard limit so gcsfuse checkpoint writes don't hit + # EMFILE ("Too many open files") when publishing the DCP directory. + cfg = _config() + job, _, _ = build_batch_job(cfg, cfg["run"]["name"]) + container = job["taskGroups"][0]["taskSpec"]["runnables"][0]["container"] + assert container["options"] == "--ipc=host --ulimit nofile=1048576:1048576" + + # An explicit shm_size sizes an isolated /dev/shm instead. + cfg = _config() + cfg["resources"]["scheduler_options"]["shm_size"] = "16g" + job, _, _ = build_batch_job(cfg, cfg["run"]["name"]) + container = job["taskGroups"][0]["taskSpec"]["runnables"][0]["container"] + assert container["options"] == "--shm-size=16g --ulimit nofile=1048576:1048576" + + +def test_gcs_file_cache_mount_options(): + # On by default: the bucket is mounted with the gcsfuse file cache enabled + # (a `--cache-dir` turns it on) so the dataset's random reads hit local disk + # after first touch. + cfg = _config() + job, _, _ = build_batch_job(cfg, cfg["run"]["name"]) + volume = job["taskGroups"][0]["taskSpec"]["volumes"][0] + assert volume["mountOptions"] == [ + "--cache-dir /mnt/disks/gcsfuse-cache", + "--file-cache-max-size-mb -1", + "--max-conns-per-host 100", + ] + + # Overridable dir/size. + cfg = _config() + cfg["resources"]["scheduler_options"].update( + gcs_cache_dir="/mnt/disks/cache", gcs_cache_max_size_mb=100000 + ) + job, _, _ = build_batch_job(cfg, cfg["run"]["name"]) + volume = job["taskGroups"][0]["taskSpec"]["volumes"][0] + assert volume["mountOptions"] == [ + "--cache-dir /mnt/disks/cache", + "--file-cache-max-size-mb 100000", + "--max-conns-per-host 100", + ] + + # Disable the file cache: only the connection cap remains. + cfg = _config() + cfg["resources"]["scheduler_options"]["gcs_file_cache"] = False + job, _, _ = build_batch_job(cfg, cfg["run"]["name"]) + volume = job["taskGroups"][0]["taskSpec"]["volumes"][0] + assert volume["mountOptions"] == ["--max-conns-per-host 100"] + + # Disable both: no mountOptions on the volume. + cfg = _config() + cfg["resources"]["scheduler_options"].update( + gcs_file_cache=False, gcs_max_conns_per_host=0 + ) + job, _, _ = build_batch_job(cfg, cfg["run"]["name"]) + volume = job["taskGroups"][0]["taskSpec"]["volumes"][0] + assert "mountOptions" not in volume + + +def test_explicit_accelerator_for_non_a2(): + cfg = _config() + cfg["resources"]["scheduler_options"].update( + machine_type="n1-standard-8", + accelerator_type="nvidia-tesla-a100", + accelerator_count=1, + ) + job, _, _ = build_batch_job(cfg, cfg["run"]["name"]) + accel = job["allocationPolicy"]["instances"][0]["policy"]["accelerators"][0] + assert accel == {"type": "nvidia-tesla-a100", "count": 1} + + +def _runnable(job): + return job["taskGroups"][0]["taskSpec"]["runnables"][0] + + +def test_wandb_key_exported_in_script(): + # The W&B key (finalize_config moves --run.wandb-api-key into env) is + # exported in the rendered script; no Secret Manager secretVariables exist. + cfg = _config() + cfg["env"]["WANDB_API_KEY"] = "abc123" + job, script, _ = build_batch_job(cfg, cfg["run"]["name"]) + assert "environment" not in _runnable(job) + assert "export WANDB_API_KEY=abc123" in script + + +@pytest.mark.parametrize("value", ["", " ", None]) +def test_missing_wandb_key_rejected(value): + # A gcloud run cannot authenticate to W&B without a real key, and the VM + # never sees the local .env -- so an empty/unset key fails fast at build. + cfg = _config() + if value is None: + cfg["env"].pop("WANDB_API_KEY", None) + else: + cfg["env"]["WANDB_API_KEY"] = value + with pytest.raises(SystemExit, match="WANDB_API_KEY"): + build_batch_job(cfg, cfg["run"]["name"]) + + +def test_non_gs_exp_root_rejected(): + cfg = _config() + cfg["paths"]["exp_root"] = "/local/exp" + with pytest.raises(SystemExit, match="gs://"): + validate_launch_config(cfg) + + +def test_missing_project_rejected(): + cfg = _config() + del cfg["resources"]["scheduler_options"]["project"] + with pytest.raises(SystemExit, match="project"): + validate_launch_config(cfg) From 15a175023366a13ca733b86d0ecd9bec8559cb7c Mon Sep 17 00:00:00 2001 From: Gregor Krzmanc Date: Fri, 24 Jul 2026 22:01:58 -0700 Subject: [PATCH 2/3] docs: document the gcloud launch site in launch/README.md Add a "Managed Google Cloud Batch Submission" section (prerequisites, `pimm submit --site gcloud` example, and gs:///staging/WANDB notes), list the site under File Ownership, and note the new `gcloud` scheduler value. Drop a stale gcloud.yaml comment referencing personal submit scripts. Co-Authored-By: Claude Opus 4.8 --- launch/README.md | 45 +++++++++++++++++++++++++++++++++++++++- launch/sites/gcloud.yaml | 3 +-- 2 files changed, 45 insertions(+), 3 deletions(-) diff --git a/launch/README.md b/launch/README.md index c09aa3e..d430bb6 100644 --- a/launch/README.md +++ b/launch/README.md @@ -72,6 +72,46 @@ 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 +export WANDB_API_KEY=... # the Batch VM does not see the local .env + +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" +``` + +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 and must be passed at submit time (the VM cannot + read the local `.env`); submission fails fast if it is unset. +- 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 +205,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 +216,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. diff --git a/launch/sites/gcloud.yaml b/launch/sites/gcloud.yaml index 02e3327..5959e9f 100644 --- a/launch/sites/gcloud.yaml +++ b/launch/sites/gcloud.yaml @@ -95,8 +95,7 @@ env: WANDB_CACHE_DIR: /tmp/wandb/cache WANDB_ARTIFACT_DIR: /tmp/wandb/artifacts # WANDB_API_KEY is a secret -- do NOT commit it here. Pass it at submit time: - # uv run pimm submit --site gcloud ... --run.wandb-api-key "$WANDB_API_KEY" - # (the submit scripts' --gcloud branch already forwards $WANDB_API_KEY). + # pimm submit --site gcloud ... --run.wandb-api-key "$WANDB_API_KEY" # # It is REQUIRED: the gcloud VM does not see the local .env, so submission # fails fast if WANDB_API_KEY is empty/unset (a run cannot authenticate to From b3a327b553d11cfae024aaa43773a1295d825023 Mon Sep 17 00:00:00 2001 From: Gregor Krzmanc Date: Tue, 28 Jul 2026 11:40:57 -0700 Subject: [PATCH 3/3] gcloud: auto-read WANDB_API_KEY from repo .env finalize_config now falls back to a WANDB_API_KEY=... line in the repo .env when neither --run.wandb-api-key nor an env-block value is set, so a gcloud submit authenticates to W&B via EITHER the .env line or the flag (the flag still wins). The .env file itself is never staged to the VM; the submit host reads the single key and injects it into the rendered Batch job env. Update gcloud.yaml, gcloud.py, and launch/README.md to document both paths. Co-Authored-By: Claude Opus 4.8 --- launch/README.md | 13 ++++++++----- launch/sites/gcloud.yaml | 15 ++++++++------- pimm/launch/config.py | 36 ++++++++++++++++++++++++++++++++++-- pimm/launch/gcloud.py | 7 ++++--- 4 files changed, 54 insertions(+), 17 deletions(-) diff --git a/launch/README.md b/launch/README.md index d430bb6..69b1b25 100644 --- a/launch/README.md +++ b/launch/README.md @@ -87,13 +87,13 @@ Edit `launch/sites/gcloud.yaml` for your project (project, location, machine type, `gs://` exp_root, and data paths), then submit: ```bash -export WANDB_API_KEY=... # the Batch VM does not see the local .env - +# 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" + --run.wandb-api-key "$WANDB_API_KEY" # optional; overrides the `.env` value ``` Notes: @@ -107,8 +107,11 @@ Notes: - `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 and must be passed at submit time (the VM cannot - read the local `.env`); submission fails fast if it is unset. +- `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. diff --git a/launch/sites/gcloud.yaml b/launch/sites/gcloud.yaml index 5959e9f..fbd8785 100644 --- a/launch/sites/gcloud.yaml +++ b/launch/sites/gcloud.yaml @@ -94,13 +94,14 @@ env: 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. Pass it at submit time: - # pimm submit --site gcloud ... --run.wandb-api-key "$WANDB_API_KEY" - # - # It is REQUIRED: the gcloud VM does not see the local .env, so submission - # fails fast if WANDB_API_KEY is empty/unset (a run cannot authenticate to - # W&B without it). Export it before submitting: - # export WANDB_API_KEY=$(grep '^WANDB_API_KEY=' .env | cut -d= -f2-) + # 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. diff --git a/pimm/launch/config.py b/pimm/launch/config.py index a8f173b..dfe41f5 100644 --- a/pimm/launch/config.py +++ b/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", {}) diff --git a/pimm/launch/gcloud.py b/pimm/launch/gcloud.py index c1503ec..d41bf3b 100644 --- a/pimm/launch/gcloud.py +++ b/pimm/launch/gcloud.py @@ -251,9 +251,10 @@ def build_batch_job( "maxRetryCount": 0, } - # W&B key: the key (--run.wandb-api-key, moved into env by finalize_config) - # is rendered as an env export in the script. It must be a non-empty value: - # the gcloud VM does not see the local .env, so a missing key means the run + # W&B key: the key (from --run.wandb-api-key or the repo `.env`, moved into + # env by finalize_config) is rendered as an env export in the script. It must + # be a non-empty value: the gcloud VM does not see the local .env, so a + # missing key means the run # cannot authenticate to W&B. Fail loudly at build time rather than letting # the job start and die. build_batch_job is the single chokepoint every # submit (and --dry-run render) goes through, so this guards them all.