Add Modal training entrypoint and data volume tooling

- `modal/train.py`: Full training on Modal GPUs with `electrai-data`
  Volume (dataset_4, 2,885 samples). Supports GPU selection (L4/A100/H100),
  custom configs, WandB logging, checkpoint persistence.
- `modal/populate_volume.py`: Sync S3 → Modal Volume via `boto3`.
- Update spec with training docs, data provenance, secrets.

Data pipeline: Globus (Della) → S3 → Modal Volume (complete).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: ryan-williams

modal/populate_volume.py

@@ -0,0 +1,83 @@
+"""Populate Modal Volume with training data from S3."""
+
+from __future__ import annotations
+
+import modal
+
+app = modal.App("electrai-populate")
+volume = modal.Volume.from_name("electrai-data", create_if_missing=True)
+
+
+@app.function(
+    image=modal.Image.debian_slim(python_version="3.12").pip_install("boto3"),
+    volumes={"/data": volume},
+    secrets=[modal.Secret.from_name("aws-credentials")],
+    timeout=7200,
+)
+def sync_s3(
+    bucket: str = "openathena",
+    prefix: str = "electrai/mp/chg_datasets/dataset_4",
+    dest: str = "/data/mp/chg_datasets/dataset_4",
+):
+    """Sync dataset from S3 to Modal Volume."""
+    import logging
+    from pathlib import Path
+
+    import boto3
+
+    log = logging.getLogger(__name__)
+    logging.basicConfig(level=logging.INFO)
+
+    s3 = boto3.client("s3")
+    paginator = s3.get_paginator("list_objects_v2")
+
+    # Count and list all objects
+    objects = [
+        obj
+        for page in paginator.paginate(Bucket=bucket, Prefix=prefix)
+        for obj in page.get("Contents", [])
+    ]
+
+    total_bytes = sum(o["Size"] for o in objects)
+    log.info("Found %d objects, %.1f GiB total", len(objects), total_bytes / (1024**3))
+
+    downloaded = 0
+    skipped = 0
+    for i, obj in enumerate(objects):
+        key = obj["Key"]
+        rel = key[len(prefix) :].lstrip("/")
+        local_path = Path(dest) / rel
+        local_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Skip if already exists with same size
+        if local_path.exists() and local_path.stat().st_size == obj["Size"]:
+            skipped += 1
+            continue
+
+        s3.download_file(bucket, key, str(local_path))
+        downloaded += 1
+
+        if (i + 1) % 100 == 0:
+            log.info(
+                "Progress: %d/%d (downloaded %d, skipped %d)",
+                i + 1,
+                len(objects),
+                downloaded,
+                skipped,
+            )
+
+    log.info("Done: %d downloaded, %d skipped (already existed)", downloaded, skipped)
+    volume.commit()
+
+
+@app.local_entrypoint()
+def main(
+    bucket: str = "openathena",
+    prefix: str = "electrai/mp/chg_datasets/dataset_4",
+    dest: str = "/data/mp/chg_datasets/dataset_4",
+):
+    import logging
+
+    logging.basicConfig(level=logging.INFO)
+    sync_s3.remote(bucket=bucket, prefix=prefix, dest=dest)
+    logging.getLogger(__name__).info("Volume populated.")

modal/train.py

