diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index bb3be81..b35034d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -20,6 +20,10 @@ jobs: run: ruff check src/ - name: Check formatting run: ruff format --check src/ + - name: Check MLIP tests and recipes + run: | + ruff check tests/test_mlip*.py tests/test_cli.py examples/mlip_gpu.py alcf/polaris/mlip/smoke.py + ruff format --check tests/test_mlip*.py tests/test_cli.py examples/mlip_gpu.py alcf/polaris/mlip/smoke.py test: runs-on: ubuntu-latest diff --git a/CHANGELOG.md b/CHANGELOG.md index a4213e9..cbc6fc7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- MLIP numerical validation, strict CLI outcomes, incremental atomic batch + persistence, and an explicit experimental GPU validation recipe +- Separate PR hardening and future development handoff plans in `docs/plans` +- Runtime-selectable, agent-free MLIP APIs and CLI commands for direct ASE + MACE, Rootstock, and NVIDIA ALCHEMI MACE calculations +- Ordered MLIP batch execution with per-item JSON results, calculator/model + reuse, partial-failure handling, and a persistent batch manifest +- Polaris installation and live smoke-test recipes for the three MLIP paths - Zeo++ module (`matkit.zeopp`) for pore geometry analysis: pore diameters (Di/Df/Dif), accessible surface area, accessible volume, pore size distribution, and channel identification - CLI `matkit zeopp run` and `matkit zeopp analyze` subcommands with support for high accuracy mode (`-ha`), custom radii files (`-r UFF.rad`), and configurable probe parameters - CLI interface (`matkit` command) with subcommands for graspa, graspa_sycl, raspa2, and tobacco diff --git a/README.md b/README.md index 4c00825..4b491aa 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ - **RASPA2** -- Classical GCMC simulations - **RASPA3** -- Force field format conversion from RASPA2 - **Zeo++** -- Pore geometry analysis (pore diameters, surface area, volume, channels) -- **MACE-MP** -- ML interatomic potential geometry/cell optimization +- **MLIPs** -- experimental direct MACE, Rootstock, and NVIDIA ALCHEMI adapters - **ORCA** -- Quantum chemistry (planned) ## Features @@ -44,6 +44,12 @@ pip install -e ".[rdkit]" # For ML interatomic potentials (MACE) pip install -e ".[mlip]" +# Lightweight access to cluster-managed Rootstock models +pip install -e ".[rootstock]" + +# NVIDIA ALCHEMI MACE support (install a matching CUDA extra too) +pip install -e ".[nvalchemi_mace]" + # All optional dependencies pip install -e ".[all]" @@ -75,8 +81,88 @@ matkit zeopp run --cif structure.cif --analysis res --analysis sa --radii UFF.ra # Parse existing Zeo++ output files matkit zeopp analyze --path output_dir/ + +# Run MACE directly through ASE +matkit mlip run --input structure.cif --backend ase-mace \ + --checkpoint medium --device cuda --dtype float32 + +# Run a Rootstock checkpoint already deployed on Polaris +matkit mlip run --input structure.cif --backend rootstock \ + --checkpoint mace-mp-0-medium --cluster polaris --device cuda + +# Run a native NVIDIA ALCHEMI batch +matkit mlip run-batch --input-dir cifs --backend nvalchemi-mace \ + --checkpoint medium --device cuda --batch-size 16 +``` + +### GPU examples + +[`examples/mlip_gpu.py`](examples/mlip_gpu.py) runs one backend per Python +process so GPU runtime state is isolated. It accepts one or more ASE-readable +structure files and writes a manifest plus one JSON result per input. + +```bash +# Direct MACE calculator through ASE +python examples/mlip_gpu.py --backend ase-mace structure.cif + +# Rootstock-managed MACE checkpoint on Polaris +python examples/mlip_gpu.py --backend rootstock \ + --cluster polaris --checkpoint mace-mp-0-medium structure.cif + +# NVIDIA ALCHEMI MACE with native GPU batching +python examples/mlip_gpu.py --backend nvalchemi-mace \ + --checkpoint medium --batch-size 16 structures/*.cif ``` +All three commands force `device="cuda"` and use `float32` where the backend +exposes a dtype. Pass `--driver opt` for a fixed-cell geometry optimization. + +The new MLIP adapters remain **experimental** until a real execution is +recorded for the backend, capability, and environment. CPU tests exercise +validation and adapter contracts; they do not establish GPU compatibility or +scientific accuracy. See the [Polaris smoke recipe](alcf/polaris/mlip/README.md) +for opt-in energy, optimization, and native batch checks. + +### MLIP outcomes and batch files + +`success` means a calculation produced valid numerical results. Optimization +convergence is reported separately as `converged`; usable unconverged results +are retained. The CLI and GPU example exit **1** if any item fails or any +requested optimization remains unconverged, **2** for invalid arguments, and +**0** only when the requested calculations succeed. JSON summaries go to stdout; +MatKit's failure diagnostics go to stderr. Optional engines may also emit logs. + +Batch result files and `batch_manifest.json` are replaced atomically as each +item completes (after each native ALCHEMI chunk). The manifest starts `running` +with pending items, then finishes `completed`, `partial`, or `failure`. It +includes `pending`, `unconverged`, per-item `converged`, and a run-level `error`. +An execution-complete batch can still contain unconverged optimizations. +Catchable orchestration failures record `interrupted`; abrupt termination can +leave `running`. Previously committed result files remain available. An +unrecoverable filesystem failure may also prevent the last manifest update. + +Use a **fresh output directory for every batch**. Existing manifests or active +batch locks are refused; automatic resume and retries are deferred. The GPU +example accepts `--output-dir` and the CLI accepts `--outdir` for reruns. + +Explicit options that do not apply to the selected backend/driver are errors: + +- `--optimizer`, `--fmax`, and `--steps` require `--driver opt`. +- `--dt`, `--compile-model`, `--enable-cueq`, `--batch-size`, and `--max-atoms` + belong to ALCHEMI; `--dt` also requires optimization. `max_atoms` limits chunk + grouping; a single larger structure is still processed alone. +- Rootstock precision uses a model-supported `--setup-kwarg`, not `--dtype`. + Rootstock workers own their GPU environment; the caller does not need CUDA. +- The `mace_anicc` factory controls precision and rejects an explicit CLI dtype. + The configuration records requested settings; comprehensive resolved model + provenance is future work. + +MLIP optional packages may require a newer Python than MatKit's core. ALCHEMI +0.2 requires Python 3.11–3.13 and a matching CUDA stack; the Polaris recipe uses +Python 3.12/CUDA 12. Its MACE extra pins 0.3.15, which conflicts with ChemGraph's +currently required MACE >=0.3.16. Use a MatKit-specific environment rather than +combining these stacks. The `all` extra does not include Rootstock or ALCHEMI. + ## Python API ```python @@ -111,8 +197,34 @@ print(result["results"]["sa"]) # {'ASA': 4004.7, 'ASA_m2_g': 3918.3, ...} # Parse existing Zeo++ output files result = get_output_data("output_dir/") + +# Agent-free, runtime-selectable MLIP execution +from matkit.mlip import ( + ASEMACEConfig, + MLIPCalculationConfig, + run_mlip, +) + +result = run_mlip( + "structure.cif", + ASEMACEConfig( + checkpoint="medium", + device="cuda", + dtype="float32", + ), + MLIPCalculationConfig(driver="energy"), + output_file="mace_result.json", +) ``` +## Development roadmap + +The [PR #14 handoff](docs/plans/pr14-hardening.md) records immediate MLIP fixes +and validation. The separate [MatKit roadmap](docs/plans/matkit-roadmap.md) +starts with GCMC correctness and capability status, followed by common result +contracts, reproducibility/benchmarks, a porous-material workflow, ChemGraph +integration, and tested skills. + ## License MIT License - Copyright 2025 Thang Pham (Argonne National Laboratory) diff --git a/alcf/polaris/mlip/README.md b/alcf/polaris/mlip/README.md new file mode 100644 index 0000000..2910cbe --- /dev/null +++ b/alcf/polaris/mlip/README.md @@ -0,0 +1,90 @@ +# MLIP playground on Polaris + +This directory installs and smoke-tests the three agent-free MLIP paths in +MatKit: + +- direct MACE through ASE; +- a cluster-managed MACE checkpoint through Rootstock; +- native batched MACE through NVIDIA ALCHEMI Toolkit. + +The adapters and these recipes are **experimental**. Preparing or testing the +scripts with doubles does not verify a live installation. The scripts collect +integration evidence when explicitly run on Polaris; they are not performance +benchmarks and do not establish model parity or scientific accuracy. + +## Prerequisites + +Rootstock is deployed on Polaris, but ALCF users need access to its shared +installation. Follow the current Polaris instructions in the +[Matter Model Almanac](https://garden-ai.github.io/almanac/clusters/) before +running the Rootstock smoke test. + +Run the installer from the MatKit checkout on a Polaris login node: + +```bash +export MATKIT_MLIP_ENV=/lus/eagle/projects/PROJECT/USER/envs/matkit-mlip +bash alcf/polaris/mlip/install.sh +``` + +The installer creates an isolated Python 3.12 environment, installs the CUDA +12 and MACE extras for ALCHEMI, installs Rootstock, and installs this checkout +in editable mode. Override `MATKIT_MLIP_ENV`; the default is `.venv` in the +repository. + +This is a MatKit-specific environment. ALCHEMI 0.2 supports Python 3.11–3.13; +its MACE extra pins `mace-torch==0.3.15`. ChemGraph currently requires MACE +`>=0.3.16`, so installing its full dependency stack here is incompatible. +Different environments are required to support those versions simultaneously. +Rootstock uses a separate deployment-managed model environment. Its caller +does not require a local CUDA-enabled PyTorch installation. + +## Smoke test + +Edit the `#PBS -A` project in `smoke.pbs`, then submit it while the checkout is +your working directory: + +```bash +qsub -v MATKIT_MLIP_ENV="$MATKIT_MLIP_ENV" \ + alcf/polaris/mlip/smoke.pbs +``` + +The defaults use the small periodic structure in `tests/data`, direct and +ALCHEMI checkpoint alias `medium`, and Rootstock checkpoint +`mace-mp-0-medium`. Override paths or checkpoint names when submitting: + +```bash +qsub -v MATKIT_MLIP_ENV="$MATKIT_MLIP_ENV",\ +MATKIT_SMOKE_INPUT=/path/to/input.cif,\ +MACE_CHECKPOINT=/path/to/model.pt,\ +ROOTSTOCK_CHECKPOINT=mace-mp-0-medium \ + alcf/polaris/mlip/smoke.pbs +``` + +Results are written under `projects/mlip_smoke_$PBS_JOBID` by default. Set +`MATKIT_SMOKE_OUTPUT` to choose a **new** persistent directory. Reusing an +existing output directory is refused. + +The PBS script invokes `smoke.py`, which runs six separate Python processes: +energy and fixed-cell optimization for each backend. ALCHEMI evaluates a +two-structure native batch for both drivers. Inputs are the original structure +and a reproducibly perturbed copy. Override `MATKIT_SMOKE_STEPS` (default 1000) +or `MATKIT_SMOKE_FMAX` (default 0.01 eV/angstrom) when submitting. + +Every case retains stdout, stderr, calculation results, and its exit code. +`smoke_report.json` records the caller's package versions, GPU information, +Rootstock deployment resolution, input hash, settings, commands, observed +convergence, and validation outcome. It is updated after each case; one failure +does not prevent other cases from running. A case passes only if its CLI exits +zero, results are finite and correctly shaped, the cell is unchanged within +serialization precision, and requested optimizations converge. The overall +runner exits nonzero if any case fails. + +Keep the full output directory with any capability-validation record, including +the MatKit commit used. Promote only the backend/capability/environment actually +verified by that record. Caller package versions do not identify all packages +inside a Rootstock worker; retain the deployment information and obtain worker +versions when assessing reproducibility or parity. + +For compute-node downloads, the PBS script exports the ALCF HTTP proxy. Model +weights should be allowed to finish downloading before treating later timings +as performance measurements. diff --git a/alcf/polaris/mlip/install.sh b/alcf/polaris/mlip/install.sh new file mode 100644 index 0000000..1ceebcf --- /dev/null +++ b/alcf/polaris/mlip/install.sh @@ -0,0 +1,44 @@ +#!/bin/bash -l + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +MLIP_ENV="${MATKIT_MLIP_ENV:-${REPO_ROOT}/.venv}" + +module use /soft/modulefiles +module load conda/2025-09-25 + +if [[ ! -x "${MLIP_ENV}/bin/python" ]]; then + python -m venv "${MLIP_ENV}" +fi + +source "${MLIP_ENV}/bin/activate" +python -m pip install --upgrade pip setuptools wheel + +python -m pip install \ + --extra-index-url https://download.pytorch.org/whl/cu126 \ + --extra-index-url https://pypi.nvidia.com \ + 'nvalchemi-toolkit[cu12,mace]>=0.2,<0.3' + +python -m pip install -e \ + "${REPO_ROOT}[mlip,rootstock,nvalchemi_mace]" + +python - <<'PY' +from importlib.metadata import version + +for package in ( + "matkit", + "ase", + "mace-torch", + "rootstock", + "nvalchemi-toolkit", + "torch", +): + print(f"{package}=={version(package)}") +PY + +echo +echo "Environment installed at ${MLIP_ENV}" +echo "Check Rootstock access with: rootstock resolve --cluster polaris --json" +echo "Submit alcf/polaris/mlip/smoke.pbs from the MatKit checkout next." diff --git a/alcf/polaris/mlip/smoke.pbs b/alcf/polaris/mlip/smoke.pbs new file mode 100644 index 0000000..c936669 --- /dev/null +++ b/alcf/polaris/mlip/smoke.pbs @@ -0,0 +1,38 @@ +#!/bin/bash -l +#PBS -N matkit-mlip-smoke +#PBS -l select=1:system=polaris +#PBS -l place=scatter +#PBS -l walltime=00:30:00 +#PBS -l filesystems=home:eagle +#PBS -q debug +#PBS -A PROJECT + +set -euo pipefail + +: "${MATKIT_MLIP_ENV:?Submit with -v MATKIT_MLIP_ENV=/path/to/env}" + +MATKIT_REPO="${MATKIT_REPO:-${PBS_O_WORKDIR}}" +INPUT_FILE="${MATKIT_SMOKE_INPUT:-${MATKIT_REPO}/tests/data/test_structure.cif}" +OUTPUT_DIR="${MATKIT_SMOKE_OUTPUT:-${MATKIT_REPO}/projects/mlip_smoke_${PBS_JOBID}}" +MACE_CHECKPOINT="${MACE_CHECKPOINT:-medium}" +ROOTSTOCK_CHECKPOINT="${ROOTSTOCK_CHECKPOINT:-mace-mp-0-medium}" + +module use /soft/modulefiles +module load conda/2025-09-25 +source "${MATKIT_MLIP_ENV}/bin/activate" + +export HTTP_PROXY="http://proxy.alcf.anl.gov:3128" +export HTTPS_PROXY="http://proxy.alcf.anl.gov:3128" +export http_proxy="${HTTP_PROXY}" +export https_proxy="${HTTPS_PROXY}" +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" +export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}" + +cd "${MATKIT_REPO}" +python alcf/polaris/mlip/smoke.py \ + --input "${INPUT_FILE}" \ + --output-dir "${OUTPUT_DIR}" \ + --checkpoint "${MACE_CHECKPOINT}" \ + --rootstock-checkpoint "${ROOTSTOCK_CHECKPOINT}" \ + --steps "${MATKIT_SMOKE_STEPS:-1000}" \ + --fmax "${MATKIT_SMOKE_FMAX:-0.01}" diff --git a/alcf/polaris/mlip/smoke.py b/alcf/polaris/mlip/smoke.py new file mode 100644 index 0000000..7ccee5c --- /dev/null +++ b/alcf/polaris/mlip/smoke.py @@ -0,0 +1,260 @@ +#!/usr/bin/env python3 +"""Opt-in live MLIP validation; each backend/driver uses a separate process.""" + +import argparse +import hashlib +from importlib.metadata import PackageNotFoundError, version +import json +import math +import os +from pathlib import Path +import platform +import subprocess +import sys +from datetime import datetime, timezone + +import numpy as np +from ase.io import read, write + + +def _probe(command): + try: + result = subprocess.run( + command, capture_output=True, text=True, timeout=30 + ) + return { + "return_code": result.returncode, + "stdout": result.stdout, + "stderr": result.stderr, + } + except (OSError, subprocess.TimeoutExpired) as exc: + return {"error": str(exc)} + + +def _record_report(path, report): + temporary = path.with_suffix(".tmp") + temporary.write_text( + json.dumps(report, indent=2, allow_nan=False), encoding="utf-8" + ) + temporary.replace(path) + + +def _validate_outputs(case_dir, structures, driver, fmax, batched): + if batched: + manifest = json.loads((case_dir / "batch_manifest.json").read_text()) + assert ( + manifest["pending"] + == manifest["failed"] + == manifest["unconverged"] + == 0 + ) + assert [item["index"] for item in manifest["items"]] == list( + range(len(structures)) + ) + paths = [Path(item["result_file"]) for item in manifest["items"]] + else: + paths = [case_dir / "result.json"] + summaries = [] + for path, original in zip(paths, structures): + result = json.loads(path.read_text()) + assert result["success"], result["error"] + assert math.isfinite(result["energy"]) + assert result["energy_unit"] == "eV" + assert result["force_unit"] == "eV/angstrom" + forces = np.asarray(result["forces"]) + assert forces.shape == (len(original), 3) and np.isfinite(forces).all() + final = result["final_structure"] + assert final["atomic_numbers"] == original.numbers.tolist() + assert final["pbc"] == original.pbc.tolist() + assert np.allclose( + final["cell"], original.cell.array, rtol=1e-6, atol=1e-6 + ) + positions = np.asarray(final["positions"]) + assert ( + positions.shape == (len(original), 3) + and np.isfinite(positions).all() + ) + if result["stress"] is not None: + stress = np.asarray(result["stress"]) + assert stress.shape == (3, 3) and np.isfinite(stress).all() + if driver == "opt": + assert result["converged"], "optimization did not converge" + assert np.linalg.norm(forces, axis=1).max() <= fmax + 1e-8 + else: + assert np.allclose(positions, original.positions, rtol=0, atol=1e-6) + summaries.append( + { + "result_file": str(path), + "energy": result["energy"], + "converged": result["converged"], + "n_steps": result["n_steps"], + "max_force": float(np.linalg.norm(forces, axis=1).max()), + } + ) + return summaries + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input", required=True, type=Path) + parser.add_argument("--output-dir", required=True, type=Path) + parser.add_argument("--checkpoint", default="medium") + parser.add_argument("--rootstock-checkpoint", default="mace-mp-0-medium") + parser.add_argument("--cluster", default="polaris") + parser.add_argument("--steps", type=int, default=1000) + parser.add_argument("--fmax", type=float, default=0.01) + args = parser.parse_args(argv) + if args.steps < 1 or not math.isfinite(args.fmax) or args.fmax <= 0: + parser.error("steps and fmax must be positive; fmax must be finite") + out = args.output_dir.resolve() + if out.exists(): + parser.error("--output-dir must be a new directory") + original = read(args.input) + perturbed = original.copy() + perturbed.rattle(stdev=0.01, seed=7) + inputs_dir = out / "inputs" + inputs_dir.mkdir(parents=True) + input_paths = [ + inputs_dir / "original.extxyz", + inputs_dir / "perturbed.extxyz", + ] + structures = [original, perturbed] + for path, atoms in zip(input_paths, structures): + write(path, atoms) + packages = {} + for name in ( + "matkit", + "ase", + "mace-torch", + "rootstock", + "nvalchemi-toolkit", + "torch", + ): + try: + packages[name] = version(name) + except PackageNotFoundError: + packages[name] = None + report = { + "started_utc": datetime.now(timezone.utc).isoformat(), + "python": sys.version, + "platform": platform.platform(), + "matkit_revision": _probe( + [ + "git", + "-C", + str(Path(__file__).resolve().parents[3]), + "rev-parse", + "HEAD", + ] + ), + "packages": packages, + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "gpu": _probe(["nvidia-smi"]), + "rootstock_deployment": _probe( + ["rootstock", "resolve", "--cluster", args.cluster, "--json"] + ), + "input_file": str(args.input.resolve()), + "input_sha256": hashlib.sha256(args.input.read_bytes()).hexdigest(), + "settings": { + "checkpoint": args.checkpoint, + "rootstock_checkpoint": args.rootstock_checkpoint, + "cluster": args.cluster, + "steps": args.steps, + "fmax": args.fmax, + }, + "cases": [], + "note": ( + "Integration evidence only; this does not establish parity, " + "performance, or scientific accuracy." + ), + } + report_path = out / "smoke_report.json" + _record_report(report_path, report) + for backend in ("ase-mace", "rootstock", "nvalchemi-mace"): + for driver in ("energy", "opt"): + batched = backend == "nvalchemi-mace" + case_dir = out / backend / driver + case_dir.mkdir(parents=True) + command = [ + sys.executable, + "-m", + "matkit.cli", + "mlip", + "run-batch" if batched else "run", + "--backend", + backend, + "--driver", + driver, + "--device", + "cuda", + ] + selected = list(range(2)) if batched else [int(driver == "opt")] + for index in selected: + command.extend(["--input", str(input_paths[index])]) + command.extend( + [ + "--outdir" if batched else "--output", + str(case_dir if batched else case_dir / "result.json"), + ] + ) + if backend == "rootstock": + command.extend( + [ + "--checkpoint", + args.rootstock_checkpoint, + "--cluster", + args.cluster, + "--timeout", + "1200", + ] + ) + else: + command.extend( + ["--checkpoint", args.checkpoint, "--dtype", "float32"] + ) + if batched: + command.extend(["--batch-size", "2"]) + if driver == "opt": + command.extend( + ["--steps", str(args.steps), "--fmax", str(args.fmax)] + ) + case = { + "backend": backend, + "driver": driver, + "command": command, + "status": "running", + } + report["cases"].append(case) + _record_report(report_path, report) + try: + with ( + (case_dir / "stdout.log").open("w") as stdout, + (case_dir / "stderr.log").open("w") as stderr, + ): + process = subprocess.run( + command, stdout=stdout, stderr=stderr + ) + case["return_code"] = process.returncode + case["results"] = _validate_outputs( + case_dir, + [structures[i] for i in selected], + driver, + args.fmax, + batched, + ) + assert process.returncode == 0, ( + f"CLI exited {process.returncode}" + ) + case["status"] = "passed" + except Exception as exc: + case.update( + status="failed", error=str(exc) or type(exc).__name__ + ) + _record_report(report_path, report) + print(f"{backend} {driver}: {case['status']}") + print(f"Evidence report: {report_path}") + return int(any(case["status"] != "passed" for case in report["cases"])) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/docs/plans/matkit-roadmap.md b/docs/plans/matkit-roadmap.md new file mode 100644 index 0000000..76003c0 --- /dev/null +++ b/docs/plans/matkit-roadmap.md @@ -0,0 +1,107 @@ +# MatKit development roadmap + +## Intent and boundaries + +MatKit has two equal goals: an independently useful scientific toolkit and an +experimental testing ground for ChemGraph. Keep an explicit stable/experimental +boundary, and promote capabilities using evidence. + +- MatKit owns scientific preparation, validation, adapters, parsing, and results. +- MatKit experiments establish numerical correctness and performance. +- ChemGraph owns agent planning, memory, tool selection, and distributed jobs. +- Optional MCP adapters expose MatKit functions; workflow skills explain their + appropriate use. Neither agent framework is a core MatKit dependency. +- Compatible environments can use Python directly; incompatible ML stacks use + separate worker/service environments. Separate processes alone do not resolve + conflicting package requirements. + +## Prerequisite and next-session entry point + +Read [PR #14 hardening](pr14-hardening.md), check its final revision and merge +state, and refresh repository instructions before beginning. Start with +milestone 1. Each milestone is separate work from PR #14 and should use focused +PRs; preserve the user's existing work and public compatibility. + +## Ordered milestones + +### 1. Reliable foundations + +- Propagate custom cutoffs into unit-cell replication in gRASPA, gRASPA SYCL, + RASPA2, and pygRASPA setup, including cached batch calculations. +- Set RASPA2 parsing success when valid results are obtained. +- Document support per capability (energy, forces, stress, geometry/cell + optimization), including environment requirements and evidence. +- Graduation requires documented interfaces, reference fixtures, failure tests, + reproducible installation, and a recorded real execution. + +Acceptance: regression cases reproduce and fix the custom-cutoff mismatch and +false RASPA2 failure. Every advertised capability has a support status. + +### 2. Shared contracts and provenance + +- Introduce a versioned common result envelope with operation-specific payloads. +- Separate scientific model identity, calculator adapter, and execution location. +- Standardize `potential_energy` with compatibility adapters for existing fields; + preserve legacy defaults through wrappers. +- Record requested/resolved settings, input and model hashes where available, + versions, precision, units, convergence, and artifact references. + +Acceptance: representative results validate and round-trip; legacy APIs remain +usable; MatKit's core remains independent of ChemGraph. + +### 3. Reproducibility, resume, and benchmarks + +- Add a licensed reference corpus of molecules, crystals, MOFs, varied sizes, + and intentionally invalid inputs. +- Establish matching-checkpoint numerical parity before measuring performance. + Use independent reference data for scientific accuracy claims. +- Measure startup, warm execution, throughput, peak memory, and failures + separately, recording the environment and repeated observations. +- Add explicit resume based on input/configuration/model identity. Reuse only + compatible completed results; rerun failures and unconverged optimizations. + +Acceptance: reproducible benchmark artifacts and tests that reject changed +resume inputs/settings and recover interrupted work without losing results. + +### 4. Flagship porous-material workflow + +Provide a scripted pipeline: validate structure, optionally desolvate, assign +charges, analyze pores, prepare/run adsorption, and analyze isotherms. Record +artifact lineage, scientific settings, equilibration, and uncertainty methods. + +Acceptance: fixture-driven execution in CPU CI and an opt-in real-engine run +with inspectable artifacts at every stage. + +### 5. Narrow ChemGraph integration + +- Start with pore analysis, adsorption preparation, and MLIP execution. +- Use Python for compatible environments and optional MCP services for isolated + environments; reuse ChemGraph's execution and job tracking. +- Keep large tensors/trajectories in retrievable artifacts. +- Consolidate duplicated scientific implementations after parity tests pass. + +Acceptance: tool discovery, invocation, structured results, artifact retrieval, +and failure/nonconvergence handling tested with deterministic fixtures. + +### 6. Tested skills and agent evaluations + +Add workflow recipes with use conditions, inputs, scientific choices, expected +outputs, and recovery. Connect them to an explicit ChemGraph skill-loading path. +Turn successful scripted workflows into agent evaluation tasks. + +Acceptance: recipe commands execute against tested APIs; evaluations check +numerical results, artifacts, tool selection, and failure handling. Agent +failures become new scientific/integration regression cases. + +## CI throughout the roadmap + +Test installed wheels and bundled templates, then selected optional-dependency +environments. Keep ordinary CI fast, deterministic, and CPU-only. Gate model +downloads, external engines, live endpoints, and GPUs explicitly. + +## Status + +Roadmap implementation has not started. PR #14 hardening is complete at +implementation commit `28e764f`; its handoff records 259 passing CPU tests and +one skip. Check the PR's current head and merge state before beginning milestone +1. No GPU capability has been promoted on the basis of this implementation. diff --git a/docs/plans/pr14-hardening.md b/docs/plans/pr14-hardening.md new file mode 100644 index 0000000..6c614df --- /dev/null +++ b/docs/plans/pr14-hardening.md @@ -0,0 +1,93 @@ +# PR #14 hardening and session handoff + +PR: https://github.com/tdpham2/MatKit/pull/14 + +Starting revision: `7a58171ee92513047517c0fee63eaa4f4e4f8573`. +Branch: `mlip-playground-backends`. This plan covers the new MLIP interface; +the independent growth plan is [matkit-roadmap.md](matkit-roadmap.md). + +## Agreed decisions + +- Persist completed work immediately; automatic resume is deferred. +- CLI exit 0 requires valid calculations and converged requested optimizations. + Calculation failure/nonconvergence exits 1; invalid arguments exit 2. +- `success` records numerically valid completion; `converged` separately records + optimization convergence. Keep usable unconverged results. +- Keep existing public entry points, configuration classes, `energy`, units, + and fixed-cell optimization scope. Full contract migration belongs later. +- Backends/capabilities without real execution evidence remain experimental. + CPU tests and prepared GPU recipes do not count as GPU validation. +- Update the existing PR; merging is a separate action. + +## Implementation checklist + +- [x] Validate finite configuration, structures, energy, forces, and stress; + distinguish unavailable stress from calculator failure; use strict JSON. +- [x] Reject explicitly unsupported CLI/example options; implement strict + exits and expose unconverged counts. +- [x] Persist ordered per-item results and manifests atomically during execution; + retain completed work after interruptions, teardown, and persistence errors. +- [x] Make mocked GPU tests hermetic; cover adapter boundaries; extend opt-in + Polaris energy/optimization/batch checks and record their environment. +- [x] Run the complete CPU suite, lint/format checks, package build, example + help checks, and shell syntax checks. + +## Persistence contract + +An initial manifest has ordered pending items and status `running`. A completed +item's result is committed before its manifest entry. Final execution statuses +remain `completed`, `partial`, and `failure`. Catchable orchestration failures +mark the manifest `interrupted` and propagate; an uncatchable termination can +leave `running`. Completed JSON files remain readable in either case. + +Manifests expose pending and unconverged counts, per-item convergence, and a +run-level error. Success counts include valid but unconverged results. Reruns +must use a fresh batch directory; existing manifests are not overwritten. + +## Acceptance scenarios + +- Reject NaN/Inf inputs/outputs and malformed tensors before reporting success. +- Preserve valid converged and unconverged optimization results. +- Forward applicable options and reject explicit unsupported options. +- Preserve ordering, duplicate basenames, calculator reuse, and partial failures. +- Preserve completed records after interruption, write failure, and teardown. +- Mocked tests require neither CUDA nor downloads; real GPU checks are opt-in. + +## Verification and next session + +Implementation is complete at `28e764f` (four focused implementation commits; +this handoff is a subsequent documentation-only change). PR #14's commit list +is the source of truth for the final published branch head and GitHub checks. + +| Commit | Change | +| --- | --- | +| `4346208` | Scientific input/result validation and strict JSON | +| `065d3cc` | Explicit option validation and strict CLI/example exits | +| `c646505` | Incremental atomic persistence and Rootstock worker isolation | +| `28e764f` | Adapter boundary tests, GPU recipe, CI checks, and documentation | + +Local verification on Python 3.12: + +- `PYTHONPATH=src pytest tests/ -q`: **259 passed, 1 skipped**. +- Ruff lint and format checks pass for `src/`, MLIP tests, CLI tests, + `examples/mlip_gpu.py`, and `alcf/polaris/mlip/smoke.py`. +- `python -m build --no-isolation --outdir /private/tmp/matkit-pr14-dist`: + source distribution and wheel built successfully. +- Both GPU Python entry points pass `--help`; both Polaris shell scripts + pass `bash -n`; `git diff --check` passes. + +No real MACE model, Rootstock deployment, ALCHEMI GPU kernel, or cluster job was +run during this implementation. Those capabilities remain experimental. The +opt-in smoke runner records six cases, environment/commit information, logs, +results, and convergence; it continues after individual case failures. + +Known limits: no automatic resume/retry, no parity/performance evidence, no +common provenance migration, and no GCMC fixes in this PR. ALCHEMI native +optimization step counts remain unknown (`n_steps=null`); per-item calculation +times for native batches include shared chunk work and are not throughput +measurements. Inputs/results remain accumulated in memory. Abrupt termination +can leave a `running` manifest and a batch lock; use a fresh directory. + +After PR #14 is reviewed, begin roadmap milestone 1: cutoff propagation, +RASPA2 success reporting, and explicit capability support status. Do not begin +resume, new backends, benchmarking, or ChemGraph changes as part of this PR. diff --git a/examples/mlip_gpu.py b/examples/mlip_gpu.py new file mode 100644 index 0000000..eee5687 --- /dev/null +++ b/examples/mlip_gpu.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +"""Run one MatKit MLIP backend on one or more structures using a GPU.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +from matkit.mlip import ( + ASEMACEConfig, + MLIPCalculationConfig, + NVAlchemiMACEConfig, + RootstockConfig, + run_mlip_batch, +) +from matkit.mlip.config import _validate_explicit_options + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Run direct ASE MACE, Rootstock, or NVIDIA ALCHEMI MACE on a GPU. " + "Use one backend per process to keep GPU runtime state isolated." + ) + ) + parser.add_argument( + "inputs", + nargs="+", + type=Path, + help="Structure files readable by ASE.", + ) + parser.add_argument( + "--backend", + choices=("ase-mace", "rootstock", "nvalchemi-mace"), + required=True, + ) + parser.add_argument( + "--checkpoint", + help=( + "Model alias or checkpoint path. Defaults to 'medium' for MACE and " + "ALCHEMI, or 'mace-mp-0-medium' for Rootstock." + ), + ) + parser.add_argument( + "--output-dir", + type=Path, + help="Result directory (default: mlip_gpu_results/).", + ) + parser.add_argument("--driver", choices=("energy", "opt"), default="energy") + parser.add_argument( + "--fmax", type=float, help="Optimization only (default: 0.01)." + ) + parser.add_argument( + "--steps", type=int, help="Optimization only (default: 1000)." + ) + parser.add_argument( + "--batch-size", + type=int, + help="Native batch size for ALCHEMI only (default: 16).", + ) + parser.add_argument( + "--cluster", + help="Rootstock cluster name (defaults to 'polaris').", + ) + parser.add_argument( + "--root", + type=Path, + help="Rootstock deployment root instead of a named cluster.", + ) + parser.add_argument( + "--compile-model", + action="store_true", + default=None, + help="Enable model compilation for NVIDIA ALCHEMI MACE.", + ) + parser.add_argument( + "--enable-cueq", + action="store_true", + default=None, + help="Enable CuEquivariance for NVIDIA ALCHEMI MACE.", + ) + return parser + + +def main() -> int: + parser = build_parser() + args = parser.parse_args() + + if args.cluster and args.root: + parser.error("--cluster and --root are mutually exclusive") + + provided = { + "root_path" if name == "root" else name + for name, value in vars(args).items() + if value is not None + } + try: + _validate_explicit_options( + args.backend, args.driver, "mace_mp", provided + ) + calculation = MLIPCalculationConfig( + driver=args.driver, + fmax=args.fmax if args.fmax is not None else 0.01, + steps=args.steps if args.steps is not None else 1000, + ) + if args.batch_size is not None and args.batch_size < 1: + raise ValueError("--batch-size must be at least 1") + except ValueError as exc: + parser.error(str(exc)) + + if args.backend == "ase-mace": + backend = ASEMACEConfig( + checkpoint=args.checkpoint or "medium", + device="cuda", + dtype="float32", + ) + elif args.backend == "rootstock": + backend = RootstockConfig( + checkpoint=args.checkpoint or "mace-mp-0-medium", + cluster=None if args.root else (args.cluster or "polaris"), + root=str(args.root) if args.root else None, + device="cuda", + ) + else: + backend = NVAlchemiMACEConfig( + checkpoint=args.checkpoint or "medium", + device="cuda", + dtype="float32", + compile_model=bool(args.compile_model), + enable_cueq=bool(args.enable_cueq), + ) + + output_dir = args.output_dir or Path("mlip_gpu_results") / args.backend + try: + summary = run_mlip_batch( + args.inputs, + backend, + calculation, + output_dir=output_dir, + batch_size=args.batch_size if args.batch_size is not None else 16, + ) + except Exception as exc: + print(json.dumps({"status": "failure", "error": str(exc)}, indent=2)) + print(str(exc), file=sys.stderr) + return 1 + unconverged = ( + sum( + result["success"] and not result["converged"] + for result in summary["results"] + ) + if calculation.driver == "opt" + else 0 + ) + + print( + json.dumps( + { + "status": summary["status"], + "backend": args.backend, + "total": summary["total"], + "succeeded": summary["succeeded"], + "failed": summary["failed"], + "unconverged": unconverged, + "manifest_file": summary["manifest_file"], + }, + indent=2, + ) + ) + if summary["failed"] or unconverged: + print( + f"{summary['failed']} failed, {unconverged} unconverged; " + "available results retained.", + file=sys.stderr, + ) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/pyproject.toml b/pyproject.toml index 711c5bd..bce10c1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,11 +18,13 @@ dependencies = [ [project.optional-dependencies] rdkit = ["rdkit"] mlip = ["mace-torch"] +rootstock = ["rootstock>=1.6,<2"] +nvalchemi_mace = ["nvalchemi-toolkit[mace]>=0.2,<0.3"] plot = ["matplotlib>=3.5"] -pacmof2 = ["pacmof2"] +pacmof2 = ["pacmof2 @ git+https://github.com/snurr-group/pacmof2.git"] graspa = ["pyyaml>=6.0"] pygraspa = ["pyyaml>=6.0"] -all = ["rdkit", "mace-torch", "matplotlib>=3.5", "pacmof2", "pyyaml>=6.0"] +all = ["rdkit", "mace-torch", "matplotlib>=3.5", "pacmof2 @ git+https://github.com/snurr-group/pacmof2.git", "pyyaml>=6.0"] dev = ["pytest>=7.0", "ruff>=0.4"] [project.scripts] diff --git a/skills.md b/skills.md index fff7ab2..66eaf17 100644 --- a/skills.md +++ b/skills.md @@ -19,7 +19,7 @@ src/matkit/ raspa3/ # RASPA2 -> RASPA3 format conversion zeopp/ # Zeo++ pore geometry analysis (wraps network binary) tobacco/ # SMILES -> CIF linker generation for ToBaCCo - mlip/ # MACE-MP ML interatomic potential optimization + mlip/ # Direct, Rootstock, and ALCHEMI MLIP execution orca/ # ORCA quantum chemistry (stub) io/ # File format converters (SMILES, PubChem JSON) utils/ # Shared utilities (unit cell calc, solvent removal, CIF sampling) @@ -50,7 +50,8 @@ matkit - **Language**: Python >= 3.10 - **Core deps**: ase (atomic simulation), click (CLI), networkx (graph analysis), numpy -- **Optional deps**: rdkit (SMILES), mace-torch (MLIP), openbabel CLI (obabel) +- **Optional deps**: rdkit, mace-torch, rootstock, nvalchemi-toolkit, + openbabel CLI (obabel) - **Build**: setuptools via pyproject.toml (PEP 621) - **Linting**: ruff (E, F rules, 80 char line length) - **Testing**: pytest (tests/ directory) diff --git a/src/matkit/cli.py b/src/matkit/cli.py index 6449d4b..bc0eb62 100644 --- a/src/matkit/cli.py +++ b/src/matkit/cli.py @@ -1012,6 +1012,379 @@ def mlip_cli(): pass +def _mlip_options(function): + """Add runtime-neutral MLIP options to a Click command.""" + decorators = [ + click.option( + "--backend", + required=True, + type=click.Choice(["ase-mace", "rootstock", "nvalchemi-mace"]), + help="MLIP execution backend.", + ), + click.option( + "--checkpoint", + required=True, + help="Model alias, canonical Rootstock ID, or checkpoint path.", + ), + click.option( + "--device", default=None, help="Device such as cpu or cuda." + ), + click.option( + "--dtype", + default=None, + type=click.Choice(["float32", "float64"]), + help="Floating-point precision.", + ), + click.option( + "--driver", + default="energy", + show_default=True, + type=click.Choice(["energy", "opt"]), + ), + click.option( + "--optimizer", + default="fire", + show_default=True, + type=click.Choice(["bfgs", "lbfgs", "gpmin", "fire", "mdmin"]), + ), + click.option("--fmax", default=0.01, show_default=True, type=float), + click.option("--steps", default=1000, show_default=True, type=int), + click.option( + "--calculator-type", + default="mace_mp", + show_default=True, + type=click.Choice(["mace_mp", "mace_off", "mace_anicc"]), + help="Direct ASE MACE calculator factory.", + ), + click.option( + "--dispersion/--no-dispersion", + default=False, + help="Enable MACE-MP D3 dispersion.", + ), + click.option("--cluster", default=None, help="Rootstock cluster ID."), + click.option( + "--root", + "root_path", + default=None, + type=click.Path(), + help="Custom Rootstock installation root.", + ), + click.option( + "--cache-root", + default=None, + type=click.Path(), + help="Custom Rootstock cache root.", + ), + click.option( + "--setup-kwarg", + multiple=True, + metavar="KEY=JSON", + help="Rootstock setup keyword; repeat as needed.", + ), + click.option( + "--timeout", + default=600.0, + show_default=True, + type=float, + help="Rootstock worker startup timeout.", + ), + click.option( + "--weights", + default=None, + type=click.Path(), + help="Rootstock custom checkpoint weights.", + ), + click.option( + "--dt", + default=0.1, + show_default=True, + type=float, + help="ALCHEMI FIRE timestep.", + ), + click.option( + "--compile-model", + is_flag=True, + help="Compile the ALCHEMI MACE model.", + ), + click.option( + "--enable-cueq", + is_flag=True, + help="Enable cuEquivariance in ALCHEMI MACE.", + ), + ] + for decorator in reversed(decorators): + function = decorator(function) + return function + + +def _parse_setup_kwargs(values): + parsed = {} + for value in values: + if "=" not in value: + raise click.BadParameter( + "must use KEY=JSON syntax", param_hint="--setup-kwarg" + ) + key, raw = value.split("=", 1) + if not key: + raise click.BadParameter( + "key must not be empty", param_hint="--setup-kwarg" + ) + try: + parsed[key] = json.loads(raw) + except json.JSONDecodeError: + parsed[key] = raw + return parsed + + +def _build_mlip_configs(options): + from matkit.mlip import ( + ASEMACEConfig, + MLIPCalculationConfig, + NVAlchemiMACEConfig, + RootstockConfig, + ) + from matkit.mlip.config import _validate_explicit_options + + ctx = click.get_current_context() + provided = { + name + for name in ctx.params + if ctx.get_parameter_source(name) + not in ( + None, + click.core.ParameterSource.DEFAULT, + ) + } + _validate_explicit_options( + options["backend"], + options["driver"], + options["calculator_type"], + provided, + ) + + backend_name = options["backend"] + checkpoint = options["checkpoint"] + device = options["device"] + dtype = options["dtype"] + if backend_name == "ase-mace": + backend = ASEMACEConfig( + checkpoint=checkpoint, + device=device or "cpu", + dtype=dtype or "float64", + calculator_type=options["calculator_type"], + dispersion=options["dispersion"], + ) + elif backend_name == "rootstock": + backend = RootstockConfig( + checkpoint=checkpoint, + cluster=options["cluster"], + root=options["root_path"], + cache_root=options["cache_root"], + setup_kwargs=_parse_setup_kwargs(options["setup_kwarg"]), + timeout=options["timeout"], + weights=options["weights"], + device=device or "cpu", + ) + else: + backend = NVAlchemiMACEConfig( + checkpoint=checkpoint, + device=device or "cuda", + dtype=dtype or "float32", + dt=options["dt"], + compile_model=options["compile_model"], + enable_cueq=options["enable_cueq"], + ) + calculation = MLIPCalculationConfig( + driver=options["driver"], + optimizer=options["optimizer"], + fmax=options["fmax"], + steps=options["steps"], + ) + if ( + backend_name == "nvalchemi-mace" + and calculation.driver == "opt" + and calculation.optimizer != "fire" + ): + raise ValueError("NVIDIA ALCHEMI supports only the FIRE optimizer") + return backend, calculation + + +def _mlip_configs_or_usage_error(options): + try: + return _build_mlip_configs(options) + except ValueError as exc: + raise click.BadParameter(str(exc)) from exc + + +@mlip_cli.command("run") +@click.option( + "--input", + "input_file", + required=True, + type=click.Path(exists=True, dir_okay=False), + help="Input structure readable by ASE.", +) +@click.option( + "--output", + default="output.json", + show_default=True, + type=click.Path(dir_okay=False), + help="JSON result file.", +) +@_mlip_options +def mlip_run_cmd(input_file, output, **options): + """Run one MLIP energy calculation or fixed-cell optimization.""" + from matkit.mlip import run_mlip + + backend, calculation = _mlip_configs_or_usage_error(options) + try: + result = run_mlip( + input_file, + backend, + calculation=calculation, + output_file=output, + ) + except Exception as exc: + click.echo( + json.dumps({"status": "failure", "error": str(exc)}, indent=2) + ) + raise click.ClickException(str(exc)) from exc + unconverged = ( + result["success"] + and calculation.driver == "opt" + and not result["converged"] + ) + click.echo( + json.dumps( + { + "status": ( + "failure" + if not result["success"] + else "unconverged" + if unconverged + else "success" + ), + "energy": result["energy"], + "unit": result["energy_unit"], + "converged": result["converged"], + "output_results_file": str(output), + "error": result.get("error", ""), + }, + indent=2, + ) + ) + if not result["success"]: + raise click.ClickException(result["error"]) + if unconverged: + raise click.ClickException( + "Optimization did not converge; results retained." + ) + + +@mlip_cli.command("run-batch") +@click.option( + "--input", + "input_files", + multiple=True, + type=click.Path(exists=True, dir_okay=False), + help="Input structure; repeat for an ordered list.", +) +@click.option( + "--input-dir", + default=None, + type=click.Path(exists=True, file_okay=False), + help="Directory containing input structures.", +) +@click.option( + "--pattern", + default="*.cif", + show_default=True, + help="Glob used with --input-dir.", +) +@click.option( + "--outdir", + default="mlip_results", + show_default=True, + type=click.Path(file_okay=False), +) +@click.option( + "--batch-size", default=16, show_default=True, type=click.IntRange(min=1) +) +@click.option("--max-atoms", default=None, type=click.IntRange(min=1)) +@_mlip_options +def mlip_run_batch_cmd( + input_files, + input_dir, + pattern, + outdir, + batch_size, + max_atoms, + **options, +): + """Run an ordered MLIP batch and write a JSON manifest.""" + from pathlib import Path + + from matkit.mlip import run_mlip_batch + + if bool(input_files) == bool(input_dir): + raise click.UsageError("Specify exactly one of --input or --input-dir.") + if input_dir: + files = [ + str(path) + for path in sorted(Path(input_dir).glob(pattern)) + if path.is_file() + ] + if not files: + raise click.UsageError( + f"No files matching {pattern!r} in {input_dir}." + ) + else: + files = list(input_files) + + backend, calculation = _mlip_configs_or_usage_error(options) + try: + summary = run_mlip_batch( + files, + backend, + calculation=calculation, + output_dir=outdir, + batch_size=batch_size, + max_atoms=max_atoms, + ) + except Exception as exc: + click.echo( + json.dumps({"status": "failure", "error": str(exc)}, indent=2) + ) + raise click.ClickException(str(exc)) from exc + unconverged = ( + sum( + result["success"] and not result["converged"] + for result in summary["results"] + ) + if calculation.driver == "opt" + else 0 + ) + click.echo( + json.dumps( + { + "status": summary["status"], + "manifest_file": summary["manifest_file"], + "total": summary["total"], + "succeeded": summary["succeeded"], + "failed": summary["failed"], + "unconverged": unconverged, + }, + indent=2, + ) + ) + if summary["failed"] or unconverged: + raise click.ClickException( + f"MLIP batch: {summary['failed']} failed, " + f"{unconverged} unconverged; " + "available results retained." + ) + + @mlip_cli.command("mace-opt") @click.option( "--fname", diff --git a/src/matkit/mlip/__init__.py b/src/matkit/mlip/__init__.py index 83231a6..897e32e 100644 --- a/src/matkit/mlip/__init__.py +++ b/src/matkit/mlip/__init__.py @@ -1,4 +1,24 @@ -__all__ = [] +from matkit.mlip.config import ( + ASEMACEConfig, + MLIPBackendConfig, + MLIPCalculationConfig, + NVAlchemiMACEConfig, + RootstockConfig, +) +from matkit.mlip.runner import run_mlip, run_mlip_batch +from matkit.types import MLIPBatchSummary, MLIPResult + +__all__ = [ + "ASEMACEConfig", + "MLIPBackendConfig", + "MLIPBatchSummary", + "MLIPCalculationConfig", + "MLIPResult", + "NVAlchemiMACEConfig", + "RootstockConfig", + "run_mlip", + "run_mlip_batch", +] try: from matkit.mlip.mace_opt import run_opt_mace diff --git a/src/matkit/mlip/config.py b/src/matkit/mlip/config.py new file mode 100644 index 0000000..9499594 --- /dev/null +++ b/src/matkit/mlip/config.py @@ -0,0 +1,187 @@ +"""Configuration objects for runtime-selectable MLIP calculations.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +import json +import math +from numbers import Integral, Real +from typing import Any, Literal, TypeAlias + + +_DTYPES = {"float32", "float64"} +_OPTIMIZERS = {"bfgs", "lbfgs", "gpmin", "fire", "mdmin"} + + +def _positive_number(name: str, value: Any) -> None: + if ( + isinstance(value, bool) + or not isinstance(value, Real) + or not math.isfinite(value) + or value <= 0 + ): + raise ValueError(f"{name} must be positive and finite") + + +def _positive_integer(name: str, value: Any) -> None: + if isinstance(value, bool) or not isinstance(value, Integral) or value < 1: + raise ValueError(f"{name} must be at least 1 and an integer") + + +def _validate_explicit_options( + backend: str, driver: str, calculator_type: str, provided: set[str] +) -> None: + """Reject explicit frontend options the selected calculation cannot use.""" + backend_options = { + "ase-mace": {"calculator_type", "dispersion"}, + "rootstock": { + "cluster", + "root_path", + "cache_root", + "setup_kwarg", + "timeout", + "weights", + }, + "nvalchemi-mace": { + "dt", + "compile_model", + "enable_cueq", + "batch_size", + "max_atoms", + }, + } + for owner, names in backend_options.items(): + for name in sorted(provided & names): + if backend != owner: + flag = "root" if name == "root_path" else name.replace("_", "-") + raise ValueError(f"--{flag} requires --backend {owner}") + if "dtype" in provided: + if backend == "rootstock": + raise ValueError( + "--dtype is not supported by Rootstock; use --setup-kwarg " + "with precision settings supported by the deployed model" + ) + if backend == "ase-mace" and calculator_type == "mace_anicc": + raise ValueError("--dtype is controlled by the mace_anicc factory") + if "dispersion" in provided and calculator_type != "mace_mp": + raise ValueError("--dispersion requires --calculator-type mace_mp") + if driver == "energy": + opt_only = provided & {"optimizer", "fmax", "steps", "dt"} + if opt_only: + flag = sorted(opt_only)[0].replace("_", "-") + raise ValueError(f"--{flag} requires --driver opt") + + +@dataclass(frozen=True) +class ASEMACEConfig: + """Run a MACE calculator directly through ASE.""" + + checkpoint: str = "medium" + device: str = "cpu" + dtype: Literal["float32", "float64"] = "float64" + calculator_type: Literal["mace_mp", "mace_off", "mace_anicc"] = "mace_mp" + dispersion: bool = False + damping: str = "bj" + dispersion_xc: str = "pbe" + dispersion_cutoff: float = 21.167088422553647 + type: Literal["ase-mace"] = field(default="ase-mace", init=False) + + def __post_init__(self) -> None: + if not self.checkpoint: + raise ValueError("checkpoint must not be empty") + if self.dtype not in _DTYPES: + raise ValueError(f"Unsupported dtype: {self.dtype}") + if self.calculator_type not in { + "mace_mp", + "mace_off", + "mace_anicc", + }: + raise ValueError( + f"Unsupported MACE calculator type: {self.calculator_type}" + ) + _positive_number("dispersion_cutoff", self.dispersion_cutoff) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class RootstockConfig: + """Run a Rootstock-managed checkpoint through its ASE calculator.""" + + checkpoint: str + cluster: str | None = None + root: str | None = None + cache_root: str | None = None + setup_kwargs: dict[str, Any] = field(default_factory=dict) + timeout: float = 600.0 + weights: str | None = None + device: str = "cpu" + type: Literal["rootstock"] = field(default="rootstock", init=False) + + def __post_init__(self) -> None: + if not self.checkpoint: + raise ValueError("checkpoint must not be empty") + if self.cluster is not None and self.root is not None: + raise ValueError("Rootstock cannot specify both cluster and root") + _positive_number("timeout", self.timeout) + try: + json.dumps(self.setup_kwargs, allow_nan=False) + except (TypeError, ValueError) as exc: + raise ValueError( + "setup_kwargs must contain finite JSON data" + ) from exc + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class NVAlchemiMACEConfig: + """Run MACE using NVIDIA ALCHEMI Toolkit.""" + + checkpoint: str + device: str = "cuda" + dtype: Literal["float32", "float64"] = "float32" + dt: float = 0.1 + compile_model: bool = False + enable_cueq: bool = False + type: Literal["nvalchemi-mace"] = field( + default="nvalchemi-mace", init=False + ) + + def __post_init__(self) -> None: + if not self.checkpoint: + raise ValueError("checkpoint must not be empty") + if self.dtype not in _DTYPES: + raise ValueError(f"Unsupported dtype: {self.dtype}") + _positive_number("dt", self.dt) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +MLIPBackendConfig: TypeAlias = ( + ASEMACEConfig | RootstockConfig | NVAlchemiMACEConfig +) + + +@dataclass(frozen=True) +class MLIPCalculationConfig: + """Calculation settings shared by all MLIP backends.""" + + driver: Literal["energy", "opt"] = "energy" + optimizer: Literal["bfgs", "lbfgs", "gpmin", "fire", "mdmin"] = "fire" + fmax: float = 0.01 + steps: int = 1000 + + def __post_init__(self) -> None: + if self.driver not in {"energy", "opt"}: + raise ValueError(f"Unsupported MLIP driver: {self.driver}") + if self.optimizer not in _OPTIMIZERS: + raise ValueError(f"Unsupported ASE optimizer: {self.optimizer}") + _positive_number("fmax", self.fmax) + _positive_integer("steps", self.steps) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) diff --git a/src/matkit/mlip/runner.py b/src/matkit/mlip/runner.py new file mode 100644 index 0000000..3ad94dc --- /dev/null +++ b/src/matkit/mlip/runner.py @@ -0,0 +1,806 @@ +"""Plain-Python execution core for runtime-selectable MLIPs.""" + +from __future__ import annotations + +import json +import logging +import os +import tempfile +import time +from contextlib import ExitStack, contextmanager +from pathlib import Path +from typing import Any, Callable, Iterator, Sequence + +import numpy as np +from ase.calculators.calculator import PropertyNotImplementedError +from ase.io import read as ase_read + +from matkit.mlip.config import ( + ASEMACEConfig, + MLIPBackendConfig, + MLIPCalculationConfig, + NVAlchemiMACEConfig, + RootstockConfig, + _positive_integer, +) +from matkit.types import MLIPBatchSummary, MLIPResult + +logger = logging.getLogger(__name__) +_ResultCallback = Callable[[int, dict[str, Any]], None] + + +def _finite_array(name: str, value, shape: tuple[int, ...]) -> np.ndarray: + array = np.asarray(value, dtype=float) + if array.shape != shape or not np.isfinite(array).all(): + raise ValueError(f"{name} must have shape {shape} and finite values") + return array + + +def _validate_atoms(atoms) -> None: + if not len(atoms): + raise ValueError("Structure must contain at least one atom") + _finite_array("positions", atoms.positions, (len(atoms), 3)) + cell = _finite_array("cell", atoms.cell.array, (3, 3)) + periodic = cell[atoms.pbc] + if len(periodic) and np.linalg.matrix_rank(periodic) != len(periodic): + raise ValueError( + "Periodic cell vectors must be nonzero and independent" + ) + + +def _atoms_payload(atoms) -> dict[str, Any]: + """Convert the portable portion of an ASE Atoms object to JSON data.""" + return { + "atomic_numbers": atoms.get_atomic_numbers().tolist(), + "positions": atoms.get_positions().tolist(), + "cell": atoms.cell.array.tolist(), + "pbc": atoms.pbc.tolist(), + } + + +def _optional_stress(atoms): + if not atoms.pbc.any() or atoms.cell.rank != 3: + return None + try: + return atoms.get_stress(voigt=False) + except PropertyNotImplementedError: + return None + + +def _synchronize_device(device: str) -> None: + if not device.startswith("cuda"): + return + import torch + + if not torch.cuda.is_available(): + raise RuntimeError(f"CUDA is unavailable for requested device {device}") + torch.cuda.synchronize(torch.device(device)) + + +def _create_mace_calculator(config: ASEMACEConfig): + try: + import mace.calculators as mace_calculators + except ImportError as exc: + raise ImportError( + "Direct MACE requires the 'mlip' extra: pip install matkit[mlip]" + ) from exc + + try: + factory = getattr(mace_calculators, config.calculator_type) + except AttributeError as exc: + raise ImportError( + f"Installed mace-torch does not provide {config.calculator_type}." + ) from exc + + if config.calculator_type == "mace_anicc": + if config.dispersion: + raise ValueError( + "Dispersion options are supported only by " + "calculator_type='mace_mp'." + ) + return factory(device=config.device, model_path=config.checkpoint) + + kwargs: dict[str, Any] = { + "model": config.checkpoint, + "device": config.device, + "default_dtype": config.dtype, + } + if config.calculator_type == "mace_mp": + kwargs["dispersion"] = config.dispersion + if config.dispersion: + kwargs.update( + { + "damping": config.damping, + "dispersion_xc": config.dispersion_xc, + "dispersion_cutoff": config.dispersion_cutoff, + } + ) + elif config.dispersion: + raise ValueError( + "Dispersion options are supported only by calculator_type=" + "'mace_mp'." + ) + return factory(**kwargs) + + +@contextmanager +def _ase_backend_context( + config: ASEMACEConfig | RootstockConfig, +) -> Iterator[Any]: + """Create one ASE calculator and retain it for the whole request.""" + if isinstance(config, ASEMACEConfig): + yield _create_mace_calculator(config) + return + + try: + from rootstock import RootstockCalculator + except ImportError as exc: + raise ImportError( + "Rootstock requires its optional extra: " + "pip install matkit[rootstock]" + ) from exc + + kwargs = { + "checkpoint": config.checkpoint, + "cluster": config.cluster, + "root": config.root, + "cache_root": config.cache_root, + "device": config.device, + "setup_kwargs": config.setup_kwargs, + "timeout": config.timeout, + "weights": config.weights, + } + kwargs = {key: value for key, value in kwargs.items() if value is not None} + with RootstockCalculator(**kwargs) as calculator: + yield calculator + + +def _optimizer_class(name: str): + from ase.optimize import BFGS, FIRE, GPMin, LBFGS, MDMin + + return { + "bfgs": BFGS, + "lbfgs": LBFGS, + "gpmin": GPMin, + "fire": FIRE, + "mdmin": MDMin, + }[name] + + +def _success_result( + input_file: str, + backend: MLIPBackendConfig, + calculation: MLIPCalculationConfig, + atoms, + energy: float, + forces, + stress, + converged: bool, + n_steps: int | None, + calculation_time: float, +) -> dict[str, Any]: + _validate_atoms(atoms) + energy = _finite_array("energy", energy, ()).item() + forces = _finite_array("forces", forces, (len(atoms), 3)) + if stress is not None: + stress = _finite_array("stress", stress, (3, 3)) + return { + "schema_version": 1, + "success": True, + "error": "", + "input_structure_file": input_file, + "backend_info": backend.to_dict(), + "calculation_input": calculation.to_dict(), + "energy": float(energy), + "energy_unit": "eV", + "forces": None if forces is None else forces.tolist(), + "force_unit": "eV/angstrom", + "stress": None if stress is None else stress.tolist(), + "stress_unit": "eV/angstrom^3", + "converged": bool(converged), + "n_steps": n_steps, + "final_structure": _atoms_payload(atoms), + "calculation_time_s": calculation_time, + } + + +def _failure_result( + input_file: str, + backend: MLIPBackendConfig, + calculation: MLIPCalculationConfig, + error: Exception | str, + calculation_time: float = 0.0, +) -> dict[str, Any]: + return { + "schema_version": 1, + "success": False, + "error": str(error), + "input_structure_file": input_file, + "backend_info": backend.to_dict(), + "calculation_input": calculation.to_dict(), + "energy": None, + "energy_unit": "eV", + "forces": None, + "force_unit": "eV/angstrom", + "stress": None, + "stress_unit": "eV/angstrom^3", + "converged": False, + "n_steps": None, + "final_structure": None, + "calculation_time_s": calculation_time, + } + + +def _run_ase_item( + input_file: str, + atoms, + calculator, + backend: ASEMACEConfig | RootstockConfig, + calculation: MLIPCalculationConfig, +) -> dict[str, Any]: + atoms.calc = calculator + # Rootstock's synchronous worker owns the GPU and its dependencies. + if isinstance(backend, ASEMACEConfig): + _synchronize_device(backend.device) + started = time.perf_counter() + converged = True + n_steps = 0 + if calculation.driver == "opt": + optimizer = _optimizer_class(calculation.optimizer)(atoms, logfile=None) + converged = bool( + optimizer.run(fmax=calculation.fmax, steps=calculation.steps) + ) + n_steps = optimizer.nsteps + + energy = atoms.get_potential_energy() + forces = atoms.get_forces() + stress = _optional_stress(atoms) + if isinstance(backend, ASEMACEConfig): + _synchronize_device(backend.device) + elapsed = time.perf_counter() - started + return _success_result( + input_file, + backend, + calculation, + atoms, + energy, + forces, + stress, + converged, + n_steps, + elapsed, + ) + + +def _load_nvalchemi_model(config: NVAlchemiMACEConfig): + try: + import torch + from nvalchemi.models.mace import MACEWrapper + except ImportError as exc: + raise ImportError( + "NVIDIA ALCHEMI MACE requires the 'nvalchemi_mace' " + "extra and a matching CUDA extra." + ) from exc + + model = MACEWrapper.from_checkpoint( + config.checkpoint, + device=torch.device(config.device), + dtype=getattr(torch, config.dtype), + enable_cueq=config.enable_cueq, + compile_model=config.compile_model, + ) + model.eval() + return model + + +def _atoms_to_nvalchemi_data(atoms, config: NVAlchemiMACEConfig): + try: + import torch + from nvalchemi.data import AtomicData + except ImportError as exc: + raise ImportError( + "NVIDIA ALCHEMI MACE requires the 'nvalchemi_mace' " + "extra and a matching CUDA extra." + ) from exc + + dtype = getattr(torch, config.dtype) + data = AtomicData.from_atoms( + atoms, + device=torch.device(config.device), + dtype=dtype, + ) + data.forces = torch.zeros( + data.num_nodes, 3, device=data.device, dtype=dtype + ) + data.energy = torch.zeros(1, 1, device=data.device, dtype=dtype) + data.velocities = torch.zeros( + data.num_nodes, 3, device=data.device, dtype=dtype + ) + return data + + +def _nvalchemi_result( + input_file: str, + original_atoms, + data, + backend: NVAlchemiMACEConfig, + calculation: MLIPCalculationConfig, + converged: bool, + started: float, +) -> dict[str, Any]: + final_atoms = original_atoms.copy() + final_atoms.positions = data.positions.detach().cpu().numpy() + if data.cell is not None: + final_atoms.cell = data.cell.squeeze(0).detach().cpu().numpy() + if data.pbc is not None: + final_atoms.pbc = data.pbc.squeeze(0).detach().cpu().numpy() + + energy_values = data.energy.detach().cpu().numpy() + if energy_values.size != 1: + raise ValueError("ALCHEMI must return one energy per structure") + energy = energy_values.reshape(()).item() + forces = None + if data.forces is not None: + forces = data.forces.detach().cpu().numpy() + stress = None + if data.stress is not None: + stress = data.stress.detach().cpu().numpy() + if stress.shape == (1, 3, 3): + stress = stress[0] + return _success_result( + input_file, + backend, + calculation, + final_atoms, + energy, + forces, + stress, + converged, + None if calculation.driver == "opt" else 0, + time.perf_counter() - started, + ) + + +def _run_nvalchemi_chunk( + model, + entries: Sequence[tuple[int, str, Any]], + backend: NVAlchemiMACEConfig, + calculation: MLIPCalculationConfig, +) -> list[tuple[int, dict[str, Any]]]: + try: + from nvalchemi.data import Batch + from nvalchemi.dynamics import BaseDynamics, ConvergenceHook, FIRE + except ImportError as exc: + raise ImportError( + "NVIDIA ALCHEMI MACE requires the 'nvalchemi_mace' " + "extra and a matching CUDA extra." + ) from exc + + started = time.perf_counter() + data_list = [ + _atoms_to_nvalchemi_data(atoms, backend) for _, _, atoms in entries + ] + batch = Batch.from_data_list(data_list) + hooks = model.make_neighbor_hooks() + convergence = None + if calculation.driver == "energy": + dynamics = BaseDynamics(model=model, hooks=hooks, n_steps=1) + else: + convergence = ConvergenceHook.from_fmax(calculation.fmax) + dynamics = FIRE( + model=model, + dt=backend.dt, + hooks=hooks, + convergence_hook=convergence, + n_steps=calculation.steps, + ) + + _synchronize_device(backend.device) + with dynamics: + batch = dynamics.run(batch) + _synchronize_device(backend.device) + + converged_indices = set(range(len(entries))) + if convergence is not None: + indices = convergence.evaluate(batch) + converged_indices = ( + set() if indices is None else set(indices.detach().cpu().tolist()) + ) + + data_items = batch.to_data_list() + if len(data_items) != len(entries): + raise RuntimeError("ALCHEMI returned a different number of structures") + results = [] + for chunk_index, ((original_index, input_file, atoms), data) in enumerate( + zip(entries, data_items) + ): + try: + result = _nvalchemi_result( + input_file, + atoms, + data, + backend, + calculation, + chunk_index in converged_indices, + started, + ) + except Exception as exc: + result = _failure_result( + input_file, + backend, + calculation, + exc, + time.perf_counter() - started, + ) + results.append((original_index, result)) + return results + + +def _chunks_by_capacity( + entries: Sequence[tuple[int, str, Any]], + batch_size: int, + max_atoms: int | None, +) -> Iterator[list[tuple[int, str, Any]]]: + chunk: list[tuple[int, str, Any]] = [] + atom_count = 0 + for entry in entries: + n_atoms = len(entry[2]) + exceeds_atoms = ( + max_atoms is not None + and bool(chunk) + and atom_count + n_atoms > max_atoms + ) + if len(chunk) >= batch_size or exceeds_atoms: + yield chunk + chunk = [] + atom_count = 0 + chunk.append(entry) + atom_count += n_atoms + if chunk: + yield chunk + + +def _read_inputs( + input_files: Sequence[str | Path], + backend: MLIPBackendConfig, + calculation: MLIPCalculationConfig, + on_result: _ResultCallback | None = None, +) -> tuple[list[tuple[int, str, Any]], list[dict[str, Any] | None]]: + prepared = [] + results: list[dict[str, Any] | None] = [None] * len(input_files) + for index, value in enumerate(input_files): + input_file = str(Path(value).expanduser().resolve()) + try: + if not Path(input_file).is_file(): + raise FileNotFoundError( + f"Input structure file does not exist: {value}" + ) + atoms = ase_read(input_file) + _validate_atoms(atoms) + prepared.append((index, input_file, atoms)) + except Exception as exc: + results[index] = _failure_result( + input_file, backend, calculation, exc + ) + results[index]["setup_time_s"] = 0.0 + # Persistence errors must escape, not become calculation failures. + if results[index] is not None and on_result is not None: + on_result(index, results[index]) + return prepared, results + + +def _validate_execution_request( + input_files: Sequence[str | Path], + backend: MLIPBackendConfig, + calculation: MLIPCalculationConfig, + batch_size: int, + max_atoms: int | None, +) -> None: + if not input_files: + raise ValueError("At least one input structure file is required") + _positive_integer("batch_size", batch_size) + if max_atoms is not None: + _positive_integer("max_atoms", max_atoms) + if ( + isinstance(backend, NVAlchemiMACEConfig) + and calculation.driver == "opt" + and calculation.optimizer != "fire" + ): + raise ValueError("NVIDIA ALCHEMI supports only the FIRE optimizer") + + +def _execute_inputs( + input_files: Sequence[str | Path], + backend: MLIPBackendConfig, + calculation: MLIPCalculationConfig, + batch_size: int, + max_atoms: int | None, + on_result: _ResultCallback | None = None, +) -> tuple[list[dict[str, Any]], float]: + _validate_execution_request( + input_files, backend, calculation, batch_size, max_atoms + ) + prepared, results = _read_inputs( + input_files, backend, calculation, on_result + ) + if not prepared: + return [result for result in results if result is not None], 0.0 + + def complete(index, result): + result["setup_time_s"] = setup_time + results[index] = result + if on_result is not None: + on_result(index, result) + + setup_started = time.perf_counter() + if isinstance(backend, (ASEMACEConfig, RootstockConfig)): + # Only startup failures apply to every prepared item. Body/teardown + # errors must not replace already committed results. + with ExitStack() as stack: + try: + calculator = stack.enter_context(_ase_backend_context(backend)) + if isinstance(backend, ASEMACEConfig): + _synchronize_device(backend.device) + except Exception as exc: + setup_time = time.perf_counter() - setup_started + for index, input_file, _ in prepared: + complete( + index, + _failure_result(input_file, backend, calculation, exc), + ) + else: + setup_time = time.perf_counter() - setup_started + for index, input_file, atoms in prepared: + item_started = time.perf_counter() + try: + result = _run_ase_item( + input_file, + atoms, + calculator, + backend, + calculation, + ) + except Exception as exc: + result = _failure_result( + input_file, + backend, + calculation, + exc, + time.perf_counter() - item_started, + ) + complete(index, result) + else: + try: + model = _load_nvalchemi_model(backend) + _synchronize_device(backend.device) + except Exception as exc: + setup_time = time.perf_counter() - setup_started + for index, input_file, _ in prepared: + complete( + index, + _failure_result(input_file, backend, calculation, exc), + ) + else: + setup_time = time.perf_counter() - setup_started + for chunk in _chunks_by_capacity(prepared, batch_size, max_atoms): + chunk_started = time.perf_counter() + try: + chunk_results = _run_nvalchemi_chunk( + model, chunk, backend, calculation + ) + if [index for index, _ in chunk_results] != [ + index for index, _, _ in chunk + ]: + raise RuntimeError( + "ALCHEMI returned inconsistent result indices." + ) + except Exception as exc: + chunk_results = [ + ( + index, + _failure_result( + input_file, + backend, + calculation, + exc, + time.perf_counter() - chunk_started, + ), + ) + for index, input_file, _ in chunk + ] + for index, result in chunk_results: + complete(index, result) + + final_results = [] + for result in results: + if result is None: + raise RuntimeError("Internal MLIP result ordering error") + final_results.append(result) + return final_results, setup_time + + +def _write_json(path: Path, data: dict[str, Any]) -> str: + """Commit strict JSON with atomic replacement on the same filesystem.""" + encoded = json.dumps(data, indent=2, allow_nan=False) + path.parent.mkdir(parents=True, exist_ok=True) + temporary = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + delete=False, + ) as stream: + temporary = Path(stream.name) + stream.write(encoded) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + finally: + if temporary is not None: + temporary.unlink(missing_ok=True) + return str(path.resolve()) + + +def run_mlip( + input_file: str | Path, + backend: MLIPBackendConfig, + calculation: MLIPCalculationConfig | None = None, + output_file: str | Path | None = None, +) -> MLIPResult: + """Run one energy calculation or fixed-cell optimization.""" + calculation = calculation or MLIPCalculationConfig() + + def persist(_index, result): + path = Path(output_file).expanduser().resolve() + result["output_results_file"] = str(path) + _write_json(path, result) + + results, _ = _execute_inputs( + [input_file], + backend, + calculation, + batch_size=1, + max_atoms=None, + on_result=persist if output_file is not None else None, + ) + result = results[0] + return result + + +@contextmanager +def _fresh_batch_directory(output_path: Path): + """Claim a batch directory without racing another MatKit writer.""" + output_path.mkdir(parents=True, exist_ok=True) + lock = output_path / ".matkit_batch.lock" + with lock.open("x"): + pass + try: + if (output_path / "batch_manifest.json").exists(): + raise FileExistsError( + f"Batch manifest already exists in {output_path}; " + "use a fresh directory (resume is not supported yet)." + ) + yield + finally: + lock.unlink(missing_ok=True) + + +def run_mlip_batch( + input_files: Sequence[str | Path], + backend: MLIPBackendConfig, + calculation: MLIPCalculationConfig | None = None, + output_dir: str | Path = "mlip_results", + batch_size: int = 16, + max_atoms: int | None = None, +) -> MLIPBatchSummary: + """Persist each completed item; require a fresh batch directory. + + Execution status and optimization convergence are independent. A completed + batch can have unconverged results. Catchable orchestration errors persist + an interrupted manifest and propagate; completed results remain available. + """ + calculation = calculation or MLIPCalculationConfig() + _validate_execution_request( + input_files, backend, calculation, batch_size, max_atoms + ) + started = time.perf_counter() + output_path = Path(output_dir).expanduser().resolve() + items = [] + for index, value in enumerate(input_files): + path = Path(value).expanduser().resolve() + items.append( + { + "index": index, + "input_structure_file": str(path), + "status": "pending", + "converged": None, + "result_file": str( + output_path / f"{index:05d}_{path.stem}.json" + ), + "error": "", + } + ) + manifest = { + "schema_version": 1, + "status": "running", + "backend_info": backend.to_dict(), + "calculation_input": calculation.to_dict(), + "setup_time_s": 0.0, + "wall_time_s": 0.0, + "total": len(items), + "succeeded": 0, + "failed": 0, + "pending": len(items), + "unconverged": 0, + "error": "", + "items": items, + } + manifest_path = output_path / "batch_manifest.json" + + def checkpoint(status): + manifest["status"] = status + for key, state in ( + ("succeeded", "success"), + ("failed", "failure"), + ("pending", "pending"), + ): + manifest[key] = sum(item["status"] == state for item in items) + manifest["unconverged"] = ( + sum( + item["status"] == "success" and item["converged"] is False + for item in items + ) + if calculation.driver == "opt" + else 0 + ) + manifest["wall_time_s"] = time.perf_counter() - started + _write_json(manifest_path, manifest) + + def persist(index, result): + item = items[index] + result["output_results_file"] = item["result_file"] + _write_json(Path(item["result_file"]), result) + item.update( + status="success" if result["success"] else "failure", + converged=result["converged"] if result["success"] else None, + error=result["error"], + ) + manifest["setup_time_s"] = max( + manifest["setup_time_s"], result["setup_time_s"] + ) + checkpoint("running") + + with _fresh_batch_directory(output_path): + try: + checkpoint("running") + results, setup_time = _execute_inputs( + input_files, + backend, + calculation, + batch_size, + max_atoms, + on_result=persist, + ) + manifest["setup_time_s"] = setup_time + if manifest["failed"] == 0: + status = "completed" + elif manifest["succeeded"]: + status = "partial" + else: + status = "failure" + checkpoint(status) + except BaseException as exc: + manifest["error"] = str(exc) or type(exc).__name__ + try: + checkpoint("interrupted") + except Exception: + logger.exception("Could not persist interrupted batch manifest") + raise + return { + **manifest, + "manifest_file": str(manifest_path), + "results": results, + } diff --git a/src/matkit/types.py b/src/matkit/types.py index b463428..d72fb07 100644 --- a/src/matkit/types.py +++ b/src/matkit/types.py @@ -73,3 +73,49 @@ class UMABatchResult(TypedDict): final_energy: Optional[float] n_steps: Optional[int] error_message: Optional[str] + + +class _MLIPArtifacts(TypedDict, total=False): + output_results_file: str + + +class MLIPResult(_MLIPArtifacts): + """Runtime-neutral result from ``matkit.mlip.run_mlip``.""" + + schema_version: int + success: bool + error: str + input_structure_file: str + backend_info: dict + calculation_input: dict + energy: Optional[float] + energy_unit: str + forces: Optional[list[list[float]]] + force_unit: str + stress: Optional[list[list[float]]] + stress_unit: str + converged: bool + n_steps: Optional[int] + final_structure: Optional[dict] + calculation_time_s: float + setup_time_s: float + + +class MLIPBatchSummary(TypedDict): + """Persistent summary from ``matkit.mlip.run_mlip_batch``.""" + + schema_version: int + status: str + backend_info: dict + calculation_input: dict + setup_time_s: float + wall_time_s: float + total: int + succeeded: int + failed: int + pending: int + unconverged: int + error: str + items: list[dict] + manifest_file: str + results: list[MLIPResult] diff --git a/tests/test_cli.py b/tests/test_cli.py index 342a4d2..5754214 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,5 +1,12 @@ """Tests for matkit CLI.""" +import json +import importlib.util +from pathlib import Path +import sys + +import pytest + from click.testing import CliRunner from matkit.cli import main @@ -59,3 +66,310 @@ def test_unknown_command(self): runner = CliRunner() result = runner.invoke(main, ["nonexistent"]) assert result.exit_code != 0 + + def test_mlip_group_exposes_runtime_neutral_commands(self): + runner = CliRunner() + result = runner.invoke(main, ["mlip", "--help"]) + assert result.exit_code == 0 + assert "run" in result.output + assert "run-batch" in result.output + + def test_mlip_run_delegates_to_python_api( + self, sample_cif, tmp_path, monkeypatch + ): + import matkit.mlip + + received = {} + + def fake_run(input_file, backend, calculation, output_file): + received.update( + { + "input": input_file, + "backend": backend, + "calculation": calculation, + "output": output_file, + } + ) + return { + "success": True, + "energy": -1.25, + "energy_unit": "eV", + "converged": True, + } + + monkeypatch.setattr(matkit.mlip, "run_mlip", fake_run) + output = tmp_path / "result.json" + runner = CliRunner() + result = runner.invoke( + main, + [ + "mlip", + "run", + "--input", + sample_cif, + "--output", + str(output), + "--backend", + "rootstock", + "--checkpoint", + "mace-mp-0-medium", + "--cluster", + "polaris", + "--device", + "cuda", + "--setup-kwarg", + 'default_dtype="float32"', + ], + ) + + assert result.exit_code == 0 + assert json.loads(result.output)["energy"] == -1.25 + assert received["backend"].cluster == "polaris" + assert received["backend"].setup_kwargs == {"default_dtype": "float32"} + assert received["output"] == str(output) + + def test_mlip_batch_requires_one_input_source(self): + runner = CliRunner() + result = runner.invoke( + main, + [ + "mlip", + "run-batch", + "--backend", + "ase-mace", + "--checkpoint", + "medium", + ], + ) + assert result.exit_code != 0 + assert "exactly one" in result.output + + +@pytest.mark.parametrize( + "backend,flags,message", + [ + ("rootstock", ["--dtype", "float32"], "--setup-kwarg"), + ("ase-mace", ["--cluster", "polaris"], "requires --backend rootstock"), + ("ase-mace", ["--enable-cueq"], "requires --backend nvalchemi-mace"), + ("nvalchemi-mace", ["--dispersion"], "requires --backend ase-mace"), + ( + "ase-mace", + ["--calculator-type", "mace_anicc", "--dtype", "float64"], + "mace_anicc factory", + ), + ( + "ase-mace", + ["--calculator-type", "mace_off", "--dispersion"], + "requires --calculator-type mace_mp", + ), + ("ase-mace", ["--steps", "10"], "requires --driver opt"), + ("ase-mace", ["--fmax", "0.01"], "requires --driver opt"), + ("nvalchemi-mace", ["--dt", "0.1"], "requires --driver opt"), + ( + "nvalchemi-mace", + ["--driver", "opt", "--optimizer", "bfgs"], + "only the FIRE optimizer", + ), + ("ase-mace", ["--driver", "opt", "--fmax", "nan"], "finite"), + ], +) +def test_mlip_rejects_unsupported_explicit_options( + sample_cif, + monkeypatch, + backend, + flags, + message, +): + import matkit.mlip + + monkeypatch.setattr( + matkit.mlip, + "run_mlip", + lambda *a, **kw: pytest.fail("invalid request ran"), + ) + result = CliRunner().invoke( + main, + [ + "mlip", + "run", + "--input", + sample_cif, + "--backend", + backend, + "--checkpoint", + "medium", + *flags, + ], + ) + assert result.exit_code == 2 + assert message in result.output + + +@pytest.mark.parametrize( + "success,converged,code", + [(True, True, 0), (True, False, 1), (False, False, 1)], +) +def test_mlip_single_exit_and_summary( + sample_cif, monkeypatch, success, converged, code +): + import matkit.mlip + + monkeypatch.setattr( + matkit.mlip, + "run_mlip", + lambda *a, **kw: { + "success": success, + "converged": converged, + "energy": 1.0 if success else None, + "energy_unit": "eV", + "error": "" if success else "calculation failed", + }, + ) + result = CliRunner().invoke( + main, + [ + "mlip", + "run", + "--input", + sample_cif, + "--backend", + "ase-mace", + "--checkpoint", + "medium", + "--driver", + "opt", + ], + ) + assert result.exit_code == code + summary = json.loads(result.stdout) + assert summary["converged"] is converged + assert summary["status"] == ( + "failure" if not success else "success" if converged else "unconverged" + ) + assert bool(result.stderr) is bool(code) + + +@pytest.mark.parametrize("failed,unconverged", [(0, 0), (1, 0), (0, 1), (1, 1)]) +def test_mlip_batch_exit_and_summary( + sample_cif, monkeypatch, failed, unconverged +): + import matkit.mlip + + def fake_batch(*args, **kwargs): + return { + "status": "partial" if failed else "completed", + "total": 2, + "succeeded": 2 - failed, + "failed": failed, + "manifest_file": "batch.json", + "results": [ + {"success": True, "converged": not unconverged}, + {"success": not failed, "converged": not failed}, + ], + } + + monkeypatch.setattr(matkit.mlip, "run_mlip_batch", fake_batch) + result = CliRunner().invoke( + main, + [ + "mlip", + "run-batch", + "--input", + sample_cif, + "--backend", + "ase-mace", + "--checkpoint", + "medium", + "--driver", + "opt", + ], + ) + assert result.exit_code == int(bool(failed or unconverged)) + assert json.loads(result.stdout)["unconverged"] == unconverged + assert bool(result.stderr) is bool(failed or unconverged) + + +def test_mlip_runtime_error_summary(sample_cif, monkeypatch): + import matkit.mlip + + def fail(*args, **kwargs): + raise OSError("disk full") + + monkeypatch.setattr(matkit.mlip, "run_mlip", fail) + result = CliRunner().invoke( + main, + [ + "mlip", + "run", + "--input", + sample_cif, + "--backend", + "ase-mace", + "--checkpoint", + "medium", + ], + ) + assert result.exit_code == 1 + assert json.loads(result.stdout) == { + "status": "failure", + "error": "disk full", + } + assert "disk full" in result.stderr + + +@pytest.fixture +def gpu_example(): + path = Path(__file__).parents[1] / "examples/mlip_gpu.py" + spec = importlib.util.spec_from_file_location("mlip_gpu_example", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_gpu_example_rejects_unused_options(gpu_example, monkeypatch): + monkeypatch.setattr( + sys, + "argv", + [ + "mlip_gpu.py", + "--backend", + "ase-mace", + "--cluster", + "polaris", + "in.cif", + ], + ) + with pytest.raises(SystemExit) as error: + gpu_example.main() + assert error.value.code == 2 + + +def test_gpu_example_unconverged_exit(gpu_example, monkeypatch, capsys): + monkeypatch.setattr( + sys, + "argv", + [ + "mlip_gpu.py", + "--backend", + "ase-mace", + "--driver", + "opt", + "in.cif", + ], + ) + monkeypatch.setattr( + gpu_example, + "run_mlip_batch", + lambda *a, **kw: { + "status": "completed", + "total": 1, + "succeeded": 1, + "failed": 0, + "manifest_file": "batch.json", + "results": [{"success": True, "converged": False}], + }, + ) + assert gpu_example.main() == 1 + output = capsys.readouterr() + assert json.loads(output.out)["unconverged"] == 1 + assert "unconverged" in output.err diff --git a/tests/test_mlip_nvalchemi.py b/tests/test_mlip_nvalchemi.py new file mode 100644 index 0000000..97dce69 --- /dev/null +++ b/tests/test_mlip_nvalchemi.py @@ -0,0 +1,279 @@ +"""Exercise ALCHEMI adapter boundaries with CPU-only protocol doubles. + +These tests verify MatKit's conversions and dispatch, not GPU kernel behavior +or compatibility with a live installation. That evidence comes from smoke.pbs. +""" + +import sys +from types import ModuleType, SimpleNamespace + +import numpy as np +import pytest +from ase import Atoms +from ase.io import write + +from matkit.mlip import MLIPCalculationConfig, NVAlchemiMACEConfig, run_mlip +from matkit.mlip import runner + + +class Tensor: + def __init__(self, values): + self.values = np.asarray(values) + + def detach(self): + return self + + def cpu(self): + return self + + def numpy(self): + return self.values.copy() + + def squeeze(self, axis): + return Tensor(self.values.squeeze(axis)) + + def tolist(self): + return self.values.tolist() + + +@pytest.fixture +def alchemi(monkeypatch): + modules = {} + for name in ( + "torch", + "nvalchemi", + "nvalchemi.data", + "nvalchemi.dynamics", + "nvalchemi.models", + "nvalchemi.models.mace", + ): + module = ModuleType(name) + module.__path__ = [] + monkeypatch.setitem(sys.modules, name, module) + modules[name] = module + if "." in name: + parent, leaf = name.rsplit(".", 1) + setattr(modules[parent], leaf, module) + state = SimpleNamespace( + events=[], model_loads=[], corrupt=None, converged=[0] + ) + torch = modules["torch"] + torch.device = str + torch.float32, torch.float64 = "float32", "float64" + torch.cuda = SimpleNamespace( + is_available=lambda: True, + synchronize=lambda device: state.events.append(("sync", device)), + ) + + def zeros(*shape, device, dtype): + state.events.append(("zeros", shape, device, dtype)) + return Tensor(np.zeros(shape, dtype=dtype)) + + torch.zeros = zeros + + class AtomicData: + @classmethod + def from_atoms(cls, atoms, *, device, dtype): + state.events.append(("from_atoms", len(atoms), device, dtype)) + return SimpleNamespace( + num_nodes=len(atoms), + device=device, + positions=Tensor(atoms.positions.astype(dtype)), + cell=Tensor(atoms.cell.array[None].astype(dtype)), + pbc=Tensor(atoms.pbc[None]), + stress=None, + ) + + class Batch: + @classmethod + def from_data_list(cls, data): + result = cls() + result.data = data + return result + + def to_data_list(self): + return self.data + + class Model: + def eval(self): + state.events.append(("eval",)) + + def make_neighbor_hooks(self): + return ["neighbor-hook"] + + class MACEWrapper: + @classmethod + def from_checkpoint(cls, checkpoint, **kwargs): + state.model_loads.append((checkpoint, kwargs)) + return Model() + + class BaseDynamics: + def __init__(self, **kwargs): + state.events.append((type(self).__name__, kwargs)) + + def __enter__(self): + state.events.append(("enter",)) + return self + + def __exit__(self, *args): + state.events.append(("exit",)) + + def run(self, batch): + for index, data in enumerate(batch.data): + assert data.energy.numpy().shape == (1, 1) + assert data.forces.numpy().shape == (data.num_nodes, 3) + assert data.velocities.numpy().shape == (data.num_nodes, 3) + data.energy = Tensor([[index + 1.25]]) + data.forces = Tensor(np.zeros((data.num_nodes, 3))) + data.stress = Tensor(np.eye(3)[None]) + if isinstance(self, FIRE): + data.positions = Tensor(data.positions.numpy() + 0.1) + if state.corrupt and index == 1: + state.corrupt(data) + return batch + + class FIRE(BaseDynamics): + pass + + class ConvergenceHook: + @classmethod + def from_fmax(cls, value): + state.events.append(("fmax", value)) + return cls() + + def evaluate(self, batch): + return None if state.converged is None else Tensor(state.converged) + + modules["nvalchemi.data"].AtomicData = AtomicData + modules["nvalchemi.data"].Batch = Batch + modules["nvalchemi.models.mace"].MACEWrapper = MACEWrapper + for cls in (BaseDynamics, FIRE, ConvergenceHook): + setattr(modules["nvalchemi.dynamics"], cls.__name__, cls) + state.model = Model() + return state + + +def entries(): + return [ + ( + 3, + "first.xyz", + Atoms("Cu", positions=[[1, 2, 3]], cell=[8] * 3, pbc=True), + ), + ( + 7, + "second.xyz", + Atoms("Cu2", positions=[[0, 0, 0], [2, 0, 0]], cell=[9] * 3), + ), + ] + + +@pytest.mark.parametrize("dtype", ["float32", "float64"]) +def test_model_loading_and_atomic_conversion(alchemi, dtype): + backend = NVAlchemiMACEConfig( + "weights.pt", dtype=dtype, compile_model=True, enable_cueq=True + ) + runner._load_nvalchemi_model(backend) + assert alchemi.model_loads == [ + ( + "weights.pt", + { + "device": "cuda", + "dtype": dtype, + "compile_model": True, + "enable_cueq": True, + }, + ) + ] + assert ("eval",) in alchemi.events + atoms = entries()[1][2] + data = runner._atoms_to_nvalchemi_data(atoms, backend) + assert np.array_equal(data.positions.numpy(), atoms.positions) + assert data.positions.numpy().dtype == np.dtype(dtype) + assert np.array_equal(data.cell.numpy()[0], atoms.cell.array) + assert np.array_equal(data.pbc.numpy()[0], atoms.pbc) + assert data.forces.numpy().shape == (2, 3) + assert ("zeros", (2, 3), "cuda", dtype) in alchemi.events + + +@pytest.mark.parametrize("driver", ["energy", "opt"]) +def test_dynamics_and_result_mapping(alchemi, driver): + inputs = entries() + backend = NVAlchemiMACEConfig("medium", dt=0.2) + calculation = MLIPCalculationConfig(driver=driver, fmax=0.02, steps=7) + outputs = runner._run_nvalchemi_chunk( + alchemi.model, inputs, backend, calculation + ) + assert [index for index, _ in outputs] == [3, 7] + for (_, result), (_, _, original) in zip(outputs, inputs): + assert result["success"] + assert ( + result["final_structure"]["atomic_numbers"] + == original.numbers.tolist() + ) + assert result["final_structure"]["cell"] == original.cell.array.tolist() + assert result["final_structure"]["pbc"] == original.pbc.tolist() + expected = original.positions + (0.1 if driver == "opt" else 0) + assert np.allclose(result["final_structure"]["positions"], expected) + assert np.array_equal(result["stress"], np.eye(3)) + assert result["force_unit"] == "eV/angstrom" + name = "FIRE" if driver == "opt" else "BaseDynamics" + settings = next(event[1] for event in alchemi.events if event[0] == name) + assert settings["hooks"] == ["neighbor-hook"] + assert settings["n_steps"] == (7 if driver == "opt" else 1) + assert [result["energy"] for _, result in outputs] == [1.25, 2.25] + assert ("enter",) in alchemi.events and ("exit",) in alchemi.events + assert [result["converged"] for _, result in outputs] == [ + True, + driver == "energy", + ] + if driver == "opt": + assert settings["dt"] == 0.2 + assert ("fmax", 0.02) in alchemi.events + + +@pytest.mark.parametrize( + "field,value", + [ + ("energy", [[1, 2]]), + ("energy", [[float("nan")]]), + ("forces", [[0, 0, 0]]), + ("stress", np.zeros((2, 3, 3))), + ], +) +def test_one_invalid_native_result_preserves_other_items(alchemi, field, value): + alchemi.corrupt = lambda data: setattr(data, field, Tensor(value)) + outputs = runner._run_nvalchemi_chunk( + alchemi.model, + entries(), + NVAlchemiMACEConfig("medium"), + MLIPCalculationConfig(), + ) + assert outputs[0][1]["success"] + assert not outputs[1][1]["success"] + assert outputs[1][1]["error"] + + +def test_no_native_convergence(alchemi): + alchemi.converged = None + outputs = runner._run_nvalchemi_chunk( + alchemi.model, + entries(), + NVAlchemiMACEConfig("medium"), + MLIPCalculationConfig(driver="opt"), + ) + assert all( + result["success"] and not result["converged"] for _, result in outputs + ) + + +def test_energy_does_not_apply_optimizer_restriction(alchemi, tmp_path): + path = tmp_path / "structure.xyz" + write(path, entries()[0][2]) + result = run_mlip( + path, + NVAlchemiMACEConfig("medium"), + MLIPCalculationConfig(optimizer="bfgs"), + ) + assert result["success"] + assert any(event[0] == "BaseDynamics" for event in alchemi.events) diff --git a/tests/test_mlip_persistence.py b/tests/test_mlip_persistence.py new file mode 100644 index 0000000..8bfd68a --- /dev/null +++ b/tests/test_mlip_persistence.py @@ -0,0 +1,261 @@ +"""Durability regressions for completed results and interrupted batches.""" + +from contextlib import contextmanager +import json +from pathlib import Path + +import numpy as np +import pytest +from ase import Atoms +from ase.calculators.emt import EMT +from ase.io import write + +from matkit.mlip import ( + ASEMACEConfig, + MLIPCalculationConfig, + NVAlchemiMACEConfig, + RootstockConfig, + run_mlip_batch, +) +from matkit.mlip import runner + + +@pytest.fixture +def inputs(tmp_path, monkeypatch): + paths = [] + for directory in ("a", "b"): + path = tmp_path / directory / "same.xyz" + path.parent.mkdir() + write( + path, Atoms("Cu2", positions=[[0, 0, 0], [3, 0, 0]], cell=[8] * 3) + ) + paths.append(path) + + @contextmanager + def context(_backend): + yield EMT() + + monkeypatch.setattr(runner, "_ase_backend_context", context) + return paths + + +def read_manifest(out): + return json.loads((out / "batch_manifest.json").read_text()) + + +def test_interrupt_preserves_completed_item(inputs, tmp_path, monkeypatch): + out = tmp_path / "results" + original = runner._run_ase_item + + def run_item(input_file, *args): + manifest = read_manifest(out) + assert manifest["status"] == "running" + if input_file == str(inputs[1]): + assert manifest["succeeded"] == 1 + assert manifest["pending"] == 1 + stored = json.loads( + Path(manifest["items"][0]["result_file"]).read_text() + ) + assert stored["success"] + raise KeyboardInterrupt() + assert manifest["pending"] == 2 + return original(input_file, *args) + + monkeypatch.setattr(runner, "_run_ase_item", run_item) + with pytest.raises(KeyboardInterrupt): + run_mlip_batch(inputs, ASEMACEConfig(), output_dir=out) + manifest = read_manifest(out) + assert manifest["status"] == "interrupted" + assert manifest["error"] == "KeyboardInterrupt" + assert manifest["succeeded"] == manifest["pending"] == 1 + assert not (out / "00001_same.json").exists() + + +def test_teardown_failure_preserves_all_results(inputs, tmp_path, monkeypatch): + @contextmanager + def broken_context(_backend): + yield EMT() + raise RuntimeError("worker cleanup failed") + + monkeypatch.setattr(runner, "_ase_backend_context", broken_context) + out = tmp_path / "results" + with pytest.raises(RuntimeError, match="worker cleanup failed"): + run_mlip_batch(inputs, RootstockConfig("test"), output_dir=out) + manifest = read_manifest(out) + assert manifest["status"] == "interrupted" + assert manifest["succeeded"] == 2 + assert manifest["failed"] == manifest["pending"] == 0 + assert manifest["error"] == "worker cleanup failed" + for item in manifest["items"]: + assert item["status"] == "success" + assert json.loads(Path(item["result_file"]).read_text())["success"] + + +@pytest.mark.parametrize("failure", ["result", "manifest"]) +def test_persistence_failure_does_not_become_calculation_failure( + inputs, + tmp_path, + monkeypatch, + failure, +): + out = tmp_path / "results" + replace = runner.os.replace + failed = False + + def fail_once(source, destination): + nonlocal failed + destination = Path(destination) + if failure == "result": + trigger = destination.name == "00001_same.json" + else: + trigger = ( + destination.name == "batch_manifest.json" + and (out / "00001_same.json").exists() + ) + if trigger and not failed: + failed = True + raise OSError("disk write failed") + return replace(source, destination) + + monkeypatch.setattr(runner.os, "replace", fail_once) + with pytest.raises(OSError, match="disk write failed"): + run_mlip_batch(inputs, ASEMACEConfig(), output_dir=out) + manifest = read_manifest(out) + assert manifest["status"] == "interrupted" + assert manifest["failed"] == 0 + assert manifest["succeeded"] == (1 if failure == "result" else 2) + assert manifest["pending"] == (1 if failure == "result" else 0) + assert json.loads((out / "00000_same.json").read_text())["success"] + assert not list(out.glob("*.tmp")) + + +def test_atomic_replace_failure_retains_previous_json(tmp_path, monkeypatch): + path = tmp_path / "result.json" + path.write_text('{"previous": true}') + + def fail(*args): + raise OSError("replace failed") + + monkeypatch.setattr(runner.os, "replace", fail) + with pytest.raises(OSError, match="replace failed"): + runner._write_json(path, {"new": True}) + assert json.loads(path.read_text()) == {"previous": True} + assert list(tmp_path.iterdir()) == [path] + + +def test_fresh_directory_and_complete_payloads(inputs, tmp_path, monkeypatch): + out = tmp_path / "results" + summary = run_mlip_batch(inputs, ASEMACEConfig(), output_dir=out) + manifest = read_manifest(out) + assert manifest == { + k: v + for k, v in summary.items() + if k not in ("manifest_file", "results") + } + assert manifest["pending"] == manifest["unconverged"] == 0 + assert [Path(item["result_file"]).name for item in manifest["items"]] == [ + "00000_same.json", + "00001_same.json", + ] + for item, result in zip(manifest["items"], summary["results"]): + assert json.loads(Path(item["result_file"]).read_text()) == result + original_manifest = (out / "batch_manifest.json").read_bytes() + monkeypatch.setattr( + runner, + "_ase_backend_context", + lambda _: pytest.fail("existing batch was rerun"), + ) + with pytest.raises(FileExistsError, match="fresh directory"): + run_mlip_batch(inputs, ASEMACEConfig(), output_dir=out) + assert (out / "batch_manifest.json").read_bytes() == original_manifest + + +def test_concurrent_writer_cannot_claim_directory(inputs, tmp_path): + out = tmp_path / "results" + with runner._fresh_batch_directory(out): + with pytest.raises(FileExistsError): + run_mlip_batch(inputs, ASEMACEConfig(), output_dir=out) + assert (out / ".matkit_batch.lock").exists() + + +def test_valid_unconverged_batch_is_counted(inputs, tmp_path): + summary = run_mlip_batch( + inputs, + ASEMACEConfig(), + MLIPCalculationConfig(driver="opt", steps=1), + output_dir=tmp_path / "results", + ) + assert summary["status"] == "completed" + assert summary["succeeded"] == summary["unconverged"] == 2 + assert all(item["converged"] is False for item in summary["items"]) + + +def test_startup_failure_retains_input_errors(inputs, tmp_path, monkeypatch): + @contextmanager + def broken_start(_backend): + raise ImportError("missing calculator") + yield + + monkeypatch.setattr(runner, "_ase_backend_context", broken_start) + summary = run_mlip_batch( + [inputs[0], tmp_path / "missing.xyz"], + ASEMACEConfig(), + output_dir=tmp_path / "results", + ) + assert summary["status"] == "failure" + assert summary["failed"] == 2 + assert "missing calculator" in summary["items"][0]["error"] + assert "does not exist" in summary["items"][1]["error"] + + +def test_rootstock_gpu_worker_does_not_require_parent_cuda( + inputs, tmp_path, monkeypatch +): + monkeypatch.setattr( + runner, + "_synchronize_device", + lambda _: pytest.fail("Rootstock CUDA belongs to the worker"), + ) + result = run_mlip_batch( + inputs, + RootstockConfig("medium", device="cuda"), + output_dir=tmp_path / "results", + ) + assert result["succeeded"] == 2 + + +def test_native_chunk_is_saved_before_next_chunk(inputs, tmp_path, monkeypatch): + out = tmp_path / "results" + monkeypatch.setattr(runner, "_load_nvalchemi_model", lambda _: object()) + monkeypatch.setattr(runner, "_synchronize_device", lambda _: None) + + def chunk(model, entries, backend, calculation): + index, input_file, atoms = entries[0] + if index: + assert read_manifest(out)["succeeded"] == 1 + assert json.loads((out / "00000_same.json").read_text())["success"] + raise KeyboardInterrupt() + return [ + ( + index, + runner._success_result( + input_file, + backend, + calculation, + atoms, + 1.0, + np.zeros((len(atoms), 3)), + None, + True, + 0, + 0.1, + ), + ) + ] + + monkeypatch.setattr(runner, "_run_nvalchemi_chunk", chunk) + with pytest.raises(KeyboardInterrupt): + run_mlip_batch( + inputs, NVAlchemiMACEConfig("medium"), output_dir=out, batch_size=1 + ) + assert read_manifest(out)["status"] == "interrupted" diff --git a/tests/test_mlip_runner.py b/tests/test_mlip_runner.py new file mode 100644 index 0000000..87cc712 --- /dev/null +++ b/tests/test_mlip_runner.py @@ -0,0 +1,314 @@ +"""Tests for runtime-selectable, agent-free MLIP execution.""" + +import json +import sys +from contextlib import contextmanager +from types import SimpleNamespace + +import numpy as np +import pytest +from ase import Atoms +from ase.calculators.emt import EMT +from ase.io import write + +from matkit.mlip import ( + ASEMACEConfig, + MLIPCalculationConfig, + NVAlchemiMACEConfig, + RootstockConfig, + run_mlip, + run_mlip_batch, +) +from matkit.mlip import runner as mlip_runner + + +def _write_copper(path, distance=2.5): + atoms = Atoms( + "Cu2", + positions=[[0.0, 0.0, 0.0], [distance, 0.0, 0.0]], + cell=[8.0, 8.0, 8.0], + pbc=False, + ) + write(path, atoms) + return atoms + + +def _emt_context(counter=None): + @contextmanager + def calculator_context(_config): + if counter is not None: + counter["entered"] += 1 + yield EMT() + + return calculator_context + + +def test_backend_config_validation(): + with pytest.raises(ValueError, match="both cluster and root"): + RootstockConfig( + checkpoint="mace-mp-0-medium", + cluster="polaris", + root="/shared/rootstock", + ) + with pytest.raises(ValueError, match="dt must be positive"): + NVAlchemiMACEConfig(checkpoint="medium", dt=0) + with pytest.raises(ValueError, match="steps must be at least 1"): + MLIPCalculationConfig(steps=0) + + +def test_mace_anicc_uses_model_path_signature(monkeypatch): + received = {} + + def mace_anicc(**kwargs): + received.update(kwargs) + return "calculator" + + monkeypatch.setitem( + sys.modules, + "mace.calculators", + SimpleNamespace(mace_anicc=mace_anicc), + ) + monkeypatch.setitem( + sys.modules, + "mace", + SimpleNamespace(calculators=sys.modules["mace.calculators"]), + ) + + calculator = mlip_runner._create_mace_calculator( + ASEMACEConfig( + checkpoint="ani.model", + calculator_type="mace_anicc", + device="cuda", + ) + ) + + assert calculator == "calculator" + assert received == {"device": "cuda", "model_path": "ani.model"} + + +def test_run_mlip_energy_writes_runtime_neutral_result(tmp_path, monkeypatch): + input_file = tmp_path / "copper.xyz" + output_file = tmp_path / "result.json" + _write_copper(input_file) + monkeypatch.setattr(mlip_runner, "_ase_backend_context", _emt_context()) + + result = run_mlip( + input_file, + ASEMACEConfig(checkpoint="unused"), + output_file=output_file, + ) + + assert result["success"] is True + assert result["energy"] is not None + assert len(result["forces"]) == 2 + assert result["stress"] is None + assert result["n_steps"] == 0 + stored = json.loads(output_file.read_text()) + assert stored == result + assert stored["backend_info"]["type"] == "ase-mace" + assert stored["final_structure"]["atomic_numbers"] == [29, 29] + + +def test_periodic_energy_includes_full_stress(tmp_path, monkeypatch): + input_file = tmp_path / "periodic.xyz" + atoms = Atoms( + "Cu", + positions=[[0.0, 0.0, 0.0]], + cell=[3.6, 3.6, 3.6], + pbc=True, + ) + write(input_file, atoms) + monkeypatch.setattr(mlip_runner, "_ase_backend_context", _emt_context()) + + result = run_mlip( + input_file, + ASEMACEConfig(checkpoint="unused"), + ) + + assert result["success"] is True + assert np.asarray(result["stress"]).shape == (3, 3) + + +def test_run_mlip_fixed_cell_optimization(tmp_path, monkeypatch): + input_file = tmp_path / "copper.xyz" + initial = _write_copper(input_file, distance=3.0) + monkeypatch.setattr(mlip_runner, "_ase_backend_context", _emt_context()) + + result = run_mlip( + input_file, + ASEMACEConfig(checkpoint="unused"), + MLIPCalculationConfig(driver="opt", steps=2), + ) + + assert result["success"] is True + assert result["n_steps"] <= 2 + assert result["final_structure"]["cell"] == initial.cell.tolist() + + +def test_batch_reuses_calculator_and_preserves_failures(tmp_path, monkeypatch): + first = tmp_path / "first.xyz" + second = tmp_path / "second.xyz" + missing = tmp_path / "missing.xyz" + _write_copper(first) + _write_copper(second, distance=2.7) + counter = {"entered": 0} + monkeypatch.setattr( + mlip_runner, + "_ase_backend_context", + _emt_context(counter), + ) + + summary = run_mlip_batch( + [first, missing, second], + ASEMACEConfig(checkpoint="unused"), + output_dir=tmp_path / "results", + ) + + assert summary["status"] == "partial" + assert summary["succeeded"] == 2 + assert summary["failed"] == 1 + assert counter["entered"] == 1 + assert [item["index"] for item in summary["items"]] == [0, 1, 2] + assert [item["status"] for item in summary["items"]] == [ + "success", + "failure", + "success", + ] + manifest = json.loads( + (tmp_path / "results" / "batch_manifest.json").read_text() + ) + assert manifest["status"] == "partial" + assert all( + (tmp_path / "results" / name).exists() + for name in ( + "00000_first.json", + "00001_missing.json", + "00002_second.json", + ) + ) + + +def test_rootstock_context_forwards_options_and_closes(monkeypatch): + events = [] + + class FakeRootstockCalculator: + def __init__(self, **kwargs): + events.append(("init", kwargs)) + + def __enter__(self): + events.append(("enter", None)) + return "calculator" + + def __exit__(self, exc_type, exc, traceback): + events.append(("exit", None)) + + monkeypatch.setitem( + sys.modules, + "rootstock", + SimpleNamespace(RootstockCalculator=FakeRootstockCalculator), + ) + config = RootstockConfig( + checkpoint="mace-mp-0-medium", + cluster="polaris", + setup_kwargs={"default_dtype": "float32"}, + timeout=1200, + device="cuda", + ) + + with mlip_runner._ase_backend_context(config) as calculator: + assert calculator == "calculator" + + assert [event[0] for event in events] == ["init", "enter", "exit"] + kwargs = events[0][1] + assert kwargs["checkpoint"] == "mace-mp-0-medium" + assert kwargs["cluster"] == "polaris" + assert kwargs["setup_kwargs"] == {"default_dtype": "float32"} + assert kwargs["timeout"] == 1200 + + +def test_nvalchemi_loads_once_and_chunks_in_order(tmp_path, monkeypatch): + input_files = [] + for index in range(3): + input_file = tmp_path / f"input_{index}.xyz" + _write_copper(input_file, distance=2.5 + index * 0.1) + input_files.append(input_file) + + loaded = [] + chunks = [] + sentinel_model = object() + + def load_model(config): + loaded.append(config.checkpoint) + return sentinel_model + + def run_chunk(model, entries, backend, calculation): + assert model is sentinel_model + chunks.append([entry[1] for entry in entries]) + return [ + ( + index, + mlip_runner._success_result( + input_file, + backend, + calculation, + atoms, + 1.25, + np.zeros((len(atoms), 3)), + None, + True, + 0, + 0.01, + ), + ) + for index, input_file, atoms in entries + ] + + monkeypatch.setattr(mlip_runner, "_load_nvalchemi_model", load_model) + monkeypatch.setattr(mlip_runner, "_run_nvalchemi_chunk", run_chunk) + monkeypatch.setattr(mlip_runner, "_synchronize_device", lambda _: None) + summary = run_mlip_batch( + input_files, + NVAlchemiMACEConfig(checkpoint="medium"), + output_dir=tmp_path / "results", + batch_size=3, + max_atoms=4, + ) + + assert summary["status"] == "completed" + assert loaded == ["medium"] + assert [len(chunk) for chunk in chunks] == [2, 1] + assert chunks[0] + chunks[1] == [ + str(path.resolve()) for path in input_files + ] + + +def test_nvalchemi_missing_backend_is_persisted(tmp_path, monkeypatch): + input_file = tmp_path / "input.xyz" + output_file = tmp_path / "failure.json" + _write_copper(input_file) + + def missing_backend(_config): + raise ImportError("Install the nvalchemi_mace extra") + + monkeypatch.setattr(mlip_runner, "_load_nvalchemi_model", missing_backend) + result = run_mlip( + input_file, + NVAlchemiMACEConfig(checkpoint="medium"), + output_file=output_file, + ) + + assert result["success"] is False + assert "nvalchemi_mace" in result["error"] + assert json.loads(output_file.read_text())["success"] is False + + +def test_nvalchemi_rejects_non_fire_optimizer(tmp_path): + input_file = tmp_path / "input.xyz" + _write_copper(input_file) + + with pytest.raises(ValueError, match="only the FIRE optimizer"): + run_mlip( + input_file, + NVAlchemiMACEConfig(checkpoint="medium"), + MLIPCalculationConfig(driver="opt", optimizer="bfgs"), + ) diff --git a/tests/test_mlip_smoke.py b/tests/test_mlip_smoke.py new file mode 100644 index 0000000..e7a6f72 --- /dev/null +++ b/tests/test_mlip_smoke.py @@ -0,0 +1,80 @@ +"""Check the live smoke runner's orchestration without launching engines.""" + +import importlib.util +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest +from ase import Atoms +from ase.io import write + + +@pytest.fixture +def smoke(): + path = Path(__file__).parents[1] / "alcf/polaris/mlip/smoke.py" + spec = importlib.util.spec_from_file_location("mlip_smoke", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_failed_smoke_case_does_not_skip_other_backends( + smoke, tmp_path, monkeypatch +): + input_file = tmp_path / "input.xyz" + write( + input_file, Atoms("Cu2", positions=[[0, 0, 0], [2, 0, 0]], cell=[8] * 3) + ) + commands = [] + + def fake_run(command, **kwargs): + commands.append(command) + return SimpleNamespace(returncode=1 if len(commands) == 1 else 0) + + monkeypatch.setattr(smoke, "_probe", lambda command: {"mocked": True}) + monkeypatch.setattr(smoke, "subprocess", SimpleNamespace(run=fake_run)) + monkeypatch.setattr(smoke, "_validate_outputs", lambda *args: []) + out = tmp_path / "evidence" + assert ( + smoke.main(["--input", str(input_file), "--output-dir", str(out)]) == 1 + ) + report = json.loads((out / "smoke_report.json").read_text()) + assert len(commands) == len(report["cases"]) == 6 + assert report["cases"][0]["status"] == "failed" + assert all(case["status"] == "passed" for case in report["cases"][1:]) + for command in commands[-2:]: + assert "run-batch" in command + assert command.count("--input") == 2 + assert command[command.index("--batch-size") + 1] == "2" + assert len(report["input_sha256"]) == 64 + + +def test_smoke_rejects_unconverged_results(smoke, tmp_path): + atoms = Atoms("Cu", cell=[8] * 3) + data = { + "success": True, + "error": "", + "energy": 1.0, + "energy_unit": "eV", + "force_unit": "eV/angstrom", + "forces": [[0, 0, 0]], + "stress": None, + "converged": False, + "n_steps": 1, + "final_structure": { + "positions": atoms.positions.tolist(), + "atomic_numbers": atoms.numbers.tolist(), + "cell": atoms.cell.array.tolist(), + "pbc": atoms.pbc.tolist(), + }, + } + (tmp_path / "result.json").write_text(json.dumps(data)) + with pytest.raises(AssertionError, match="did not converge"): + smoke._validate_outputs(tmp_path, [atoms], "opt", 0.01, False) + assert ( + smoke._validate_outputs(tmp_path, [atoms], "energy", 0.01, False)[0][ + "energy" + ] + == 1.0 + ) diff --git a/tests/test_mlip_validation.py b/tests/test_mlip_validation.py new file mode 100644 index 0000000..82a7d36 --- /dev/null +++ b/tests/test_mlip_validation.py @@ -0,0 +1,216 @@ +"""Scientific validation and execution outcome regressions.""" + +from contextlib import contextmanager +import json +import sys +from types import SimpleNamespace + +import numpy as np +import pytest +from ase import Atoms +from ase.calculators.calculator import ( + Calculator, + PropertyNotImplementedError, + all_changes, +) +from ase.calculators.emt import EMT +from ase.io import write + +from matkit.mlip import ( + ASEMACEConfig, + MLIPCalculationConfig, + NVAlchemiMACEConfig, + RootstockConfig, + run_mlip, + run_mlip_batch, +) +from matkit.mlip import runner + + +@pytest.fixture +def copper(tmp_path): + path = tmp_path / "copper.xyz" + write(path, Atoms("Cu2", positions=[[0, 0, 0], [3, 0, 0]], cell=[8] * 3)) + return path + + +def use_calculator(monkeypatch, calculator): + @contextmanager + def context(_backend): + yield calculator + + monkeypatch.setattr(runner, "_ase_backend_context", context) + + +@pytest.mark.parametrize( + "value", [float("nan"), float("inf"), -float("inf"), True] +) +@pytest.mark.parametrize( + "field", ["fmax", "dt", "timeout", "dispersion_cutoff"] +) +def test_nonfinite_config_rejected(field, value): + factory = { + "fmax": MLIPCalculationConfig, + "dt": lambda **kw: NVAlchemiMACEConfig("medium", **kw), + "timeout": lambda **kw: RootstockConfig("medium", **kw), + "dispersion_cutoff": ASEMACEConfig, + }[field] + with pytest.raises(ValueError, match=field): + factory(**{field: value}) + + +@pytest.mark.parametrize("value", [0, -1, 1.5, True, float("nan")]) +def test_integer_limits(copper, tmp_path, value): + with pytest.raises(ValueError, match="steps"): + MLIPCalculationConfig(steps=value) + for field in ("batch_size", "max_atoms"): + with pytest.raises(ValueError, match=field): + run_mlip_batch( + [copper], + ASEMACEConfig(), + output_dir=tmp_path / field, + **{field: value}, + ) + + +def test_rootstock_kwargs_must_be_serializable(): + for value in (float("nan"), object()): + with pytest.raises(ValueError, match="finite JSON"): + RootstockConfig("medium", setup_kwargs={"nested": [value]}) + + +@pytest.mark.parametrize("case", ["empty", "positions", "cell", "periodic"]) +def test_invalid_structure_fails_before_model_load(copper, monkeypatch, case): + atoms = Atoms("Cu", cell=[3] * 3, pbc=True) + if case == "empty": + atoms = Atoms() + elif case == "positions": + atoms.positions[0, 0] = float("nan") + elif case == "cell": + atoms.cell[0, 0] = float("inf") + else: + atoms.cell[2] = atoms.cell[1] + monkeypatch.setattr(runner, "ase_read", lambda _: atoms) + monkeypatch.setattr( + runner, + "_ase_backend_context", + lambda _: pytest.fail("Invalid input must not load a model"), + ) + result = run_mlip(copper, ASEMACEConfig()) + assert not result["success"] + assert result["final_structure"] is None + assert result["error"] + + +def test_partial_periodic_cell_and_molecule_without_cell(): + runner._validate_atoms(Atoms("Cu", cell=[3, 3, 0], pbc=[True, True, False])) + runner._validate_atoms(Atoms("H")) + + +@pytest.mark.parametrize( + "field,value", + [ + ("energy", float("nan")), + ("energy", float("inf")), + ("forces", np.full((2, 3), float("nan"))), + ("forces", np.zeros((1, 3))), + ("stress", np.full((3, 3), float("inf"))), + ], +) +def test_invalid_calculator_output_is_persisted_failure( + copper, + tmp_path, + monkeypatch, + field, + value, +): + class InvalidCalculator(Calculator): + implemented_properties = ["energy", "forces", "stress"] + + def calculate( + self, atoms=None, properties=None, system_changes=all_changes + ): + super().calculate(atoms, properties, system_changes) + self.results = { + "energy": 1.0, + "forces": np.zeros((2, 3)), + "stress": np.zeros((3, 3)), + field: value, + } + + atoms = Atoms("Cu2", cell=[8] * 3, pbc=True) + monkeypatch.setattr(runner, "ase_read", lambda _: atoms) + use_calculator(monkeypatch, InvalidCalculator()) + output = tmp_path / "result.json" + result = run_mlip(copper, ASEMACEConfig(), output_file=output) + assert not result["success"] + assert result["energy"] is None + assert json.loads(output.read_text()) == result + assert "NaN" not in output.read_text() + assert "Infinity" not in output.read_text() + + +@pytest.mark.parametrize("unsupported", [True, False]) +def test_stress_unsupported_differs_from_calculator_failure( + copper, + monkeypatch, + unsupported, +): + class StressCalculator(EMT): + def get_stress(self, atoms=None): + if unsupported: + raise PropertyNotImplementedError("stress unavailable") + raise RuntimeError("stress calculation crashed") + + atoms = Atoms( + "Cu2", positions=[[0, 0, 0], [3, 0, 0]], cell=[8] * 3, pbc=True + ) + monkeypatch.setattr(runner, "ase_read", lambda _: atoms) + use_calculator(monkeypatch, StressCalculator()) + result = run_mlip(copper, ASEMACEConfig()) + assert result["success"] is unsupported + assert result["stress"] is None + if not unsupported: + assert "stress calculation crashed" in result["error"] + + +@pytest.mark.parametrize("steps,converged", [(1, False), (300, True)]) +def test_optimization_retains_valid_results( + copper, monkeypatch, steps, converged +): + use_calculator(monkeypatch, EMT()) + result = run_mlip( + copper, + ASEMACEConfig(), + MLIPCalculationConfig(driver="opt", steps=steps), + ) + assert result["success"] + assert result["converged"] is converged + assert np.isfinite(result["energy"]) + assert result["n_steps"] <= steps + + +def test_strict_json_preserves_previous_file(tmp_path): + path = tmp_path / "result.json" + path.write_text('{"old": true}') + with pytest.raises(ValueError): + runner._write_json(path, {"energy": float("nan")}) + assert json.loads(path.read_text()) == {"old": True} + + +def test_cuda_availability_and_faults(monkeypatch): + cuda = SimpleNamespace(is_available=lambda: False) + monkeypatch.setitem( + sys.modules, "torch", SimpleNamespace(cuda=cuda, device=str) + ) + with pytest.raises(RuntimeError, match="CUDA is unavailable"): + runner._synchronize_device("cuda") + runner._synchronize_device("cpu") + + def broken_sync(device): + raise RuntimeError("device fault") + + cuda.is_available = lambda: True + cuda.synchronize = broken_sync + with pytest.raises(RuntimeError, match="device fault"): + runner._synchronize_device("cuda")