Add Modal GPU CI workflow and app

- `modal/ci.py`: Modal app that runs `run_training()` on an L4 GPU,
  checks val_loss against expected values. Tested locally — produces
  identical results to EC2 L4 (val_loss=0.364269).
- `.github/workflows/gpu-e2e-modal.yml`: GHA workflow that calls
  `modal run` from `ubuntu-latest`. Requires `MODAL_TOKEN_ID` and
  `MODAL_TOKEN_SECRET` secrets.

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

Author: ryan-williams

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

@@ -0,0 +1,32 @@
+name: GPU E2E (Modal)
+
+on:
+  pull_request:
+    branches: [main]
+  workflow_dispatch:
+    inputs:
+      epochs:
+        description: 'Number of training epochs'
+        default: '5'
+        type: string
+
+jobs:
+  test:
+    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 e2e test on Modal GPU
+        env:
+          EPOCHS: ${{ inputs.epochs || '5' }}
+        run: modal run modal/ci.py --epochs "$EPOCHS"

modal/ci.py

@@ -0,0 +1,103 @@
+"""Modal GPU CI for electrai e2e training test."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import modal
+
+ROOT = Path(__file__).parent.parent
+
+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",
+    )
+    # 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
+    )
+    .add_local_dir(str(ROOT / "tests"), remote_path="/root/electrai/tests", copy=True)
+    .add_local_dir(str(ROOT / "data"), remote_path="/root/electrai/data", 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-ci", image=image)
+
+
+@app.function(gpu="L4", timeout=600)
+def run_e2e_test(epochs: int = 5, check: bool = True):
+    """Run e2e training test on GPU."""
+    import json
+    import logging
+    import sys
+
+    log = logging.getLogger(__name__)
+
+    sys.path.insert(0, "/root/electrai/scripts")
+    from e2e_train import run_training
+
+    results = run_training(epochs=epochs, gpu=True, verbose=True)
+    log.info("Platform: %s", results["platform"])
+    log.info("Final val_loss: %.6f", results["final_val_loss"])
+    if results["final_train_loss"] is not None:
+        log.info("Final train_loss: %.6f", results["final_train_loss"])
+    log.info("Wallclock: %.1fs", results["wallclock_s"])
+
+    if check:
+        expected_file = Path("/root/electrai/tests/expected_values.json")
+        expected_values = json.loads(expected_file.read_text())
+        platform = results["platform"]
+        if platform not in expected_values:
+            raise ValueError(
+                f"No expected values for platform {platform!r}, "
+                f"available: {list(expected_values.keys())}"
+            )
+        expected = expected_values[platform]
+        if expected.get("final_val_loss") is None:
+            raise ValueError(f"Expected values for {platform!r} are null")
+        expected_val_loss = expected["final_val_loss"]
+        diff = abs(results["final_val_loss"] - expected_val_loss)
+        tolerance = 0.001
+        if diff > tolerance:
+            raise AssertionError(
+                f"val_loss {results['final_val_loss']:.6f} differs from expected "
+                f"{expected_val_loss:.6f} by {diff:.6f} (tolerance: {tolerance})"
+            )
+        log.info(
+            "PASS: val_loss matches expected within tolerance (%.6f <= %f)",
+            diff,
+            tolerance,
+        )
+
+    return results
+
+
+@app.local_entrypoint()
+def main(epochs: int = 5, check: bool = True):
+    import logging
+
+    logging.basicConfig(level=logging.INFO)
+    results = run_e2e_test.remote(epochs=epochs, check=check)
+    logging.getLogger(__name__).info(
+        "Results: val_loss=%.6f in %.1fs",
+        results["final_val_loss"],
+        results["wallclock_s"],
+    )

specs/modal-ci.md

@@ -39,7 +39,9 @@ image = (
     .add_local_dir(str(ROOT / "tests"), remote_path="/root/electrai/tests")
     .add_local_dir(str(ROOT / "scripts"), remote_path="/root/electrai/scripts")
     .add_local_dir(str(ROOT / "data"), remote_path="/root/electrai/data")
-    .add_local_file(str(ROOT / "pyproject.toml"), remote_path="/root/electrai/pyproject.toml")
+    .add_local_file(
+        str(ROOT / "pyproject.toml"), remote_path="/root/electrai/pyproject.toml"
+    )
 )
 
 app = modal.App("electrai-ci", image=image)
@@ -49,10 +51,14 @@ app = modal.App("electrai-ci", image=image)
 def run_e2e_test(epochs: int = 5, check: bool = True):
     """Run e2e training test on GPU."""
     import subprocess
+
     cmd = [
-        "python", "scripts/e2e_train.py",
-        "--gpu", "--verbose",
-        "--epochs", str(epochs),
+        "python",
+        "scripts/e2e_train.py",
+        "--gpu",
+        "--verbose",
+        "--epochs",
+        str(epochs),
     ]
     if not check:
         cmd.append("--no-check")