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/ci.py

@@ -8,23 +8,13 @@
 
 ROOT = Path(__file__).parent.parent
 
+# Dependencies read from pyproject.toml (shared with train.py, populate_volume.py)
 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",
-        "click",
+    .pip_install_from_pyproject(
+        str(ROOT / "pyproject.toml"), optional_dependencies=["dev"]
     )
-    # Source code + test data (copy=True so pip install -e works after)
     .add_local_dir(str(ROOT / "src"), remote_path="/root/electrai/src", copy=True)
     .add_local_dir(
         str(ROOT / "scripts"), remote_path="/root/electrai/scripts", copy=True
@@ -36,13 +26,13 @@
         remote_path="/root/electrai/pyproject.toml",
         copy=True,
     )
-    .run_commands("cd /root/electrai && pip install -e .")
+    .run_commands("cd /root/electrai && pip install --no-deps -e .")
 )
 
 app = modal.App("electrai-ci", image=image)
 
 
-@app.function(gpu="L4", timeout=600)
+@app.function(gpu="L4", timeout=600, retries=0)
 def run_e2e_test(epochs: int = 5, check: bool = True):
     """Run e2e training test on GPU."""
     import json

modal/populate_volume.py

@@ -0,0 +1,84 @@
+"""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,
+    retries=0,
+)
+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,215 @@
+"""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 last checkpoint
+    modal run modal/train.py --gpu A100 --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)
+
+# Dependencies read from pyproject.toml (shared with ci.py, populate_volume.py)
+image = (
+    modal.Image.debian_slim(python_version="3.12")
+    .apt_install("git")
+    .pip_install_from_pyproject(
+        str(ROOT / "pyproject.toml"), optional_dependencies=["dev"]
+    )
+    .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 --no-deps -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
+    retries=0,
+)
+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)
+
+    # Check for existing checkpoint to resume from
+    ckpt_path = Path(cfg.get("ckpt_path", CKPT_ROOT))
+    last_ckpt = ckpt_path / "last.ckpt"
+    if last_ckpt.exists():
+        log.info("Found checkpoint: %s", last_ckpt)
+    else:
+        log.info("No checkpoint found, starting from scratch")
+
+    # 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",
+    resume: bool = False,
+):
+    import json
+    import logging
+
+    logging.basicConfig(level=logging.INFO)
+    log = logging.getLogger(__name__)
+
+    if config:
+        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:
+        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
+
+    if not resume:
+        cfg.pop("resume_from_checkpoint", None)
+
+    config_json = json.dumps(cfg, indent=2)
+    log.info("Config:\n%s", config_json)
+
+    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.")