Stop MD before a non-finite state reaches a frame or a checkpoint

docs/exact-restart.md

@@ -87,7 +87,10 @@ run keeps a single runtime checkpoint.
 `state.json` records the phase, generation, global counters, pending pathway occurrence IDs,
 detector state, configuration, any durable-state failure, and the active checkpoint. Each
 generation writes one frame per observation to `segments/<generation>/trajectory.traj`; a frame is
-written before its checkpoint, so an interrupted attempt never leaves a gap.
+written before its checkpoint, so an interrupted attempt never leaves a gap. An MD state with a
+non-finite position, momentum, cell, or integrator array gets neither: the job stops with a non-zero
+exit, and the next job repeats from the last checkpoint, or starts over if there is none. A frame or
+checkpoint an earlier version wrote from such a state is not repaired.
 
 ```python
 run = ReactionRun.create("trajectory-000", config=config)
@@ -108,7 +111,8 @@ summary = run.run(
 ```
 
 A run interrupted before its first checkpoint reopens in phase `new`; pass the initial atoms to
-`run()` again, as the CLI does.
+`run()` again, as the CLI does. After `run()` raises, reopen the run before calling it again, because
+the object that raised can hold state ahead of its checkpoint.
 
 The runner refuses a provider whose step counter or atomic state disagrees with the durable
 boundary; there is no inexact fallback. The built-in model providers are optional and

src/reactionflow/run.py

@@ -108,6 +108,23 @@ def _same_atomic_state(first: Atoms, second: Atoms) -> bool:
     )
 
 
+def _require_finite(snapshot: ExactRestartSnapshot) -> None:
+    """Refuse a non-finite MD state, so that it reaches neither a frame nor a checkpoint.
+
+    The integrator's arrays count too: NaN forces or stress on the first evaluation, or NaN stress
+    on the last step of an interval, leave the atoms finite and would poison every later step.
+    """
+
+    atoms = snapshot.atoms
+    arrays = [atoms.positions, atoms.get_momenta(), atoms.cell.array]
+    arrays += [a for a in snapshot.dynamics.arrays.values() if np.issubdtype(a.dtype, np.inexact)]
+    if not all(np.isfinite(array).all() for array in arrays):
+        raise ValueError(
+            "MD state is non-finite, usually from NaN or inf in the model's output or a blown-up "
+            "trajectory; nothing was saved for it, so the next job repeats it"
+        )
+
+
 def _write_result(directory: Path, outcome: PathwayOutcome, identity: Mapping[str, Any]) -> None:
     with Trajectory(directory / "images.traj", "w") as trajectory:
         for image in outcome.images:
@@ -381,6 +398,7 @@ def _run_segment(self, total_steps: int, runtime_provider: ExactRuntimeProvider)
         )
         with manager as runtime:
             if self._snapshot is None:
+                _require_finite(runtime.snapshot())
                 self._checkpoint(runtime)
             elif int(runtime.nsteps) != self._global_step:
                 raise ValueError("exact runtime step counter does not match the durable run")
@@ -403,6 +421,7 @@ def _run_segment(self, total_steps: int, runtime_provider: ExactRuntimeProvider)
                     steps = min(self.config.observation_interval, total_steps - self._global_step)
                     before = int(runtime.nsteps)
                     runtime.run(steps)
