Add Modal GPU benchmark (modal/benchmark.py)

Mirrors `gpu-benchmark.yml` but runs on Modal instead of EC2:
- Uses `electrai-data` Volume (2,885 samples from dataset_4)
- Configurable GPU type, model size, sample count, epochs
- WandB logging with same tags/config as EC2 benchmark
- Subsample support for quick benchmarks vs full dataset

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

Author: ryan-williams

.github/workflows/gpu-benchmark-modal.yml

@@ -0,0 +1,68 @@
+name: GPU Benchmark (Modal)
+
+on:
+  workflow_dispatch:
+    inputs:
+      gpu:
+        description: 'GPU type (L4, A10G, A100, H100)'
+        default: 'L4'
+        type: string
+      epochs:
+        description: 'Number of training epochs'
+        default: '5'
+        type: string
+      channels:
+        description: 'Number of model channels (8=tiny/L4, 32=prod/A100, 64=large)'
+        default: '8'
+        type: string
+      residual_blocks:
+        description: 'Number of residual blocks (2=tiny/L4, 16=prod/A100)'
+        default: '2'
+        type: string
+      samples:
+        description: 'Number of samples from Volume (0=all eligible)'
+        default: '50'
+        type: string
+      max_file_size:
+        description: 'Skip samples larger than N MB (0=no limit, 100=safe for L4)'
+        default: '100'
+        type: string
+      wandb_project:
+        description: 'WandB project name'
+        default: 'elf-net-ci'
+        type: string
+
+jobs:
+  benchmark:
+    runs-on: ubuntu-latest
+    env:
+      MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
+      MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v5
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Install Modal
+        run: uv pip install --system modal
+
+      - name: Run benchmark
+        env:
+          GPU: ${{ inputs.gpu || 'L4' }}
+          EPOCHS: ${{ inputs.epochs || '5' }}
+          CHANNELS: ${{ inputs.channels || '8' }}
+          RESIDUAL_BLOCKS: ${{ inputs.residual_blocks || '2' }}
+          SAMPLES: ${{ inputs.samples || '50' }}
+          MAX_FILE_SIZE: ${{ inputs.max_file_size || '100' }}
+          WANDB_PROJECT: ${{ inputs.wandb_project || 'elf-net-ci' }}
+        run: |
+          modal run modal/benchmark.py \
+            --gpu "$GPU" \
+            --epochs "$EPOCHS" \
+            --channels "$CHANNELS" \
+            --residual-blocks "$RESIDUAL_BLOCKS" \
+            --samples "$SAMPLES" \
+            --max-file-size "$MAX_FILE_SIZE" \
+            --wandb-project "$WANDB_PROJECT"

modal/benchmark.py

