Skip to content

Add an Aurora Quantum ESPRESSO long-run cap-to-resume example - #164

Draft
reowszer wants to merge 2 commits into
argonne-lcf:mainfrom
reowszer:feature/aurora-qe-example
Draft

Add an Aurora Quantum ESPRESSO long-run cap-to-resume example#164
reowszer wants to merge 2 commits into
argonne-lcf:mainfrom
reowszer:feature/aurora-qe-example

Conversation

@reowszer

Copy link
Copy Markdown
Contributor

Summary

Adds a runnable example under examples/aurora_qe_longrun/ that exercises
ChemGraph'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 resume continues from that
partial. 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:

  • Route 1 (calc layer), qe_cap_driver.py. run_ase_core calls QE directly
    with no LLM. Proves the cap fires at a step boundary, writes a resumable partial,
    and a resume continues to the same final energy.
  • Route 2 (full agent), qe_agent_e2e.py. A real LLM drives LangGraph
    single_agent to run_ase to run_ase_core to QE, plus the manifest hook and
    chemgraph resume. Covers the layers Route 1 bypasses: tool-arg unwrapping, JSON
    tool-message parsing, tool_call_id correlation, clear-pending-on-success (M2),
    and log_dir adoption on resume (M3).
  • Bulk Si (pbc=[T,T,T]) uses the configured mesh verbatim; an H2O molecule
    (cell-less) exercises the molecule path (is_nonperiodic then center(vacuum)
    then K_POINTS gamma), validating the calculator PR's molecule fix on live
    pw.x.

IRI_INTEGRATION.md is a design note for a future cross-allocation auto-resubmit
layer 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_version and is read permissively, so older readers can
continue 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 to
adopt 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:

  1. feature/dft-calculators: the VASP and Quantum ESPRESSO calculators.
  2. feature/longrun-cap-resume: the wall-clock cap, run manifest, and chemgraph resume.
  3. feature/aurora-qe-example (this PR): a runnable Aurora example that
    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

  • Bug fix
  • New feature
  • Docs
  • Chore / refactor / CI

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-process
EMT/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
assert every 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. The
geometry band is a coarse pass/fail guard against a broken run.

  • Bulk Si, Route 1 and Route 2: cap fired at a step boundary, left a resumable
    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.
  • H2O molecule, Route 1 and Route 2: all molecule
    stages passed. The generated espresso.pwi carries K_POINTS gamma plus a
    finite CELL_PARAMETERS, and the resume converged to -473.2196 eV with a
    relaxed geometry of r(O-H) = 0.971 A, angle(H-O-H) = 104.4 deg, within the
    reference 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, 28
    skipped, 2 deselected).

Checklist

  • Branched off the latest main and targets main
  • PR is focused on a single logical change (split if it grew large)
  • ruff check . passes
  • pytest tests/ -k "not tblite" passes (plus extras tests if Academy/backends touched)
  • Added/updated tests for the change (n/a: example only; the CI unit tests for this seam ship in the cap-to-resume PR. The drivers self-check via assert.)
  • Updated docs / README for any user-facing change

reowszer added 2 commits July 31, 2026 21:36
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.
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