+                    _require_finite(runtime.snapshot())
                     advanced = int(runtime.nsteps) - before
                     if advanced != steps:
                         raise RuntimeError(

tests/test_model_resume.py

@@ -1,23 +1,25 @@
 from __future__ import annotations
 
 import json
-from contextlib import nullcontext
+from contextlib import contextmanager, nullcontext
 from typing import ClassVar
 
 import numpy as np
 import pytest
 from ase import Atoms
 from ase.calculators.calculator import Calculator, all_changes
+from ase.io import read
 
 from reactionflow import (
     BondDetectorConfig,
+    ComponentState,
     ExactRestartSnapshot,
     PathwayOutcome,
     ReactionRun,
     ReactionRunConfig,
 )
 from reactionflow.adapters import _torch, uma
-from reactionflow.adapters.ase import _ASERuntime
+from reactionflow.adapters.ase import ASELangevinBAOABAdapter, _ASERuntime
 from reactionflow.campaign import TrajectorySpec
 
 
@@ -185,3 +187,81 @@ def refine(candidate, **kwargs):
     assert record["step"] == 1
     old, new = record["changed"]["package_versions"]
     assert set(old.values()) == {"test-version"} and set(new.values()) == {"newer-version"}
+
+
+class NaNFrom(Harmonic):
+    """Turns one output NaN from evaluation ``at`` on, and does not mind NaN input."""
+
+    def __init__(self, at=None, output="forces"):
+        super().__init__(1.0)
+        self.at, self.output, self.evaluations = at, output, 0
+
+    def calculate(self, atoms=None, properties=("energy",), system_changes=all_changes):
+        super().calculate(atoms, properties, system_changes)
+        self.evaluations += 1
+        if self.at is not None and self.evaluations >= self.at:
+            self.results[self.output] = np.full_like(self.results[self.output], np.nan)
+
+
+class LeasedModel(ASELangevinBAOABAdapter):
+    def __init__(self, pressure_GPa, model):
+        super().__init__(
+            trajectory=TrajectorySpec(
+                id="t",
+                total_steps=8,
+                timestep_fs=0.1,
+                temperature_K=300,
+                pressure_GPa=pressure_GPa,
+                seed=19,
+            )
+        )
+        self.model = model
+
+    @contextmanager
+    def _lease(self):
+        yield self.model, ComponentState(kind="test.model")
+
+
+@pytest.mark.parametrize(
+    ("pressure_GPa", "output"), [(None, "forces"), (0.0, "forces"), (0.0, "stress")]
+)
+@pytest.mark.parametrize(
+    ("at", "saved_step"),
+    [(1, 0), (5, 0), (7, 4)],
+    ids=["first_evaluation", "last_of_an_interval", "inside_an_interval"],
+)
+def test_non_finite_md_state_stops_the_job_before_any_frame_or_checkpoint(
+    tmp_path, pressure_GPa, output, at, saved_step
+):
+    atoms = Atoms("H2", positions=[[0, 0, 0], [0.7, 0, 0]], cell=[9, 9, 9], pbc=True)
+    config = ReactionRunConfig(
+        observation_interval=4,
+        detector=BondDetectorConfig(pair_thresholds={"H-H": (0.8, 1.2)}),
+    )
+
+    def run_md(root, model):
+        run = ReactionRun.open(root) if root.exists() else ReactionRun.create(root, config=config)
+        adapter = LeasedModel(pressure_GPa, model)
+        run.run(
+            atoms if run.phase == "new" else None,
+            runtime_provider=adapter,
+            pathway_calculator_provider=adapter.calculator,
+            total_steps=8,
+        )
+        return run
+
+    run_md(tmp_path / "control", NaNFrom())
+    root = tmp_path / "run"
+    with pytest.raises(ValueError, match="non-finite"):
+        run_md(root, NaNFrom(at, output))
+    state = json.loads((root / "state.json").read_text())
+    assert (state["phase"], state["global_step"]) == ("new" if at == 1 else "running", saved_step)
+
+    # The next job, on a good model, gives exactly the run that never met the NaN.
+    assert run_md(root, NaNFrom()).phase == "completed"
+    frames = read(root / "segments/0000/trajectory.traj", ":")
+    control = read(tmp_path / "control/segments/0000/trajectory.traj", ":")
+    assert len(frames) == len(control) == 3
+    for frame, expected in zip(frames, control, strict=True):
+        assert np.array_equal(frame.positions, expected.positions)
+        assert np.array_equal(frame.get_momenta(), expected.get_momenta())