@@ -0,0 +1,193 @@
+"""Modal GPU benchmark for electrai.
+
+Runs a configurable training benchmark on Modal GPUs with data from the
+electrai-data Volume, logs metrics to WandB, and reports wall-clock time.
+
+Usage:
+    # Default: 50 samples, 5 epochs, 8 channels, 2 blocks, L4
+    modal run modal/benchmark.py
+
+    # Production model on A100 (32ch/16blk needs >=40GB VRAM)
+    modal run modal/benchmark.py --gpu A100 --channels 32 --residual-blocks 16
+
+    # All 2,885 Volume samples
+    modal run modal/benchmark.py --samples 0
+
+    # Quick smoke test
+    modal run modal/benchmark.py --samples 10 --epochs 2 --channels 8 --residual-blocks 2
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import modal
+
+ROOT = Path(__file__).parent.parent
+
+data_volume = modal.Volume.from_name("electrai-data")
+
+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 / "scripts"), remote_path="/root/electrai/scripts", 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-benchmark", image=image)
+
+DATA_ROOT = "/data/mp/chg_datasets/dataset_4"
+
+
+@app.function(
+    gpu="L4",
+    volumes={"/data": data_volume},
+    secrets=[modal.Secret.from_name("wandb-credentials")],
+    timeout=7200,
+    retries=0,
+)
+def run_benchmark(
+    epochs: int = 5,
+    channels: int = 8,
+    residual_blocks: int = 2,
+    samples: int = 50,
+    max_file_size: float = 100,
+    seed: int = 42,
+    wandb_project: str = "elf-net-ci",
+    gpu_type: str = "L4",
+):
+    """Run benchmark and return results."""
+    import logging
+    import sys
+
+    log = logging.getLogger(__name__)
+    logging.basicConfig(level=logging.INFO)
+
+    sys.path.insert(0, "/root/electrai/scripts")
+    from e2e_train import run_training
+
+    # Filter and subsample from volume filelist
+    filelist = Path(DATA_ROOT) / "mp_filelist.txt"
+    all_ids = [s for s in filelist.read_text().strip().splitlines() if s]
+
+    # Filter by file size first (avoid OOM on large grids)
+    if max_file_size > 0:
+        max_bytes = int(max_file_size * 1024 * 1024)
+        eligible = [
+            sid
+            for sid in all_ids
+            if (Path(DATA_ROOT) / "data" / f"{sid}.CHGCAR").stat().st_size <= max_bytes
+        ]
+        log.info(
+            "File size filter: %d/%d eligible (<=%.0fMB)",
+            len(eligible),
+            len(all_ids),
+            max_file_size,
+        )
+    else:
+        eligible = all_ids
+
+    # Subsample from eligible
+    if 0 < samples < len(eligible):
+        import random
+
+        random.seed(seed)
+        subset = sorted(random.sample(eligible, samples))
+        log.info("Subsampled %d/%d eligible samples", samples, len(eligible))
+    else:
+        subset = eligible
+        samples = len(subset)
+        log.info("Using all %d eligible samples", samples)
+
+    if not subset:
+        raise ValueError(
+            f"No eligible samples (total={len(all_ids)}, max_file_size={max_file_size}MB)"
+        )
+
+    # Create temp data root with symlinks + clean filelist (no trailing newline)
+    data_root = "/tmp/benchmark_data"
+    Path(data_root).mkdir(parents=True, exist_ok=True)
+    for subdir in ["data", "label"]:
+        link = Path(data_root) / subdir
+        if not link.exists():
+            link.symlink_to(Path(DATA_ROOT) / subdir)
+    Path(data_root, "mp_filelist.txt").write_text("\n".join(subset))
+
+    log.info(
+        "Benchmark: gpu=%s, epochs=%d, channels=%d, blocks=%d, samples=%d",
+        gpu_type,
+        epochs,
+        channels,
+        residual_blocks,
+        samples,
+    )
+
+    results = run_training(
+        channels=channels,
+        residual_blocks=residual_blocks,
+        epochs=epochs,
+        seed=seed,
+        gpu=True,
+        gradient_checkpoint=True,  # save VRAM for large grids
+        data_root=data_root,
+        max_file_size=0,  # already filtered above
+        wandb_project=wandb_project,
+        verbose=True,
+    )
+
+    log.info("val_loss: %.6f", results["final_val_loss"])
+    log.info("train_loss: %.6f", results["final_train_loss"])
+    log.info("Wallclock: %.1fs", results["wallclock_s"])
+    log.info("GPU: %s", gpu_type)
+
+    return results
+
+
+@app.local_entrypoint()
+def main(
+    gpu: str = "L4",
+    epochs: int = 5,
+    channels: int = 8,
+    residual_blocks: int = 2,
+    samples: int = 50,
+    max_file_size: float = 100,
+    seed: int = 42,
+    wandb_project: str = "elf-net-ci",
+):
+    import logging
+
+    logging.basicConfig(level=logging.INFO)
+    log = logging.getLogger(__name__)
+
+    benchmark_fn = run_benchmark
+    if gpu != "L4":
+        benchmark_fn = run_benchmark.with_options(gpu=gpu)
+
+    results = benchmark_fn.remote(
+        epochs=epochs,
+        channels=channels,
+        residual_blocks=residual_blocks,
+        samples=samples,
+        max_file_size=max_file_size,
+        seed=seed,
+        wandb_project=wandb_project,
+        gpu_type=gpu,
+    )
+
+    log.info(
+        "Benchmark complete: val_loss=%.6f, wallclock=%.1fs on %s",
+        results["final_val_loss"],
+        results["wallclock_s"],
+        gpu,
+    )