From 420ad256871a981423c2f1d0c1fcaee76febcb57 Mon Sep 17 00:00:00 2001 From: Austin Ellis Date: Thu, 5 Nov 2026 08:43:24 -0500 Subject: [PATCH 1/2] Stop MD before a non-finite state reaches a frame or a checkpoint --- docs/exact-restart.md | 5 ++- src/reactionflow/run.py | 19 +++++++++ tests/test_model_resume.py | 84 +++++++++++++++++++++++++++++++++++++- 3 files changed, 105 insertions(+), 3 deletions(-) diff --git a/docs/exact-restart.md b/docs/exact-restart.md index 7337dc2..b0fc6be 100644 --- a/docs/exact-restart.md +++ b/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//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) diff --git a/src/reactionflow/run.py b/src/reactionflow/run.py index f4f7c9a..82d11af 100644 --- a/src/reactionflow/run.py +++ b/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 += snapshot.dynamics.arrays.values() + if not all(np.isfinite(array).all() for array in arrays): + raise ValueError( + "MD state is non-finite, 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( diff --git a/tests/test_model_resume.py b/tests/test_model_resume.py index e8d7f3f..7315d7d 100644 --- a/tests/test_model_resume.py +++ b/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()) From 3fcaab64115aabbc3c11c91af347dc740125d327 Mon Sep 17 00:00:00 2001 From: Austin Ellis Date: Thu, 5 Nov 2026 10:43:12 -0500 Subject: [PATCH 2/2] Check only inexact integrator arrays, and say to reopen a run after run() raises --- docs/exact-restart.md | 3 ++- src/reactionflow/run.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/docs/exact-restart.md b/docs/exact-restart.md index b0fc6be..1346c80 100644 --- a/docs/exact-restart.md +++ b/docs/exact-restart.md @@ -111,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 diff --git a/src/reactionflow/run.py b/src/reactionflow/run.py index 82d11af..b82af45 100644 --- a/src/reactionflow/run.py +++ b/src/reactionflow/run.py @@ -117,10 +117,10 @@ def _require_finite(snapshot: ExactRestartSnapshot) -> None: atoms = snapshot.atoms arrays = [atoms.positions, atoms.get_momenta(), atoms.cell.array] - arrays += snapshot.dynamics.arrays.values() + 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, from NaN or inf in the model's output or a blown-up " + "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" )