Add resumable execution to the screening framework (#33)

ThermoScreening/cli/thermo.py

@@ -122,6 +122,11 @@ def _command_parser():
         "--method", default="GFN2-xTB", choices=["GFN2-xTB", "GFN1-xTB"],
         help="GFN-xTB parametrisation for the xtb engines (default 'GFN2-xTB').",
     )
+    screen_parser.add_argument(
+        "--resume", action="store_true",
+        help="Reuse successful results from a prior <out>.json and only (re)run "
+        "the missing/failed molecules.",
+    )
 
     return parser
 
@@ -207,6 +212,7 @@ def run_screen(parser_args):
         quasi_rrho=parser_args.quasi_rrho,
         engine=parser_args.engine,
         method=parser_args.method,
+        resume=parser_args.resume,
     )
 
     failed = sum(1 for record in results if record["status"] != "ok")

ThermoScreening/thermo/screening.py

@@ -10,6 +10,7 @@
 import csv
 import json
 import logging
+import os
 from dataclasses import dataclass
 from pathlib import Path
 
@@ -114,20 +115,56 @@ def _thermo_summary(thermo):
     }
 
 
+def _load_completed(out):
+    """
+    Load the successfully-completed records from a prior ``<out>.json``.
+
+    Returns a ``{name: record}`` mapping of the records whose ``status`` is
+    ``"ok"``, so a resumed screen can skip them. A missing or unreadable file
+    yields an empty mapping (nothing to resume).
+    """
+    json_path = Path(str(out)).with_suffix(".json")
+    if not json_path.exists():
+        return {}
+    try:
+        prior = json.loads(json_path.read_text(encoding="utf-8"))
+    except (json.JSONDecodeError, OSError):
+        return {}
+    if not isinstance(prior, list):
+        return {}
+    return {
+        record["name"]: record
+        for record in prior
+        if isinstance(record, dict) and record.get("status") == "ok" and "name" in record
+    }
+
+
+def _atomic_write(path, write_fn):
+    """
+    Write via a temp file + ``os.replace`` so an interrupted write never leaves a
+    truncated (unresumable) file behind.
+    """
+    tmp = path.with_name(path.name + ".tmp")
+    with open(tmp, "w", newline="", encoding="utf-8") as handle:
+        write_fn(handle)
+    os.replace(tmp, path)
+
+
 def _write_results(results, out):
     stem = Path(str(out)).with_suffix("")
     if stem.parent != Path(""):
         stem.parent.mkdir(parents=True, exist_ok=True)
 
-    csv_path = stem.with_suffix(".csv")
-    with open(csv_path, "w", newline="", encoding="utf-8") as handle:
+    def write_csv(handle):
         writer = csv.DictWriter(handle, fieldnames=_RESULT_FIELDS)
         writer.writeheader()
         for record in results:
             writer.writerow({field: record.get(field, "") for field in _RESULT_FIELDS})
 
+    csv_path = stem.with_suffix(".csv")
     json_path = stem.with_suffix(".json")
-    json_path.write_text(json.dumps(results, indent=2), encoding="utf-8")
+    _atomic_write(csv_path, write_csv)
+    _atomic_write(json_path, lambda handle: handle.write(json.dumps(results, indent=2)))
 
     return csv_path, json_path
 