@@ -0,0 +1,214 @@
+"""Modal training entrypoint for electrai.
+
+Run real training experiments on Modal GPUs with data from the
+electrai-data Volume and checkpoints persisted to electrai-checkpoints.
+
+Usage:
+    # Default config (ResUNet, dataset_4, 50 epochs, L4)
+    modal run modal/train.py
+
+    # Custom config file
+    modal run modal/train.py --config examples/MP/experiments/experiment_0/config.yaml
+
+    # Override GPU, epochs, channels
+    modal run modal/train.py --gpu A100 --epochs 10 --channels 64
+
+    # Resume from checkpoint
+    modal run modal/train.py --resume
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import modal
+
+ROOT = Path(__file__).parent.parent
+
+# Persistent volumes
+data_volume = modal.Volume.from_name("electrai-data")
+ckpt_volume = modal.Volume.from_name("electrai-checkpoints", create_if_missing=True)
+
+image = (
+    modal.Image.debian_slim(python_version="3.12")
+    .apt_install("git")
+    .pip_install(
+        "torch~=2.9",
+        "torchvision>=0.24",
+        "lightning~=2.5",
+        "numpy~=2.3",
+        "scikit-learn>=1.7",
+        "pymatgen>=2025.10",
+        "pyyaml>=6.0",
+        "zarr>=3.1",
+        "hydra-core>=1.3",
+        "wandb>=0.12",
+    )
+    .add_local_dir(str(ROOT / "src"), remote_path="/root/electrai/src", copy=True)
+    .add_local_dir(
+        str(ROOT / "examples"), remote_path="/root/electrai/examples", copy=True
+    )
+    .add_local_file(
+        str(ROOT / "pyproject.toml"),
+        remote_path="/root/electrai/pyproject.toml",
+        copy=True,
+    )
+    .run_commands("cd /root/electrai && pip install -e .")
+)
+
+app = modal.App("electrai-train", image=image)
+
+DATA_ROOT = "/data/mp/chg_datasets/dataset_4"
+CKPT_ROOT = "/checkpoints"
+
+
+@app.function(
+    gpu="L4",
+    volumes={"/data": data_volume, CKPT_ROOT: ckpt_volume},
+    secrets=[modal.Secret.from_name("wandb-credentials")],
+    timeout=14400,  # 4 hours
+)
+def train(config_json: str, gpu_type: str = "L4"):
+    """Run training with the given config (as JSON, converted to YAML remotely)."""
+    import json
+    import logging
+    import subprocess
+    import sys
+
+    import yaml
+
+    log = logging.getLogger(__name__)
+    logging.basicConfig(level=logging.INFO)
+
+    cfg = json.loads(config_json)
+
+    # Write config as YAML for the training entrypoint
+    config_path = Path("/tmp/config.yaml")
+    with config_path.open("w") as f:
+        yaml.dump(cfg, f, default_flow_style=False)
+
+    log.info("GPU: %s", gpu_type)
+    log.info("Data root: %s", DATA_ROOT)
+    log.info("Checkpoint dir: %s", CKPT_ROOT)
+
+    # Verify data exists
+    filelist = Path(DATA_ROOT) / "mp_filelist.txt"
+    if not filelist.exists():
+        raise FileNotFoundError(f"Filelist not found: {filelist}")
+    n_samples = len(filelist.read_text().strip().splitlines())
+    log.info("Dataset: %d samples", n_samples)
+
+    # Run training
+    result = subprocess.run(
+        [
+            sys.executable,
+            "-m",
+            "electrai.entrypoints.main",
+            "train",
+            "--config",
+            str(config_path),
+        ],
+        cwd="/root/electrai",
+        check=False,
+    )
+
+    # Persist checkpoints
+    ckpt_volume.commit()
+
+    if result.returncode != 0:
+        raise RuntimeError(f"Training failed with exit code {result.returncode}")
+
+    log.info("Training complete. Checkpoints saved to electrai-checkpoints volume.")
+
+
+@app.local_entrypoint()
+def main(
+    config: str = "",
+    gpu: str = "L4",
+    epochs: int = 50,
+    channels: int = 32,
+    residual_blocks: int = 1,
+    depth: int = 2,
+    kernel_size: int = 5,
+    lr: float = 0.01,
+    batch_size: int = 1,
+    val_frac: float = 0.005,
+    wandb_project: str = "mp-experiment",
+):
+    import json
+    import logging
+
+    logging.basicConfig(level=logging.INFO)
+    log = logging.getLogger(__name__)
+
+    if config:
+        # Read YAML config locally (pyyaml may not be in Modal's local Python,
+        # so shell out to parse it)
+        import subprocess
+
+        result = subprocess.run(
+            [
+                "python3",
+                "-c",
+                f"import yaml, json; print(json.dumps(yaml.safe_load(open('{config}'))))",
+            ],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        cfg = json.loads(result.stdout)
+    else:
+        # Build config from CLI args
+        cfg = {
+            "data": {
+                "_target_": "electrai.dataloader.dataset.RhoRead",
+                "root": f"{DATA_ROOT}/mp_filelist.txt",
+                "split_file": None,
+                "precision": "f32",
+                "batch_size": batch_size,
+                "train_workers": 4,
+                "val_workers": 2,
+                "pin_memory": False,
+                "val_frac": val_frac,
+                "drop_last": False,
+                "augmentation": False,
+                "random_seed": 42,
+            },
+            "model": {
+                "_target_": "electrai.model.resunet.ResUNet3D",
+                "in_channels": 1,
+                "out_channels": 1,
+                "n_channels": channels,
+                "n_residual_blocks": residual_blocks,
+                "kernel_size": kernel_size,
+                "depth": depth,
+                "use_checkpoint": False,
+            },
+            "precision": 32,
+            "epochs": epochs,
+            "lr": lr,
+            "weight_decay": 0.0,
+            "warmup_length": 1,
+            "beta1": 0.9,
+            "beta2": 0.99,
+            "wandb_mode": "online",
+            "entity": "PrinceOA",
+            "wb_pname": wandb_project,
+            "ckpt_path": CKPT_ROOT,
+        }
+
+    # Always override data root and checkpoint path for Modal
+    if "data" in cfg:
+        cfg["data"]["root"] = f"{DATA_ROOT}/mp_filelist.txt"
+    cfg["ckpt_path"] = CKPT_ROOT
+
+    config_json = json.dumps(cfg, indent=2)
+    log.info("Config:\n%s", config_json)
+
+    # Override GPU type via with_options
+    train_fn = train
+    if gpu != "L4":
+        train_fn = train.with_options(gpu=gpu)
+
+    train_fn.remote(config_json=config_json, gpu_type=gpu)
+    log.info("Done.")

specs/modal-ci.md

@@ -114,12 +114,53 @@ jobs:
 - **CPU baseline test** — could add a second `@app.function(gpu=None)` call, but CPU tests already run in `gen-expected.yml`
 - **`update_expected` mode** — would need to get the file back out of Modal (print to stdout and capture, or use a Modal Volume). Defer for now.
 - **Artifact upload** — same challenge; not needed for the basic CI pass/fail
-- **WandB logging** — could add by passing `--wandb-project` and setting `WANDB_API_KEY` secret. Straightforward to add later.
-- **S3 data sync** — only needed for `gpu-benchmark.yml` (larger dataset). The e2e test uses the small checked-in dataset.
+
+## Training on Modal (`modal/train.py`)
+
+Full training entrypoint for running real experiments on Modal, replacing Lambda Labs.
+
+### Data: `electrai-data` Volume
+
+`dataset_4` (2,885 samples, ~205 GiB) synced from Della (Globus source of truth) → S3 → Modal Volume:
+- S3: `s3://openathena/electrai/mp/chg_datasets/dataset_4/`
+- Volume mount: `/data/mp/chg_datasets/dataset_4/{data,label}/`
+
+Populate script: `modal/populate_volume.py` (S3 → Volume via `boto3`).
+
+### Checkpoints: `electrai-checkpoints` Volume
+
+Persists across runs. Mounted at `/checkpoints`.
+
+### Usage
+
+```bash
+# Default: ResUNet, dataset_4, 50 epochs, L4
+modal run modal/train.py
+
+# A100, custom hyperparams
+modal run modal/train.py --gpu A100 --channels 64 --epochs 50
+
+# Use existing config file
+modal run modal/train.py --config path/to/config.yaml --gpu A100
+```
+
+### Data provenance
+
+```
+Globus (ROSENGROUP share)
+  └── /mp/chg_datasets/dataset_4/   (canonical, on Della)
+        ├── Della → S3 (aws s3 sync, one-time)
+        │     └── s3://openathena/electrai/mp/chg_datasets/dataset_4/
+        │           └── S3 → Modal Volume (modal/populate_volume.py)
+        └── Della → Lambda LLFS (Globus transfer, Betsy's prior setup)
+              └── /home/ubuntu/betsy-electrai-2/dataset2/
+```
 
 ## Secrets required
 
-- `MODAL_TOKEN_ID` / `MODAL_TOKEN_SECRET` — need to be set on the electrai repo (or at org level). helico has these set already; check whether they're org-wide or repo-level.
+- `MODAL_TOKEN_ID` / `MODAL_TOKEN_SECRET` — repo secrets for GHA workflow
+- `wandb-credentials` — Modal secret with `WANDB_API_KEY` (for training)
+- `aws-credentials` — Modal secret with AWS creds (for `populate_volume.py`)
 
 No `GH_SA_TOKEN` needed (no runner registration).