diff --git a/quantui/backends/worker.py b/quantui/backends/worker.py index 0144556..5325ae2 100644 --- a/quantui/backends/worker.py +++ b/quantui/backends/worker.py @@ -378,12 +378,39 @@ def _run_frequency( options = request.options or {} scf_rescue = bool(options.get("scf_rescue", True)) _write_progress(staging_dir, "running", "Running frequency analysis", 15.0) + + # M-CHECKPOINT CHK.4 — this is the calc type the real production pain + # (roadmap 34's "Real-world cost data" note) was about: a killed + # frequency job used to discard the entire 6N-displacement Hessian + # computation on every resubmission, unlike geometry_opt/pes_scan + # (CL2.8) on the same batch path. Resume diffs a *set* of already-banked + # displacement ids (CHK.4.1/.2), not a prefix — see freq_calc.py and + # freq_ir_workers.py/freq_raman_workers.py for why that matters under + # the QUANTUI_FREQ_PARALLEL opt-in. + ckpt, resumable = _begin_worker_checkpoint( + molecule, + calc_type="frequency", + method=request.method, + basis=request.basis, + staging_dir=staging_dir, + log_stream=log_stream, + ) + if resumable: + _append_log( + staging_dir, + "[checkpoint] Resuming frequency analysis — some " + "finite-difference displacement SCFs are already banked from a " + "previous attempt.", + ) + result = run_freq_calc( molecule=molecule, method=request.method, basis=request.basis, progress_stream=log_stream, scf_rescue=scf_rescue, + checkpoint=ckpt, + resume=resumable, ) return result, molecule diff --git a/quantui/checkpoint.py b/quantui/checkpoint.py index cccae4b..691722b 100644 --- a/quantui/checkpoint.py +++ b/quantui/checkpoint.py @@ -30,6 +30,9 @@ opt.traj ASE trajectory, appended per step (CHK.2) opt.restart BFGS Hessian state (CHK.2) points.jsonl one line per completed scan point (CHK.3) + items//.json one file per completed named-set item + (CHK.4.1) — e.g. frequency's per- + displacement SCF results Checkpoints live outside the results directory on purpose: a result directory is created when a calculation *succeeds*, and the runs that most need a @@ -51,6 +54,7 @@ import os import shutil import time +import uuid from dataclasses import dataclass from pathlib import Path from typing import Any, Iterable, Optional, Sequence @@ -200,6 +204,67 @@ def describe(self) -> str: # --------------------------------------------------------------------------- +def mark_item_done_at(items_root: Path, item_id: str, payload: dict) -> bool: + """Atomically write one named-item-set record — by directory path alone. + + The path-only sibling of :meth:`Checkpoint.mark_item_done`, for callers + that have only a directory path, not a full ``Checkpoint`` object — the + exact situation a ``ProcessPoolExecutor`` worker is in: it receives the + checkpoint's item-set directory as a plain string through ``initargs`` + (CHK.4.4), and a ``Checkpoint`` with a live log stream attached is not + something safe to pickle across that process boundary. + + Safe under concurrent callers writing *different* item ids: each item + gets its own file, written via temp-file-then-``os.replace`` (atomic on + POSIX), so there is nothing to lock and no shared file to race on. The + temp name includes the writer's pid plus a random token and is always + suffixed ``.tmp`` (never ``.json``), so it can never be mistaken for a + completed item even if two writers raced on the same id. + + Returns ``True`` on success, ``False`` on any failure — the caller + decides what (if anything) to log; this never raises. + """ + try: + root = Path(items_root) + root.mkdir(parents=True, exist_ok=True) + target = root / f"{item_id}.json" + tmp = root / f"{item_id}.{os.getpid()}.{uuid.uuid4().hex[:8]}.tmp" + tmp.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8") + os.replace(tmp, target) + except Exception as exc: # noqa: BLE001 — never break the calculation + logger.debug( + "checkpoint item write failed for %s/%s: %s", items_root, item_id, exc + ) + return False + return True + + +def completed_items_at(items_root: Path) -> dict: + """Read every complete named-item-set record — by directory path alone. + + The path-only sibling of :meth:`Checkpoint.completed_items`. Skips any + file that fails to parse (a corrupt item should cost re-running that one + item, not the whole resume) and any ``.tmp`` leftover from a writer that + died between the write and the rename. + """ + out: dict = {} + try: + entries = list(Path(items_root).iterdir()) + except OSError: + return out + for entry in entries: + if entry.suffix != ".json": + continue + try: + record = json.loads(entry.read_text(encoding="utf-8")) + except (OSError, UnicodeDecodeError, json.JSONDecodeError): + logger.debug("discarding corrupt checkpoint item %s", entry) + continue + if isinstance(record, dict): + out[entry.stem] = record + return out + + def _atomic_write_json(path: Path, payload: dict) -> None: """Write *payload* to *path* so a crash can't leave a half-file. @@ -460,6 +525,8 @@ def has_progress(self) -> bool: return True except OSError: pass + if self._has_item_set_progress(): + return True return self._has_leg_progress() def _has_leg_progress(self) -> bool: @@ -556,6 +623,78 @@ def completed_points(self) -> list[dict]: points.append(record) return points + # ── Named item sets (CHK.4.1) ─────────────────────────────────────────── + # + # CHK.3's ``points.jsonl`` assumes completion is a *prefix*: point 0, then + # 1, then 2, one writer, one file, appended in order. That breaks down for + # a frequency job's displaced-geometry SCFs, which `freq_ir_workers.py` + # can run **concurrently** across worker processes — completion order is + # non-deterministic, so resume has to diff a *set* of finished item ids + # against the full required set, not trim a prefix. + # + # The design that makes this safe under concurrent writers: one file per + # item, named by its id, written via temp-file + ``os.replace`` — never a + # shared file. Two displacements never share a filename (each is assigned + # to exactly one worker), so there is nothing to lock, and a worker dying + # mid-write never leaves a *renamed* (i.e. visible-as-done) file behind — + # that item simply stays "not yet complete" on the next resume. This is + # deliberately more general than "frequency displacements": *name* groups + # unrelated item sets under one checkpoint (e.g. IR vs. Raman + # displacements) so they can never collide with each other. + + def items_dir(self, name: str) -> Path: + """Directory holding one named set of per-item checkpoint records.""" + return self.dir / "items" / name + + def mark_item_done(self, name: str, item_id: str, payload: dict) -> None: + """Atomically record one item of the named set *name* as complete. + + Safe to call from any process that knows this checkpoint's root path + — including a ``ProcessPoolExecutor`` worker — because each item gets + its own file and the write is temp-file-then-rename. Delegates to + :func:`mark_item_done_at`, which a worker process can also call + directly given just the directory path (see its docstring) — a + ``Checkpoint`` with a live log stream attached is not something to + pickle across a process boundary. + """ + if mark_item_done_at(self.items_dir(name), item_id, payload): + self._log(f"saved — {name} item {item_id}") + + def completed_items(self, name: str) -> dict: + """Return every complete item's stored payload, keyed by item id. + + Skips any file that fails to parse rather than raising — the same + defensive posture as :meth:`completed_points`. A corrupt single item + should cost re-running that one item, not the whole resume. Delegates + to :func:`completed_items_at`. + """ + return completed_items_at(self.items_dir(name)) + + def completed_item_ids(self, name: str) -> set: + """Return the ids of every complete item in the named set *name*. + + Derived from :meth:`completed_items` rather than merely "a ``.json`` + file with this name exists" — the resume-time question is "do we + have a *usable* result for this item?", not "does this filename + exist?". A corrupted item file (e.g. a crash that landed between the + temp write and the rename leaving a truncated ``.json`` — which + should not happen given the atomic-rename contract, but "never break + a calculation" means not trusting that) is therefore never mistaken + for a completed one, and simply gets recomputed on resume. + """ + return set(self.completed_items(name).keys()) + + def _has_item_set_progress(self) -> bool: + """True when any named item set (see above) has a completed item.""" + try: + names = list((self.dir / "items").iterdir()) + except OSError: + return False + for name_dir in names: + if name_dir.is_dir() and self.completed_items(name_dir.name): + return True + return False + # --------------------------------------------------------------------------- # Discovery + retention diff --git a/quantui/freq_calc.py b/quantui/freq_calc.py index 2bbd278..6b9c282 100644 --- a/quantui/freq_calc.py +++ b/quantui/freq_calc.py @@ -283,6 +283,8 @@ def run_freq_calc( basis: str = "STO-3G", progress_stream: Optional[IO[str]] = None, scf_rescue: bool = True, + checkpoint: Optional[Any] = None, + resume: bool = False, ) -> FreqResult: """Run SCF + analytical Hessian to obtain vibrational frequencies. @@ -306,6 +308,19 @@ def run_freq_calc( intensities) automatically retries through the shared rescue helper on non-convergence (M-SCF-ROBUST, see :mod:`quantui.scf_robust`). Default ``True``. + checkpoint: Optional :class:`~quantui.checkpoint.Checkpoint` + (M-CHECKPOINT CHK.4). When given, each completed + finite-difference displacement SCF (the ``6N``-solve numerical + IR-intensity step) is durably recorded — see + :mod:`quantui.freq_displacement_ids` — so a resumed run can skip + displacements already banked from an earlier attempt, whether + the calc ran serially or through + :mod:`quantui.freq_ir_workers`'s parallel worker pool. + resume: Skip displacements already completed in *checkpoint* rather + than recomputing every one from scratch. Has no effect if + *checkpoint* is ``None`` or has no banked displacements — the + calc still runs, it just starts from nothing, same as if resume + were never requested. Returns: :class:`FreqResult` with frequencies, ZPVE, and SCF properties. @@ -355,6 +370,8 @@ def run_freq_calc( basis=basis, progress_stream=progress_stream, scf_rescue=scf_rescue, + checkpoint=checkpoint, + resume=resume, _dft=dft, _gto=gto, _scf=scf, @@ -370,6 +387,8 @@ def _run_freq_calc_body( basis: str, progress_stream: Optional[IO[str]], scf_rescue: bool = True, + checkpoint: Optional[Any] = None, + resume: bool = False, _dft: Any, _gto: Any, _scf: Any, @@ -573,13 +592,16 @@ def _status(msg: str) -> None: try: from .config import BOHR_TO_ANGSTROM as _BOHR_TO_ANG + from .freq_displacement_ids import ( + displacement_id, + parse_displacement_id, + ) _DELTA = 0.01 # Bohr _KM_MOL_FAC = 42.255 # (D/Å)²/amu → km/mol _n_ir = mol.natm _ir_total_solves = _n_ir * 3 * 2 - _ir_done_solves = 0 _coords0 = mol.atom_coords().copy() _dpdx = _np_ir.zeros((_n_ir * 3, 3)) _xc = getattr(mf, "xc", None) @@ -596,6 +618,46 @@ def _status(msg: str) -> None: # raised a shape-mismatch ValueError inside PySCF and # silently dropped IR intensities for the whole calc (caught # by the broad except below). + + # --- M-CHECKPOINT CHK.4: resume already-banked displacements --- + # ``_dipoles`` is the single source of truth for every + # displacement's result, whichever path produced it (resumed + # from an earlier attempt, computed in this run's parallel + # pool, or computed in this run's serial fallback) — the + # final dpdx assembly below reads only from this dict, never + # from a path-specific variable, which is what makes the + # assembly step itself a proper CHK.4.5 gate: it only runs + # once every required id is present here, regardless of how + # it got there. + _ITEM_SET_NAME = "freq_displacements" + _dipoles: dict = {} + if checkpoint is not None and resume: + for _item_id, _payload in checkpoint.completed_items( + _ITEM_SET_NAME + ).items(): + try: + _key = parse_displacement_id(_item_id) + _dipoles[_key] = _np_ir.asarray( + _payload["dipole"], dtype=float + ) + except (ValueError, KeyError, TypeError): + # Malformed/foreign record — never trust it, just + # recompute this one displacement (checkpoint.py's + # "never break a calculation" rule applies here + # too: a bad record costs one displacement, not + # the whole resume). + continue + _ir_done_solves = len(_dipoles) + if _ir_done_solves and checkpoint is not None: + try: + checkpoint.log_resumed( + f"{_ir_done_solves}/{_ir_total_solves} finite-difference " + "displacement SCFs already banked from an earlier attempt" + ) + except ( + Exception + ): # noqa: BLE001 — provenance is never worth a crash + pass _status( "Numerical IR intensities: " f"{_ir_done_solves}/{_ir_total_solves} finite-difference displacement SCFs done (6 per atom) " @@ -655,11 +717,35 @@ def _displaced_scf_dipole() -> _np_ir.ndarray: displacement_count=_ir_total_solves, ) + # Checkpoint items directory as a plain string, for workers + # (CHK.4.4) — a ``Checkpoint`` with a live log stream attached + # is not something safe to pickle across the + # ``ProcessPoolExecutor`` process boundary, so workers get + # only the directory path and write via + # ``checkpoint.mark_item_done_at`` (see freq_ir_workers.py). + _ckpt_items_dir = ( + str(checkpoint.items_dir(_ITEM_SET_NAME)) + if checkpoint is not None + else None + ) + _mol_v = mol.verbose mol.verbose = 0 _parallel_failed = False try: - if _use_parallel: + # Build the task list once, up front, skipping anything + # already in ``_dipoles`` (resumed from an earlier + # attempt) — shared between the parallel and (if it + # falls back) serial paths below, so neither ever + # recomputes a displacement the other already has. + _remaining: list[tuple[int, int, int]] = [ + (_I, _ax, _sign) + for _I in range(_n_ir) + for _ax in range(3) + for _sign in (1, -1) + if (_I, _ax, _sign) not in _dipoles + ] + if _use_parallel and _remaining: try: # Stash dm0 once on disk so workers can map-load it # via initargs (avoids per-task pickling). @@ -669,23 +755,19 @@ def _displaced_scf_dipole() -> _np_ir.ndarray: import tempfile as _tempfile _n_workers = _ir_par.pick_worker_count( - _cpu_count, _ir_total_solves + _cpu_count, len(_remaining) ) _threads_each = _ir_par.threads_per_worker( _cpu_count, _n_workers ) - # Build all 6N task arguments first; pickling-safe - # flat lists per-displacement. + # Pickling-safe flat lists per-displacement, one + # task per still-required id. _tasks: list[tuple[int, int, int, list[float]]] = [] - for _I in range(_n_ir): - for _ax in range(3): - _cp = _coords0.copy() - _cp[_I, _ax] += _DELTA - _tasks.append((_I, _ax, +1, _cp.flatten().tolist())) - _cm = _coords0.copy() - _cm[_I, _ax] -= _DELTA - _tasks.append((_I, _ax, -1, _cm.flatten().tolist())) + for _I, _ax, _sign in _remaining: + _c = _coords0.copy() + _c[_I, _ax] += _DELTA * _sign + _tasks.append((_I, _ax, _sign, _c.flatten().tolist())) _dm0_handle = _tempfile.NamedTemporaryFile( delete=False, suffix=".dm0.pkl" @@ -711,19 +793,28 @@ def _displaced_scf_dipole() -> _np_ir.ndarray: _xc, _dm0_handle.name, _threads_each, + _ckpt_items_dir, ), ) as _pool: # Submit all and store futures keyed by task # index so we can assemble +/- per (I, ax). + # Each task also carries its own + # displacement id (CHK.4.2) so the worker + # can durably record it (CHK.4.4) the + # moment it finishes — before this parent + # process ever calls ``.result()``, so a + # parent crash mid-wave never loses a + # sibling that already completed. _futs = { _pool.submit( - _ir_par.run_displaced_scf, _task[3] + _ir_par.run_displaced_scf, + displacement_id( + _task[0], _task[1], _task[2] + ), + _task[3], ): _task for _task in _tasks } - # Accumulate results into a temporary map - # ``(I, ax, sign) -> dipole_array``. - _dipoles: dict = {} for _fut in _cf.as_completed(_futs): _I, _ax, _sign, _coords_done = _futs[_fut] _dipoles[(_I, _ax, _sign)] = _fut.result() @@ -741,13 +832,6 @@ def _displaced_scf_dipole() -> _np_ir.ndarray: os.unlink(_dm0_handle.name) except OSError: pass - - # Assemble dpdx now that all dipoles are in hand. - for _I in range(_n_ir): - for _ax in range(3): - _mu_p = _dipoles[(_I, _ax, +1)] - _mu_m = _dipoles[(_I, _ax, -1)] - _dpdx[3 * _I + _ax] = (_mu_p - _mu_m) / (2 * _DELTA) except Exception as _par_exc: logger.warning( "Parallel IR-intensity computation failed (%s); falling back to serial.", @@ -757,42 +841,54 @@ def _displaced_scf_dipole() -> _np_ir.ndarray: "Parallel IR intensities failed; falling back to serial computation." ) _parallel_failed = True - # Reset so the serial loop's progress messages - # below start clean rather than continuing from - # wherever the failed parallel attempt left off. - _ir_done_solves = 0 + # Reflect whatever is actually in ``_dipoles`` — + # resumed items, plus anything the parallel pool + # completed (and durably banked, worker-side) + # before it hit the exception — rather than + # discarding that progress. The serial fallback + # below skips anything already in ``_dipoles``, + # so this is never recomputed. + _ir_done_solves = len(_dipoles) if not _use_parallel or _parallel_failed: - for _I in range(_n_ir): - for _ax in range(3): - # +Δ displacement - _cp = _coords0.copy() - _cp[_I, _ax] += _DELTA - mol.set_geom_(_cp, unit="Bohr") - _mu_p = _displaced_scf_dipole() - _ir_done_solves += 1 - _status( - "Numerical IR intensities: " - f"{_ir_done_solves}/{_ir_total_solves} " - "finite-difference displacement SCFs done (6 per atom) " - f"({_ir_total_solves - _ir_done_solves} " - "remaining)" - ) - - # -Δ displacement - _cm = _coords0.copy() - _cm[_I, _ax] -= _DELTA - mol.set_geom_(_cm, unit="Bohr") - _mu_m = _displaced_scf_dipole() - _ir_done_solves += 1 - _status( - "Numerical IR intensities: " - f"{_ir_done_solves}/{_ir_total_solves} " - "finite-difference displacement SCFs done (6 per atom) " - f"({_ir_total_solves - _ir_done_solves} " - "remaining)" + for _I, _ax, _sign in _remaining: + _key = (_I, _ax, _sign) + if _key in _dipoles: + # Already resumed, or already completed by a + # parallel pool that failed only partway + # through — never recompute either. + continue + _c = _coords0.copy() + _c[_I, _ax] += _DELTA * _sign + mol.set_geom_(_c, unit="Bohr") + _mu = _displaced_scf_dipole() + _dipoles[_key] = _mu + if checkpoint is not None: + checkpoint.mark_item_done( + _ITEM_SET_NAME, + displacement_id(_I, _ax, _sign), + {"dipole": _mu.tolist()}, ) + _ir_done_solves += 1 + _status( + "Numerical IR intensities: " + f"{_ir_done_solves}/{_ir_total_solves} " + "finite-difference displacement SCFs done (6 per atom) " + f"({_ir_total_solves - _ir_done_solves} " + "remaining)" + ) - _dpdx[3 * _I + _ax] = (_mu_p - _mu_m) / (2 * _DELTA) + # --- CHK.4.5: assembly gate --- + # Runs once, here, only after every required id is in + # ``_dipoles`` — whichever path (resume, parallel, + # serial fallback) put it there. Reading from one shared + # dict rather than a path-specific variable is what + # makes this a proper gate rather than three separate, + # possibly-inconsistent assembly steps. + for _I in range(_n_ir): + for _ax in range(3): + _mu_p = _dipoles[(_I, _ax, 1)] + _mu_m = _dipoles[(_I, _ax, -1)] + _dpdx[3 * _I + _ax] = (_mu_p - _mu_m) / (2 * _DELTA) finally: mol.set_geom_(_coords0, unit="Bohr") mol.verbose = _mol_v @@ -833,6 +929,8 @@ def _displaced_scf_dipole() -> _np_ir.ndarray: hessian=h, atom_str=molecule.to_pyscf_format(), scf_rescue=scf_rescue, + checkpoint=checkpoint, + resume=resume, ) if len(_raman) == len(frequencies_cm1): raman_activities = _raman @@ -906,6 +1004,18 @@ def _tv(v): logger.warning("Thermochemistry failed: %s", _exc) _status("Thermochemistry failed; frequency backend complete.") + # M-CHECKPOINT CHK.4 — the Hessian/frequency step (this whole outer + # try) succeeded, so nothing here is worth resuming from anymore. + # A thermochemistry failure just above does not change that: thermo + # is best-effort enrichment, not the core deliverable this + # checkpoint tracks. Left uncalled on the outer except below — + # a genuinely failed Hessian keeps the checkpoint resumable. + if checkpoint is not None: + try: + checkpoint.mark_complete() + except Exception: # noqa: BLE001 — provenance is never worth a crash + pass + except Exception as exc: logger.warning("Hessian/frequency computation failed: %s", exc) _status("Hessian/frequency step failed.") diff --git a/quantui/freq_displacement_ids.py b/quantui/freq_displacement_ids.py new file mode 100644 index 0000000..b5f7923 --- /dev/null +++ b/quantui/freq_displacement_ids.py @@ -0,0 +1,79 @@ +"""Deterministic displacement-id scheme for frequency checkpointing (CHK.4.2). + +A Frequency job's numerical IR-intensity step needs one SCF per Cartesian +displacement of each atom (``+Δ`` and ``-Δ``) — ``6N`` total for an +``N``-atom molecule. Unlike Geometry-Opt's optimizer steps or PES-Scan's scan +points, these displacements have no inherent order once +``freq_ir_workers.py``'s ``QUANTUI_FREQ_PARALLEL`` opt-in runs them +concurrently across worker processes — completion order is +non-deterministic. So CHK.4 (see +``QuantUI-development-tracking/TODO/roadmaps/34-m-checkpoint-calc-restart-roadmap.md``) +resumes a *set* of completed displacement ids rather than a *prefix*, and +this module is the single source of truth for what those ids are — shared +by :mod:`quantui.freq_calc`'s serial loop, its parallel dispatch, and +:mod:`quantui.freq_ir_workers` / :mod:`quantui.freq_raman_workers`'s worker +processes — so none of them can ever disagree about what "required" means. + +``atom_index`` is 0-based and follows ``mol.atom_coords()`` ordering — the +same atom-ordering stability :class:`~quantui.checkpoint.CalcIdentity` +already depends on elsewhere in the checkpoint layer. +""" + +from __future__ import annotations + +from typing import List, Tuple + +_AXIS_LABELS = ("x", "y", "z") +_SIGN_LABELS = {1: "+", -1: "-"} +_SIGN_FROM_LABEL = {"+": 1, "-": -1} + + +def displacement_id(atom_index: int, axis: int, sign: int) -> str: + """Stable string id for one Cartesian displacement. + + Args: + atom_index: 0-based atom index. + axis: 0, 1, or 2 for x, y, or z. + sign: ``+1`` or ``-1``. + + Returns: + e.g. ``"d000_x_+"`` for atom 0, x-axis, ``+Δ``. + """ + if axis not in (0, 1, 2): + raise ValueError(f"axis must be 0, 1, or 2 (got {axis!r})") + if sign not in (1, -1): + raise ValueError(f"sign must be +1 or -1 (got {sign!r})") + return f"d{atom_index:03d}_{_AXIS_LABELS[axis]}_{_SIGN_LABELS[sign]}" + + +def parse_displacement_id(item_id: str) -> Tuple[int, int, int]: + """Inverse of :func:`displacement_id` — returns ``(atom_index, axis, sign)``. + + Raises: + ValueError: *item_id* is not a well-formed displacement id. + """ + try: + if not item_id.startswith("d"): + raise ValueError + atom_part, axis_label, sign_label = item_id[1:].split("_") + atom_index = int(atom_part) + axis = _AXIS_LABELS.index(axis_label) + sign = _SIGN_FROM_LABEL[sign_label] + except (ValueError, IndexError, KeyError) as exc: + raise ValueError(f"not a displacement id: {item_id!r}") from exc + return atom_index, axis, sign + + +def required_displacement_ids(n_atoms: int) -> List[str]: + """Every displacement id for an *n_atoms*-atom molecule (``6N`` total). + + Deterministic ordering — atom, then axis, then ``+``/``-`` — purely so + the id list is reproducible for logging/debugging; resume itself treats + this as a set (see the module docstring), never a prefix. + """ + ids: List[str] = [] + for atom_index in range(n_atoms): + for axis in range(3): + ids.append(displacement_id(atom_index, axis, 1)) + ids.append(displacement_id(atom_index, axis, -1)) + return ids diff --git a/quantui/freq_ir_workers.py b/quantui/freq_ir_workers.py index f2518e8..d42d8d8 100644 --- a/quantui/freq_ir_workers.py +++ b/quantui/freq_ir_workers.py @@ -55,6 +55,7 @@ def init_worker( xc: str | None, dm0_pickle_path: str, omp_threads: int, + checkpoint_items_dir: str | None = None, ) -> None: """ProcessPoolExecutor worker initializer. @@ -85,6 +86,14 @@ def init_worker( omp_threads: BLAS thread budget for this worker. Set as ``OMP_NUM_THREADS`` / ``MKL_NUM_THREADS`` / ``OPENBLAS_NUM_THREADS`` / ``PYSCF_NUM_THREADS``. + checkpoint_items_dir: + M-CHECKPOINT CHK.4.4. Directory path (plain string — a + ``Checkpoint`` object with a live log stream is not safe to pickle + across this process boundary) where each completed displacement's + result is durably recorded via + :func:`quantui.checkpoint.mark_item_done_at`. ``None`` when the run + has no checkpoint (checkpointing is always optional — see + :mod:`quantui.checkpoint`'s "never break a calculation" rule). """ # Order matters: set env vars before any NumPy / PySCF import. threads = str(int(omp_threads)) @@ -105,16 +114,20 @@ def init_worker( spin=int(spin), xc=xc, dm0=dm0, + checkpoint_items_dir=checkpoint_items_dir, ) -def run_displaced_scf(coords_bohr_flat) -> Any: +def run_displaced_scf(item_id: str, coords_bohr_flat) -> Any: """Run one SCF at the displaced geometry; return the dipole as ndarray. Called by :class:`concurrent.futures.ProcessPoolExecutor` once per submitted displacement task. ``coords_bohr_flat`` is the displaced geometry packed as a flat Python list (``[x0, y0, z0, x1, y1, z1, ...]``) for cheap pickling — reshaped to ``(N_atoms, 3)`` inside the worker. + ``item_id`` is the displacement's stable id (CHK.4.2, e.g. ``"d000_x_+"``) + — used only to durably record completion (CHK.4.4) when this run has a + checkpoint; it plays no role in the SCF itself. Uses ``_WORKER_STATE`` populated by :func:`init_worker` for the invariant inputs (atom string, basis, etc.) + the shared initial-guess @@ -130,6 +143,10 @@ def run_displaced_scf(coords_bohr_flat) -> Any: Any exception raised here propagates to the parent via the ``Future.result()`` call. The freq_calc driver catches such failures and falls back to the serial loop so the user's calc still completes. + The checkpoint write happens *before* this function returns — if this + worker process is killed immediately after (e.g. the whole SLURM job is + preempted right as this task finishes), the durable record already + exists regardless of whether the parent ever collects this Future. """ import numpy as np from pyscf import dft, gto, scf @@ -169,7 +186,18 @@ def run_displaced_scf(coords_bohr_flat) -> Any: from .scf_robust import run_scf_with_rescue run_scf_with_rescue(mf, dm0=dm0) - return np.array(mf.dip_moment(verbose=0)) + dipole = np.array(mf.dip_moment(verbose=0)) + + items_dir = state.get("checkpoint_items_dir") + if items_dir: + from quantui.checkpoint import mark_item_done_at + + # Best-effort by construction — mark_item_done_at never raises. A + # failed write here just means this displacement isn't resumable + # from disk; the calc itself is unaffected either way. + mark_item_done_at(items_dir, item_id, {"dipole": dipole.tolist()}) + + return dipole def freq_parallel_opt_in() -> bool: diff --git a/quantui/freq_raman_workers.py b/quantui/freq_raman_workers.py index 6aed4c3..f088561 100644 --- a/quantui/freq_raman_workers.py +++ b/quantui/freq_raman_workers.py @@ -23,8 +23,17 @@ def init_raman_worker( omp_threads: int, dm0_is_unrestricted: bool, density_fit_used: bool, + checkpoint_items_dir: str | None = None, ) -> None: - """Worker initializer — same threading discipline as IR workers.""" + """Worker initializer — same threading discipline as IR workers. + + ``checkpoint_items_dir`` (M-CHECKPOINT CHK.4.4): directory path (plain + string — see :func:`quantui.freq_ir_workers.init_worker`'s docstring for + why not a ``Checkpoint`` object) where each completed displacement's + polarizability is durably recorded via + :func:`quantui.checkpoint.mark_item_done_at`. ``None`` when the run has + no checkpoint. + """ import os import pickle @@ -46,6 +55,7 @@ def init_raman_worker( dm0=dm0, dm0_is_unrestricted=bool(dm0_is_unrestricted), density_fit_used=bool(density_fit_used), + checkpoint_items_dir=checkpoint_items_dir, ) @@ -63,8 +73,14 @@ def _polarizability_module(mol: Any, dm0_is_unrestricted: bool): return pol_mod -def run_displaced_polarizability(coords_bohr_flat) -> list[list[float]]: - """Run one SCF at a displaced geometry; return α as a nested 3×3 list.""" +def run_displaced_polarizability(item_id: str, coords_bohr_flat) -> list[list[float]]: + """Run one SCF at a displaced geometry; return α as a nested 3×3 list. + + ``item_id`` (CHK.4.2, e.g. ``"d000_x_+"``) is used only to durably + record completion (CHK.4.4) when this run has a checkpoint — see + :func:`quantui.freq_ir_workers.run_displaced_scf`'s docstring for the + same pattern and its crash-safety rationale. + """ import numpy as np from pyscf import dft, gto, scf @@ -102,4 +118,16 @@ def run_displaced_polarizability(coords_bohr_flat) -> list[list[float]]: pol_mod = _polarizability_module(mol, dm0_is_unrestricted) alpha = np.asarray(pol_mod.polarizability(pol_mod.Polarizability(mf)), dtype=float) reshaped = alpha.reshape(3, 3) - return [[float(x) for x in row] for row in reshaped.tolist()] + nested = [[float(x) for x in row] for row in reshaped.tolist()] + + items_dir = state.get("checkpoint_items_dir") + if items_dir: + from quantui.checkpoint import mark_item_done_at + + # Best-effort by construction — never raises. See + # freq_ir_workers.run_displaced_scf for why this write happens here, + # inside the worker, rather than only after the parent collects the + # Future. + mark_item_done_at(items_dir, item_id, {"alpha": nested}) + + return nested diff --git a/quantui/raman_calc.py b/quantui/raman_calc.py index 86c409c..2f83b3f 100644 --- a/quantui/raman_calc.py +++ b/quantui/raman_calc.py @@ -151,20 +151,59 @@ def _cpu_raman_activities_fd( status: Callable[[str], None], atom_str: str | None = None, scf_rescue: bool = True, + checkpoint: Optional[Any] = None, + resume: bool = False, ) -> List[float]: - """CPU Raman via pyscf-properties polarizability + geometry FD.""" + """CPU Raman via pyscf-properties polarizability + geometry FD. + + ``checkpoint``/``resume`` (M-CHECKPOINT CHK.4.4) mirror + :func:`quantui.freq_calc.run_freq_calc`'s displacement checkpointing, + under a separate named item set (``"raman_displacements"``) so this + loop's per-displacement records can never collide with the IR loop's + ``"freq_displacements"`` set, even though both are keyed on the exact + same ``(atom_idx, axis, sign)`` displacements of the same molecule. + """ import os pol_mod = _polarizability_module(mf, dm0_is_unrestricted) from quantui.density_fitting import try_density_fit as _try_density_fit + from quantui.freq_displacement_ids import ( + displacement_id as _disp_id, + ) + from quantui.freq_displacement_ids import ( + parse_displacement_id as _parse_disp_id, + ) from quantui.gpu_offload import try_to_gpu as _try_to_gpu_inner + _ITEM_SET_NAME = "raman_displacements" _xc = getattr(mf, "xc", None) _n_atoms = mol.natm _coords0 = mol.atom_coords().copy() _total = _n_atoms * 3 * 2 - _done = 0 + + # --- M-CHECKPOINT CHK.4: resume already-banked displacements --- + # Single source of truth for every displacement's polarizability, + # whichever path produced it (resumed, parallel pool, serial fallback) — + # the assembly step below reads only from this dict, which is what + # makes it a proper CHK.4.5 gate. + _alphas: dict = {} + if checkpoint is not None and resume: + for _item_id, _payload in checkpoint.completed_items(_ITEM_SET_NAME).items(): + try: + _key = _parse_disp_id(_item_id) + _alphas[_key] = np.asarray(_payload["alpha"], dtype=float) + except (ValueError, KeyError, TypeError): + continue + _done = len(_alphas) + if _done and checkpoint is not None: + try: + checkpoint.log_resumed( + f"{_done}/{_total} finite-difference polarizability " + "evaluations already banked from an earlier attempt" + ) + except Exception: # noqa: BLE001 — provenance is never worth a crash + pass status( "Numerical Raman activities (CPU): " @@ -210,27 +249,33 @@ def _displaced_alpha(atom_idx: int, ax: int, sign: int) -> np.ndarray: cpu_count=_cpu_count, displacement_count=_total, ) + _ckpt_items_dir = ( + str(checkpoint.items_dir(_ITEM_SET_NAME)) if checkpoint is not None else None + ) _parallel_failed = False try: - if _use_parallel: + _remaining: list[tuple[int, int, int]] = [ + (atom_idx, ax, sign) + for atom_idx in range(_n_atoms) + for ax in range(3) + for sign in (1, -1) + if (atom_idx, ax, sign) not in _alphas + ] + if _use_parallel and _remaining: try: import concurrent.futures as _cf import multiprocessing as _mp import pickle as _pickle import tempfile as _tempfile - _n_workers = _ir_par.pick_worker_count(_cpu_count, _total) + _n_workers = _ir_par.pick_worker_count(_cpu_count, len(_remaining)) _threads_each = _ir_par.threads_per_worker(_cpu_count, _n_workers) _tasks: list[tuple[int, int, int, list[float]]] = [] - for atom_idx in range(_n_atoms): - for ax in range(3): - cp = _coords0.copy() - cp[atom_idx, ax] += _DELTA_BOHR - _tasks.append((atom_idx, ax, +1, cp.flatten().tolist())) - cm = _coords0.copy() - cm[atom_idx, ax] -= _DELTA_BOHR - _tasks.append((atom_idx, ax, -1, cm.flatten().tolist())) + for atom_idx, ax, sign in _remaining: + cp = _coords0.copy() + cp[atom_idx, ax] += sign * _DELTA_BOHR + _tasks.append((atom_idx, ax, sign, cp.flatten().tolist())) _dm0_handle = _tempfile.NamedTemporaryFile( delete=False, suffix=".dm0.pkl" @@ -267,12 +312,14 @@ def _displaced_alpha(atom_idx: int, ax: int, sign: int) -> np.ndarray: _threads_each, dm0_is_unrestricted, density_fit_used, + _ckpt_items_dir, ), ) as _pool: - _alphas: dict = {} _futs = { _pool.submit( - _ram_par.run_displaced_polarizability, task[3] + _ram_par.run_displaced_polarizability, + _disp_id(task[0], task[1], task[2]), + task[3], ): task for task in _tasks } @@ -294,12 +341,6 @@ def _displaced_alpha(atom_idx: int, ax: int, sign: int) -> np.ndarray: os.unlink(_dm0_handle.name) except OSError: pass - - for atom_idx in range(_n_atoms): - for ax in range(3): - ap = _alphas[(atom_idx, ax, +1)] - am = _alphas[(atom_idx, ax, -1)] - dalpha[3 * atom_idx + ax] = (ap - am) / (2.0 * _DELTA_BOHR) except Exception as _par_exc: logger.warning( "Parallel Raman computation failed (%s); falling back to serial.", @@ -309,14 +350,29 @@ def _displaced_alpha(atom_idx: int, ax: int, sign: int) -> np.ndarray: "Parallel Raman activities failed; falling back to serial computation." ) _parallel_failed = True - _done = 0 + # Preserve resumed + partially-completed progress rather + # than discarding it — mirrors freq_calc.py's IR loop. + _done = len(_alphas) if not _use_parallel or _parallel_failed: - for atom_idx in range(_n_atoms): - for ax in range(3): - ap = _displaced_alpha(atom_idx, ax, +1) - am = _displaced_alpha(atom_idx, ax, -1) - dalpha[3 * atom_idx + ax] = (ap - am) / (2.0 * _DELTA_BOHR) + for atom_idx, ax, sign in _remaining: + if (atom_idx, ax, sign) in _alphas: + continue # resumed, or already done by a partial parallel attempt + alpha = _displaced_alpha(atom_idx, ax, sign) + _alphas[(atom_idx, ax, sign)] = alpha + if checkpoint is not None: + checkpoint.mark_item_done( + _ITEM_SET_NAME, + _disp_id(atom_idx, ax, sign), + {"alpha": alpha.tolist()}, + ) + + # --- CHK.4.5: assembly gate --- runs once, from the shared dict. + for atom_idx in range(_n_atoms): + for ax in range(3): + ap = _alphas[(atom_idx, ax, 1)] + am = _alphas[(atom_idx, ax, -1)] + dalpha[3 * atom_idx + ax] = (ap - am) / (2.0 * _DELTA_BOHR) finally: mol.set_geom_(_coords0, unit="Bohr") mol.verbose = _mol_v @@ -360,6 +416,8 @@ def compute_raman_activities( hessian: Any = None, atom_str: str | None = None, scf_rescue: bool = True, + checkpoint: Optional[Any] = None, + resume: bool = False, ) -> List[float]: """Compute static Raman activities (Å⁴/amu) per normal mode. @@ -401,6 +459,8 @@ def compute_raman_activities( status=status, atom_str=atom_str, scf_rescue=scf_rescue, + checkpoint=checkpoint, + resume=resume, ) except ImportError as exc: logger.warning("pyscf-properties polarizability unavailable: %s", exc) diff --git a/tests/test_backends_worker.py b/tests/test_backends_worker.py index 151b0f2..72ad658 100644 --- a/tests/test_backends_worker.py +++ b/tests/test_backends_worker.py @@ -472,3 +472,77 @@ def test_preopt_checkpoint_result_is_cached_across_attempts( assert outcome2.status == "success" mock_opt.assert_called_once() # still just the one call from attempt 1 assert "Reusing saved preopt geometry" in (staging / "live.log").read_text() + + @patch("quantui.freq_calc.run_freq_calc") + def test_frequency_first_attempt_gets_a_fresh_checkpoint(self, mock_freq, staging): + """M-CHECKPOINT CHK.4 — the calc type roadmap 34's real production + cost data was about: a killed frequency job used to discard the + entire 6N-displacement Hessian on every resubmission.""" + from quantui.checkpoint import Checkpoint + + mock_freq.return_value = SimpleNamespace( + energy_hartree=-1.12, + homo_lumo_gap_ev=10.0, + converged=True, + n_iterations=5, + method="RHF", + basis="STO-3G", + formula="H2", + frequencies_cm1=[4400.0], + ir_intensities=[1.0], + raman_activities=[0.2], + zpve_hartree=0.01, + displacements=[[[0.0, 0.0, 1.0], [0.0, 0.0, -1.0]]], + ) + data = json.loads((staging / "request.json").read_text()) + data["calc_type"] = "frequency" + (staging / "request.json").write_text(json.dumps(data)) + + outcome = run_worker_request(staging / "request.json") + assert outcome.status == "success" + _args, kwargs = mock_freq.call_args + assert kwargs["resume"] is False + assert isinstance(kwargs["checkpoint"], Checkpoint) + assert (staging / ".checkpoint").is_dir() + + @patch("quantui.freq_calc.run_freq_calc") + def test_frequency_resumes_when_a_prior_attempt_left_progress( + self, mock_freq, staging + ): + """A prior attempt that banked some displacements (see + checkpoint.py's ``items/freq_displacements/`` — CHK.4.1) before + getting killed must resume, not restart the whole Hessian.""" + from quantui.checkpoint import Checkpoint + + # Simulate a previous attempt that got killed mid-run: open the same + # checkpoint identity the worker will compute, and bank one + # displacement — status stays "running" (never marked complete). + ckpt = Checkpoint( + self._identity(calc_type="frequency"), root=staging / ".checkpoint" + ) + ckpt.begin() + ckpt.mark_item_done("freq_displacements", "d000_x_+", {"dipole": [0, 0, 0]}) + + mock_freq.return_value = SimpleNamespace( + energy_hartree=-1.12, + homo_lumo_gap_ev=10.0, + converged=True, + n_iterations=5, + method="RHF", + basis="STO-3G", + formula="H2", + frequencies_cm1=[4400.0], + ir_intensities=[1.0], + raman_activities=[0.2], + zpve_hartree=0.01, + displacements=[[[0.0, 0.0, 1.0], [0.0, 0.0, -1.0]]], + ) + data = json.loads((staging / "request.json").read_text()) + data["calc_type"] = "frequency" + (staging / "request.json").write_text(json.dumps(data)) + + outcome = run_worker_request(staging / "request.json") + assert outcome.status == "success" + _args, kwargs = mock_freq.call_args + assert kwargs["resume"] is True + assert "Resuming frequency analysis" in (staging / "live.log").read_text() diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py index 8699148..4e54667 100644 --- a/tests/test_checkpoint.py +++ b/tests/test_checkpoint.py @@ -475,6 +475,139 @@ def test_append_creates_the_directory_if_needed(self, root): assert ckpt.completed_points() +# ══ Named item sets (CHK.4.1) ════════════════════════════════════════════════ +# +# The set-oriented sibling of TestPoints above: frequency's displaced-geometry +# SCFs can complete in any order (freq_ir_workers.py's QUANTUI_FREQ_PARALLEL +# opt-in runs them concurrently), so resume has to diff a *set* of finished +# item ids, not trim a *prefix* the way CHK.3's scan points do. + + +class TestNamedItemSets: + def test_round_trip(self, root): + ckpt = C.Checkpoint(_identity()) + ckpt.begin() + ckpt.mark_item_done("freq_displacements", "d000_x_+", {"energy": -1.1}) + ckpt.mark_item_done("freq_displacements", "d000_x_-", {"energy": -1.2}) + assert ckpt.completed_item_ids("freq_displacements") == {"d000_x_+", "d000_x_-"} + items = ckpt.completed_items("freq_displacements") + assert items["d000_x_+"]["energy"] == -1.1 + assert items["d000_x_-"]["energy"] == -1.2 + + def test_missing_set_reads_as_empty(self, root): + ckpt = C.Checkpoint(_identity()) + assert ckpt.completed_item_ids("freq_displacements") == set() + assert ckpt.completed_items("freq_displacements") == {} + + def test_mark_item_done_creates_the_directory_if_needed(self, root): + ckpt = C.Checkpoint(_identity()) + ckpt.mark_item_done("freq_displacements", "d000_x_+", {"energy": -1.1}) + assert ckpt.completed_item_ids("freq_displacements") == {"d000_x_+"} + + def test_rewriting_an_item_replaces_it(self, root): + """A resubmitted worker recomputing the same id must not duplicate.""" + ckpt = C.Checkpoint(_identity()) + ckpt.begin() + ckpt.mark_item_done("freq_displacements", "d000_x_+", {"energy": -1.1}) + ckpt.mark_item_done("freq_displacements", "d000_x_+", {"energy": -1.15}) + assert ckpt.completed_item_ids("freq_displacements") == {"d000_x_+"} + assert ckpt.completed_items("freq_displacements")["d000_x_+"]["energy"] == -1.15 + + def test_no_stray_tmp_file_left_behind_after_a_write(self, root): + ckpt = C.Checkpoint(_identity()) + ckpt.mark_item_done("freq_displacements", "d000_x_+", {"energy": -1.1}) + leftovers = list(ckpt.items_dir("freq_displacements").iterdir()) + assert leftovers == [ckpt.items_dir("freq_displacements") / "d000_x_+.json"] + + def test_a_corrupt_item_file_is_excluded_from_both_ids_and_payloads(self, root): + """A corrupted item must never be mistaken for a completed one — + the resume-time question is "do we have a usable result?", not + "does this filename exist?". Recomputing one displacement is cheap; + silently assembling a Hessian from a phantom result is not.""" + ckpt = C.Checkpoint(_identity()) + ckpt.begin() + ckpt.mark_item_done("freq_displacements", "d000_x_+", {"energy": -1.1}) + d = ckpt.items_dir("freq_displacements") + (d / "d001_y_-.json").write_text('{"energy": -1.', encoding="utf-8") + assert ckpt.completed_item_ids("freq_displacements") == {"d000_x_+"} + assert set(ckpt.completed_items("freq_displacements").keys()) == {"d000_x_+"} + + def test_a_dot_tmp_file_never_counts_as_complete(self, root): + ckpt = C.Checkpoint(_identity()) + d = ckpt.items_dir("freq_displacements") + d.mkdir(parents=True) + (d / "d000_x_+.12345.abcd1234.tmp").write_text("{}", encoding="utf-8") + assert ckpt.completed_item_ids("freq_displacements") == set() + + def test_two_named_sets_never_collide(self, root): + """IR vs. Raman displacement sets under one checkpoint (CHK.4.4).""" + ckpt = C.Checkpoint(_identity()) + ckpt.mark_item_done("ir_displacements", "d000_x_+", {"kind": "ir"}) + ckpt.mark_item_done("raman_displacements", "d000_x_+", {"kind": "raman"}) + assert ckpt.completed_items("ir_displacements")["d000_x_+"]["kind"] == "ir" + assert ( + ckpt.completed_items("raman_displacements")["d000_x_+"]["kind"] == "raman" + ) + + def test_has_progress_true_once_an_item_is_marked_done(self, root): + ckpt = C.Checkpoint(_identity()) + ckpt.begin() + assert not ckpt.has_progress() + ckpt.mark_item_done("freq_displacements", "d000_x_+", {"energy": -1.1}) + assert ckpt.has_progress() + + def test_concurrent_writers_never_collide_on_different_ids(self, root): + """Simulates freq_ir_workers.py's ProcessPoolExecutor writers: each + one only ever touches its own item id's filename, so nothing here + needs a lock. Not a true multiprocess test (that belongs in + CHK.4.4's integration tests) — just the filesystem-contract half.""" + ckpt = C.Checkpoint(_identity()) + for i in range(20): + ckpt.mark_item_done("freq_displacements", f"d{i:03d}_x_+", {"i": i}) + ids = ckpt.completed_item_ids("freq_displacements") + assert len(ids) == 20 + items = ckpt.completed_items("freq_displacements") + assert all(items[f"d{i:03d}_x_+"]["i"] == i for i in range(20)) + + +class TestPathOnlyItemFunctions: + """The cross-process sibling API (CHK.4.4): a ProcessPoolExecutor worker + only has a directory path (via ``initargs``), not a picklable + ``Checkpoint`` object — these free functions let it write/read the same + on-disk records ``Checkpoint.mark_item_done``/``completed_items`` do.""" + + def test_write_then_read_by_path_alone(self, root, tmp_path): + items_root = tmp_path / "items" / "freq_displacements" + assert C.mark_item_done_at(items_root, "d000_x_+", {"energy": -1.1}) is True + assert C.completed_items_at(items_root) == {"d000_x_+": {"energy": -1.1}} + + def test_interoperates_with_the_checkpoint_object(self, root): + """A worker writing by path and the parent reading via Checkpoint + (or vice versa) must see the exact same records — same directory, + same file format, either API.""" + ckpt = C.Checkpoint(_identity()) + C.mark_item_done_at( + ckpt.items_dir("freq_displacements"), "d000_x_+", {"energy": -1.1} + ) + assert ckpt.completed_items("freq_displacements") == { + "d000_x_+": {"energy": -1.1} + } + ckpt.mark_item_done("freq_displacements", "d000_x_-", {"energy": -1.2}) + assert C.completed_items_at(ckpt.items_dir("freq_displacements")) == { + "d000_x_+": {"energy": -1.1}, + "d000_x_-": {"energy": -1.2}, + } + + def test_missing_directory_reads_as_empty(self, tmp_path): + assert C.completed_items_at(tmp_path / "nope") == {} + + def test_returns_false_rather_than_raising_on_an_unwritable_path(self, tmp_path): + # A file where a directory is expected: mkdir(parents=True) raises. + blocker = tmp_path / "blocker" + blocker.write_text("not a directory", encoding="utf-8") + assert C.mark_item_done_at(blocker / "sub", "d000_x_+", {}) is False + + # ══ Warm start discovery ═════════════════════════════════════════════════════ diff --git a/tests/test_freq_calc.py b/tests/test_freq_calc.py index 1b59174..e278a2c 100644 --- a/tests/test_freq_calc.py +++ b/tests/test_freq_calc.py @@ -350,11 +350,83 @@ def test_freq_ir_workers_dispatches_on_dm0_shape_not_spin(self): tmp.close() init_worker(atom_str, "sto-3g", 0, 0, None, tmp.name, 1) coords = mol.atom_coords(unit="Bohr").flatten().tolist() - dip = run_displaced_scf(coords) # must not raise + dip = run_displaced_scf("d000_x_+", coords) # must not raise assert len(dip) == 3 finally: os.unlink(tmp.name) + def test_worker_writes_its_own_checkpoint_record(self, tmp_path): + """M-CHECKPOINT CHK.4.4's crash-safety claim, at the unit level: the + worker itself durably records completion — not just the parent + after collecting the Future — so a parent crash mid-wave never + loses a sibling that already finished. Exercised in-process here + (no real ProcessPoolExecutor), since ``mark_item_done_at`` only + needs a directory path — the same one a real spawned worker would + get via ``initargs``. + """ + pytest.importorskip("pyscf") + import os + import pickle + import tempfile + + from pyscf import gto, scf + + from quantui.checkpoint import completed_items_at + from quantui.freq_ir_workers import init_worker, run_displaced_scf + + atom_str = "O 0 0 0.119; H 0 0.763 -0.477; H 0 -0.763 -0.477" + mol = gto.M(atom=atom_str, basis="sto-3g", spin=0, charge=0, verbose=0) + mf = scf.RHF(mol) + mf.kernel() + dm0 = mf.make_rdm1() + + items_dir = tmp_path / "items" / "freq_displacements" + tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") + try: + pickle.dump(dm0, tmp) + tmp.close() + init_worker(atom_str, "sto-3g", 0, 0, None, tmp.name, 1, str(items_dir)) + coords = mol.atom_coords(unit="Bohr").flatten().tolist() + dip = run_displaced_scf("d000_x_+", coords) + + recorded = completed_items_at(items_dir) + assert set(recorded) == {"d000_x_+"} + assert recorded["d000_x_+"]["dipole"] == pytest.approx(list(dip), abs=1e-9) + finally: + os.unlink(tmp.name) + + def test_worker_with_no_checkpoint_dir_writes_nothing(self, tmp_path): + """checkpointing must stay fully optional — a worker with + ``checkpoint_items_dir=None`` (the default) must not create + anything on disk.""" + pytest.importorskip("pyscf") + import os + import pickle + import tempfile + + from pyscf import gto, scf + + from quantui.freq_ir_workers import init_worker, run_displaced_scf + + atom_str = "O 0 0 0.119; H 0 0.763 -0.477; H 0 -0.763 -0.477" + mol = gto.M(atom=atom_str, basis="sto-3g", spin=0, charge=0, verbose=0) + mf = scf.RHF(mol) + mf.kernel() + dm0 = mf.make_rdm1() + + tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") + try: + pickle.dump(dm0, tmp) + tmp.close() + init_worker( + atom_str, "sto-3g", 0, 0, None, tmp.name, 1 + ) # no checkpoint dir + coords = mol.atom_coords(unit="Bohr").flatten().tolist() + run_displaced_scf("d000_x_+", coords) + assert not (tmp_path / "items").exists() + finally: + os.unlink(tmp.name) + @pyscf_only @pytest.mark.slow def test_parallel_failure_falls_back_to_serial(self, monkeypatch): @@ -401,5 +473,203 @@ def __exit__(self, *a): ) +# ============================================================================ +# M-CHECKPOINT CHK.4 — displacement-level checkpoint/resume +# ============================================================================ + + +@pyscf_only +@pytest.mark.slow +class TestChk4DisplacementCheckpointing: + """The concurrency-safe, set-based design (roadmap 34's CHK.4 section): + resume diffs a *set* of completed displacement ids, not a prefix. + + Raman is disabled in every test here (``QUANTUI_RAMAN=0``) so the call + count asserted below is exactly the IR loop's — otherwise it would also + depend on whether pyscf-properties happens to be installed. + """ + + def _checkpoint(self, tmp_path, molecule): + from quantui.checkpoint import CalcIdentity, Checkpoint + + identity = CalcIdentity.from_molecule( + molecule, calc_type="frequency", method="RHF", basis="STO-3G" + ) + # Matches production: backends/worker.py's _begin_worker_checkpoint + # always calls .begin() before handing a checkpoint to a calc + # function — a checkpoint object is never "live" (has a meta.json + # to update) until this runs. + ckpt = Checkpoint(identity, root=tmp_path / "ckpt") + ckpt.begin() + return ckpt + + def _count_rescue_calls(self, monkeypatch): + """Count real run_scf_with_rescue calls without changing behavior. + + Patched at the source (quantui.scf_robust), not at + quantui.freq_calc — the call sites use a *local* ``from .scf_robust + import run_scf_with_rescue`` re-executed on every call to + _run_freq_calc_body, so patching the source module is what actually + takes effect on the next run. + """ + import quantui.scf_robust as scf_robust_mod + + real_rescue = scf_robust_mod.run_scf_with_rescue + calls: list = [] + + def _counting(*args, **kwargs): + calls.append(1) + return real_rescue(*args, **kwargs) + + monkeypatch.setattr(scf_robust_mod, "run_scf_with_rescue", _counting) + return calls + + def test_every_displacement_is_recorded(self, tmp_path, monkeypatch): + from quantui.freq_calc import run_freq_calc + + monkeypatch.setenv("QUANTUI_RAMAN", "0") + molecule = _water() + ckpt = self._checkpoint(tmp_path, molecule) + run_freq_calc(molecule, method="RHF", basis="STO-3G", checkpoint=ckpt) + # 3 atoms x 3 axes x 2 signs = 18. + assert len(ckpt.completed_item_ids("freq_displacements")) == 18 + + def test_successful_run_marks_the_checkpoint_complete(self, tmp_path, monkeypatch): + """A successful frequency run must not linger forever in the + "unfinished calculations" listing — resumable_checkpoints() filters + on STATUS_COMPLETE, so this has to actually be set on success.""" + from quantui.checkpoint import STATUS_COMPLETE + from quantui.freq_calc import run_freq_calc + + monkeypatch.setenv("QUANTUI_RAMAN", "0") + molecule = _water() + ckpt = self._checkpoint(tmp_path, molecule) + run_freq_calc(molecule, method="RHF", basis="STO-3G", checkpoint=ckpt) + assert ckpt.load_state()["status"] == STATUS_COMPLETE + assert ckpt.resumable_state() is None + + def test_fully_banked_resume_recomputes_nothing_but_the_reference_scf( + self, tmp_path, monkeypatch + ): + """The strong CHK.4 claim, proven by call count, not just a + plausible-looking answer: every displacement already banked means + zero new displacement SCFs on resume.""" + from quantui.freq_calc import run_freq_calc + + monkeypatch.setenv("QUANTUI_RAMAN", "0") + molecule = _water() + ckpt = self._checkpoint(tmp_path, molecule) + + calls = self._count_rescue_calls(monkeypatch) + baseline = run_freq_calc( + molecule, method="RHF", basis="STO-3G", checkpoint=ckpt + ) + assert len(calls) == 1 + 18 # reference SCF + all 18 displacements + + calls.clear() + # A real resubmission calls .begin() again (backends/worker.py's + # _begin_worker_checkpoint runs on every attempt) — it resets status + # to "running" but must not touch the already-banked item files. + ckpt.begin() + resumed = run_freq_calc( + molecule, method="RHF", basis="STO-3G", checkpoint=ckpt, resume=True + ) + assert len(calls) == 1 # only the reference SCF — zero displacement recompute + assert resumed.ir_intensities == pytest.approx( + baseline.ir_intensities, abs=1e-9 + ) + + def test_partial_resume_recomputes_only_the_missing_displacements( + self, tmp_path, monkeypatch + ): + """A checkpoint interrupted partway through: resume must recompute + exactly the missing ids, reuse the rest, and land on the same + answer as an uninterrupted run.""" + from quantui.freq_calc import run_freq_calc + from quantui.freq_displacement_ids import required_displacement_ids + + monkeypatch.setenv("QUANTUI_RAMAN", "0") + molecule = _water() + ckpt = self._checkpoint(tmp_path, molecule) + + baseline = run_freq_calc( + molecule, method="RHF", basis="STO-3G", checkpoint=ckpt + ) + all_ids = required_displacement_ids(3) + assert ckpt.completed_item_ids("freq_displacements") == set(all_ids) + + # Simulate an interrupted run: keep only the first 3 banked. + keep = set(all_ids[:3]) + items_dir = ckpt.items_dir("freq_displacements") + for item_id in all_ids: + if item_id not in keep: + (items_dir / f"{item_id}.json").unlink() + assert ckpt.completed_item_ids("freq_displacements") == keep + + ckpt.begin() # a real resubmission calls .begin() again — must not erase items + calls = self._count_rescue_calls(monkeypatch) + resumed = run_freq_calc( + molecule, method="RHF", basis="STO-3G", checkpoint=ckpt, resume=True + ) + # reference SCF + the 15 displacements that were NOT kept. + assert len(calls) == 1 + (len(all_ids) - len(keep)) + assert ckpt.completed_item_ids("freq_displacements") == set(all_ids) + assert resumed.ir_intensities == pytest.approx( + baseline.ir_intensities, abs=1e-6 + ) + + def test_resume_false_ignores_a_populated_checkpoint(self, tmp_path, monkeypatch): + """resume=False must behave like no checkpoint was ever passed for + *reading* progress — even though the run still writes into it, so a + later resume has something to build on. Mirrors optimizer.py's + CHK.2 convention (resume is opt-in per call, not implied by merely + passing a checkpoint object).""" + from quantui.freq_calc import run_freq_calc + + monkeypatch.setenv("QUANTUI_RAMAN", "0") + molecule = _water() + ckpt = self._checkpoint(tmp_path, molecule) + run_freq_calc(molecule, method="RHF", basis="STO-3G", checkpoint=ckpt) + assert len(ckpt.completed_item_ids("freq_displacements")) == 18 + + ckpt.begin() + calls = self._count_rescue_calls(monkeypatch) + run_freq_calc( + molecule, method="RHF", basis="STO-3G", checkpoint=ckpt, resume=False + ) + # Started fresh despite 18 already banked — same as the first run. + assert len(calls) == 1 + 18 + + def test_a_corrupt_banked_item_is_recomputed_not_trusted( + self, tmp_path, monkeypatch + ): + """checkpoint.py's "never break a calculation" rule, exercised + through the actual freq_calc resume path rather than checkpoint.py + in isolation.""" + from quantui.freq_calc import run_freq_calc + from quantui.freq_displacement_ids import required_displacement_ids + + monkeypatch.setenv("QUANTUI_RAMAN", "0") + molecule = _water() + ckpt = self._checkpoint(tmp_path, molecule) + run_freq_calc(molecule, method="RHF", basis="STO-3G", checkpoint=ckpt) + + all_ids = required_displacement_ids(3) + items_dir = ckpt.items_dir("freq_displacements") + corrupt_id = all_ids[0] + (items_dir / f"{corrupt_id}.json").write_text("{not json", encoding="utf-8") + assert corrupt_id not in ckpt.completed_item_ids("freq_displacements") + + ckpt.begin() + calls = self._count_rescue_calls(monkeypatch) + resumed = run_freq_calc( + molecule, method="RHF", basis="STO-3G", checkpoint=ckpt, resume=True + ) + # reference SCF + the one corrupted (and therefore recomputed) displacement. + assert len(calls) == 2 + assert resumed.ir_intensities, "must still produce a usable result" + assert corrupt_id in ckpt.completed_item_ids("freq_displacements") + + if __name__ == "__main__": pytest.main([__file__, "-v", "--tb=short"]) diff --git a/tests/test_freq_displacement_ids.py b/tests/test_freq_displacement_ids.py new file mode 100644 index 0000000..fdbbd58 --- /dev/null +++ b/tests/test_freq_displacement_ids.py @@ -0,0 +1,74 @@ +"""Tests for quantui.freq_displacement_ids (M-CHECKPOINT CHK.4.2). + +Pure functions, no PySCF, no checkpoint I/O — just the id scheme that +CHK.4.3-.5's serial/parallel wiring and resume logic build on. +""" + +from __future__ import annotations + +import pytest + +from quantui.freq_displacement_ids import ( + displacement_id, + parse_displacement_id, + required_displacement_ids, +) + + +class TestDisplacementId: + def test_format(self): + assert displacement_id(0, 0, 1) == "d000_x_+" + assert displacement_id(0, 0, -1) == "d000_x_-" + assert displacement_id(12, 2, 1) == "d012_z_+" + + def test_rejects_bad_axis(self): + with pytest.raises(ValueError): + displacement_id(0, 3, 1) + + def test_rejects_bad_sign(self): + with pytest.raises(ValueError): + displacement_id(0, 0, 0) + + def test_ids_are_unique_across_atoms_axes_signs(self): + ids = { + displacement_id(a, ax, s) + for a in range(5) + for ax in range(3) + for s in (1, -1) + } + assert len(ids) == 5 * 3 * 2 + + +class TestParseDisplacementId: + def test_round_trips(self): + for atom_index, axis, sign in [(0, 0, 1), (12, 2, -1), (3, 1, 1)]: + item_id = displacement_id(atom_index, axis, sign) + assert parse_displacement_id(item_id) == (atom_index, axis, sign) + + @pytest.mark.parametrize( + "bad_id", ["", "not_a_displacement", "d000_w_+", "d000_x_0", "d0a0_x_+"] + ) + def test_rejects_malformed_ids(self, bad_id): + with pytest.raises(ValueError): + parse_displacement_id(bad_id) + + +class TestRequiredDisplacementIds: + def test_count_is_6n(self): + assert len(required_displacement_ids(3)) == 18 + assert len(required_displacement_ids(19)) == 114 + + def test_zero_atoms_is_empty(self): + assert required_displacement_ids(0) == [] + + def test_every_id_parses_back_and_covers_the_grid(self): + n_atoms = 4 + ids = required_displacement_ids(n_atoms) + parsed = {parse_displacement_id(i) for i in ids} + expected = { + (a, ax, s) for a in range(n_atoms) for ax in range(3) for s in (1, -1) + } + assert parsed == expected + + def test_deterministic_order(self): + assert required_displacement_ids(2) == required_displacement_ids(2) diff --git a/tests/test_raman_calc.py b/tests/test_raman_calc.py index e835cd7..6b063d1 100644 --- a/tests/test_raman_calc.py +++ b/tests/test_raman_calc.py @@ -259,3 +259,82 @@ def test_h2o_bend_weaker_ir_than_raman(self, monkeypatch): assert ( ir_low < ram_low ), f"Expected bend-like mode IR ({ir_low:.2f}) < Raman ({ram_low:.2f})" + + +class TestChk4RamanDisplacementCheckpointing: + """M-CHECKPOINT CHK.4.4 — Raman shares the exact same (atom, axis, sign) + displacement grid as the IR loop, so it needs its own named item set + ("raman_displacements") to avoid colliding with IR's + ("freq_displacements") — see roadmap 34's CHK.4 section.""" + + def _checkpoint(self, tmp_path, molecule): + from quantui.checkpoint import CalcIdentity, Checkpoint + + identity = CalcIdentity.from_molecule( + molecule, calc_type="frequency", method="RHF", basis="STO-3G" + ) + # Matches production: backends/worker.py's _begin_worker_checkpoint + # always calls .begin() before handing a checkpoint to a calc + # function. + ckpt = Checkpoint(identity, root=tmp_path / "ckpt") + ckpt.begin() + return ckpt + + def _count_rescue_calls(self, monkeypatch): + import quantui.scf_robust as scf_robust_mod + + real_rescue = scf_robust_mod.run_scf_with_rescue + calls: list = [] + + def _counting(*args, **kwargs): + calls.append(1) + return real_rescue(*args, **kwargs) + + monkeypatch.setattr(scf_robust_mod, "run_scf_with_rescue", _counting) + return calls + + @pyscf_only + @pytest.mark.slow + def test_raman_displacements_recorded_under_their_own_item_set( + self, tmp_path, monkeypatch + ): + monkeypatch.setenv("QUANTUI_RAMAN", "1") + from quantui.freq_calc import run_freq_calc + + molecule = _water() + ckpt = self._checkpoint(tmp_path, molecule) + run_freq_calc(molecule, method="RHF", basis="STO-3G", checkpoint=ckpt) + + ir_ids = ckpt.completed_item_ids("freq_displacements") + raman_ids = ckpt.completed_item_ids("raman_displacements") + assert len(ir_ids) == 18 + assert len(raman_ids) == 18 + # Same displacement grid, but the two sets never collide — reading + # one back must not be affected by the other's records. + assert ir_ids == raman_ids + + @pyscf_only + @pytest.mark.slow + def test_fully_banked_resume_skips_all_raman_recompute(self, tmp_path, monkeypatch): + monkeypatch.setenv("QUANTUI_RAMAN", "1") + from quantui.freq_calc import run_freq_calc + + molecule = _water() + ckpt = self._checkpoint(tmp_path, molecule) + + calls = self._count_rescue_calls(monkeypatch) + baseline = run_freq_calc( + molecule, method="RHF", basis="STO-3G", checkpoint=ckpt + ) + # reference SCF + 18 IR displacements + 18 Raman displacements. + assert len(calls) == 1 + 18 + 18 + + calls.clear() + ckpt.begin() # a real resubmission calls .begin() again + resumed = run_freq_calc( + molecule, method="RHF", basis="STO-3G", checkpoint=ckpt, resume=True + ) + assert len(calls) == 1 # only the reference SCF + assert resumed.raman_activities == pytest.approx( + baseline.raman_activities, abs=1e-6 + )