Skip to content

Refine a trajectory's barriers after its MD when the campaign asks for it - #82

Merged
Austin243 merged 6 commits into
mainfrom
refine-after-md
Nov 5, 2026
Merged

Austin243 merged 6 commits into
mainfrom
refine-after-md

Conversation

@Austin243

@Austin243 Austin243 commented Nov 5, 2026 •

Copy link
Copy Markdown
Owner

Before: a trajectory stops MD at each confirmed reaction, refines it, and then continues, and nothing turns that off. A slow refinement holds up its trajectory's MD. In the 288-atom acetonitrile T-P series, one event held a trajectory for 12.7 hours and another trajectory made about 3 ps of MD in 17 hours. Running MD alone later took a subclass that overrides private methods.

After: with "refine_after_md": true at the top of a campaign, MD runs straight through to total_steps in one segment. Confirmed reactions wait in the queue state.json already keeps, and the same process refines them when MD ends, in the order MD found them. A job that ends partway resumes as usual. With the option off, results are unchanged.

The key is not part of the trajectory contract, so trajectories already running keep their hash and can resume, and it can change on resubmission. A trajectory already paused for refinement finishes that refinement first. status needs no change: the phase is running with events counting up during MD, then refining.

A queue can hold many events and results are never overwritten, so three failures are handled:

  • The model is lost. After a refinement that gives no barrier, the last checkpoint is restored, which loads and evaluates the model. If that raises or gives NaN forces, the job exits non-zero with the rest of the queue intact, and the next job refines it. A lost GPU costs one event, not the whole queue. An event that fails on its own does not stop the queue.
  • MD cannot go on. The model raises or returns NaN at some geometry, so the job never reaches MD's end. Set the option to false and resubmit: the queue is refined before MD is touched, and MD then stops as it did.
  • The run is marked failed during MD or the drain (a locked database, a failed publish). A resubmission refines the queue, then reports the first failure again.

Three small changes also reach the option-off path. A run marked failed with reactions still queued refines them on resubmission, where main raises at once. An event with no barrier and more events queued runs the model check before going on, which costs one model load and one evaluation and changes no result unless the model is gone. And that check refuses a model giving NaN forces, where main publishes every event left in that queue first and stops only if MD resumes.

How: _observe enters the refining phase only when the option is off or at MD's last step, and _checkpoint saves that phase together with that step. _refine_due decides whether the queue runs now: it waits only in the running phase with the option on. The existing model check on resume, moved into _check_model, still comes before any result is published. That is about 45 changed lines in run.py, plus the campaign key and one line in cli.py.

Tests: two, written before the code. One runs a campaign through to completion (one segment, MD's lease closed before refinement). The other interrupts during MD and before refinement (queue kept, each event refined once, model checked once, MD not repeated). 342 passed, 25 skipped; ruff is clean.

Checked outside the test suite, with real reactionflow run processes and BAOAB MD, where "main" is the commit before #84:

  • Stateless double-well potential (301 frames, 33 events, 29 validated saddles): frames, events, statuses, and barriers are identical to main with the option off, with it on, and with it on through 4 SIGKILLs. The frames match bitwise.
  • Every durable boundary interrupted by SIGKILL or KeyboardInterrupt, at the first, middle, and last call, with the option flipped both ways on resume: 1874 cases, none differ from main's uninterrupted run.
  • A stateless periodic double well at constant pressure under SSNEB, and a chain model with stepwise events: frames, events, statuses, and barriers are identical to main with the option absent, false, and true.
  • A lost GPU, simulated by a calculator that raises on every evaluation, and a deleted model file, with 11 events queued: the job exits 1 with one failed and 10 queued, another job on the bad GPU loses nothing more, and a job on a good one finishes with those 10 equal to main's.
  • The same stop, one event lost and the rest queued, when the loss starts in the Hessian, in the connectivity check, or inside a stepwise event; when the calculator gives NaN instead of raising, at constant volume and at constant pressure; and through TorchModelAdapter over a stub torch. Each model check costs one evaluation, measured at constant volume and constant pressure.
  • A model that raises or returns NaN at step 1930 of 3000: the 23 queued events wait through resubmissions, and option false gives main's 23 results before MD stops as before.
  • reactions.sqlite3 locked past its timeout during an observation, with no code patched: main keeps 5 results and fails, and here the resubmission gives those same 5 and reports the failure. A failed publish during the drain: all 11 results equal main's. A second failure during the drain leaves the first in state.json and status, and goes to that job's log and last-error.json.
  • A changed model file on these paths is refused before anything is refined, with no run file changed.
  • A trajectory started on main and killed resumes here with its contract unchanged and finishes with main's frames.
  • EMT, on acetonitrile and water with no metal, and on Cu with H at 1 GPa: events and statuses match. The acetonitrile frames differ by about 1e-13 Å, because EMT's neighbor list is internal state. Main's own frames differ by 1e-7 Å after a kill.

Not run: a real MLIP on a GPU. The lost GPU above is simulated, and a real one may hang or kill the process instead of raising; a killed process keeps the queue. Whether GPU memory is freed when MD's lease closes has not been measured. Calculators stay alive until garbage collection on main too.

Limits: the model check is one evaluation, so a fault that comes and goes can pass it and fail events across the queue. With one evaluation in 50 raising, 9 of 11 events ended failed and the job exited 0, where main stops at its next MD failure. The saved runs show no such fault: none of their 94 distinct results is failed. And a calculator that cannot evaluate the last checkpoint when freshly loaded, as an SCF without its restart guess might, fails the check and holds the queue with no way through. Simulated, with MD ending at such a checkpoint: exit 1 after 2 of 24 results, where main finishes all 24. A stateless model cannot reach this, because MD already evaluated that geometry. docs/exact-restart.md limits the generic adapter to stateless inference, though examples/xtb runs tblite through it; tblite was not run.

#84, now merged, stops MD before NaN or inf from the model can leave a bad frame or a checkpoint no later job can use. This merged with it: 351 passed, 25 skipped; ruff is clean.

Refinement stays serial, so a long queue keeps the job busy after MD ends. Merges with #83 in either order with no conflicts.

…eturns NaN

A lost GPU can end an event as saddle_unchecked or stepwise, and a model returning NaN as
relaxation_failed, so checking only after a failed status let the next event fail too or let the
whole queue finish with no results. Also keep the failure that stopped MD when a later one follows.
@Austin243
Austin243 marked this pull request as ready for review November 5, 2026 18:30
@Austin243
Austin243 merged commit a8d8453 into main Nov 5, 2026
4 checks passed
@Austin243
Austin243 deleted the refine-after-md branch November 5, 2026 18:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant