fep-freesolv: --resume + incremental CSV write so multi-day CPU runs survive crashes

Previously freesolv_validate wrote the CSV only at the end of the
whole bench, so a 30-h CPU run dying mid-loop (lid close, kernel
panic, power) lost every completed compound. The friend's pilot-1
lost 3 compounds the same way on the M5 Max.

Changes:
  - Incremental atomic CSV write after each compound completes
    (write to .tmp sibling, then rename). Idempotent with the
    final summary write.
  - --resume flag: read existing CSV, skip compounds with non-
    empty dG_pred_kcalmol. Compounds that were attempted but
    didn't complete (pred=None / blank) re-run cleanly.
  - run_freesolv_m5max.sh passes --resume by default; honors
    OUT_DIR env override so resuming a crashed run is:
        OUT_DIR=run/fep/20260523_100401 bash scripts/run_freesolv_m5max.sh
    Switches `tee` to `tee -a` so resume doesn't truncate the
    log. New runs without OUT_DIR still get a fresh timestamped
    dir; --resume on an empty CSV is a no-op.

Regression: 7/7 covering round-trip, mid-compound-crash exclusion,
missing-file, atomic-tmp-rename, schema column pinning, partial-
run realistic case, fep-report compatibility.

This is precisely the lid-close protection I offered to ship at
the start of the laptop CPU run; current in-flight run (PID 91412,
started before this commit) doesn't benefit, but ANY rerun from
the same OUT_DIR will pick up where the prior run died.

Author: Rockman6

.github/workflows/smoke.yml

@@ -190,6 +190,9 @@ jobs:
       - name: fep-binding bench --resume regression
         run: python -u tests/fep/test_bench_resume_smoke.py
 
+      - name: fep-freesolv --resume + incremental CSV (7/7)
+        run: python -u tests/fep/test_freesolv_resume_smoke.py
+
       - name: bridge — Campaign-1 → Campaign-2 rate-law emitter
         run: python -u tests/bridge/test_binding_to_hill_smoke.py
 

scripts/run_freesolv_m5max.sh

@@ -34,10 +34,18 @@ fi
 source "$(conda info --base)/etc/profile.d/conda.sh"
 conda activate cellsim
 
-# Output dir, timestamped so repeat runs don't collide
-STAMP=$(date +%Y%m%d_%H%M%S)
-OUT_DIR="run/fep/${STAMP}"
+# Output dir, timestamped so repeat runs don't collide. Override
+# via env: `OUT_DIR=run/fep/20260523_100401 bash <this>` resumes
+# a crashed run from where it left off (any compound with a
+# non-empty dG_pred_kcalmol in freesolv_results.csv is skipped).
+if [ -z "${OUT_DIR}" ]; then
+    STAMP=$(date +%Y%m%d_%H%M%S)
+    OUT_DIR="run/fep/${STAMP}"
+fi
 mkdir -p "${OUT_DIR}"
+if [ -f "${OUT_DIR}/freesolv_results.csv" ]; then
+    echo "[run_freesolv] resuming from existing CSV in ${OUT_DIR}"
+fi
 
 # Header block — mirror to env.log so `cellsim fep-report` can
 # extract the `git commit:` line for provenance. Earlier versions
@@ -103,7 +111,8 @@ time ./scripts/cellsim fep-freesolv \
     --equilibration-steps 5000 \
     --sample-stride 250 \
     --out-csv "${OUT_DIR}/freesolv_results.csv" \
-    2>&1 | tee "${OUT_DIR}/run.log"
+    --resume \
+    2>&1 | tee -a "${OUT_DIR}/run.log"
 
 EXIT_CODE=${PIPESTATUS[0]}
 

src/fep/freesolv_validate.py

@@ -98,6 +98,73 @@ def as_dict(self) -> dict:
         return asdict(self)
 
 
