Add an Aurora Quantum ESPRESSO long-run cap-to-resume example - #164
Draft
reowszer wants to merge 2 commits into
Draft
Add an Aurora Quantum ESPRESSO long-run cap-to-resume example#164reowszer wants to merge 2 commits into
reowszer wants to merge 2 commits into
Conversation
Provide a runnable example that exercises the wall-clock cap and `chemgraph resume` path on real Quantum ESPRESSO on Aurora, covering both the calculator layer and the full agent stack. Two drivers: - qe_cap_driver.py drives the calculator layer directly: run under a cap until it fires at an ASE optimizer-step boundary, leave a resumable partial (.traj + restart + xyz), then resume from that partial to completion. - qe_agent_e2e.py drives the full agent stack end to end: run1 caps and records status='capped' with a pending step in run_manifest.json under the agent's own log_dir; a second agent resumes with resume_from=<session_id>, adopts the prior log_dir from the session DB, clears the pending step, and continues. Four PBS scripts submit these on Aurora: run_qe_cap.pbs / run_qe_mol.pbs for the calculator driver (Si bulk and a centered H2O molecule) and run_agent_e2e.pbs / run_agent_mol.pbs for the agent driver. The README explains the milestones each stage proves and the environment they need (pseudopotentials, the prebuilt pw.x launch, and the optional LLM proxy for firewalled compute nodes). The gitignore rules keep the tracked drivers, PBS scripts, and README while ignoring the per-run output directories these scripts create.
Document how the cap->resume mechanism validated in the example would extend into unattended, cross-allocation auto-resubmission through the DOE IRI Facility API, and record why that extension is purely additive on top of the shipped work. The note (IRI_INTEGRATION.md) is a design note, not shipped code. It maps the Facility API compute endpoints and JobSpec fields onto this example's PBS scripts, explains the two-layer split (within-allocation enforcement, which is implemented and validated here, versus the future across-allocation launcher), and gives three reasons the launcher is strictly additive: the manifest carries a schema_version and is read permissively so new fields never break an older reader; the cap's self-computed CHEMGRAPH_ALLOCATION_DEADLINE stays authoritative because JobStatus exposes no live seconds-remaining; and the join key is already session_id, which resume uses to adopt the prior log_dir. It also lists the live-facility prerequisites (a non-demo adapter, a separate iri-api Globus scope, a persistent login-node service) that keep it a deferred effort. The README gains a short "Future direction" section pointing at the note so a reader sees the current example as a deliberate first layer.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a runnable example under
examples/aurora_qe_longrun/that exercisesChemGraph's long-running-calculation support using real Quantum ESPRESSO DFT
on ALCF Aurora. A calculation that would exceed the job's wall-clock allocation
self-terminates at an ASE optimizer-step boundary, writes a durable partial
geometry and run manifest, and a later
chemgraph resumecontinues from thatpartial. This is the real-DFT, HPC counterpart to the in-process EMT/MACE unit
tests: it proves the same cap-to-resume seam end to end on a subprocess DFT engine.
This PR is example and documentation only. It adds no library code and changes no
shipped behavior; it is the validation harness for the calculator and
cap-to-resume PRs.
The example drives the seam at two layers, each over a periodic and a
non-periodic system:
qe_cap_driver.py.run_ase_corecalls QE directlywith no LLM. Proves the cap fires at a step boundary, writes a resumable partial,
and a resume continues to the same final energy.
qe_agent_e2e.py. A real LLM drives LangGraphsingle_agenttorun_asetorun_ase_coreto QE, plus the manifest hook andchemgraph resume. Covers the layers Route 1 bypasses: tool-arg unwrapping, JSONtool-message parsing,
tool_call_idcorrelation, clear-pending-on-success (M2),and log_dir adoption on resume (M3).
pbc=[T,T,T]) uses the configured mesh verbatim; an H2O molecule(cell-less) exercises the molecule path (
is_nonperiodicthencenter(vacuum)then
K_POINTS gamma), validating the calculator PR's molecule fix on livepw.x.IRI_INTEGRATION.mdis a design note for a future cross-allocation auto-resubmitlayer built on the DOE IRI Facility API. It is additive on top of the shipped
manifest and is not implemented here. Three properties keep it additive: the
manifest carries a
schema_versionand is read permissively, so older readers cancontinue to parse manifests that contain new fields; the cap's self-computed
deadline stays authoritative because the Facility API exposes no live
seconds-remaining; and the join key is already
session_id, which resume uses toadopt the prior
log_dir.Part of a 3-PR series
This is the third of three PRs that together add long-running-calculation support
for subprocess DFT. They are meant to land in order:
chemgraph resume.validates the seam on real QE DFT. Depends on both earlier PRs, so it opens as
Draft until they merge.
Related issues
None
Type of change
How was this tested?
The portable CI coverage of this seam lives in the repo unit tests
(
tests/test_cap*.py,tests/test_manifest.py,tests/test_allocation_cap.py,tests/test_resume_injection.py,tests/test_mace_cap.py) on the in-processEMT/MACE calculators and stays green. This directory is the HPC counterpart and
cannot run in CI (it needs Aurora plus QE plus an LLM endpoint); the drivers
assertevery invariant, so a non-zero exit signals a real regression.Live-run scope: these runs validate the flow (the cap-to-resume seam, and for H2O
the non-periodic Gamma/centering path on live
pw.x), not physical accuracy. Thegeometry band is a coarse pass/fail guard against a broken run.
partial, and every resume converged to the same energy as a single uncapped opt
(
-308.187965 eV), including a re-run under the langgraph 1.x dependency bump.stages passed. The generated
espresso.pwicarriesK_POINTS gammaplus afinite
CELL_PARAMETERS, and the resume converged to-473.2196 eVwith arelaxed geometry of
r(O-H) = 0.971 A, angle(H-O-H) = 104.4 deg, within thereference range observed across plane-wave PBE codes (0.96-1.00 A, 102-106 deg).
Full per-job numbers and porting notes are in the example README.
ruff check .passes.pytest tests/ -k "not tblite"passes (293 passed, 28skipped, 2 deselected).
Checklist
mainand targetsmainruff check .passespytest tests/ -k "not tblite"passes (plus extras tests if Academy/backends touched)assert.)