@@ -146,6 +183,7 @@ def screen(
     quasi_rrho=False,
     engine="dftb+",
     method="GFN2-xTB",
+    resume=False,
 ):
     """
     Run a thermochemistry screen over a set of molecules.
@@ -192,6 +230,11 @@ def screen(
         applies only to DFTB+; ``solvent`` applies to DFTB+ and xtb-cli.
     method : str
         GFN-xTB parametrisation for the xtb engines (``"GFN2-xTB"`` default).
+    resume : bool
+        If True, reuse the successful (``status="ok"``) records from a prior
+        ``<out>.json`` and skip those molecules; molecules that previously failed
+        (or were never run) are (re-)run. Results are written incrementally after
+        every molecule regardless, so an interrupted screen can be resumed.
 
     Returns
     -------
@@ -210,8 +253,16 @@ def screen(
     jobs = _load_jobs(source, charge, spin)
     root = Path(directory)
 
+    completed = _load_completed(out) if resume else {}
+
     results = []
+    csv_path = json_path = None
     for job in jobs:
+        if job.name in completed:
+            logger.info(f"Skipping {job.name} (already completed)")
+            results.append(completed[job.name])
+            csv_path, json_path = _write_results(results, out)
+            continue
         record = {
             "name": job.name,
             "path": str(job.path),
@@ -267,7 +318,8 @@ def screen(
             # we must keep screening the remaining molecules
             logger.warning(f"Screening failed for {job.name}: {exc}")
         results.append(record)
+        # write after every molecule so an interrupted run stays resumable
+        csv_path, json_path = _write_results(results, out)
 
-    csv_path, json_path = _write_results(results, out)
     logger.info(f"Wrote {csv_path} and {json_path}")
     return results

tests/thermo/test_screening.py

@@ -174,6 +174,70 @@ def fake_dftb(*args, **kwargs):
     assert captured["method"] == "GFN1-xTB"
 
 
+def test_screen_resume_skips_completed_and_reruns_failed(monkeypatch, tmp_path):
+    from pathlib import Path
+
+    _write_xyz(tmp_path / "mol_a.xyz")
+    _write_xyz(tmp_path / "mol_b.xyz")
+    out = tmp_path / "out"
+
+    calls = []
+
+    def thermo_fail_b(atoms, directory=None, **kwargs):
+        name = Path(directory).name
+        calls.append(name)
+        if name == "mol_b":
+            raise RuntimeError("boom")
+        return _FakeThermo()
+
+    # first run: mol_a ok, mol_b fails
+    monkeypatch.setattr(screening, "dftbplus_thermo", thermo_fail_b)
+    first = screening.screen(str(tmp_path), out=str(out), directory=str(tmp_path / "r1"))
+    assert {r["name"]: r["status"] for r in first} == {"mol_a": "ok", "mol_b": "error"}
+
+    # resume: mol_a (ok) skipped, only mol_b (was error) re-run
+    calls.clear()
+    monkeypatch.setattr(
+        screening, "dftbplus_thermo",
+        lambda atoms, directory=None, **kw: (calls.append(Path(directory).name), _FakeThermo())[1],
+    )
+    second = screening.screen(
+        str(tmp_path), out=str(out), directory=str(tmp_path / "r2"), resume=True
+    )
+    assert calls == ["mol_b"]  # only the previously-failed molecule re-run
+    assert {r["name"]: r["status"] for r in second} == {"mol_a": "ok", "mol_b": "ok"}
+
+
+def test_load_completed_handles_missing_and_corrupt(tmp_path):
+    # no prior file -> nothing to resume
+    assert screening._load_completed(str(tmp_path / "nope")) == {}
+    # unreadable/corrupt json -> empty, does not crash
+    (tmp_path / "bad.json").write_text("{ not valid json", encoding="utf-8")
+    assert screening._load_completed(str(tmp_path / "bad")) == {}
+    # valid JSON but not a list of records -> empty, does not crash
+    (tmp_path / "scalar.json").write_text("42", encoding="utf-8")
+    assert screening._load_completed(str(tmp_path / "scalar")) == {}
+
+
+def test_screen_writes_results_incrementally(monkeypatch, tmp_path):
+    for name in ("mol_a", "mol_b", "mol_c"):
+        _write_xyz(tmp_path / f"{name}.xyz")
+
+    writes = []
+    real_write = screening._write_results
+
+    def spy_write(results, out):
+        writes.append(len(results))
+        return real_write(results, out)
+
+    monkeypatch.setattr(screening, "_write_results", spy_write)
+    monkeypatch.setattr(screening, "dftbplus_thermo", lambda atoms, **kw: _FakeThermo())
+
+    screening.screen(str(tmp_path), out=str(tmp_path / "out"), directory=str(tmp_path / "r"))
+    # written after each molecule with a growing result set (durable/resumable)
+    assert writes == [1, 2, 3]
+
+
 def test_screen_dispatches_to_xtb_cli_engine(monkeypatch, tmp_path):
     _write_xyz(tmp_path / "mol.xyz")
 
@@ -345,6 +409,9 @@ def test_cli_parse_args_routes_screen():
     cli_args = cli.parse_args(["screen", "molecules.csv", "--engine", "xtb-cli"])
     assert cli_args.engine == "xtb-cli"
 
+    assert args.resume is False  # default
+    assert cli.parse_args(["screen", "molecules.csv", "--resume"]).resume is True
+
 
 def test_cli_run_screen_returns_failure_count(monkeypatch):
     import ThermoScreening.cli.thermo as cli
@@ -357,6 +424,7 @@ def test_cli_run_screen_returns_failure_count(monkeypatch):
         source="x", out="res", charge=0.0, temperature=298.15,
         pressure=101325.0, directory="screening", parameter_set="3ob",
         solvent=None, quasi_rrho=False, engine="dftb+", method="GFN2-xTB",
+        resume=False,
     )
 
     assert cli.run_screen(args) == 1  # one molecule failed