+_CSV_COLS = [
+    "name", "smiles", "dG_expt_kcalmol",
+    "dG_pred_kcalmol", "uncertainty_kcalmol",
+    "residual_kcalmol", "ghmc_accept_mean",
+    "ghmc_accept_min", "wall_seconds", "ok", "reason",
+]
+
+
+def _write_csv(csv_path: Path, entries: list) -> None:
+    """Atomic CSV write: write to a .tmp sibling, then rename. Avoids
+    leaving a half-written file if the process is killed mid-write."""
+    csv_path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = csv_path.with_suffix(csv_path.suffix + ".tmp")
+    with tmp.open("w", newline="", encoding="utf-8-sig") as fo:
+        w = csv.DictWriter(fo, fieldnames=_CSV_COLS,
+                           extrasaction="ignore")
+        w.writeheader()
+        for p in entries:
+            w.writerow({k: getattr(p, k, "") for k in _CSV_COLS})
+    tmp.replace(csv_path)
+
+
+def _load_resume_rows(csv_path: Path) -> tuple[list, set]:
+    """Load previously-completed FreeSolvPoint rows from a partial
+    CSV. Returns (rows, names_completed). Used by --resume so a
+    crashed multi-day CPU run picks up where it left off without
+    re-doing compounds with non-empty dG_pred_kcalmol."""
+    rows: list = []
+    names: set = set()
+    if not csv_path.exists():
+        return rows, names
+    with csv_path.open("r", encoding="utf-8-sig", newline="") as fi:
+        reader = csv.DictReader(fi)
+        for r in reader:
+            pred = (r.get("dG_pred_kcalmol") or "").strip()
+            if not pred or pred == "None":
+                # Compound was attempted but didn't complete; re-run
+                # it fresh by not adding to names_completed.
+                continue
+            def _f(key):
+                v = (r.get(key) or "").strip()
+                if not v or v == "None":
+                    return None
+                try:
+                    return float(v)
+                except ValueError:
+                    return None
+            p = FreeSolvPoint(
+                name=r.get("name", "").strip(),
+                smiles=r.get("smiles", "").strip(),
+                dG_expt_kcalmol=float(
+                    r.get("dG_expt_kcalmol") or "nan"),
+                dG_pred_kcalmol=_f("dG_pred_kcalmol"),
+                uncertainty_kcalmol=_f("uncertainty_kcalmol"),
+                residual_kcalmol=_f("residual_kcalmol"),
+                wall_seconds=_f("wall_seconds"),
+                ok=((r.get("ok", "") or "").strip().lower()
+                    in ("true", "1", "yes")),
+                reason=r.get("reason", "").strip(),
+                ghmc_accept_mean=_f("ghmc_accept_mean"),
+                ghmc_accept_min=_f("ghmc_accept_min"),
+            )
+            rows.append(p)
+            names.add(p.name)
+    return rows, names
+
+
 def run_freesolv_validation(
     yaml_path: str | Path,
     *,
@@ -106,8 +173,22 @@ def run_freesolv_validation(
     n_equilibration_steps: int = 2500,
     sample_stride: int = 250,
     seed: int = 1,
+    out_csv: Optional[Path] = None,
+    resume: bool = False,
 ) -> FreeSolvResult:
-    """Run the FreeSolv subset and compute MAE + correlations."""
+    """Run the FreeSolv subset and compute MAE + correlations.
+
+    If `out_csv` is provided, the CSV is written **after every
+    compound completes** (not just at the end), so a 30+ hour CPU
+    run that dies mid-way leaves all completed compounds intact for
+    the next `--resume` invocation. The legacy main() also writes a
+    final summary CSV when out_csv is given.
+
+    If `resume=True` and `out_csv` exists, compounds with a non-
+    empty `dG_pred_kcalmol` in that file are kept and skipped in
+    the bench loop. Mid-compound crashes (no pred written) re-run
+    cleanly because only completed rows are loaded.
+    """
     import yaml as pyyaml
     import numpy as np
 
@@ -117,7 +198,18 @@ def run_freesolv_validation(
     result = FreeSolvResult(yaml_path=str(yaml_path))
     t0 = time.time()
 
+    completed_names: set = set()
+    if resume and out_csv:
+        prior_rows, completed_names = _load_resume_rows(out_csv)
+        result.entries.extend(prior_rows)
+        if completed_names:
+            print(f"[freesolv] --resume: {len(completed_names)} "
+                  f"compound(s) already in {out_csv}, skipping: "
+                  f"{sorted(completed_names)}", flush=True)
+
     for entry in data.get("entries", []):
+        if entry["name"] in completed_names:
+            continue
         p = FreeSolvPoint(
             name=entry["name"],
             smiles=entry["smiles"],
@@ -160,6 +252,16 @@ def run_freesolv_validation(
             p.reason = r.reason
             print(f"  FAIL: {r.reason}", flush=True)
         result.entries.append(p)
+        # Atomic incremental write — if a multi-day CPU run dies
+        # mid-loop (lid close / kernel panic / power), every
+        # completed compound survives in the CSV for --resume.
+        if out_csv:
+            try:
+                _write_csv(out_csv, result.entries)
+            except OSError as e:
+                logger.warning(
+                    "incremental CSV write failed: %s "
+                    "(continuing run)", e)
 
     result.n_total = len(result.entries)
     ok = [p for p in result.entries if p.ok]
@@ -197,30 +299,33 @@ def main(argv: Optional[list] = None) -> int:
     ap.add_argument("--sample-stride", type=int, default=250)
     ap.add_argument("--seed", type=int, default=1)
     ap.add_argument("--out-csv", type=Path, default=None,
-                    help="write per-compound results as CSV")
+                    help="write per-compound results as CSV "
+                         "(written incrementally after each "
+                         "compound, so a crashed run preserves "
+                         "completed compounds for --resume)")
+    ap.add_argument("--resume", action="store_true",
+                    help="if --out-csv exists, skip compounds "
+                         "with a non-empty dG_pred_kcalmol already "
+                         "in it. For restarting a crashed multi-day "
+                         "CPU run without losing completed work.")
     args = ap.parse_args(argv)
 
     r = run_freesolv_validation(
         args.yaml,
         n_windows=args.n_windows,
         n_production_steps=args.production_steps,
         n_equilibration_steps=args.equilibration_steps,
-        sample_stride=args.sample_stride, seed=args.seed)
+        sample_stride=args.sample_stride, seed=args.seed,
+        out_csv=args.out_csv, resume=args.resume)
     print()
     print(r.summary())
     if args.out_csv:
-        args.out_csv.parent.mkdir(parents=True, exist_ok=True)
-        cols = ["name", "smiles", "dG_expt_kcalmol",
-                "dG_pred_kcalmol", "uncertainty_kcalmol",
-                "residual_kcalmol", "ghmc_accept_mean",
-                "ghmc_accept_min", "wall_seconds", "ok", "reason"]
-        with args.out_csv.open(
-                "w", newline="", encoding="utf-8-sig") as fo:
-            w = csv.DictWriter(
-                fo, fieldnames=cols, extrasaction="ignore")
-            w.writeheader()
-            for p in r.entries:
-                w.writerow({k: getattr(p, k, "") for k in cols})
+        # Final write — same atomic helper. Idempotent with the
+        # per-compound incremental writes above; this ensures the
+        # last row reaches the file even if the loop body's writer
+        # didn't run for the last compound (e.g. all skipped via
+        # --resume).
+        _write_csv(args.out_csv, r.entries)
         print(f"\nwrote {args.out_csv}")
 
     if r.mae_kcalmol is not None and r.mae_kcalmol <= 1.5:

tests/fep/test_freesolv_resume_smoke.py

@@ -0,0 +1,175 @@
+"""freesolv_validate --resume regression.
+
+Pins:
+  - Incremental CSV write after each compound (so a kill mid-run
+    preserves completed work)
+  - --resume reads back the CSV and skips compounds with non-empty
+    dG_pred_kcalmol
+  - Mid-compound crash (no pred written) doesn't poison the resume
+  - Final write is idempotent with the incremental writes
+
+Multi-day CPU FreeSolv runs can die on lid-close / power / kernel
+panic. Without these properties the user loses ~30h of work; with
+them, --resume picks up exactly where the crash happened.
+"""
+from __future__ import annotations
+
+import csv
+import sys
+import tempfile
+from pathlib import Path
+
+REPO = Path(__file__).resolve().parents[2]
+sys.path.insert(0, str(REPO))
+
+from src.fep.freesolv_validate import (
+    FreeSolvPoint,
+    _CSV_COLS,
+    _load_resume_rows,
+    _write_csv,
+)
+
+
+def _make_point(name, pred=None, ok=True):
+    return FreeSolvPoint(
+        name=name, smiles="C" if name == "methane" else "CCO",
+        dG_expt_kcalmol=2.00 if name == "methane" else -5.01,
+        dG_pred_kcalmol=pred,
+        uncertainty_kcalmol=0.3 if pred is not None else None,
+        residual_kcalmol=(pred - 2.0) if pred is not None else None,
+        wall_seconds=420.0 if pred is not None else None,
+        ok=ok, reason="" if ok else "sampling failed: foo",
+        ghmc_accept_mean=0.85 if pred is not None else None,
+        ghmc_accept_min=0.72 if pred is not None else None,
+    )
+
+
+def test_write_csv_round_trip_via_load_resume():
+    with tempfile.TemporaryDirectory(prefix="freesolv_resume_") as tmp:
+        p = Path(tmp) / "results.csv"
+        rows = [_make_point("methane", pred=+1.80),
+                _make_point("ethanol", pred=-5.30)]
+        _write_csv(p, rows)
+        assert p.exists()
+        loaded, names = _load_resume_rows(p)
+        assert names == {"methane", "ethanol"}
+        assert len(loaded) == 2
+        assert loaded[0].name == "methane"
+        assert loaded[0].dG_pred_kcalmol == 1.80
+        assert loaded[1].dG_pred_kcalmol == -5.30
+
+
+def test_resume_skips_incomplete_compound_so_it_reruns():
+    """Compound that started but didn't finish writes a row with
+    pred=None. --resume must NOT mark it complete; the loop must
+    re-run that compound from scratch."""
+    with tempfile.TemporaryDirectory(prefix="freesolv_resume_") as tmp:
+        p = Path(tmp) / "results.csv"
+        rows = [
+            _make_point("methane", pred=+1.80),
+            _make_point("ethane", pred=None, ok=False),  # crashed
+        ]
+        _write_csv(p, rows)
+        loaded, names = _load_resume_rows(p)
+        # methane resumes-as-done; ethane is missing → will re-run
+        assert names == {"methane"}, names
+        assert len(loaded) == 1
+
+
+def test_load_resume_handles_missing_file():
+    """First run, no CSV exists → empty lists, no crash."""
+    with tempfile.TemporaryDirectory(prefix="freesolv_resume_") as tmp:
+        p = Path(tmp) / "does_not_exist.csv"
+        loaded, names = _load_resume_rows(p)
+        assert loaded == []
+        assert names == set()
+
+
+def test_write_csv_is_atomic_via_tmp_rename():
+    """If the .tmp file exists and full CSV doesn't, a previous
+    write was interrupted; current code overwrites tmp cleanly."""
+    with tempfile.TemporaryDirectory(prefix="freesolv_resume_") as tmp:
+        p = Path(tmp) / "results.csv"
+        tmp_path = p.with_suffix(p.suffix + ".tmp")
+        # Plant a stale tmp file from a "killed" prior write
+        tmp_path.write_text("stale,partial,write,leftover\n")
+        # New write should clobber and produce a valid CSV
+        _write_csv(p, [_make_point("methane", pred=+1.80)])
+        assert p.exists()
+        # Stale tmp was renamed away (rename clobbers target)
+        assert not tmp_path.exists()
+        # And the CSV is readable
+        loaded, names = _load_resume_rows(p)
+        assert names == {"methane"}
+
+
+def test_csv_schema_columns_pinned():
+    """Adding a column without thinking can break fep-report. Pin
+    the exact set + order."""
+    assert _CSV_COLS == [
+        "name", "smiles", "dG_expt_kcalmol",
+        "dG_pred_kcalmol", "uncertainty_kcalmol",
+        "residual_kcalmol", "ghmc_accept_mean",
+        "ghmc_accept_min", "wall_seconds", "ok", "reason",
+    ]
+
+
+def test_resume_loads_partial_run_three_done_one_pending():
+    """End-to-end realistic case: 12 compounds, 3 completed, 9
+    waiting. --resume keeps the 3, the next bench run runs the 9."""
+    with tempfile.TemporaryDirectory(prefix="freesolv_resume_") as tmp:
+        p = Path(tmp) / "results.csv"
+        rows = [
+            _make_point("methane", pred=+1.80),
+            _make_point("ethane", pred=+1.55),
+            _make_point("propane", pred=+2.31),
+        ]
+        _write_csv(p, rows)
+        loaded, names = _load_resume_rows(p)
+        assert names == {"methane", "ethane", "propane"}
+        assert len(loaded) == 3
+        # The three preds round-tripped exactly
+        preds = {r.name: r.dG_pred_kcalmol for r in loaded}
+        assert preds == {"methane": 1.80, "ethane": 1.55,
+                         "propane": 2.31}
+
+
+def test_csv_compatible_with_fep_report():
+    """The incremental writer must produce a CSV that cellsim
+    fep-report can consume. Smoke-check via a DictReader and
+    field-presence."""
+    with tempfile.TemporaryDirectory(prefix="freesolv_resume_") as tmp:
+        p = Path(tmp) / "results.csv"
+        _write_csv(p, [_make_point("methane", pred=+1.80)])
+        with p.open("r", encoding="utf-8-sig") as fi:
+            r = next(csv.DictReader(fi))
+        for col in ("name", "smiles", "dG_expt_kcalmol",
+                    "dG_pred_kcalmol", "wall_seconds", "ok"):
+            assert col in r, f"missing column {col}"
+
+
+if __name__ == "__main__":
+    funcs = [
+        test_write_csv_round_trip_via_load_resume,
+        test_resume_skips_incomplete_compound_so_it_reruns,
+        test_load_resume_handles_missing_file,
+        test_write_csv_is_atomic_via_tmp_rename,
+        test_csv_schema_columns_pinned,
+        test_resume_loads_partial_run_three_done_one_pending,
+        test_csv_compatible_with_fep_report,
+    ]
+    fails = []
+    for f in funcs:
+        try:
+            f()
+            print(f"[PASS] {f.__name__}")
+        except AssertionError as e:
+            print(f"[FAIL] {f.__name__}: {e}")
+            fails.append(f.__name__)
+        except Exception as e:
+            import traceback
+            traceback.print_exc()
+            print(f"[ERROR] {f.__name__}: {e}")
+            fails.append(f.__name__)
+    print(f"{len(funcs) - len(fails)}/{len(funcs)} PASS")
+    sys.exit(0 if not fails else 1)