diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b35034d..6aa7728 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -24,6 +24,10 @@ jobs: 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 + - name: Check unified API tests and recipes + run: | + ruff check tests/test_api*.py tests/test_foundations.py tests/test_operation_cli.py tests/test_mcp_api.py tests/test_packaged_resources.py tests/test_unified_smoke.py tests/fixtures/fake_engine.py examples/unified_smoke.py + ruff format --check tests/test_api*.py tests/test_foundations.py tests/test_operation_cli.py tests/test_mcp_api.py tests/test_packaged_resources.py tests/test_unified_smoke.py tests/fixtures/fake_engine.py examples/unified_smoke.py test: runs-on: ubuntu-latest @@ -36,9 +40,30 @@ jobs: with: python-version: ${{ matrix.python-version }} - name: Install dependencies - run: pip install -e ".[dev]" + run: | + pip install build pytest + python -m build --wheel + pip install dist/*.whl - name: Run tests - run: pytest tests/ -v + env: + MATKIT_WHEEL_TEST: "1" + run: | + mkdir -p "$RUNNER_TEMP/matkit-wheel-tests" + cp -R tests examples alcf "$RUNNER_TEMP/matkit-wheel-tests/" + cd "$RUNNER_TEMP/matkit-wheel-tests" + pytest tests/ -v + + mcp: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install optional MCP environment + run: pip install ".[dev,mcp]" + - name: Test local stdio and worker cancellation + run: pytest tests/test_mcp_api.py -v build: runs-on: ubuntu-latest diff --git a/CHANGELOG.md b/CHANGELOG.md index cbc6fc7..488bd25 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- Experimental unified scientific API, versioned requests/results, structure + lineage, relocatable run bundles, and supervised local execution +- Shared operation CLI and optional bounded stdio MCP tools with artifact + retrieval, plus installed-wheel and deterministic transport tests +- Single-component gRASPA CUDA prepare/run/analyze path, capability inventory, + and opt-in execution evidence recorder - 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` @@ -28,6 +34,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `conftest.py` for pytest configuration ### Fixed +- Custom cutoff propagation through GCMC unit-cell replication and cached + batch setup, and successful RASPA2 parsing incorrectly reporting failure - **MACE optimizer bug**: `geo_opt_cell_opt` mode called `dyn1.run()` instead of `dyn.run()` for cell optimization step, meaning the cell was never actually optimized - **Missing f-string**: `raspa2.py` error message `"Unit {unit} is not supported"` was missing `f` prefix - **Unreachable code**: Removed dead `return result` after `raise ValueError` in `graspa.py` and `graspa_sycl.py` diff --git a/README.md b/README.md index 4b491aa..fe714b6 100644 --- a/README.md +++ b/README.md @@ -165,6 +165,12 @@ combining these stacks. The `all` extra does not include Rootstock or ALCHEMI. ## Python API +The experimental [unified operation API](docs/unified-api.md) provides shared +Python/CLI requests, relocatable calculation bundles, and optional MCP tools. +It covers MLIP evaluation/relaxation, Zeo++ analysis, and single-component +gRASPA CUDA execution. See the [capability inventory](docs/capabilities.md) +for implementation status, environment requirements, and validation limits. + ```python from matkit.graspa import setup_simulation, get_output_data from matkit.utils import calculate_cell_size, remove_solvent, sample_cifs diff --git a/alcf/polaris/unified/run.pbs b/alcf/polaris/unified/run.pbs new file mode 100644 index 0000000..ce428c9 --- /dev/null +++ b/alcf/polaris/unified/run.pbs @@ -0,0 +1,16 @@ +#!/bin/bash +#PBS -N matkit-unified +#PBS -l select=1:system=polaris +#PBS -l place=scatter +#PBS -l walltime=01:00:00 +#PBS -q debug + +# Supply your account with qsub -A ACCOUNT and export the three variables below. +# This script executes an already prepared bundle inside the allocated node. +set -euo pipefail +: "${MATKIT_PYTHON:?Set an absolute MatKit environment Python path}" +: "${MATKIT_BUNDLE:?Set the absolute prepared-bundle path}" +: "${MATKIT_EXECUTION_PROFILE:?Set the absolute execution-profile JSON path}" +cd "${PBS_O_WORKDIR:?}" +exec "$MATKIT_PYTHON" -m matkit.cli execute "$MATKIT_BUNDLE" \ + --execution "$MATKIT_EXECUTION_PROFILE" diff --git a/docs/capabilities.md b/docs/capabilities.md new file mode 100644 index 0000000..2c7b9f2 --- /dev/null +++ b/docs/capabilities.md @@ -0,0 +1,38 @@ +# Capability inventory + +This inventory distinguishes implemented interfaces from scientific validation. +No GPU or external-engine capability has been promoted by CPU fixtures. The +unified API remains experimental; record real execution for each capability, +model, and environment before promotion. + +| Interface | Implemented capability | Environment | Evidence/status | +| --- | --- | --- | --- | +| Unified direct MACE | Energy, requested forces/stress, fixed-cell relaxation, sequential batches | MACE + compatible CPU/CUDA stack; model-dependent properties/species | CPU adapter contracts; experimental | +| Unified Rootstock | Energy, requested forces/stress, fixed-cell relaxation, sequential batches | Rootstock client and separately installed deployment | CPU adapter contracts; worker evidence required | +| Unified ALCHEMI MACE | Energy, requested forces/stress, fixed-cell FIRE, native batches | Compatible ALCHEMI/CUDA environment | Mocked native contracts; GPU evidence required | +| Unified Zeo++ | Diameter, area, volume, PSD, channels | `network` binary and radii definitions | Parser/subprocess fixtures; real execution required | +| Unified gRASPA CUDA | Pure-component preparation, execution, absolute uptake and heat parsing | Charged periodic CIF; templates; CUDA executable for execution | Synthetic output/subprocess fixtures; real execution required | +| Legacy gRASPA/pygRASPA | Pure/mixture and grid setup; existing parsers; pygRASPA reference-energy helper | Core for preparation; engine-specific environment for execution | Setup/parser fixtures; mixtures outside unified result contract | +| Legacy gRASPA SYCL | Setup and parsing | Core for preparation; Intel GPU environment for execution | Cutoff regressions; Aurora recipe; execution evidence required | +| Legacy RASPA2 | Setup and parsing | Core for preparation; RASPA2 for execution | Cutoff and success-reporting regressions; execution evidence required | +| RASPA3 | Force-field conversion only | Core | Conversion fixtures; simulation execution not implemented | +| Legacy MACE optimization | Geometry, cell, sequential geometry/cell optimization | MACE with required forces/stress | Existing interface; outside unified validation contract | +| Legacy UMA | Single point, geometry/cell optimization, MD and batch optimization | FAIRChem/UMA installation | Existing interfaces; outside unified validation contract | +| PACMOF2 | Charge-prediction wrapper | PACMOF2 installation | Existing interface; per-output scientific validation not yet unified | +| Structure utilities/ToBaCCo | Solvent removal, sampling, linker/conversion helpers | Core; optional RDKit/Open Babel as applicable | Legacy compatibility; future structural operations belong in MOFforge | +| Isotherm plotting | Single/mixture plots and selectivity from existing data | Matplotlib extra | Parser/plot fixtures; does not establish input scientific accuracy | +| ORCA | Stub | Not applicable | No supported execution capability | +| CLI/MCP | Unified requests, results, artifacts; bounded stdio tools | Core CLI; optional MCP SDK 2 | Deterministic subprocess and local stdio integration tests | + +Availability in `matkit capabilities` is a caller-side installation check, not a +model suitability or GPU compatibility claim. Preparation needs no engine +binary. Generic energy evaluation needs only its requested properties; +relaxation needs forces, and cell optimization is not advertised by the unified +API. Native ALCHEMI rejects unsupported atom arrays/constraints. Model aliases +whose content cannot be resolved are identified as such in provenance. + +Promotion requires interface documentation, licensed reference fixtures, +failure tests, reproducible installation, and reviewed real execution. Numerical +parity uses matching checkpoints/settings; scientific accuracy requires +independent reference data. Record benchmark startup, warm execution, throughput, +memory, and failures separately. diff --git a/docs/unified-api.md b/docs/unified-api.md new file mode 100644 index 0000000..3b6aff7 --- /dev/null +++ b/docs/unified-api.md @@ -0,0 +1,278 @@ +# Unified scientific operations + +The experimental `matkit.api` interface runs without an agent, MOFforge, or +ChemGraph. Python, the operation CLI, and optional MCP tools share validated +requests, scientific results, and persistent run bundles. Existing engine APIs, +CLI commands, defaults, and result layouts remain available. + +## Install and select an environment + +```bash +pip install . # core API, preparation, parsing, CLI +pip install '.[mlip]' # direct MACE +pip install '.[rootstock]' # Rootstock client +pip install '.[mcp]' # optional stdio server (MCP SDK 2) +``` + +Zeo++ and gRASPA executables must be installed separately. See +[Polaris MLIP environments](../alcf/polaris/mlip/README.md) for the experimental +GPU adapters. Incompatible model packages require separate installations and +interpreters. Core imports and capability discovery do not load models or CUDA. + +```bash +matkit capabilities --json +``` + +Discovery reports caller-side availability, restrictions, and evidence +separately. A remote Rootstock deployment or a configured worker environment can +differ from the caller. Species coverage, stress support, and scientific +applicability depend on the selected model. See the +[capability inventory](capabilities.md). + +## Python + +```python +from matkit.api import ( + EvaluateRequest, ExecutionConfig, MLIPMethod, StructureRef, evaluate, +) + +request = EvaluateRequest( + structure=StructureRef(path="structure.cif"), + method=MLIPMethod(checkpoint="medium"), + properties=["potential_energy", "forces"], +) +result = evaluate(request, output_dir="runs/energy") +print(result.accepted, result.payload) +``` + +The method identifies the science. `MACEAdapter`, `RootstockAdapter`, and +`AlchemiAdapter` select the implementation; `ExecutionConfig` selects the device, +interpreter, environment overrides, and executable commands. Defaults preserve +the existing adapters: direct MACE uses CPU/float64, ALCHEMI CUDA/float32, +and Rootstock uses its deployment configuration. + +`EvaluateRequest` defaults to energy only. Forces and stress must be explicitly +requested, and unavailable requested properties fail. `RelaxRequest` uses +fixed-cell FIRE with `fmax=0.01` eV/angstrom and 1000 maximum steps by default. +ALCHEMI supports only FIRE. Cell optimization and MD remain on legacy APIs. + +```python +from matkit.api import RelaxRequest, relax, run_batch + +relaxation = RelaxRequest( + structure=request.structure, method=request.method, fmax=0.02, steps=500, +) +result = relax(relaxation, output_dir="runs/relaxation") +assert result.accepted # includes requested force convergence + +# Requests must share operation, method, adapter, and scientific settings. +# Repeated basenames remain distinct and results preserve input ordering. +batch = run_batch([request, request], output_dir="runs/batch") +``` + +Batches retain one calculator per request group; ALCHEMI retains native chunking. +Per-item native timings include shared chunk work and are not throughput +benchmarks. Startup is recorded separately. Inputs remain accumulated in memory +in this first release. + +## CLI and prepared calculations + +Save a specification as `pores.json`: + +```json +{ + "operation": "pores", + "structure": {"path": "structure.cif"}, + "analyses": ["res", "sa"], + "probe_radius": 1.86, + "channel_radius": 1.86, + "num_samples": 100000 +} +``` + +Paths in CLI specifications are relative to the specification file. Model names +remain identifiers unless they resolve to local checkpoint files. Python paths +are relative to the caller's working directory. + +```bash +matkit pores --spec pores.json --outdir runs/pores +matkit prepare --spec pores.json --outdir runs/prepared +# Copy the entire prepared directory to the execution environment if needed. +matkit execute runs/prepared --execution execution.json +matkit inspect runs/prepared +``` + +Example execution profile (`execution.json`): + +```json +{ + "python": "/path/to/matkit-environment/bin/python", + "device": "cuda", + "executables": { + "zeopp": ["/path/to/network"], + "graspa": ["/path/to/gRASPA/bin/simulate"] + }, + "environment": {"OMP_NUM_THREADS": "1"}, + "timeout_s": 3600 +} +``` + +Executable values are argument lists, never shell fragments. Every worker +interpreter must have a compatible MatKit installation and the selected engine +dependencies. Environment overrides travel through the process environment; +their values are not copied into worker configuration artifacts. Selected +numerical runtime settings are recorded in provenance. + +`matkit evaluate` and `matkit relax` accept their corresponding specifications. +`matkit batch --spec requests.json --outdir runs/batch` accepts a JSON list of +homogeneous requests. CLI calculations run in a subprocess so engine output +goes to logs. Scientific JSON goes to stdout, diagnostics to stderr. Exit 2 +means invalid arguments; exit 1 means a failed calculation or failed required +convergence; exit 0 means accepted requested results. Inspection can exit 0 +while reporting a failed calculation because reading the record succeeded. + +## Single-component gRASPA CUDA + +The unified path requires a fully periodic CIF with finite, atom-mapped +`_atom_site_charge` values summing to the requested net charge. It copies the +charged CIF unchanged. Disordered sites and symmetry expansions without a +provable atom mapping are rejected; provide an explicit P1 structure. + +```json +{ + "operation": "adsorption", + "structure": {"path": "charged_framework.cif"}, + "adsorbate": "CO2", + "temperature_K": 298, + "pressure_Pa": 100000, + "cutoff_angstrom": 12.8, + "initialization_cycles": 1000, + "equilibration_cycles": 1000, + "production_cycles": 10000, + "number_of_blocks": 5, + "fugacity_coefficient": "PR-EOS", + "net_charge": 0, + "unit": "mol/kg" +} +``` + +These are example settings, not evidence of equilibration or adequate sampling. +The default template and its force-field definitions are staged and hashed. +`template_dir` may select a complete custom template directory. The first +unified parser handles one component; existing mixture setup APIs are unchanged. + +```bash +matkit adsorption prepare --spec adsorption.json --outdir runs/adsorption +matkit execute runs/adsorption --execution execution.json +matkit adsorption analyze runs/adsorption +``` + +Preparation requires neither CUDA nor the gRASPA binary. Execution requires a +zero engine exit code and complete, finite requested results. Sampling quality +remains `unknown`; uncertainty is reported as supplied by the engine, with an +unknown method unless independently established. An engine random seed not +resolved by the adapter remains unknown. This does not promise bitwise replay. + +For manual engine launches, keep `work/raspa.log` and record the actual engine +return code in `exit.json` as `{"returncode": 0}` (using the actual code). +`analyze_adsorption` can then collect the prepared bundle. Prefer `matkit execute` +to capture these records reliably. The [Polaris PBS example](../alcf/polaris/unified/run.pbs) +runs inside an existing allocation; MatKit does not submit scheduler jobs. + +## Contracts and artifacts + +Requests use `schema_name="matkit.request"`, results `"matkit.run"`, and batches +`"matkit.batch"`, each at version 1. Legacy MLIP version-1 files retain their +original meaning. Request schemas are available through +`matkit.api.models.REQUEST_ADAPTER.json_schema()` and result schemas through +`RunResult.model_json_schema()`. + +`state` records execution (`prepared`, `running`, `completed`, `failed`, or +`interrupted`). `numerical_validity` and named scientific checks are separate. +An unconverged relaxation can have valid numerical results and `state=completed` +while `accepted` is false. Unknown sampling quality does not become a claim of +equilibrium. Failure information is separate from adsorption uncertainty. + +Energy uses `potential_energy` in eV with a model-specific reference; forces +use eV/angstrom; stress uses the ASE Cartesian convention in eV/angstrom³. +The shared interface does not make energies from different methods comparable. +Pore result quantities retain their documented Zeo++ units. Adsorption uptake +is absolute loading on the framework basis, with component and unit recorded. + +Each bundle contains: + +- `request.json`: staged request using bundle-relative input references. +- `inputs/`: original structure, metadata sidecar, and supporting files. +- `work/`: engine inputs, logs, and output artifacts. +- `run.json`: execution record, settings, provenance, and artifact inventory. +- `result.json`: authoritative committed result after completion/failure. +- `execution.json`, `command.json`, and `exit.json` when applicable. + +Input files, sidecars, templates, force fields, and available local model files +are hashed. Unresolved checkpoint identities/versions remain explicitly unknown. +Prepared inputs are checked before execution. Keep the whole bundle together +when moving it; all artifact inventory paths are relative to its root. + +Structure sidecars preserve species, coordinates, cell, periodicity, atom IDs, +supported arrays, labels, bonds, and constraints. Unsupported objects and +unmappable inputs fail explicitly. Changed geometries receive parent lineage +and lose inherited derived charges/energies/pore metadata. Recompute charges +before sending a relaxed geometry to charge-dependent adsorption calculations. +Do not move a generated structure without its `.metadata.json` sidecar. + +Use a fresh directory for each run/batch. Completed results are committed before +manifest updates; interrupted work remains inspectable. If later teardown or +manifest persistence fails, inspection reports that orchestration failure while +retaining the committed numerical payload. An unreadable manifest does not +prevent recovery of a valid committed result. A hard process kill +can leave a running record or stale lock; automatic resume and restart are not +implemented. Copy a completed result for inspection, and prepare a fresh bundle +for another execution. Engine-specific restart requires additional future work. + +## MCP + +```bash +matkit-mcp --run-root /scratch/matkit-runs --input-root /data/structures \ + --profiles profiles.json --timeout 60 +``` + +`profiles.json` maps profile names to execution profiles, for example +`{"default": {"executables": {"zeopp": ["/path/to/network"]}}}`. Callers select +a profile name; executable configuration stays with the server. Input paths +must lie inside a configured input root or run root. Use repeated `--input-root` +arguments for model and structure directories. + +Tools are `matkit_capabilities`, `matkit_evaluate`, `matkit_relax`, `matkit_pores`, +`matkit_prepare`, `matkit_prepare_adsorption`, and `matkit_inspect`. Select a +catalog with `--tools matkit_capabilities,matkit_pores,matkit_inspect`. + +Calls return scalar summaries, scientific checks, failures, and artifact links +such as `matkit://runs//artifacts/`. Read those MCP resources to +retrieve structures, arrays, and full result JSON; a worker-local path alone is +not the transfer mechanism. Verify `accepted` and the scientific checks. + +Calculations are synchronous and bounded (60 seconds by default, including model +startup but excluding preparation). Timeouts/cancellation terminate workers and +preserve interrupted records. Long calculations use prepared bundles and +CLI/job-script execution. The server has no persistent background queue, HTTP +transport, or scheduler integration in this release. MOFforge and ChemGraph +remain independent; MatKit does not re-export MOFforge's tool catalog. + +## Validation and promotion + +CPU fixtures test contracts and failures, including subprocesses and real stdio +MCP sessions. Synthetic fixtures are not scientific reference calculations. +Run the opt-in [execution recorder](../examples/unified_smoke.py) in the actual +engine environment. It retains per-case bundles, failures, environment details, +and hashes; it does not assert scientific accuracy or promote capabilities. + +```bash +python examples/unified_smoke.py --spec pores.json --spec adsorption.json \ + --execution execution.json --outdir evidence/first-run +``` + +Support promotion requires reviewed reference cases, failure tests, reproducible +installation, and recorded real execution for the specific capability and +environment. The next milestones are reference/parity benchmarks and safe result +reuse, optional MOFforge/charge integration, the porous-material workflow, wider +simulation support, and ChemGraph agent evaluations. diff --git a/examples/unified_smoke.py b/examples/unified_smoke.py new file mode 100644 index 0000000..8e79ebf --- /dev/null +++ b/examples/unified_smoke.py @@ -0,0 +1,61 @@ +#!/usr/bin/env python3 +"""Record opt-in execution evidence for explicitly supplied specifications.""" + +import argparse +import json +from pathlib import Path + +from matkit.api import ExecutionConfig, run +from matkit.api.bundles import atomic_json, environment_versions +from matkit.api.structures import sha256 +from matkit.operation_cli import resolve_request_paths + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--spec", action="append", required=True, type=Path) + parser.add_argument("--execution", type=Path) + parser.add_argument("--outdir", required=True, type=Path) + args = parser.parse_args(argv) + root = args.outdir.resolve() + if root.exists(): + parser.error("Use a fresh evidence directory") + profile = json.loads(args.execution.read_text()) if args.execution else {} + execution = ExecutionConfig.model_validate( + {**profile, "mode": "subprocess"} + ) + root.mkdir(parents=True) + report = { + "kind": "execution evidence; not an accuracy benchmark", + "environment": environment_versions(), + "cases": [], + } + for index, spec in enumerate(args.spec): + case = { + "spec": str(spec.resolve()), + "bundle": f"{index:05d}", + "accepted": False, + } + try: + case["spec_sha256"] = sha256(spec) + request = resolve_request_paths( + json.loads(spec.read_text()), spec.resolve().parent + ) + result = run( + request, output_dir=root / case["bundle"], execution=execution + ) + case.update( + accepted=result.accepted, + state=result.state, + failure=result.failure.model_dump() if result.failure else None, + ) + except Exception as exc: + case.update(state="failed", failure=str(exc)) + report["cases"].append(case) + atomic_json(root / "execution_report.json", report) + print(json.dumps(report, indent=2, allow_nan=False)) + return 0 if all(case["accepted"] for case in report["cases"]) else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/pyproject.toml b/pyproject.toml index bce10c1..4d1fb21 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,9 +13,11 @@ dependencies = [ "click>=8.0", "networkx>=2.6", "numpy>=1.22", + "pydantic>=2,<3", ] [project.optional-dependencies] +mcp = ["mcp>=2,<3"] rdkit = ["rdkit"] mlip = ["mace-torch"] rootstock = ["rootstock>=1.6,<2"] @@ -29,6 +31,7 @@ dev = ["pytest>=7.0", "ruff>=0.4"] [project.scripts] matkit = "matkit.cli:main" +matkit-mcp = "matkit.mcp:main" [build-system] requires = ["setuptools>=61.0", "wheel"] @@ -58,3 +61,4 @@ include-package-data = true "matkit.pygraspa.files.template_mixture" = ["*"] "matkit.pygraspa.files.template_mixture_isotherm" = ["*"] "matkit.zeopp.files" = ["*"] +"matkit" = ["**/files/**/*"] diff --git a/src/matkit/__init__.py b/src/matkit/__init__.py index fe43258..c83ae69 100644 --- a/src/matkit/__init__.py +++ b/src/matkit/__init__.py @@ -3,6 +3,7 @@ __version__ = "0.1.0" _SUBMODULES = { + "api", "graspa", "graspa_sycl", "raspa2", diff --git a/src/matkit/api/__init__.py b/src/matkit/api/__init__.py new file mode 100644 index 0000000..d4b1456 --- /dev/null +++ b/src/matkit/api/__init__.py @@ -0,0 +1,64 @@ +"""Agent-independent scientific operations with portable, validated results.""" + +from .bundles import inspect_run, prepare +from .capabilities import list_capabilities +from .models import ( + AdsorptionRequest, + AlchemiAdapter, + Artifact, + BatchResult, + EvaluateRequest, + ExecutionConfig, + MACEAdapter, + MLIPMethod, + PoreRequest, + RelaxRequest, + RootstockAdapter, + RunRequest, + RunResult, + StructureData, + StructureRef, + parse_request, +) +from .runtime import ( + analyze_adsorption, + analyze_pores, + evaluate, + execute, + prepare_adsorption, + relax, + run, + run_adsorption, + run_batch, +) + +__all__ = [ + "AdsorptionRequest", + "AlchemiAdapter", + "Artifact", + "BatchResult", + "EvaluateRequest", + "ExecutionConfig", + "MACEAdapter", + "MLIPMethod", + "PoreRequest", + "RelaxRequest", + "RootstockAdapter", + "RunRequest", + "RunResult", + "StructureData", + "StructureRef", + "analyze_adsorption", + "analyze_pores", + "evaluate", + "execute", + "inspect_run", + "list_capabilities", + "parse_request", + "prepare", + "prepare_adsorption", + "relax", + "run", + "run_adsorption", + "run_batch", +] diff --git a/src/matkit/api/adapters.py b/src/matkit/api/adapters.py new file mode 100644 index 0000000..873d9fb --- /dev/null +++ b/src/matkit/api/adapters.py @@ -0,0 +1,383 @@ +"""Small adapters over the maintained calculator and engine implementations.""" + +from __future__ import annotations + +from contextlib import contextmanager +import os +from pathlib import Path +import shutil +import signal +import subprocess + +import numpy as np + +from .models import ( + AdsorptionPayload, + AdsorptionRequest, + AlchemiAdapter, + EvaluationPayload, + EvaluateRequest, + MACEAdapter, + PorePayload, + PoreRequest, + RelaxRequest, + ScientificCheck, +) +from .structures import final_structure, load_structure, sha256 + + +def backend_config(request, execution): + from matkit.mlip.config import ( + ASEMACEConfig, + NVAlchemiMACEConfig, + RootstockConfig, + ) + + method = request.method + adapter = request.adapter + device = execution.device or ( + "cuda" if isinstance(adapter, AlchemiAdapter) else "cpu" + ) + if isinstance(adapter, MACEAdapter): + return ASEMACEConfig( + **method.model_dump(), + device=device, + dtype=adapter.dtype or "float64", + ) + if isinstance(adapter, AlchemiAdapter): + return NVAlchemiMACEConfig( + checkpoint=method.checkpoint, + device=device, + dtype=adapter.dtype, + dt=adapter.dt or 0.1, + compile_model=adapter.compile_model, + enable_cueq=adapter.enable_cueq, + ) + return RootstockConfig( + checkpoint=method.checkpoint, + device=device, + **adapter.model_dump(exclude={"type"}), + ) + + +def calculation_config(request): + from matkit.mlip.config import MLIPCalculationConfig + + if isinstance(request, RelaxRequest): + return MLIPCalculationConfig( + driver="opt", + optimizer=request.optimizer, + fmax=request.fmax, + steps=request.steps, + ) + return MLIPCalculationConfig() + + +@contextmanager +def calculator_session(request, execution): + from matkit.mlip import runner + + backend = backend_config(request, execution) + if isinstance(request.adapter, AlchemiAdapter): + model = runner._load_nvalchemi_model(backend) + yield backend, model + else: + with runner._ase_backend_context(backend) as calculator: + yield backend, calculator + + +def evaluate_items(entries, request, execution, session): + """Yield item results while keeping the calculator/model alive.""" + from matkit.mlip import runner + + backend, calculator = session + calculation = calculation_config(request) + if isinstance(request.adapter, AlchemiAdapter): + chunks = runner._chunks_by_capacity( + entries, request.adapter.batch_size, request.adapter.max_atoms + ) + for chunk in chunks: + try: + results = runner._run_nvalchemi_chunk( + calculator, chunk, backend, calculation + ) + if [i for i, _ in results] != [i for i, _, _ in chunk]: + raise ValueError( + "ALCHEMI returned inconsistent result ordering" + ) + except Exception as exc: + results = [(i, exc) for i, _, _ in chunk] + yield from results + else: + properties = ( + request.properties if isinstance(request, EvaluateRequest) else None + ) + for index, input_file, atoms in entries: + try: + result = runner._run_ase_item( + input_file, + atoms, + calculator, + backend, + calculation, + requested_properties=properties, + ) + except Exception as exc: + result = exc + yield index, result + + +def calculator_payload(root, request, legacy): + if isinstance(legacy, Exception): + raise legacy + if not legacy["success"]: + raise ValueError(legacy["error"]) + original_atoms, structure, source_hash = load_structure(request.structure) + properties = ( + request.properties + if isinstance(request, EvaluateRequest) + else ["potential_energy", "forces"] + ) + names = { + "potential_energy": "energy", + "forces": "forces", + "stress": "stress", + } + for name in properties: + if legacy.get(names[name]) is None: + raise ValueError(f"Requested property {name} is unavailable") + final = legacy["final_structure"] + atoms = original_atoms.copy() + if final["atomic_numbers"] != atoms.numbers.tolist(): + raise ValueError("Calculator changed atom correspondence") + atoms.positions = final["positions"] + atoms.cell = final["cell"] + atoms.pbc = final["pbc"] + float32 = ( + isinstance(request.adapter, AlchemiAdapter) + and request.adapter.dtype == "float32" + ) + relative_tolerance = 1e-6 if float32 else 0 + absolute_tolerance = 1e-6 if float32 else 1e-7 + if not np.allclose( + atoms.cell.array, + original_atoms.cell.array, + rtol=relative_tolerance, + atol=absolute_tolerance, + ) or not np.array_equal(atoms.pbc, original_atoms.pbc): + raise ValueError("Fixed-cell calculation changed cell or periodicity") + if isinstance(request, EvaluateRequest) and not np.allclose( + atoms.positions, + original_atoms.positions, + rtol=relative_tolerance, + atol=absolute_tolerance, + ): + raise ValueError("Property evaluation changed atomic positions") + out = root / "work" / "final_structure.extxyz" + final_structure(atoms, structure, source_hash, out) + checks = [] + converged = None + if isinstance(request, RelaxRequest): + max_force = float(np.linalg.norm(legacy["forces"], axis=1).max()) + converged = ( + bool(legacy["converged"]) + and max_force <= request.fmax * (1 + 1e-6) + 1e-12 + ) + checks.append( + ScientificCheck( + name="force_convergence", + status="passed" if converged else "failed", + required=True, + detail=( + f"Requested fmax={request.fmax} eV/angstrom; " + f"optimizer reports convergence={converged}" + ), + ) + ) + payload = EvaluationPayload( + potential_energy=legacy.get("energy") + if "potential_energy" in properties + else None, + forces=legacy.get("forces") if "forces" in properties else None, + stress=legacy.get("stress") if "stress" in properties else None, + converged=converged, + n_steps=legacy.get("n_steps"), + final_structure="work/final_structure.extxyz", + ) + return payload, checks, {"calculation_s": legacy["calculation_time_s"]} + + +def stop_process(process, *, group=True): + if group and os.name == "posix": + try: + os.killpg(process.pid, signal.SIGTERM) + except ProcessLookupError: + pass + elif process.poll() is None: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + if not (group and os.name == "posix"): + process.kill() + finally: + if group and os.name == "posix": + # A child can outlive a terminated group leader. + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + process.wait() + + +def external_command(root, execution, engine, arguments): + from .bundles import atomic_json + + command = execution.executables.get(engine) + if command is None: + binary = shutil.which( + "network" if engine == "zeopp" else "simulate", + path=execution.environment.get("PATH"), + ) + if binary is None: + raise FileNotFoundError( + f"Configure the {engine} executable in the execution profile" + ) + command = [binary] + binary = shutil.which(command[0], path=execution.environment.get("PATH")) + if binary is None: + raise FileNotFoundError(f"Executable unavailable: {command[0]}") + command = [str(Path(binary).resolve()), *command[1:], *arguments] + atomic_json( + root / "command.json", + { + "argv": command, + "cwd": "work", + "executable_sha256": sha256(Path(command[0])), + "argument_file_hashes": { + token: sha256(candidate) + for token in command + if (candidate := Path(token)).is_file() + }, + }, + ) + # Worker children share its process group so supervision covers the engine. + own_group = os.environ.get("MATKIT_WORKER_PROCESS") != "1" + log_name = "raspa.log" if engine == "graspa" else "engine.stdout.log" + with ( + (root / "work" / log_name).open("w") as stdout, + (root / "work" / "engine.stderr.log").open("w") as stderr, + ): + process = subprocess.Popen( + command, + cwd=root / "work", + env={**os.environ, **execution.environment}, + stdout=stdout, + stderr=stderr, + start_new_session=own_group and os.name == "posix", + ) + try: + code = process.wait(timeout=execution.timeout_s) + except BaseException: + stop_process(process, group=own_group) + raise + atomic_json(root / "exit.json", {"returncode": code}) + if code != 0: + raise RuntimeError( + f"{engine} exited with code {code}; see work/engine.stderr.log" + ) + + +def run_pores(root, request, execution): + args = ["-ha"] if request.high_accuracy else [] + args.extend(["-r", request.radii_file]) + for analysis in request.analyses: + args.append(f"-{analysis}") + if analysis in {"sa", "vol", "psd"}: + args.extend( + [ + str(request.probe_radius), + str(request.channel_radius), + str(request.num_samples), + ] + ) + elif analysis == "chan": + args.append(str(request.probe_radius)) + args.append("structure.cif") + external_command(root, execution, "zeopp", args) + return parse_pores(root, request) + + +def parse_pores(root, request): + from matkit.zeopp.zeopp import _PARSERS + + required = { + "res": {"Di", "Df", "Dif"}, + "sa": {"ASA", "NASA", "density", "unitcell_volume"}, + "vol": {"AV", "NAV", "density", "unitcell_volume"}, + "psd": {"bin_lower", "counts"}, + "chan": {"num_channels", "dimensionalities"}, + } + results = {} + for analysis in request.analyses: + data = _PARSERS[analysis](root / "work" / f"structure.{analysis}") + if not required[analysis] <= data.keys(): + raise ValueError(f"Incomplete requested Zeo++ analysis: {analysis}") + if analysis == "psd" and not data["counts"]: + raise ValueError("Empty pore-size distribution") + if analysis == "chan" and data["num_channels"] != len( + data["dimensionalities"] + ): + raise ValueError("Channel count does not match dimensionalities") + results[analysis] = data + return ( + PorePayload(results=results), + [ + ScientificCheck( + name="sampling_quality", + status="unknown", + detail="Execution does not establish sampling accuracy", + ) + ], + {}, + ) + + +def parse_adsorption(root, request): + from matkit.graspa import get_output_data + + data = get_output_data( + str(root / "work"), + unit=request.unit, + eos=request.fugacity_coefficient == "PR-EOS", + ) + if not data["success"]: + raise ValueError("gRASPA output is incomplete") + payload = AdsorptionPayload( + component=request.adsorbate, + uptake=data["uptake"], + unit=data["unit"], + uncertainty=data["error"], + heat_of_adsorption=data["qst"], + heat_uncertainty=data["error_qst"], + ) + checks = [ + ScientificCheck( + name="sampling_quality", + status="unknown", + detail=( + "Engine statistics do not establish equilibration " + "or independent samples" + ), + ) + ] + return payload, checks, {"engine_reported_s": data["calc_time_in_s"]} + + +def run_external(root, request, execution): + if isinstance(request, PoreRequest): + return run_pores(root, request, execution) + if isinstance(request, AdsorptionRequest): + external_command(root, execution, "graspa", []) + return parse_adsorption(root, request) + raise TypeError("Not an external-engine request") diff --git a/src/matkit/api/bundles.py b/src/matkit/api/bundles.py new file mode 100644 index 0000000..e0510e4 --- /dev/null +++ b/src/matkit/api/bundles.py @@ -0,0 +1,434 @@ +"""Relocatable calculation bundles and atomic run records.""" + +from __future__ import annotations + +from contextlib import contextmanager +from importlib import metadata, resources +import json +import mimetypes +import os +from pathlib import Path +import platform +import re +import shutil +import tempfile +from uuid import uuid4 + +import numpy as np +from ase.io import write +from ase.io.cif import parse_cif + +from .models import ( + AdsorptionRequest, + Artifact, + CalculatorRequest, + Failure, + PoreRequest, + RunResult, + StructureRef, + parse_request, +) +from .structures import load_structure, sha256 + + +def atomic_json(path: Path, value) -> None: + data = ( + value.model_dump(mode="json") if hasattr(value, "model_dump") else value + ) + encoded = json.dumps(data, allow_nan=False, indent=2) + temporary = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=path.parent, + prefix=f".{path.name}.", + 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) + + +@contextmanager +def claim(root: Path): + lock = root / ".matkit.lock" + with lock.open("x") as stream: + stream.write(str(os.getpid())) + try: + yield + finally: + lock.unlink(missing_ok=True) + + +def contained_path(root: Path, relative: str) -> Path: + path = (root / relative).resolve() + if Path(relative).is_absolute() or not path.is_relative_to(root.resolve()): + raise ValueError("Artifact paths must stay within the run bundle") + return path + + +def artifact(root: Path, path: Path, role: str) -> Artifact: + relative = str(path.relative_to(root)) + contained_path(root, relative) + return Artifact( + path=relative, + sha256=sha256(path), + size_bytes=path.stat().st_size, + media_type=mimetypes.guess_type(path.name)[0] + or "application/octet-stream", + role=role, + ) + + +def collect_artifacts(root: Path) -> list[Artifact]: + result = [] + for path in sorted(root.rglob("*")): + if path.is_symlink(): + raise ValueError(f"Symlink artifacts are unsupported: {path.name}") + if ( + path.is_file() + and not path.name.startswith(".") + and path.name not in {"run.json", "result.json"} + ): + relative = path.relative_to(root) + role = ( + "input" + if relative.parts[0] == "inputs" + or relative.name == "request.json" + else "output" + ) + result.append(artifact(root, path, role)) + return result + + +def verify_inputs(root: Path, record: RunResult) -> None: + # Verify the full prepared inventory, including generated work files. + for ref in record.artifacts: + path = contained_path(root, ref.path) + if ( + path.is_symlink() + or not path.is_file() + or sha256(path) != ref.sha256 + ): + raise ValueError( + f"Staged artifact changed or is missing: {ref.path}" + ) + + +def inspect_run(path: str | Path) -> RunResult: + root = Path(path).expanduser().resolve() + # The committed result is authoritative after a manifest-update failure. + result_path = root / "result.json" + if not result_path.exists(): + result_path = root / "run.json" + result = RunResult.model_validate_json(result_path.read_text()) + if result_path.name == "result.json" and (root / "run.json").exists(): + try: + manifest = RunResult.model_validate_json( + (root / "run.json").read_text() + ) + except (OSError, ValueError): + return result + if manifest.run_id == result.run_id and manifest.state in { + "failed", + "interrupted", + }: + return manifest + return result + + +def commit_result(root: Path, result: RunResult) -> RunResult: + # Validate updates as model_copy deliberately skips validation. + result = RunResult.model_validate(result.model_dump(mode="json")) + atomic_json(root / "result.json", result) + atomic_json(root / "run.json", result) + return result + + +def environment_versions() -> dict: + versions = {} + for package in ( + "matkit", + "ase", + "numpy", + "pydantic", + "mace-torch", + "rootstock", + "nvalchemi-toolkit", + "torch", + ): + try: + versions[package] = metadata.version(package) + except metadata.PackageNotFoundError: + versions[package] = None + return { + "python": platform.python_version(), + "platform": platform.platform(), + "packages": versions, + "runtime_settings": { + key: os.environ.get(key) + for key in ( + "OMP_NUM_THREADS", + "CUDA_VISIBLE_DEVICES", + "MKL_NUM_THREADS", + ) + }, + } + + +def _copy_tree(source, destination: Path): + destination.mkdir() + for entry in source.iterdir(): + if isinstance(entry, Path) and entry.is_symlink(): + raise ValueError("Template symlinks are unsupported") + target = destination / entry.name + if entry.is_dir(): + _copy_tree(entry, target) + else: + target.write_bytes(entry.read_bytes()) + + +def _set_input(path: Path, key: str, value) -> None: + content, count = re.subn( + rf"^{re.escape(key)}\s+.*$", + f"{key} {value}", + path.read_text(), + flags=re.M, + ) + if count != 1: + raise ValueError(f"Template must define {key} exactly once") + path.write_text(content) + + +def _prepare_adsorption(root, request, atoms): + from matkit.graspa import setup_simulation + from matkit.utils import calculate_cell_size + + source = root / request.structure.path + if ( + source.suffix.lower() != ".cif" + or not atoms.pbc.all() + or atoms.cell.rank != 3 + ): + raise ValueError("gRASPA requires a fully periodic CIF") + if atoms.constraints: + raise ValueError("gRASPA constraint conversion is unsupported") + tags = dict(list(parse_cif(str(source)))[0]) + charges = np.asarray(tags.get("_atom_site_charge", []), dtype=float) + if charges.shape != (len(atoms),) or not np.isfinite(charges).all(): + raise ValueError( + "gRASPA requires finite, atom-mapped _atom_site_charge " + "values in the CIF" + ) + if not np.isclose(charges.sum(), request.net_charge, atol=1e-4, rtol=0): + raise ValueError("CIF charges do not sum to requested net_charge") + if request.number_of_blocks > request.production_cycles: + raise ValueError("number_of_blocks exceeds production_cycles") + template = root / request.template_dir + for name in ( + "simulation.input", + "pseudo_atoms.def", + "force_field.def", + "force_field_mixing_rules.def", + f"{request.adsorbate}.def", + ): + if not (template / name).is_file(): + raise ValueError(f"Missing simulation definition: {name}") + sizes = calculate_cell_size(atoms, request.cutoff_angstrom) + setup_simulation( + str(source), + str(root / "work"), + [ + { + "MoleculeName": request.adsorbate, + "FugacityCoefficient": request.fugacity_coefficient, + } + ], + temperature=request.temperature_K, + pressure=request.pressure_Pa, + cutoff=request.cutoff_angstrom, + n_cycle=request.production_cycles, + template_dir=str(template), + cell_size=sizes, + ) + input_file = root / "work" / "simulation.input" + for key, value in { + "NumberOfInitializationCycles": request.initialization_cycles, + "NumberOfEquilibrationCycles": request.equilibration_cycles, + "NumberOfProductionCycles": request.production_cycles, + "NumberOfBlocks": request.number_of_blocks, + }.items(): + _set_input(input_file, key, value) + text = input_file.read_text() + if len(re.findall(r"^Component\s+\d+\s+MoleculeName", text, re.M)) != 1: + raise ValueError( + "Unified gRASPA execution supports exactly one component" + ) + for key, expected in { + "UseChargesFromCIFFile": "yes", + "NumberOfSimulations": "1", + "SingleSimulation": "yes", + "RestartFile": "no", + }.items(): + match = re.search(rf"^{key}\s+(\S+)", text, re.M) + if not match or match[1].lower() != expected: + raise ValueError( + f"Unified gRASPA execution requires {key} {expected}" + ) + return { + "unit_cells": sizes, + "simulation_input": "work/simulation.input", + "random_seed": None, + "uncertainty_method": "unknown; engine-reported", + } + + +def prepare(request, *, output_dir: str | Path) -> RunResult: + request = parse_request(request) + atoms, structure, digest = load_structure(request.structure) + if ( + isinstance(request, CalculatorRequest) + and request.adapter.type == "nvalchemi-mace" + ): + if atoms.constraints or structure.arrays or structure.bonds: + raise ValueError( + "Native ALCHEMI does not support this structure's " + "arrays, bonds or constraints" + ) + if isinstance(request, PoreRequest) and ( + not atoms.pbc.all() or atoms.cell.rank != 3 + ): + raise ValueError("Zeo++ requires a fully periodic cell") + root = Path(output_dir).expanduser().resolve() + root.mkdir(parents=True, exist_ok=True) + with claim(root): + if any(p.name != ".matkit.lock" for p in root.iterdir()): + raise FileExistsError( + "Use a fresh output directory; resume is not supported" + ) + record = RunResult( + run_id=uuid4().hex, + operation=request.operation, + state="prepared", + requested=request.model_dump(mode="json"), + provenance={ + "input_sha256": digest, + "model_sha256": None, + "preparation_environment": environment_versions(), + }, + ) + atomic_json(root / "run.json", record) + try: + (root / "inputs").mkdir() + (root / "work").mkdir() + source = Path(request.structure.path).expanduser().resolve() + staged = root / "inputs" / f"structure{source.suffix.lower()}" + shutil.copyfile(source, staged) + atomic_json(staged.with_suffix(".metadata.json"), structure) + values = request.model_dump(mode="json") + values["structure"] = StructureRef( + path=str(staged.relative_to(root)), sha256=digest + ).model_dump(mode="json") + if isinstance(request, CalculatorRequest): + checkpoint = Path(request.method.checkpoint).expanduser() + if checkpoint.is_file(): + target = root / "inputs" / f"model{checkpoint.suffix}" + shutil.copyfile(checkpoint, target) + values["method"]["checkpoint"] = str( + target.relative_to(root) + ) + record.provenance["model_sha256"] = sha256(target) + elif checkpoint.is_absolute(): + raise FileNotFoundError( + f"Checkpoint file does not exist: {checkpoint}" + ) + if ( + request.adapter.type == "rootstock" + and request.adapter.weights + ): + weights = Path(request.adapter.weights).expanduser() + target = root / "inputs" / f"weights{weights.suffix}" + shutil.copyfile(weights, target) + values["adapter"]["weights"] = str(target.relative_to(root)) + record.provenance["model_sha256"] = sha256(target) + if isinstance(request, PoreRequest): + radii = root / "inputs" / "radii.rad" + if request.radii_file: + shutil.copyfile( + Path(request.radii_file).expanduser(), radii + ) + else: + radii.write_bytes( + resources.files("matkit.zeopp") + .joinpath("files/UFF.rad") + .read_bytes() + ) + values["radii_file"] = str(radii.relative_to(root)) + if source.suffix.lower() == ".cif": + shutil.copyfile(source, root / "work" / "structure.cif") + else: + write(root / "work" / "structure.cif", atoms) + if isinstance(request, AdsorptionRequest): + template = ( + Path(request.template_dir).expanduser() + if request.template_dir + else resources.files("matkit.graspa").joinpath( + "files/template" + ) + ) + _copy_tree(template, root / "inputs" / "template") + values["template_dir"] = "inputs/template" + staged_request = parse_request(values) + record.resolved.update( + _prepare_adsorption(root, staged_request, atoms) + ) + atomic_json(root / "request.json", values) + record.resolved["request"] = values + record.artifacts = collect_artifacts(root) + record.provenance["prepared_hashes"] = { + a.path: a.sha256 for a in record.artifacts + } + atomic_json(root / "run.json", record) + return record + except Exception as exc: + failed = record.model_copy( + update={ + "state": "failed", + "failure": Failure( + code=type(exc).__name__, + stage="preparation", + message=str(exc), + ), + } + ) + atomic_json(root / "run.json", failed) + raise + + +def staged_request(root: Path): + request = parse_request(json.loads((root / "request.json").read_text())) + values = request.model_dump(mode="json") + values["structure"]["path"] = str( + contained_path(root, request.structure.path) + ) + if isinstance(request, CalculatorRequest): + checkpoint = request.method.checkpoint + if checkpoint.startswith("inputs/"): + values["method"]["checkpoint"] = str( + contained_path(root, checkpoint) + ) + if request.adapter.type == "rootstock" and request.adapter.weights: + values["adapter"]["weights"] = str( + contained_path(root, request.adapter.weights) + ) + if isinstance(request, PoreRequest): + values["radii_file"] = str(contained_path(root, request.radii_file)) + if isinstance(request, AdsorptionRequest): + values["template_dir"] = str(contained_path(root, request.template_dir)) + return parse_request(values) diff --git a/src/matkit/api/capabilities.py b/src/matkit/api/capabilities.py new file mode 100644 index 0000000..89e4510 --- /dev/null +++ b/src/matkit/api/capabilities.py @@ -0,0 +1,87 @@ +"""Discovery must not load optional calculators, models or GPU runtimes.""" + +from importlib.metadata import PackageNotFoundError, version +import shutil + + +_CAPABILITIES = { + "ase-mace": ( + "mace-torch", + ["evaluate", "relax"], + ["potential_energy", "forces", "stress (model-dependent)"], + "CPU/CUDA; MACE-supported Python and dependencies", + "Species, periodicity and supported properties depend on the model", + ), + "rootstock": ( + "rootstock", + ["evaluate", "relax"], + ["potential_energy", "forces", "stress (deployment-dependent)"], + "Separately installed Rootstock model environment", + "Properties and species depend on the deployed model; " + "caller discovery does not verify the worker environment", + ), + "nvalchemi-mace": ( + "nvalchemi-toolkit", + ["evaluate", "relax"], + ["potential_energy", "forces", "stress (model-dependent)"], + "Compatible Python/CUDA stack; native batches; FIRE only", + "Model-dependent properties/species; unsupported atom arrays, " + "bonds and constraints are rejected", + ), + "zeopp": ( + None, + ["pores"], + ["res", "sa", "vol", "psd", "chan"], + "Zeo++ network executable", + "Fully periodic structures; explicit radii and requested analyses", + ), + "graspa": ( + None, + ["adsorption"], + ["single-component absolute uptake", "heat of adsorption"], + "gRASPA CUDA; charged periodic CIF and force-field definitions", + "Single component; atom-mapped CIF charges; sampling quality unknown", + ), +} + + +def list_capabilities() -> list[dict]: + capabilities = [] + for adapter, ( + package, + operations, + properties, + environment, + restrictions, + ) in _CAPABILITIES.items(): + installed_version = None + if package: + try: + installed_version = version(package) + except PackageNotFoundError: + pass + available = installed_version is not None + else: + available = ( + shutil.which("network" if adapter == "zeopp" else "simulate") + is not None + ) + capabilities.append( + { + "adapter": adapter, + "operations": operations, + "properties": properties, + "available_in_caller": available, + "installed_version": installed_version, + "environment": environment, + "status": "experimental", + "evidence": "CPU fixtures; real-execution evidence not bundled", + "geometry_optimization": adapter + in {"ase-mace", "rootstock", "nvalchemi-mace"}, + "cell_optimization": False, + "restart": False, + "native_batch": adapter == "nvalchemi-mace", + "restrictions": restrictions, + } + ) + return capabilities diff --git a/src/matkit/api/models.py b/src/matkit/api/models.py new file mode 100644 index 0000000..d147a38 --- /dev/null +++ b/src/matkit/api/models.py @@ -0,0 +1,395 @@ +"""Versioned scientific requests and portable, engine-independent results.""" + +from __future__ import annotations + +import json +from typing import Annotated, Literal, Union + +from pydantic import ( + BaseModel, + ConfigDict, + Field, + JsonValue, + StrictInt, + TypeAdapter, + model_validator, +) + +Positive = Annotated[float, Field(gt=0, strict=True)] +Nonnegative = Annotated[float, Field(ge=0, strict=True)] +Count = Annotated[StrictInt, Field(ge=1)] +Cycles = Annotated[StrictInt, Field(ge=0)] +Vector = tuple[float, float, float] +Matrix = tuple[Vector, Vector, Vector] +Digest = Annotated[str, Field(pattern=r"^[a-f0-9]{64}$")] + + +class Model(BaseModel): + model_config = ConfigDict( + extra="forbid", + allow_inf_nan=False, + validate_default=True, + revalidate_instances="always", + ) + + @model_validator(mode="after") + def finite_json(self): + # Reject non-standard JSON in nested JsonValue fields too. + json.dumps(self.model_dump(mode="json"), allow_nan=False) + return self + + +class StructureData(Model): + atomic_numbers: list[Count] + positions: list[Vector] + cell: Matrix + pbc: tuple[bool, bool, bool] + atom_ids: list[str] + arrays: dict[str, JsonValue] = Field(default_factory=dict) + info: dict[str, JsonValue] = Field(default_factory=dict) + constraints: list[dict[str, JsonValue]] = Field(default_factory=list) + labels: list[str] | None = None + bonds: list[tuple[str, str]] = Field(default_factory=list) + parent_sha256: Digest | None = None + derived_from: dict[str, Digest] = Field(default_factory=dict) + + @model_validator(mode="after") + def atom_correspondence(self): + size = len(self.atomic_numbers) + if ( + not size + or len(self.positions) != size + or len(self.atom_ids) != size + ): + raise ValueError( + "Species, positions and atom_ids require equal nonzero lengths" + ) + if len(set(self.atom_ids)) != size: + raise ValueError("atom_ids must be unique") + if self.labels is not None and len(self.labels) != size: + raise ValueError("labels must correspond to every atom") + for name, array in self.arrays.items(): + if not isinstance(array, list) or len(array) != size: + raise ValueError(f"Array {name} must correspond to every atom") + if any( + a not in self.atom_ids or b not in self.atom_ids + for a, b in self.bonds + ): + raise ValueError("Bond endpoints must refer to atom_ids") + return self + + +class StructureRef(Model): + path: str + sha256: Digest | None = None + metadata: StructureData | None = None + + +class ExecutionConfig(Model): + """Execution location belongs to the caller, not the scientific method.""" + + mode: Literal["inprocess", "subprocess"] = "inprocess" + python: str | None = None + device: str | None = None + environment: dict[str, str] = Field(default_factory=dict) + executables: dict[Literal["zeopp", "graspa"], list[str]] = Field( + default_factory=dict + ) + timeout_s: Positive | None = None + + @model_validator(mode="after") + def commands_not_empty(self): + if any( + not command or any(not token for token in command) + for command in self.executables.values() + ): + raise ValueError( + "Executable commands must be nonempty argument lists" + ) + return self + + +class MLIPMethod(Model): + checkpoint: str = Field(min_length=1) + calculator_type: Literal["mace_mp", "mace_off", "mace_anicc"] = "mace_mp" + dispersion: bool = False + damping: str = "bj" + dispersion_xc: str = "pbe" + dispersion_cutoff: Positive = 21.167088422553647 + + +class MACEAdapter(Model): + type: Literal["ase-mace"] = "ase-mace" + dtype: Literal["float32", "float64"] | None = None + + +class RootstockAdapter(Model): + type: Literal["rootstock"] = "rootstock" + cluster: str | None = None + root: str | None = None + cache_root: str | None = None + setup_kwargs: dict[str, JsonValue] = Field(default_factory=dict) + timeout: Positive = 600 + weights: str | None = None + + @model_validator(mode="after") + def location(self): + if self.cluster is not None and self.root is not None: + raise ValueError("Rootstock cannot specify both cluster and root") + return self + + +class AlchemiAdapter(Model): + type: Literal["nvalchemi-mace"] = "nvalchemi-mace" + dtype: Literal["float32", "float64"] = "float32" + dt: Positive | None = None + compile_model: bool = False + enable_cueq: bool = False + batch_size: Count = 16 + max_atoms: Count | None = None + + +CalculatorAdapter = Annotated[ + Union[MACEAdapter, RootstockAdapter, AlchemiAdapter], + Field(discriminator="type"), +] + + +class RequestBase(Model): + schema_name: Literal["matkit.request"] = "matkit.request" + schema_version: Literal[1] = 1 + structure: StructureRef + + +class CalculatorRequest(RequestBase): + method: MLIPMethod + adapter: CalculatorAdapter = Field(default_factory=MACEAdapter) + + @model_validator(mode="after") + def applicable_method(self): + method = self.method + defaults = MLIPMethod(checkpoint=method.checkpoint) + if self.adapter.type != "ase-mace" and method != defaults: + raise ValueError( + "MACE factory/dispersion settings require ase-mace" + ) + if method.dispersion and method.calculator_type != "mace_mp": + raise ValueError("Dispersion requires mace_mp") + if ( + isinstance(self.adapter, MACEAdapter) + and method.calculator_type == "mace_anicc" + and self.adapter.dtype is not None + ): + raise ValueError("mace_anicc controls its own precision") + return self + + +class EvaluateRequest(CalculatorRequest): + operation: Literal["evaluate"] = "evaluate" + properties: list[Literal["potential_energy", "forces", "stress"]] = Field( + default_factory=lambda: ["potential_energy"], min_length=1 + ) + + @model_validator(mode="after") + def applicable_options(self): + if len(set(self.properties)) != len(self.properties): + raise ValueError("Requested properties must be unique") + if ( + isinstance(self.adapter, AlchemiAdapter) + and self.adapter.dt is not None + ): + raise ValueError("ALCHEMI dt requires relaxation") + return self + + +class RelaxRequest(CalculatorRequest): + operation: Literal["relax"] = "relax" + optimizer: Literal["bfgs", "lbfgs", "gpmin", "fire", "mdmin"] = "fire" + fmax: Positive = 0.01 + steps: Count = 1000 + + @model_validator(mode="after") + def native_optimizer(self): + if self.adapter.type == "nvalchemi-mace" and self.optimizer != "fire": + raise ValueError("ALCHEMI supports only FIRE") + return self + + +class PoreRequest(RequestBase): + operation: Literal["pores"] = "pores" + adapter: Literal["zeopp"] = "zeopp" + analyses: list[Literal["res", "sa", "vol", "psd", "chan"]] = Field( + default_factory=lambda: ["res"], min_length=1 + ) + probe_radius: Positive = 1.86 + channel_radius: Positive = 1.86 + num_samples: Count = 2000 + high_accuracy: bool = True + radii_file: str | None = None + + @model_validator(mode="after") + def unique_analyses(self): + if len(set(self.analyses)) != len(self.analyses): + raise ValueError("Analyses must be unique") + return self + + +class AdsorptionRequest(RequestBase): + operation: Literal["adsorption"] = "adsorption" + adapter: Literal["graspa"] = "graspa" + adsorbate: str = Field(pattern=r"^[A-Za-z0-9_-]+$") + temperature_K: Positive + pressure_Pa: Nonnegative + cutoff_angstrom: Positive = 12.8 + initialization_cycles: Cycles = 1000 + equilibration_cycles: Cycles = 0 + production_cycles: Count = 1000 + number_of_blocks: Count = 1 + fugacity_coefficient: Union[Positive, Literal["PR-EOS"]] = "PR-EOS" + unit: Literal["mol/kg", "mg/g", "g/L"] = "mol/kg" + template_dir: str | None = None + net_charge: float = 0 + + +RunRequest = Annotated[ + Union[EvaluateRequest, RelaxRequest, PoreRequest, AdsorptionRequest], + Field(discriminator="operation"), +] +REQUEST_ADAPTER = TypeAdapter(RunRequest) + + +def parse_request(value: RunRequest | dict) -> RunRequest: + return REQUEST_ADAPTER.validate_python(value) + + +class Artifact(Model): + path: str + sha256: Digest + size_bytes: Annotated[StrictInt, Field(ge=0)] + media_type: str = "application/octet-stream" + role: str + + +class ScientificCheck(Model): + name: str + status: Literal["passed", "failed", "unknown", "not_applicable"] + required: bool = False + detail: str = "" + + +class Failure(Model): + code: str + stage: str + message: str + + +class EvaluationPayload(Model): + kind: Literal["evaluation"] = "evaluation" + potential_energy: float | None = None + energy_unit: Literal["eV"] = "eV" + energy_definition: str = ( + "calculator potential energy; model-specific reference" + ) + forces: list[Vector] | None = None + force_unit: Literal["eV/angstrom"] = "eV/angstrom" + stress: Matrix | None = None + stress_unit: Literal["eV/angstrom^3"] = "eV/angstrom^3" + stress_convention: str = "ASE Cartesian tensor, positive in tension" + converged: bool | None = None + n_steps: Cycles | None = None + final_structure: str | None = None + + +class PorePayload(Model): + kind: Literal["pores"] = "pores" + results: dict[str, dict[str, JsonValue]] + + +class AdsorptionPayload(Model): + kind: Literal["adsorption"] = "adsorption" + component: str + uptake: float + uptake_basis: Literal["absolute, per framework mass/volume"] = ( + "absolute, per framework mass/volume" + ) + unit: Literal["mol/kg", "mg/g", "g/L"] + uncertainty: Nonnegative + uncertainty_method: str = "unknown; reported by engine" + heat_of_adsorption: float + heat_unit: Literal["kJ/mol"] = "kJ/mol" + heat_uncertainty: Nonnegative + heat_convention: str = "as reported by gRASPA" + + +Payload = Annotated[ + Union[EvaluationPayload, PorePayload, AdsorptionPayload], + Field(discriminator="kind"), +] + + +class RunResult(Model): + schema_name: Literal["matkit.run"] = "matkit.run" + schema_version: Literal[1] = 1 + run_id: str + operation: Literal["evaluate", "relax", "pores", "adsorption"] + state: Literal["prepared", "running", "completed", "failed", "interrupted"] + numerical_validity: Literal["valid", "invalid", "unknown"] = "unknown" + checks: list[ScientificCheck] = Field(default_factory=list) + requested: dict[str, JsonValue] + resolved: dict[str, JsonValue] = Field(default_factory=dict) + provenance: dict[str, JsonValue] = Field(default_factory=dict) + timings: dict[str, Nonnegative] = Field(default_factory=dict) + artifacts: list[Artifact] = Field(default_factory=list) + payload: Payload | None = None + failure: Failure | None = None + + @property + def accepted(self) -> bool: + return ( + self.state == "completed" + and self.numerical_validity == "valid" + and all( + not check.required + or check.status in {"passed", "not_applicable"} + for check in self.checks + ) + ) + + @model_validator(mode="after") + def consistent_outcome(self): + if self.state == "completed" and ( + self.payload is None + or self.numerical_validity != "valid" + or self.failure is not None + ): + raise ValueError( + "Completed runs require valid results and no execution failure" + ) + if self.state in {"failed", "interrupted"} and self.failure is None: + raise ValueError("Failed/interrupted runs require a failure record") + if self.payload is not None: + kind = ( + "evaluation" + if self.operation in {"evaluate", "relax"} + else self.operation + ) + if self.payload.kind != kind: + raise ValueError("Payload kind does not match operation") + return self + + +class BatchResult(Model): + schema_name: Literal["matkit.batch"] = "matkit.batch" + schema_version: Literal[1] = 1 + state: Literal["running", "completed", "partial", "failed", "interrupted"] + items: list[dict[str, JsonValue]] + failure: Failure | None = None + + @property + def accepted(self) -> bool: + return ( + self.state == "completed" + and self.failure is None + and bool(self.items) + and all(item.get("accepted") for item in self.items) + ) diff --git a/src/matkit/api/runtime.py b/src/matkit/api/runtime.py new file mode 100644 index 0000000..f142a40 --- /dev/null +++ b/src/matkit/api/runtime.py @@ -0,0 +1,573 @@ +"""Synchronous operations, supervised workers and durable batches.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +import subprocess +import sys +import time + +from . import adapters +from .bundles import ( + atomic_json, + claim, + collect_artifacts, + commit_result, + environment_versions, + inspect_run, + prepare, + staged_request, + verify_inputs, +) +from .models import ( + AdsorptionRequest, + BatchResult, + CalculatorRequest, + EvaluateRequest, + ExecutionConfig, + Failure, + PoreRequest, + RelaxRequest, + parse_request, +) +from .structures import load_structure + + +def _begin(root, record, execution): + verify_inputs(root, record) + request = staged_request(root) + record = record.model_copy(deep=True) + record.state = "running" + record.provenance["execution_environment"] = environment_versions() + record.resolved["execution"] = execution.model_dump(mode="json") + if isinstance(request, CalculatorRequest): + backend = adapters.backend_config(request, execution) + record.resolved["calculator"] = backend.to_dict() + if request.method.calculator_type == "mace_anicc": + record.resolved["calculator"]["dtype"] = None + record.provenance["precision_source"] = "calculator factory" + if request.adapter.type == "rootstock": + record.provenance["calculator_environment"] = ( + "Rootstock worker; resolved environment unavailable" + ) + # Checkpoint aliases may resolve inside a worker; do not invent a hash. + record.provenance["model_identity_evidence"] = ( + "content hash" + if record.provenance.get("model_sha256") + else "requested identifier only" + ) + atomic_json(root / "run.json", record) + return request, record + + +def _finish(root, record, output, started): + payload, checks, timings = output + finished = record.model_copy( + update={ + "state": "completed", + "numerical_validity": "valid", + "payload": payload, + "checks": checks, + "timings": {**timings, "wall_s": time.perf_counter() - started}, + "artifacts": collect_artifacts(root), + "failure": None, + } + ) + return commit_result(root, finished) + + +def _fail(root, record, exc, stage, interrupted=False): + # Never replace an already committed scientific result after a write or + # teardown failure. The orchestration error remains visible to the caller. + if (root / "result.json").exists(): + try: + committed = inspect_run(root) + failure = Failure( + code=type(exc).__name__, + stage="orchestration", + message=str(exc) or type(exc).__name__, + ) + atomic_json( + root / "run.json", + committed.model_copy( + update={"state": "interrupted", "failure": failure} + ), + ) + except Exception: + pass # The committed result still survives a filesystem failure. + raise exc + try: + artifacts = collect_artifacts(root) + except Exception: + artifacts = record.artifacts + failed = record.model_copy( + update={ + "state": "interrupted" if interrupted else "failed", + "numerical_validity": "invalid" + if stage == "calculation" + else "unknown", + "failure": Failure( + code=type(exc).__name__, + stage=stage, + message=str(exc) or type(exc).__name__, + ), + "artifacts": artifacts, + } + ) + return commit_result(root, failed) + + +def _execute_claimed(root, execution): + record = inspect_run(root) + if record.state != "prepared": + raise FileExistsError( + "Only prepared runs can execute; use a fresh bundle for reruns" + ) + started = time.perf_counter() + stage = "validation" + try: + request, record = _begin(root, record, execution) + stage = "calculation" + if isinstance(request, CalculatorRequest): + atoms, _, _ = load_structure(request.structure) + setup = time.perf_counter() + with adapters.calculator_session(request, execution) as session: + setup_s = time.perf_counter() - setup + _, legacy = next( + adapters.evaluate_items( + [(0, request.structure.path, atoms)], + request, + execution, + session, + ) + ) + output = adapters.calculator_payload(root, request, legacy) + output[2]["setup_s"] = setup_s + result = _finish(root, record, output, started) + return result + return _finish( + root, + record, + adapters.run_external(root, request, execution), + started, + ) + except Exception as exc: + return _fail( + root, + record, + exc, + stage, + interrupted=isinstance(exc, subprocess.TimeoutExpired), + ) + except BaseException as exc: + _fail(root, record, exc, stage, interrupted=True) + raise + + +def _worker_command(root, execution, batch): + python = execution.python or sys.executable + command = [python, "-m", "matkit.worker", str(root)] + if batch: + command.append("--batch") + return command + + +def _supervise_claimed(root, execution, batch=False): + atomic_json( + root / "execution.json", + execution.model_copy(update={"environment": {}}), + ) + with ( + (root / "worker.stdout.log").open("w") as stdout, + (root / "worker.stderr.log").open("w") as stderr, + ): + process = subprocess.Popen( + _worker_command(root, execution, batch), + env={ + **os.environ, + **execution.environment, + "MATKIT_WORKER_PROCESS": "1", + }, + stdout=stdout, + stderr=stderr, + start_new_session=os.name == "posix", + ) + try: + code = process.wait(timeout=execution.timeout_s) + adapters.stop_process(process) + except BaseException as exc: + adapters.stop_process(process) + if batch: + _interrupt_batch(root, exc) + elif not (root / "result.json").exists(): + _fail(root, inspect_run(root), exc, "worker", interrupted=True) + if isinstance(exc, subprocess.TimeoutExpired): + if batch: + return read_batch(root) + result = inspect_run(root) + return commit_result( + root, + result.model_copy( + update={"artifacts": collect_artifacts(root)} + ), + ) + raise + if batch: + result = read_batch(root) + if result.state == "running": + _interrupt_batch( + root, RuntimeError(f"Worker exited with code {code}") + ) + return read_batch(root) + result = inspect_run(root) + if result.state in {"prepared", "running"}: + return _fail( + root, + result, + RuntimeError( + f"Worker exited with code {code}; see worker.stderr.log" + ), + "worker", + interrupted=True, + ) + if code != (0 if result.accepted else 1): + raise RuntimeError( + f"Worker exited with code {code}; " + f"committed results remain in {root}" + ) + return commit_result( + root, result.model_copy(update={"artifacts": collect_artifacts(root)}) + ) + + +def execute( + bundle: str | Path, *, execution: ExecutionConfig | dict | None = None +): + root = Path(bundle).expanduser().resolve() + if (root / ".matkit.batch").exists() and ( + root.parent / ".matkit.lock" + ).exists(): + raise FileExistsError("This item is owned by an active batch") + config = ExecutionConfig.model_validate(execution or {}) + with claim(root): + record = inspect_run(root) + if record.state != "prepared": + raise FileExistsError("Only prepared bundles can execute") + verify_inputs(root, record) + if ( + config.mode == "subprocess" + or config.python + or config.environment + or config.timeout_s + ): + try: + return _supervise_claimed(root, config) + except OSError as exc: + return _fail(root, record, exc, "worker") + return _execute_claimed(root, config) + + +def run(request, *, output_dir, execution=None): + prepare(request, output_dir=output_dir) + return execute(output_dir, execution=execution) + + +def _operation(value, cls): + if isinstance(value, dict): + value = cls.model_validate(value) + if not isinstance(value, cls): + raise ValueError(f"Expected {cls.__name__}") + return value + + +def evaluate(request: EvaluateRequest | dict, *, output_dir, execution=None): + return run( + _operation(request, EvaluateRequest), + output_dir=output_dir, + execution=execution, + ) + + +def relax(request: RelaxRequest | dict, *, output_dir, execution=None): + return run( + _operation(request, RelaxRequest), + output_dir=output_dir, + execution=execution, + ) + + +def analyze_pores(request: PoreRequest | dict, *, output_dir, execution=None): + return run( + _operation(request, PoreRequest), + output_dir=output_dir, + execution=execution, + ) + + +def prepare_adsorption(request: AdsorptionRequest | dict, *, output_dir): + return prepare( + _operation(request, AdsorptionRequest), output_dir=output_dir + ) + + +def run_adsorption( + request: AdsorptionRequest | dict, *, output_dir, execution=None +): + return run( + _operation(request, AdsorptionRequest), + output_dir=output_dir, + execution=execution, + ) + + +def analyze_adsorption(bundle: str | Path): + """Inspect a run, or collect a manually executed prepared gRASPA bundle. + + Manual execution must save work/raspa.log and exit.json containing the + engine's returncode. Prefer ``matkit execute`` to record these reliably. + """ + root = Path(bundle).expanduser().resolve() + with claim(root): + record = inspect_run(root) + if record.operation != "adsorption": + raise ValueError("Expected an adsorption bundle") + if record.state not in {"prepared", "running"}: + return record + request = staged_request(root) + # Verify the original inventory, excluding newly created logs. + verify_inputs(root, record) + exit_record = json.loads((root / "exit.json").read_text()) + started = time.perf_counter() + try: + if exit_record.get("returncode") != 0: + raise ValueError("Engine did not report a zero exit code") + return _finish( + root, record, adapters.parse_adsorption(root, request), started + ) + except Exception as exc: + return _fail(root, record, exc, "collection") + + +def read_batch(root): + return BatchResult.model_validate_json( + (Path(root) / "batch_manifest.json").read_text() + ) + + +def _checkpoint_batch(root, state=None, failure=None): + batch = read_batch(root) + for item in batch.items: + item_dir = root / item["bundle"] + if (item_dir / "run.json").exists(): + result = inspect_run(item_dir) + item.update( + state=result.state, + accepted=result.accepted, + result_file=f"{item['bundle']}/result.json" + if (item_dir / "result.json").exists() + else None, + failure=result.failure.model_dump(mode="json") + if result.failure + else None, + ) + if state: + batch.state = state + if failure: + batch.failure = failure + atomic_json(root / "batch_manifest.json", batch) + return batch + + +def _interrupt_batch(root, exc): + failure = Failure( + code=type(exc).__name__, + stage="batch", + message=str(exc) or type(exc).__name__, + ) + batch = read_batch(root) + for item in batch.items: + item_dir = root / item["bundle"] + if (item_dir / "run.json").exists() and not ( + item_dir / "result.json" + ).exists(): + _fail( + item_dir, inspect_run(item_dir), exc, "batch", interrupted=True + ) + return _checkpoint_batch(root, "interrupted", failure) + + +def _execute_batch_claimed(root, execution): + batch = read_batch(root) + entries, records, requests = [], {}, {} + started = time.perf_counter() + try: + for item in batch.items: + item_root = root / item["bundle"] + if item["state"] != "prepared": + continue + index = item["index"] + record = inspect_run(item_root) + try: + request, record = _begin(item_root, record, execution) + atoms, _, _ = load_structure(request.structure) + requests[index], records[index] = request, record + entries.append((index, request.structure.path, atoms)) + except Exception as exc: + _fail(item_root, record, exc, "validation") + _checkpoint_batch(root) + if entries: + supporting_hashes = [ + { + path: digest + for path, digest in record.provenance[ + "prepared_hashes" + ].items() + if path.startswith("inputs/") + and not path.startswith("inputs/structure") + } + for record in records.values() + ] + if any( + hashes != supporting_hashes[0] + for hashes in supporting_hashes[1:] + ): + raise ValueError( + "Supporting inputs changed while preparing the batch" + ) + request = requests[entries[0][0]] + if isinstance(request, CalculatorRequest): + setup = time.perf_counter() + try: + with adapters.calculator_session( + request, execution + ) as session: + setup_s = time.perf_counter() - setup + for index, legacy in adapters.evaluate_items( + entries, request, execution, session + ): + item_root = root / batch.items[index]["bundle"] + try: + output = adapters.calculator_payload( + item_root, requests[index], legacy + ) + output[2]["setup_s"] = setup_s + except Exception as exc: + _fail( + item_root, + records[index], + exc, + "calculation", + ) + else: + _finish( + item_root, records[index], output, started + ) + _checkpoint_batch(root) + except Exception as exc: + # Startup failures affect pending items. Teardown and + # persistence failures must preserve completed results. + if any( + (root / item["bundle"] / "result.json").exists() + for item in batch.items + if item["index"] in records + ): + raise + for index in records: + _fail( + root / batch.items[index]["bundle"], + records[index], + exc, + "calculator_setup", + ) + else: + for index, _, _ in entries: + item_root = root / batch.items[index]["bundle"] + try: + output = adapters.run_external( + item_root, requests[index], execution + ) + except Exception as exc: + _fail(item_root, records[index], exc, "calculation") + else: + _finish(item_root, records[index], output, started) + _checkpoint_batch(root) + batch = _checkpoint_batch(root) + states = [item["state"] for item in batch.items] + state = ( + "completed" + if all(s == "completed" for s in states) + else "partial" + if "completed" in states + else "failed" + ) + return _checkpoint_batch(root, state) + except BaseException as exc: + _interrupt_batch(root, exc) + raise + + +def run_batch(requests, *, output_dir, execution=None): + requests = [parse_request(request) for request in requests] + if not requests: + raise ValueError("A batch requires at least one request") + comparison = requests[0].model_dump(exclude={"structure"}) + if any( + r.model_dump(exclude={"structure"}) != comparison for r in requests[1:] + ): + raise ValueError( + "Batches require homogeneous operations, models and settings" + ) + config = ExecutionConfig.model_validate(execution or {}) + root = Path(output_dir).expanduser().resolve() + root.mkdir(parents=True, exist_ok=True) + with claim(root): + if any(p.name != ".matkit.lock" for p in root.iterdir()): + raise FileExistsError("Use a fresh batch output directory") + items = [ + { + "index": i, + "bundle": f"{i:05d}", + "state": "pending", + "accepted": False, + "failure": None, + } + for i in range(len(requests)) + ] + atomic_json( + root / "batch_manifest.json", + BatchResult(state="running", items=items), + ) + try: + for i, request in enumerate(requests): + try: + prepare(request, output_dir=root / items[i]["bundle"]) + (root / items[i]["bundle"] / ".matkit.batch").write_text( + "1" + ) + items[i]["state"] = "prepared" + except Exception as exc: + items[i].update( + state="failed", + failure={ + "code": type(exc).__name__, + "stage": "preparation", + "message": str(exc), + }, + ) + atomic_json( + root / "batch_manifest.json", + BatchResult(state="running", items=items), + ) + if ( + config.mode == "subprocess" + or config.python + or config.environment + or config.timeout_s + ): + return _supervise_claimed(root, config, batch=True) + return _execute_batch_claimed(root, config) + except BaseException as exc: + _interrupt_batch(root, exc) + raise diff --git a/src/matkit/api/structures.py b/src/matkit/api/structures.py new file mode 100644 index 0000000..97c77c1 --- /dev/null +++ b/src/matkit/api/structures.py @@ -0,0 +1,220 @@ +"""Structure handoffs retain original files and explicit atom correspondence.""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path + +import numpy as np +from ase import Atoms +from ase.constraints import dict2constraint +from ase.io import read, write +from ase.io.cif import parse_cif +from ase.spacegroup import Spacegroup + +from .models import StructureData, StructureRef + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def portable(value): + """Convert supported scientific metadata; never silently drop an object.""" + if isinstance(value, np.ndarray): + return portable(value.tolist()) + if isinstance(value, np.generic): + return portable(value.item()) + if isinstance(value, Spacegroup): + return {"number": value.no, "symbol": value.symbol} + if isinstance(value, dict): + return {str(k): portable(v) for k, v in value.items()} + if isinstance(value, (tuple, list)): + return [portable(v) for v in value] + if value is None or isinstance(value, (str, bool, int, float)): + json.dumps(value, allow_nan=False) + return value + raise ValueError(f"Unsupported structure metadata: {type(value).__name__}") + + +def to_atoms(data: StructureData) -> Atoms: + atoms = Atoms( + numbers=data.atomic_numbers, + positions=data.positions, + cell=data.cell, + pbc=data.pbc, + ) + for key, value in data.arrays.items(): + if key in {"numbers", "positions"}: + raise ValueError( + "Species and positions cannot be overridden by arrays" + ) + atoms.new_array(key, np.asarray(value)) + atoms.info = dict(data.info) + try: + atoms.set_constraint([dict2constraint(c) for c in data.constraints]) + except Exception as exc: + raise ValueError(f"Unsupported constraint handoff: {exc}") from exc + return atoms + + +def validate_geometry(atoms: Atoms) -> None: + if ( + not len(atoms) + or (atoms.numbers < 1).any() + or (atoms.numbers > 118).any() + ): + raise ValueError("Structure must contain supported atomic species") + if ( + not np.isfinite(atoms.positions).all() + or not np.isfinite(atoms.cell.array).all() + ): + raise ValueError("Structure coordinates and cell must be finite") + periodic = atoms.cell.array[atoms.pbc] + if len(periodic) and np.linalg.matrix_rank(periodic) != len(periodic): + raise ValueError( + "Periodic cell vectors must be nonzero and independent" + ) + + +def load_structure(ref: StructureRef) -> tuple[Atoms, StructureData, str]: + path = Path(ref.path).expanduser().resolve() + digest = sha256(path) + if ref.sha256 is not None and ref.sha256 != digest: + raise ValueError("Structure hash does not match the supplied reference") + frames = read(path, index=":") + if len(frames) != 1: + raise ValueError("Provide a single structure, not a trajectory") + atoms = frames[0] + validate_geometry(atoms) + metadata = ref.metadata + sidecar = path.with_suffix(".metadata.json") + if metadata is None and sidecar.is_file(): + metadata = StructureData.model_validate_json(sidecar.read_text()) + if metadata is not None: + if ( + metadata.atomic_numbers != atoms.numbers.tolist() + or metadata.pbc != tuple(atoms.pbc) + or not np.allclose( + metadata.positions, atoms.positions, rtol=0, atol=1e-7 + ) + or not np.allclose( + metadata.cell, atoms.cell.array, rtol=0, atol=1e-7 + ) + ): + raise ValueError( + "Structure metadata does not match geometry or atom order" + ) + if any(source != digest for source in metadata.derived_from.values()): + raise ValueError( + "Derived results were invalidated by a structural change; " + "recompute them" + ) + return to_atoms(metadata), metadata, digest + + labels = None + if path.suffix.lower() == ".cif": + blocks = list(parse_cif(str(path))) + if len(blocks) != 1: + raise ValueError("Provide a single CIF data block") + tags = dict(blocks[0]) + occupancy = tags.get("_atom_site_occupancy", []) + if occupancy and not np.allclose(occupancy, 1): + raise ValueError( + "Disordered/partially occupied CIF sites are unsupported" + ) + labels = tags.get("_atom_site_label") + charges = tags.get("_atom_site_charge") + if labels is not None and len(labels) != len(atoms): + raise ValueError( + "CIF symmetry expansion cannot preserve site correspondence; " + "provide an explicit P1 structure" + ) + if charges is not None: + if len(charges) != len(atoms): + raise ValueError( + "CIF charges cannot be mapped to expanded atoms" + ) + atoms.set_initial_charges(charges) + atoms.info["cif_tags"] = tags + + constraints = [portable(c.todict()) for c in atoms.constraints] + data = StructureData( + atomic_numbers=atoms.numbers.tolist(), + positions=atoms.positions.tolist(), + cell=atoms.cell.array.tolist(), + pbc=atoms.pbc.tolist(), + atom_ids=[f"{digest[:16]}:{i}" for i in range(len(atoms))], + arrays={ + key: portable(value) + for key, value in atoms.arrays.items() + if key not in {"numbers", "positions"} + }, + info=portable(atoms.info), + constraints=constraints, + labels=labels, + derived_from={"charges": digest} + if "initial_charges" in atoms.arrays + else {}, + ) + # Round-trip constraints now, before accepting the structure. + return to_atoms(data), data, digest + + +def final_structure( + atoms: Atoms, original: StructureData, source_hash: str, path: Path +) -> StructureData: + """Persist final geometry with metadata; invalidate inherited properties.""" + validate_geometry(atoms) + if atoms.numbers.tolist() != original.atomic_numbers: + raise ValueError("Calculator changed species or atom correspondence") + changed = ( + not np.array_equal(atoms.positions, original.positions) + or not np.array_equal(atoms.cell.array, original.cell) + or tuple(atoms.pbc) != original.pbc + ) + values = original.model_dump() + values.update( + positions=atoms.positions.tolist(), + cell=atoms.cell.array.tolist(), + pbc=atoms.pbc.tolist(), + parent_sha256=source_hash, + ) + if changed: + values["derived_from"] = {} + for key in ("initial_charges", "forces", "energies", "stresses"): + values["arrays"].pop(key, None) + for key in ( + "energy", + "potential_energy", + "forces", + "stress", + "charges", + "pore_properties", + "cif_tags", + "spacegroup", + "occupancy", + ): + values["info"].pop(key, None) + # The sidecar carries arrays and constraints even when the file cannot. + geometry = Atoms( + numbers=atoms.numbers, + positions=atoms.positions, + cell=atoms.cell, + pbc=atoms.pbc, + ) + write(path, geometry, format="extxyz") + if not changed: + values["derived_from"] = { + name: sha256(path) for name in original.derived_from + } + result = StructureData.model_validate(values) + path.with_suffix(".metadata.json").write_text( + result.model_dump_json(indent=2) + ) + return result diff --git a/src/matkit/cli.py b/src/matkit/cli.py index bc0eb62..c29e33b 100644 --- a/src/matkit/cli.py +++ b/src/matkit/cli.py @@ -3,6 +3,8 @@ import click import json +from matkit.operation_cli import register_commands + @click.group() @click.option( @@ -24,6 +26,9 @@ def main(verbose): ) +register_commands(main) + + # ========================================== # GRASPA COMMANDS # ========================================== diff --git a/src/matkit/graspa/graspa.py b/src/matkit/graspa/graspa.py index 69434b3..2208bc8 100644 --- a/src/matkit/graspa/graspa.py +++ b/src/matkit/graspa/graspa.py @@ -210,7 +210,7 @@ def setup_simulation( uc_x, uc_y, uc_z = cell_size else: atoms = ase_read(cifpath) - uc_x, uc_y, uc_z = calculate_cell_size(atoms) + uc_x, uc_y, uc_z = calculate_cell_size(atoms, cutoff=cutoff) input_path = outdir / "simulation.input" render_template( @@ -252,7 +252,7 @@ def _setup_single_cif( directory for each (temperature, pressure) combination. """ atoms = ase_read(cif) - cell_size = calculate_cell_size(atoms) + cell_size = calculate_cell_size(atoms, cutoff=cutoff) safe_stem = sanitize_cif_stem(cif.stem) entries = [] diff --git a/src/matkit/graspa_sycl/graspa_sycl.py b/src/matkit/graspa_sycl/graspa_sycl.py index 4ea3b0c..9579d68 100644 --- a/src/matkit/graspa_sycl/graspa_sycl.py +++ b/src/matkit/graspa_sycl/graspa_sycl.py @@ -47,7 +47,7 @@ def setup_simulation( shutil.copy(cifpath, outdir / f"{safe_stem}.cif") atoms = ase_read(cif) - uc_x, uc_y, uc_z = calculate_cell_size(atoms) + uc_x, uc_y, uc_z = calculate_cell_size(atoms, cutoff=cutoff) render_template( outdir / "simulation.input", diff --git a/src/matkit/mcp.py b/src/matkit/mcp.py new file mode 100644 index 0000000..35dabd5 --- /dev/null +++ b/src/matkit/mcp.py @@ -0,0 +1,411 @@ +"""Optional stdio tools over MatKit's scientific API and supervised workers.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +import re +import signal +from uuid import uuid4 + +import click +from pydantic import Field, JsonValue + +from matkit.api import ( + AdsorptionRequest, + EvaluateRequest, + ExecutionConfig, + PoreRequest, + RelaxRequest, + inspect_run, + list_capabilities, + prepare, +) +from matkit.api.bundles import ( + artifact, + atomic_json, + claim, + collect_artifacts, + commit_result, + contained_path, +) +from matkit.api.models import Artifact, Failure, Model, ScientificCheck +from matkit.api.runtime import _fail, _worker_command +from matkit.api.structures import sha256 +from matkit.operation_cli import resolve_request_paths + + +class ArtifactLink(Artifact): + uri: str + + +class ToolResult(Model): + run_id: str + state: str + numerical_validity: str + accepted: bool + checks: list[ScientificCheck] + summary: dict[str, JsonValue] = Field(default_factory=dict) + artifacts: list[ArtifactLink] + failure: Failure | None = None + + +def _links(root, record): + inventory = [*record.artifacts] + for name in ("run.json", "result.json"): + if (root / name).is_file(): + inventory.append(artifact(root, root / name, "record")) + return [ + ArtifactLink( + **ref.model_dump(), + uri=f"matkit://runs/{record.run_id}/artifacts/{ref.sha256}", + ) + for ref in inventory + ] + + +def _summary(root): + record = inspect_run(root) + summary = {} + if record.payload: + values = record.payload.model_dump(mode="json") + summary = { + key: value + for key, value in values.items() + if key not in {"forces", "stress", "results"} + } + if "results" in values: + summary["results"] = { + name: { + key: value + for key, value in result.items() + if not isinstance(value, list) + } + for name, result in values["results"].items() + } + return ToolResult( + run_id=record.run_id, + state=record.state, + numerical_validity=record.numerical_validity, + accepted=record.accepted, + checks=record.checks, + summary=summary, + artifacts=_links(root, record), + failure=record.failure, + ) + + +async def _execute_bounded(root, execution): + import anyio + + with claim(root): + # Environment values travel through the process environment, not through + # artifacts retrievable by MCP clients. + atomic_json( + root / "execution.json", + execution.model_copy(update={"environment": {}}), + ) + with ( + (root / "worker.stdout.log").open("w") as stdout, + (root / "worker.stderr.log").open("w") as stderr, + ): + process = None + try: + process = await anyio.open_process( + _worker_command(root, execution, False), + env={ + **os.environ, + **execution.environment, + "MATKIT_WORKER_PROCESS": "1", + }, + stdout=stdout, + stderr=stderr, + start_new_session=os.name == "posix", + ) + with anyio.fail_after(execution.timeout_s): + code = await process.wait() + record = inspect_run(root) + if record.state in {"prepared", "running"}: + _fail( + root, + record, + RuntimeError( + f"Worker exited with code {code}; " + "see worker.stderr.log" + ), + "worker", + interrupted=True, + ) + except BaseException as exc: + with anyio.CancelScope(shield=True): + if process is not None and process.returncode is None: + if os.name == "posix": + os.killpg(process.pid, signal.SIGTERM) + else: + process.terminate() + with anyio.move_on_after(5) as grace: + await process.wait() + if grace.cancel_called: + if os.name == "posix": + os.killpg(process.pid, signal.SIGKILL) + else: + process.kill() + await process.wait() + if not (root / "result.json").exists(): + _fail( + root, + inspect_run(root), + exc, + "worker", + interrupted=True, + ) + if not isinstance(exc, (TimeoutError, OSError)): + raise + finally: + if process is not None: + with anyio.CancelScope(shield=True): + if os.name == "posix": + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + await process.aclose() + record = inspect_run(root) + if record.state not in {"prepared", "running"}: + commit_result( + root, + record.model_copy( + update={ + "artifacts": collect_artifacts(root) + } + ), + ) + record = inspect_run(root) + if record.state not in {"prepared", "running"}: + commit_result( + root, + record.model_copy( + update={"artifacts": collect_artifacts(root)} + ), + ) + return _summary(root) + + +def create_server( + *, run_root, input_roots, profiles=None, timeout_s=60.0, tools=None +): + from mcp.server import MCPServer + + run_root = Path(run_root).expanduser().resolve() + run_root.mkdir(parents=True, exist_ok=True) + roots = [Path(path).expanduser().resolve() for path in input_roots] + if not roots: + raise ValueError("At least one input root is required") + profile_map = { + name: ExecutionConfig.model_validate(value) + for name, value in (profiles or {"default": {}}).items() + } + if not profile_map: + raise ValueError("At least one execution profile is required") + ExecutionConfig(timeout_s=timeout_s) + server = MCPServer("MatKit") + + def allowed(path, *, output=False): + path = Path(path).expanduser().resolve() + locations = [run_root] if output else [*roots, run_root] + if not any(path.is_relative_to(base) for base in locations): + raise ValueError( + "Input path is outside the server's configured roots" + ) + return path + + def run_directory(run_id): + if not re.fullmatch(r"[a-f0-9]{32}", run_id): + raise ValueError("Invalid run identifier") + return contained_path(run_root, run_id) + + def prepare_tool(request): + values = resolve_request_paths( + request.model_dump(mode="json"), roots[0] + ) + allowed(values["structure"]["path"]) + sidecar = Path(values["structure"]["path"]).with_suffix( + ".metadata.json" + ) + if sidecar.exists(): + allowed(sidecar) + for key in ("radii_file", "template_dir"): + if values.get(key): + allowed(values[key]) + checkpoint = values.get("method", {}).get("checkpoint") + if checkpoint and ( + Path(checkpoint).is_absolute() or Path(checkpoint).exists() + ): + allowed(checkpoint) + adapter = values.get("adapter") + if isinstance(adapter, dict): + for key in ("root", "weights"): + if adapter.get(key): + allowed(adapter[key]) + if adapter.get("cache_root"): + allowed(adapter["cache_root"], output=True) + temporary = run_root / uuid4().hex + record = prepare(values, output_dir=temporary) + destination = run_directory(record.run_id) + temporary.rename(destination) + return destination + + async def calculate(request, profile): + if profile not in profile_map: + raise ValueError(f"Unknown execution profile: {profile}") + config = profile_map[profile] + timeout = ( + min(timeout_s, config.timeout_s) if config.timeout_s else timeout_s + ) + config = config.model_copy( + update={"mode": "subprocess", "timeout_s": timeout} + ) + root = prepare_tool(request) + return await _execute_bounded(root, config) + + async def matkit_evaluate( + request: EvaluateRequest, profile: str = "default" + ) -> ToolResult: + """Evaluate requested energy/forces/stress. + + Check model/species support. Use prepared bundles and CLI execution + for long runs. + """ + return await calculate(request, profile) + + async def matkit_relax( + request: RelaxRequest, profile: str = "default" + ) -> ToolResult: + """Relax positions at fixed cell. + + Check force_convergence before using the geometry and recompute + invalidated charges/properties. + """ + return await calculate(request, profile) + + async def matkit_pores( + request: PoreRequest, profile: str = "default" + ) -> ToolResult: + """Analyze a periodic structure using Zeo++. + + Choose radii/probe settings explicitly. Execution does not establish + sampling accuracy. + """ + return await calculate(request, profile) + + def matkit_prepare_adsorption(request: AdsorptionRequest) -> ToolResult: + """Stage a charged periodic CIF for single-component gRASPA CUDA. + + Execute the bundle with MatKit CLI inside an allocation; this tool + does not submit a job. + """ + return _summary(prepare_tool(request)) + + def matkit_prepare( + request: EvaluateRequest | RelaxRequest | PoreRequest, + ) -> ToolResult: + """Prepare a long calculation for later CLI execution. + + Preparation does not load models or execution engines. + """ + return _summary(prepare_tool(request)) + + def matkit_inspect(run_id: str) -> ToolResult: + """Inspect a prepared/executed run. + + Check scientific outcomes separately from execution state; + artifacts are retrievable MCP resources. + """ + return _summary(run_directory(run_id)) + + def matkit_capabilities() -> dict: + """Discover implementations, requirements and execution profiles. + + MOFforge owns construction and structural edits. + """ + return { + "capabilities": list_capabilities(), + "profiles": list(profile_map), + "timeout_s": timeout_s, + } + + catalog = { + function.__name__: function + for function in ( + matkit_evaluate, + matkit_relax, + matkit_pores, + matkit_prepare_adsorption, + matkit_prepare, + matkit_inspect, + matkit_capabilities, + ) + } + selected = set(tools) if tools is not None else set(catalog) + if not selected <= catalog.keys(): + raise ValueError(f"Unknown tools: {sorted(selected - catalog.keys())}") + for name in sorted(selected): + server.tool()(catalog[name]) + + @server.resource( + "matkit://runs/{run_id}/artifacts/{digest}", + mime_type="application/octet-stream", + ) + def read_artifact(run_id: str, digest: str) -> bytes: + root = run_directory(run_id) + if not re.fullmatch(r"[a-f0-9]{64}", digest): + raise ValueError("Invalid artifact hash") + for ref in _links(root, inspect_run(root)): + if ref.sha256 == digest: + path = contained_path(root, ref.path) + if sha256(path) != digest: + raise ValueError("Artifact changed since it was recorded") + return path.read_bytes() + raise ValueError("Artifact is not in the run manifest") + + return server + + +@click.command() +@click.option("--run-root", required=True, type=click.Path(path_type=Path)) +@click.option( + "--input-root", + multiple=True, + required=True, + type=click.Path(exists=True, file_okay=False, path_type=Path), +) +@click.option( + "--profiles", type=click.Path(exists=True, dir_okay=False, path_type=Path) +) +@click.option( + "--timeout", type=click.FloatRange(min=0, min_open=True), default=60.0 +) +@click.option("--tools", help="Comma-separated tool allowlist.") +def main(run_root, input_root, profiles, timeout, tools): + """Serve selected MatKit tools over stdio (requires matkit[mcp]).""" + try: + profile_map = json.loads(profiles.read_text()) if profiles else None + server = create_server( + run_root=run_root, + input_roots=input_root, + profiles=profile_map, + timeout_s=timeout, + tools=tools.split(",") if tools is not None else None, + ) + server.run(transport="stdio") + except ImportError as exc: + raise click.ClickException( + "Install matkit[mcp] to run the MCP server" + ) from exc + except (ValueError, OSError) as exc: + raise click.ClickException(str(exc)) from exc + + +if __name__ == "__main__": + main() diff --git a/src/matkit/mlip/runner.py b/src/matkit/mlip/runner.py index 3ad94dc..71cf0fa 100644 --- a/src/matkit/mlip/runner.py +++ b/src/matkit/mlip/runner.py @@ -178,10 +178,25 @@ def _success_result( converged: bool, n_steps: int | None, calculation_time: float, + requested_properties: Sequence[str] | None = None, ) -> dict[str, Any]: _validate_atoms(atoms) - energy = _finite_array("energy", energy, ()).item() - forces = _finite_array("forces", forces, (len(atoms), 3)) + required = ( + requested_properties + if requested_properties is not None + else ["potential_energy", "forces"] + ) + for name, value in ( + ("potential_energy", energy), + ("forces", forces), + ("stress", stress), + ): + if name in required and value is None: + raise ValueError(f"Requested property {name} is unavailable") + if energy is not None: + energy = _finite_array("energy", energy, ()).item() + if forces is not None: + forces = _finite_array("forces", forces, (len(atoms), 3)) if stress is not None: stress = _finite_array("stress", stress, (3, 3)) return { @@ -191,7 +206,7 @@ def _success_result( "input_structure_file": input_file, "backend_info": backend.to_dict(), "calculation_input": calculation.to_dict(), - "energy": float(energy), + "energy": None if energy is None else float(energy), "energy_unit": "eV", "forces": None if forces is None else forces.tolist(), "force_unit": "eV/angstrom", @@ -237,6 +252,7 @@ def _run_ase_item( calculator, backend: ASEMACEConfig | RootstockConfig, calculation: MLIPCalculationConfig, + requested_properties: Sequence[str] | None = None, ) -> dict[str, Any]: atoms.calc = calculator # Rootstock's synchronous worker owns the GPU and its dependencies. @@ -252,9 +268,18 @@ def _run_ase_item( ) n_steps = optimizer.nsteps - energy = atoms.get_potential_energy() - forces = atoms.get_forces() - stress = _optional_stress(atoms) + properties = ( + requested_properties + if requested_properties is not None + else ["potential_energy", "forces", "stress"] + ) + energy = ( + atoms.get_potential_energy() + if "potential_energy" in properties + else None + ) + forces = atoms.get_forces() if "forces" in properties else None + stress = _optional_stress(atoms) if "stress" in properties else None if isinstance(backend, ASEMACEConfig): _synchronize_device(backend.device) elapsed = time.perf_counter() - started @@ -269,6 +294,7 @@ def _run_ase_item( converged, n_steps, elapsed, + requested_properties=requested_properties, ) diff --git a/src/matkit/operation_cli.py b/src/matkit/operation_cli.py new file mode 100644 index 0000000..ef4834c --- /dev/null +++ b/src/matkit/operation_cli.py @@ -0,0 +1,231 @@ +"""Thin Click commands over the shared scientific API.""" + +from pathlib import Path +import json + +import click + + +def read_json(path): + try: + return json.loads(Path(path).read_text()) + except (OSError, ValueError) as exc: + raise click.UsageError( + f"Cannot read JSON specification: {exc}" + ) from exc + + +def resolve_request_paths(value, base): + """Resolve inputs relative to their spec; keep model aliases intact.""" + if not isinstance(value, dict): + raise ValueError("Expected a request object") + value = json.loads(json.dumps(value)) + + def path(container, key, optional=False): + if container.get(key): + candidate = Path(container[key]).expanduser() + if not candidate.is_absolute(): + candidate = base / candidate + if not optional or candidate.is_file(): + container[key] = str(candidate.resolve()) + + if isinstance(value.get("structure"), dict): + path(value["structure"], "path") + for key in ("radii_file", "template_dir"): + path(value, key) + if isinstance(value.get("method"), dict): + path(value["method"], "checkpoint", optional=True) + if isinstance(value.get("adapter"), dict): + for key in ("root", "cache_root", "weights"): + path(value["adapter"], key) + return value + + +def execution_profile(path): + from matkit.api import ExecutionConfig + + values = read_json(path) if path else {} + if not isinstance(values, dict): + raise click.UsageError("Execution profile must be a JSON object") + # CLI engines always run in a worker, keeping engine stdout out of JSON. + return ExecutionConfig.model_validate({**values, "mode": "subprocess"}) + + +def invoke(function, *args, prepared=False, **kwargs): + try: + result = function(*args, **kwargs) + except (ValueError, TypeError, FileNotFoundError, FileExistsError) as exc: + raise click.UsageError(str(exc)) from exc + except OSError as exc: + raise click.ClickException(str(exc)) from exc + click.echo(result.model_dump_json(indent=2)) + if prepared: + return + accepted = ( + result.accepted + if hasattr(result, "accepted") + else all(item["accepted"] for item in result.items) + ) + if not accepted: + raise click.exceptions.Exit(1) + + +def register_commands(main): + def operation_command(name, operation, preparation=False): + @click.command(name) + @click.option( + "--spec", + required=True, + type=click.Path(exists=True, dir_okay=False, path_type=Path), + ) + @click.option( + "--outdir", required=True, type=click.Path(path_type=Path) + ) + @click.option( + "--execution", + type=click.Path(exists=True, dir_okay=False, path_type=Path), + ) + def command(spec, outdir, execution): + from matkit.api import parse_request, prepare, run + + try: + data = resolve_request_paths( + read_json(spec), spec.resolve().parent + ) + data.setdefault("operation", operation) + request = parse_request(data) + if request.operation != operation: + raise ValueError( + f"This command requires operation={operation}" + ) + if preparation: + if execution is not None: + raise ValueError( + "Preparation does not use an execution profile" + ) + invoke(prepare, request, output_dir=outdir, prepared=True) + else: + invoke( + run, + request, + output_dir=outdir, + execution=execution_profile(execution), + ) + except (ValueError, TypeError) as exc: + raise click.UsageError(str(exc)) from exc + + return command + + for operation in ("evaluate", "relax", "pores"): + main.add_command(operation_command(operation, operation)) + + @main.group("adsorption") + def adsorption(): + """Prepare, run, or collect single-component gRASPA calculations.""" + + adsorption.add_command(operation_command("prepare", "adsorption", True)) + adsorption.add_command(operation_command("run", "adsorption")) + + @adsorption.command("analyze") + @click.argument( + "bundle", type=click.Path(exists=True, file_okay=False, path_type=Path) + ) + def analyze(bundle): + from matkit.api import analyze_adsorption + + invoke(analyze_adsorption, bundle) + + @main.command("prepare") + @click.option( + "--spec", + required=True, + type=click.Path(exists=True, dir_okay=False, path_type=Path), + ) + @click.option("--outdir", required=True, type=click.Path(path_type=Path)) + def prepare_command(spec, outdir): + """Stage any operation without loading its execution engine.""" + from matkit.api import prepare + + try: + request = resolve_request_paths( + read_json(spec), spec.resolve().parent + ) + invoke(prepare, request, output_dir=outdir, prepared=True) + except (ValueError, TypeError) as exc: + raise click.UsageError(str(exc)) from exc + + @main.command("execute") + @click.argument( + "bundle", type=click.Path(exists=True, file_okay=False, path_type=Path) + ) + @click.option( + "--execution", + type=click.Path(exists=True, dir_okay=False, path_type=Path), + ) + def execute_command(bundle, execution): + """Execute a prepared bundle in this allocation/environment.""" + from matkit.api import execute + + try: + invoke(execute, bundle, execution=execution_profile(execution)) + except ValueError as exc: + raise click.UsageError(str(exc)) from exc + + @main.command("batch") + @click.option( + "--spec", + required=True, + type=click.Path(exists=True, dir_okay=False, path_type=Path), + ) + @click.option("--outdir", required=True, type=click.Path(path_type=Path)) + @click.option( + "--execution", + type=click.Path(exists=True, dir_okay=False, path_type=Path), + ) + def batch_command(spec, outdir, execution): + """Execute a JSON list of homogeneous requests with calculator reuse.""" + from matkit.api import run_batch + + try: + values = read_json(spec) + if not isinstance(values, list): + raise ValueError("Batch specification must be a list") + requests = [ + resolve_request_paths(value, spec.resolve().parent) + for value in values + ] + invoke( + run_batch, + requests, + output_dir=outdir, + execution=execution_profile(execution), + ) + except (ValueError, TypeError) as exc: + raise click.UsageError(str(exc)) from exc + + @main.command("inspect") + @click.argument( + "bundle", type=click.Path(exists=True, file_okay=False, path_type=Path) + ) + def inspect_command(bundle): + """Read committed results without rerunning a calculation.""" + from matkit.api import inspect_run + from matkit.api.runtime import read_batch + + function = ( + read_batch + if (bundle / "batch_manifest.json").exists() + else inspect_run + ) + # Inspection succeeding is independent of the calculation's outcome. + invoke(function, bundle, prepared=True) + + @main.command("capabilities") + @click.option( + "--json", "as_json", is_flag=True, help="Print JSON (also the default)." + ) + def capabilities(as_json): + """List implementations, availability, restrictions and evidence.""" + from matkit.api import list_capabilities + + click.echo(json.dumps(list_capabilities(), indent=2, allow_nan=False)) diff --git a/src/matkit/pygraspa/pygraspa.py b/src/matkit/pygraspa/pygraspa.py index 37ed162..5a148a2 100644 --- a/src/matkit/pygraspa/pygraspa.py +++ b/src/matkit/pygraspa/pygraspa.py @@ -165,7 +165,7 @@ def setup_simulation( uc_x, uc_y, uc_z = cell_size else: atoms = ase_read(cifpath) - uc_x, uc_y, uc_z = calculate_cell_size(atoms) + uc_x, uc_y, uc_z = calculate_cell_size(atoms, cutoff=cutoff) input_path = outdir / "simulation.input" render_template( @@ -220,7 +220,7 @@ def _setup_single_cif( ) -> list[dict]: """Set up all T x P pygRASPA simulations for one CIF file.""" atoms = ase_read(cif) - cell_size = calculate_cell_size(atoms) + cell_size = calculate_cell_size(atoms, cutoff=cutoff) safe_stem = sanitize_cif_stem(cif.stem) entries = [] diff --git a/src/matkit/raspa2/raspa2.py b/src/matkit/raspa2/raspa2.py index 505e454..3a16200 100644 --- a/src/matkit/raspa2/raspa2.py +++ b/src/matkit/raspa2/raspa2.py @@ -1,4 +1,5 @@ import shutil +import math from pathlib import Path from ase.io import read as ase_read @@ -51,7 +52,7 @@ def setup_input_simulation( shutil.copy(cif, outdir) atoms = ase_read(cif) - uc_x, uc_y, uc_z = calculate_cell_size(atoms) + uc_x, uc_y, uc_z = calculate_cell_size(atoms, cutoff=cutoff) render_template( outdir / "simulation.input", @@ -110,6 +111,22 @@ def get_output_data( uptake_mg_g = float(mg_g_line.split()[5]) error_mg_g = float(mg_g_line.split()[7]) + if ( + not all( + math.isfinite(value) + for value in ( + uptake_mol_kg, + error_mol_kg, + density_kg_m3, + uptake_mg_g, + error_mg_g, + ) + ) + or density_kg_m3 <= 0 + or min(error_mol_kg, error_mg_g) < 0 + ): + raise ValueError("Invalid numerical adsorption results") + if unit == "mol/kg": result["uptake"] = uptake_mol_kg result["error"] = error_mol_kg @@ -137,4 +154,5 @@ def get_output_data( duration_seconds = int((end_time - start_time).total_seconds()) result["calc_time_in_s"] = duration_seconds + result["success"] = True return result diff --git a/src/matkit/utils/unitcell_calculator.py b/src/matkit/utils/unitcell_calculator.py index eb5caf4..f5cd794 100644 --- a/src/matkit/utils/unitcell_calculator.py +++ b/src/matkit/utils/unitcell_calculator.py @@ -1,5 +1,7 @@ import numpy as np import ase +import math +from numbers import Real def calculate_cell_size(atoms: ase.Atoms, cutoff: float = 12.8) -> list[int]: @@ -12,7 +14,19 @@ def calculate_cell_size(atoms: ase.Atoms, cutoff: float = 12.8) -> list[int]: Returns: list[int, int, int]: Unit cell in x, y and z """ + if ( + isinstance(cutoff, bool) + or not isinstance(cutoff, Real) + or not math.isfinite(cutoff) + or cutoff <= 0 + ): + raise ValueError("cutoff must be positive and finite") unit_cell = atoms.cell[:] + if ( + not np.isfinite(unit_cell).all() + or np.linalg.matrix_rank(unit_cell) != 3 + ): + raise ValueError("Cell replication requires three independent vectors") # Unit cell vectors a = unit_cell[0] b = unit_cell[1] diff --git a/src/matkit/worker.py b/src/matkit/worker.py new file mode 100644 index 0000000..835fcac --- /dev/null +++ b/src/matkit/worker.py @@ -0,0 +1,40 @@ +"""Internal subprocess entry point; communicate results through run bundles.""" + +import argparse +import os +from pathlib import Path +import signal + +from matkit.api import ExecutionConfig +from matkit.api.runtime import _execute_batch_claimed, _execute_claimed + + +def _interrupt(signum, frame): + raise KeyboardInterrupt(f"Worker received signal {signum}") + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("bundle", type=Path) + parser.add_argument("--batch", action="store_true") + args = parser.parse_args(argv) + root = args.bundle.resolve() + if os.environ.get("MATKIT_WORKER_PROCESS") != "1" or ( + root / ".matkit.lock" + ).read_text() != str(os.getppid()): + parser.error( + "Workers must be launched by MatKit's supervising executor" + ) + signal.signal(signal.SIGTERM, _interrupt) + execution = ExecutionConfig.model_validate_json( + (root / "execution.json").read_text() + ) + if args.batch: + result = _execute_batch_claimed(root, execution) + return 0 if result.accepted else 1 + result = _execute_claimed(root, execution) + return 0 if result.accepted else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/fixtures/fake_engine.py b/tests/fixtures/fake_engine.py new file mode 100644 index 0000000..b4860b0 --- /dev/null +++ b/tests/fixtures/fake_engine.py @@ -0,0 +1,29 @@ +"""Synthetic subprocess fixtures (MIT); no scientific accuracy claims.""" + +from pathlib import Path +import sys +import time + +mode = sys.argv[1] +if "--sleep" in sys.argv: + Path("started").write_text("started") + time.sleep(30) +if "--fail" in sys.argv: + print("intentional engine failure", file=sys.stderr) + raise SystemExit(7) +if mode == "zeopp": + fixtures = Path(__file__).parents[1] / "data" / "zeopp" + for analysis in ("res", "sa", "vol", "psd", "chan"): + if f"-{analysis}" in sys.argv and not ( + "--partial" in sys.argv and analysis == "sa" + ): + Path(f"structure.{analysis}").write_bytes( + (fixtures / f"test_structure.{analysis}").read_bytes() + ) +elif mode == "graspa": + for i in range(14): + print(f"Overall: Average: {25 if i == 0 else i + 1}, +/- 0.1") + if "--partial" not in sys.argv: + print("Work time 2.0") +else: + raise SystemExit("unknown fixture engine") diff --git a/tests/test_api_contracts.py b/tests/test_api_contracts.py new file mode 100644 index 0000000..6f0c6ad --- /dev/null +++ b/tests/test_api_contracts.py @@ -0,0 +1,150 @@ +"""Scientific request validation, structure metadata and compatibility.""" + +import json +from pathlib import Path + +import numpy as np +import pytest +from ase import Atoms +from ase.constraints import FixAtoms +from ase.io import write + +from matkit.api import ( + AlchemiAdapter, + EvaluateRequest, + MLIPMethod, + PoreRequest, + RelaxRequest, + RootstockAdapter, + RunResult, + StructureRef, + list_capabilities, + parse_request, + prepare, +) +from matkit.api.models import REQUEST_ADAPTER +from matkit.api.structures import final_structure, load_structure, to_atoms + + +def test_request_round_trip_and_schema(): + for adapter in (RootstockAdapter(), AlchemiAdapter()): + request = EvaluateRequest( + structure=StructureRef(path="input.cif"), + method=MLIPMethod(checkpoint="medium"), + adapter=adapter, + ) + assert parse_request(json.loads(request.model_dump_json())) == request + schema = REQUEST_ADAPTER.json_schema() + assert set(schema["discriminator"]["mapping"]) == { + "evaluate", + "relax", + "pores", + "adsorption", + } + + +@pytest.mark.parametrize( + "changes", + [{"fmax": float("nan")}, {"steps": True}, {"steps": 0}, {"cell_opt": True}], +) +def test_invalid_relaxation_settings(changes): + with pytest.raises(ValueError): + RelaxRequest( + structure=StructureRef(path="input.cif"), + method=MLIPMethod(checkpoint="medium"), + **changes, + ) + + +def test_unknown_result_schema_is_rejected(): + with pytest.raises(ValueError): + RunResult.model_validate( + {"schema_version": 1, "success": True, "energy": 0} + ) + + +def test_nonfinite_nested_configuration(): + with pytest.raises(ValueError): + RootstockAdapter(setup_kwargs={"nested": {"bad": float("inf")}}) + + +@pytest.mark.parametrize("value", [True, "1.0", -1, float("nan")]) +def test_scientific_numbers_require_finite_numeric_values(value): + with pytest.raises(ValueError): + PoreRequest( + structure=StructureRef(path="input.cif"), probe_radius=value + ) + + +def test_discovery_does_not_import_optional_engines(monkeypatch): + import builtins + + original = builtins.__import__ + + def guarded(name, *args, **kwargs): + if name.split(".")[0] in { + "torch", + "mace", + "rootstock", + "nvalchemi", + "mcp", + }: + raise AssertionError(f"Discovery imported {name}") + return original(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", guarded) + assert len(list_capabilities()) == 5 + assert all(c["status"] == "experimental" for c in list_capabilities()) + + +def test_structure_handoff_and_invalidation(tmp_path): + path = tmp_path / "input.extxyz" + atoms = Atoms( + "Cu2", positions=[[0, 0, 0], [3, 0, 0]], cell=[10] * 3, pbc=True + ) + atoms.set_initial_charges([0.2, -0.2]) + atoms.set_constraint(FixAtoms(indices=[0])) + write(path, atoms) + original, data, digest = load_structure(StructureRef(path=str(path))) + data.labels = ["Cu1", "Cu2"] + data.bonds = [(data.atom_ids[0], data.atom_ids[1])] + recovered = to_atoms(data) + np.testing.assert_allclose(recovered.get_initial_charges(), [0.2, -0.2]) + assert len(recovered.constraints) == 1 + original.positions[1, 0] -= 0.1 + output = tmp_path / "final.extxyz" + result = final_structure(original, data, digest, output) + assert result.atom_ids == data.atom_ids + assert result.labels == data.labels + assert result.bonds == data.bonds + assert result.parent_sha256 == digest + assert "initial_charges" not in result.arrays + assert not result.derived_from + _, restored, _ = load_structure(StructureRef(path=str(output))) + assert restored == result + with pytest.raises(ValueError, match="metadata does not match"): + load_structure(StructureRef(path=str(output), metadata=data)) + + +def test_wrong_structure_digest(sample_cif): + with pytest.raises(ValueError, match="hash"): + load_structure(StructureRef(path=sample_cif, sha256="0" * 64)) + + +def test_stale_charges_rejected(tmp_path): + path = tmp_path / "input.extxyz" + write(path, Atoms("Cu", cell=[10] * 3)) + _, data, _ = load_structure(StructureRef(path=str(path))) + data.derived_from = {"charges": "0" * 64} + with pytest.raises(ValueError, match="invalidated"): + load_structure(StructureRef(path=str(path), metadata=data)) + + +def test_preparation_is_engine_independent(sample_cif, tmp_path): + record = prepare( + PoreRequest(structure=StructureRef(path=sample_cif)), + output_dir=tmp_path / "run", + ) + assert record.state == "prepared" + assert any(a.path == "inputs/radii.rad" for a in record.artifacts) + assert all(not Path(a.path).is_absolute() for a in record.artifacts) diff --git a/tests/test_api_recovery.py b/tests/test_api_recovery.py new file mode 100644 index 0000000..27ecaa4 --- /dev/null +++ b/tests/test_api_recovery.py @@ -0,0 +1,105 @@ +"""Integrity checks for interruption, mutable inputs and worker artifacts.""" + +import json +from pathlib import Path + +import pytest +from ase import Atoms +from ase.io import write + +from matkit.api import ( + EvaluateRequest, + MLIPMethod, + PoreRequest, + StructureRef, + execute, + inspect_run, + prepare, + run_batch, +) +from matkit.api.models import BatchResult, Failure + + +def test_completed_items_do_not_hide_batch_orchestration_failure(): + result = BatchResult( + state="interrupted", + items=[{"accepted": True}], + failure=Failure( + code="RuntimeError", stage="teardown", message="worker failed" + ), + ) + assert not result.accepted + + +def test_active_batch_owns_its_prepared_items(sample_cif, tmp_path): + root = tmp_path / "batch" + root.mkdir() + leaf = root / "00000" + prepare( + PoreRequest(structure=StructureRef(path=sample_cif)), output_dir=leaf + ) + (root / ".matkit.lock").write_text("1") + (leaf / ".matkit.batch").write_text("1") + with pytest.raises(FileExistsError, match="active batch"): + execute(leaf) + + +def test_changed_model_cannot_be_reused_across_batch(tmp_path, monkeypatch): + from matkit.api import runtime + + path = tmp_path / "cu.extxyz" + write(path, Atoms("Cu", cell=[10] * 3)) + model = tmp_path / "model.pt" + model.write_bytes(b"model-one") + request = EvaluateRequest( + structure=StructureRef(path=str(path)), + method=MLIPMethod(checkpoint=str(model)), + ) + original = runtime.prepare + + def prepare_item(*args, **kwargs): + result = original(*args, **kwargs) + model.write_bytes(b"model-two") + return result + + monkeypatch.setattr(runtime, "prepare", prepare_item) + with pytest.raises(ValueError, match="Supporting inputs changed"): + run_batch([request, request], output_dir=tmp_path / "batch") + batch = json.loads((tmp_path / "batch" / "batch_manifest.json").read_text()) + assert batch["state"] == "interrupted" + assert all(not item["accepted"] for item in batch["items"]) + + +def test_atomic_committed_result_survives_corrupt_manifest( + sample_cif, tmp_path +): + from tests.test_api_runtime import fake_execution + + root = tmp_path / "run" + prepare( + PoreRequest(structure=StructureRef(path=sample_cif)), output_dir=root + ) + result = execute(root, execution=fake_execution("zeopp")) + (root / "run.json").write_text("corrupt") + assert inspect_run(root).payload == result.payload + + +def test_environment_values_are_not_retrievable_artifacts(sample_cif, tmp_path): + from tests.test_api_runtime import fake_execution + + root = tmp_path / "run" + prepare( + PoreRequest(structure=StructureRef(path=sample_cif)), output_dir=root + ) + result = execute( + root, + execution=fake_execution( + "zeopp", + mode="subprocess", + environment={"MATKIT_TEST_CREDENTIAL": "fixture-secret"}, + ), + ) + assert result.accepted + for item in result.artifacts: + if Path(item.path).suffix == ".json": + assert "fixture-secret" not in (root / item.path).read_text() diff --git a/tests/test_api_runtime.py b/tests/test_api_runtime.py new file mode 100644 index 0000000..e4f0397 --- /dev/null +++ b/tests/test_api_runtime.py @@ -0,0 +1,368 @@ +"""Cross-adapter durability and outcomes, without model downloads.""" + +from contextlib import contextmanager +import json +from pathlib import Path +import shutil +import sys + +import pytest +from ase import Atoms +from ase.calculators.calculator import Calculator, all_changes +from ase.calculators.emt import EMT +from ase.io import write + +from matkit.api import ( + AdsorptionRequest, + AlchemiAdapter, + EvaluateRequest, + ExecutionConfig, + MLIPMethod, + PoreRequest, + RelaxRequest, + StructureRef, + analyze_adsorption, + evaluate, + execute, + inspect_run, + prepare, + run, + run_batch, +) +from matkit.api import bundles + + +@pytest.fixture +def input_file(tmp_path): + path = tmp_path / "cu.extxyz" + write( + path, + Atoms("Cu2", positions=[[0, 0, 0], [3, 0, 0]], cell=[10] * 3, pbc=True), + ) + return str(path) + + +@pytest.fixture +def mock_calculator(monkeypatch): + from matkit.mlip import runner + + calls = [] + + @contextmanager + def context(backend): + calls.append(backend) + yield EMT() + + monkeypatch.setattr(runner, "_ase_backend_context", context) + return calls + + +def fake_execution(engine, *args, **kwargs): + fixture = Path(__file__).parent / "fixtures" / "fake_engine.py" + return ExecutionConfig( + executables={engine: [sys.executable, str(fixture), engine, *args]}, + **kwargs, + ) + + +def energy_request(path, **kwargs): + return EvaluateRequest( + structure=StructureRef(path=path), + method=MLIPMethod(checkpoint="fixture"), + **kwargs, + ) + + +def test_energy_only_calculator_does_not_require_forces( + input_file, tmp_path, monkeypatch +): + from matkit.mlip import runner + + class EnergyOnly(Calculator): + implemented_properties = ["energy"] + + def calculate( + self, atoms=None, properties=None, system_changes=all_changes + ): + super().calculate(atoms, properties, system_changes) + self.results = {"energy": -2.5} + + @contextmanager + def context(_): + yield EnergyOnly() + + monkeypatch.setattr(runner, "_ase_backend_context", context) + result = evaluate(energy_request(input_file), output_dir=tmp_path / "run") + assert result.accepted + assert result.payload.potential_energy == -2.5 + assert result.payload.forces is None + failed = evaluate( + energy_request(input_file, properties=["forces"]), + output_dir=tmp_path / "forces", + ) + assert not failed.accepted + assert failed.failure is not None + + +def test_relaxation_retains_unconverged_result( + input_file, tmp_path, mock_calculator +): + request = RelaxRequest( + structure=StructureRef(path=input_file), + method=MLIPMethod(checkpoint="fixture"), + steps=1, + fmax=1e-12, + ) + result = run(request, output_dir=tmp_path / "run") + assert result.state == "completed" + assert result.numerical_validity == "valid" + assert not result.payload.converged + assert not result.accepted + assert inspect_run(tmp_path / "run") == result + + +def test_homogeneous_batch_reuses_calculator( + input_file, tmp_path, mock_calculator +): + request = energy_request(input_file) + result = run_batch([request, request], output_dir=tmp_path / "batch") + assert result.state == "completed" + assert len(mock_calculator) == 1 + assert [i["index"] for i in result.items] == [0, 1] + assert all(i["accepted"] for i in result.items) + assert result.items[0]["result_file"] != result.items[1]["result_file"] + + +def test_batch_retains_success_when_input_is_missing( + input_file, tmp_path, mock_calculator +): + result = run_batch( + [energy_request(input_file), energy_request("/missing/input.xyz")], + output_dir=tmp_path / "batch", + ) + assert result.state == "partial" + assert result.items[0]["accepted"] + assert result.items[1]["state"] == "failed" + + +def test_mixed_batch_rejected_before_side_effects(input_file, tmp_path): + with pytest.raises(ValueError, match="homogeneous"): + run_batch( + [ + energy_request(input_file), + energy_request(input_file, properties=["forces"]), + ], + output_dir=tmp_path / "run", + ) + assert not (tmp_path / "run").exists() + + +@pytest.mark.parametrize("recovery_write_fails", [False, True]) +def test_completed_result_survives_manifest_write_failure( + input_file, tmp_path, mock_calculator, monkeypatch, recovery_write_fails +): + root = tmp_path / "run" + original = bundles.atomic_json + + def write_json(path, data): + if path.name == "run.json" and (root / "result.json").exists(): + raise OSError("manifest unavailable") + return original(path, data) + + monkeypatch.setattr(bundles, "atomic_json", write_json) + if recovery_write_fails: + from matkit.api import runtime + + monkeypatch.setattr(runtime, "atomic_json", write_json) + with pytest.raises(OSError, match="manifest unavailable"): + evaluate(energy_request(input_file), output_dir=root) + from matkit.api import RunResult + + committed = RunResult.model_validate_json( + (root / "result.json").read_text() + ) + assert committed.accepted + inspected = inspect_run(root) + assert inspected.payload == committed.payload + assert inspected.numerical_validity == "valid" + assert inspected.state == ( + "completed" if recovery_write_fails else "interrupted" + ) + + +def test_interrupted_batch_retains_first_result( + input_file, tmp_path, mock_calculator, monkeypatch +): + from matkit.mlip import runner + + original = runner._run_ase_item + count = 0 + + def item(*args, **kwargs): + nonlocal count + count += 1 + if count == 2: + raise KeyboardInterrupt("stopped") + return original(*args, **kwargs) + + monkeypatch.setattr(runner, "_run_ase_item", item) + root = tmp_path / "batch" + with pytest.raises(KeyboardInterrupt): + run_batch([energy_request(input_file)] * 2, output_dir=root) + batch = json.loads((root / "batch_manifest.json").read_text()) + assert batch["state"] == "interrupted" + assert inspect_run(root / "00000").accepted + assert inspect_run(root / "00001").state == "interrupted" + + +def test_native_batch_grouping_without_cuda(input_file, tmp_path, monkeypatch): + from matkit.mlip import runner + + chunks = [] + monkeypatch.setattr(runner, "_load_nvalchemi_model", lambda _: object()) + + def chunk(model, entries, backend, calculation): + chunks.append(len(entries)) + return [ + ( + index, + runner._success_result( + path, + backend, + calculation, + atoms, + -1, + [[0, 0, 0]] * len(atoms), + None, + True, + 0, + 0.1, + ), + ) + for index, path, atoms in entries + ] + + monkeypatch.setattr(runner, "_run_nvalchemi_chunk", chunk) + request = energy_request(input_file, adapter=AlchemiAdapter(batch_size=2)) + result = run_batch([request] * 3, output_dir=tmp_path / "batch") + assert result.state == "completed" + assert chunks == [2, 1] + + +def test_relocated_bundle_and_all_requested_pore_outputs(sample_cif, tmp_path): + root = tmp_path / "original" + prepare( + PoreRequest( + structure=StructureRef(path=sample_cif), analyses=["res", "sa"] + ), + output_dir=root, + ) + relocated = tmp_path / "relocated" + shutil.move(root, relocated) + result = execute(relocated, execution=fake_execution("zeopp")) + assert result.accepted + assert set(result.payload.results) == {"res", "sa"} + assert (relocated / "work" / "engine.stdout.log").exists() + + +@pytest.mark.parametrize("args", [("--fail",), ("--partial",)]) +def test_external_engine_failure_and_missing_analysis( + sample_cif, tmp_path, args +): + result = run( + PoreRequest( + structure=StructureRef(path=sample_cif), analyses=["res", "sa"] + ), + output_dir=tmp_path / "run", + execution=fake_execution("zeopp", *args), + ) + assert result.state == "failed" + assert result.failure is not None + + +def test_tampered_bundle_is_refused(sample_cif, tmp_path): + root = tmp_path / "run" + prepare( + PoreRequest(structure=StructureRef(path=sample_cif)), output_dir=root + ) + (root / "inputs" / "radii.rad").write_text("changed") + with pytest.raises(ValueError, match="changed"): + execute(root, execution=fake_execution("zeopp")) + + +def test_subprocess_execution_and_timeout(sample_cif, tmp_path): + request = PoreRequest(structure=StructureRef(path=sample_cif)) + result = run( + request, + output_dir=tmp_path / "good", + execution=fake_execution("zeopp", mode="subprocess"), + ) + assert result.accepted + interrupted = run( + request, + output_dir=tmp_path / "timeout", + execution=fake_execution( + "zeopp", "--sleep", mode="subprocess", timeout_s=2 + ), + ) + assert interrupted.state == "interrupted" + from matkit.api.structures import sha256 + + assert all( + sha256(tmp_path / "timeout" / a.path) == a.sha256 + for a in interrupted.artifacts + ) + assert not (tmp_path / "timeout" / ".matkit.lock").exists() + + +@pytest.mark.parametrize( + "unit,fugacity,expected", + [ + ("mol/kg", "PR-EOS", 12), + ("mg/g", "PR-EOS", 6), + ("g/L", "PR-EOS", 14), + ("mol/kg", 1.0, 7), + ], +) +def test_single_component_adsorption_requires_charges_and_collects( + sample_cif, tmp_path, unit, fugacity, expected +): + request = AdsorptionRequest( + structure=StructureRef(path=sample_cif), + adsorbate="CO2", + temperature_K=298, + pressure_Pa=1e5, + unit=unit, + fugacity_coefficient=fugacity, + ) + with pytest.raises(ValueError, match="_atom_site_charge"): + prepare(request, output_dir=tmp_path / "uncharged") + charged = tmp_path / "charged.cif" + text = ( + Path(sample_cif) + .read_text() + .replace( + " _atom_site_occupancy", + " _atom_site_charge\n _atom_site_occupancy", + ) + ) + text = ( + text.replace("0.00000 1.0000", "0.00000 0.0 1.0000") + .replace("0.25000 1.0000", "0.25000 0.0 1.0000") + .replace("0.75000 1.0000", "0.75000 0.0 1.0000") + ) + charged.write_text(text) + request.structure.path = str(charged) + result = run( + request, + output_dir=tmp_path / "charged_run", + execution=fake_execution("graspa"), + ) + assert result.accepted + assert result.payload.uptake == expected + assert result.payload.component == "CO2" + assert result.checks[0].status == "unknown" + assert ( + tmp_path / "charged_run" / "inputs" / "structure.cif" + ).read_bytes() == charged.read_bytes() + assert analyze_adsorption(tmp_path / "charged_run") == result diff --git a/tests/test_foundations.py b/tests/test_foundations.py new file mode 100644 index 0000000..fad6531 --- /dev/null +++ b/tests/test_foundations.py @@ -0,0 +1,113 @@ +"""Scientific regressions shared by the GCMC setup implementations.""" + +import importlib +import re + +import numpy as np +import pytest +from ase import Atoms +from ase.io import write + +from matkit.raspa2 import get_output_data +from matkit.utils import calculate_cell_size + + +@pytest.mark.parametrize( + "engine", ["graspa", "graspa_sycl", "raspa2", "pygraspa"] +) +@pytest.mark.parametrize("cutoff", [4.0, 18.0]) +def test_setup_uses_requested_cutoff(engine, cutoff, tmp_path): + atoms = Atoms("Si", cell=[[10, 0, 0], [3, 9, 0], [1, 2, 8]], pbc=True) + cif = tmp_path / "input.cif" + write(cif, atoms) + module = importlib.import_module(f"matkit.{engine}") + out = tmp_path / "sim" + if engine == "raspa2": + module.setup_input_simulation([str(cif)], str(out), cutoff=cutoff) + out = out / "input" + elif engine == "graspa_sycl": + module.setup_simulation(str(cif), str(out), cutoff=cutoff) + else: + kwargs = {} + if engine == "pygraspa": + kwargs = dict( + model_path="/model.pt", + model_type="FAIRChem-esen", + E_comps=[-1.0], + ) + module.setup_simulation( + str(cif), + str(out), + [{"MoleculeName": "CO2"}], + cutoff=cutoff, + **kwargs, + ) + text = (out / "simulation.input").read_text() + match = re.search( + r"^UnitCells\s+(?:0\s+)?(\d+)\s+(\d+)\s+(\d+)$", text, re.M + ) + sizes = [int(n) for n in match.groups()] + assert sizes == calculate_cell_size(atoms, cutoff) + cell = atoms.cell.array + heights = [ + abs(np.linalg.det(cell)) + / np.linalg.norm(np.cross(cell[(i + 1) % 3], cell[(i + 2) % 3])) + for i in range(3) + ] + assert all(n * h >= 2 * cutoff for n, h in zip(sizes, heights)) + + +@pytest.mark.parametrize("engine", ["graspa", "pygraspa"]) +def test_batch_cached_replication_uses_cutoff(engine, tmp_path): + module = importlib.import_module(f"matkit.{engine}") + inputs = tmp_path / "inputs" + inputs.mkdir() + write(inputs / "input.cif", Atoms("Si", cell=[10, 10, 10], pbc=True)) + kwargs = {} + if engine == "pygraspa": + kwargs = dict( + model_path="/model.pt", model_type="FAIRChem-esen", E_comps=[-1.0] + ) + module.setup_batch( + str(inputs), + str(tmp_path / "out"), + [{"MoleculeName": "CO2"}], + temperatures=[298], + pressures=[1e4, 1e5], + cutoff=18, + **kwargs, + ) + files = list((tmp_path / "out").rglob("simulation.input")) + assert len(files) == 2 + assert all("UnitCells 0 4 4 4" in f.read_text() for f in files) + + +@pytest.mark.parametrize("unit,expected", [("mol/kg", 2), ("g/L", 44)]) +def test_raspa2_valid_result_is_successful(tmp_path, unit, expected): + log = tmp_path / "raspa.log" + log.write_text( + "Average loading absolute [mol/kg framework] 2 +/- 0.1\n" + "Average loading absolute [milligram/gram framework] 88 +/- 4.4\n" + "Framework Density 500\n" + ) + result = get_output_data(str(log), unit=unit) + assert result["success"] is True + assert result["uptake"] == expected + + +def test_raspa2_incomplete_result_fails(tmp_path): + log = tmp_path / "bad.log" + log.write_text("incomplete\n") + with pytest.raises(ValueError, match="expected lines"): + get_output_data(str(log)) + + +@pytest.mark.parametrize("cutoff", [True, 0, -1, float("inf"), float("nan")]) +def test_invalid_cutoff_rejected(cutoff): + with pytest.raises(ValueError, match="positive and finite"): + calculate_cell_size(Atoms("Si", cell=[10] * 3), cutoff) + + +def test_singular_replication_cell_rejected(): + with pytest.raises(ValueError, match="independent"): + calculate_cell_size(Atoms("Si", cell=[10, 10, 0])) diff --git a/tests/test_mcp_api.py b/tests/test_mcp_api.py new file mode 100644 index 0000000..da08267 --- /dev/null +++ b/tests/test_mcp_api.py @@ -0,0 +1,165 @@ +"""Optional MCP tests use the real SDK and stdio, with synthetic engines.""" + +import asyncio +import base64 +import json +from pathlib import Path +import sys + +import pytest + +pytest.importorskip("mcp") +from mcp import Client +from mcp.client.stdio import StdioServerParameters + +from matkit.api import inspect_run +from matkit.mcp import create_server + + +def profile(*args): + return { + "executables": { + "zeopp": [ + sys.executable, + str(Path(__file__).parent / "fixtures" / "fake_engine.py"), + "zeopp", + *args, + ] + } + } + + +def test_stdio_discovery_execution_and_artifact_retrieval(sample_cif, tmp_path): + profiles = tmp_path / "profiles.json" + profiles.write_text(json.dumps({"default": profile()})) + params = StdioServerParameters( + command=sys.executable, + args=[ + "-m", + "matkit.mcp", + "--run-root", + str(tmp_path / "runs"), + "--input-root", + str(Path(sample_cif).parent), + "--profiles", + str(profiles), + ], + ) + + async def check(): + async with Client(params) as client: + catalog = await client.list_tools() + names = {tool.name for tool in catalog.tools} + assert { + "matkit_pores", + "matkit_evaluate", + "matkit_prepare_adsorption", + } <= names + pores = next( + tool for tool in catalog.tools if tool.name == "matkit_pores" + ) + assert pores.input_schema + response = await client.call_tool( + "matkit_pores", + { + "request": { + "structure": {"path": sample_cif}, + "analyses": ["res", "sa"], + } + }, + ) + assert not response.is_error, response + data = response.structured_content + assert data["accepted"] + assert set(data["summary"]["results"]) == {"res", "sa"} + ref = next( + a for a in data["artifacts"] if a["path"] == "result.json" + ) + resource = await client.read_resource(ref["uri"]) + content = resource.contents[0] + encoded = ( + content.text + if hasattr(content, "text") + else base64.b64decode(content.blob) + ) + result = json.loads(encoded) + assert result["run_id"] == data["run_id"] + assert result["payload"]["results"] == data["summary"]["results"] + + asyncio.run(check()) + + +def test_mcp_timeout_and_path_roots(sample_cif, tmp_path): + server = create_server( + run_root=tmp_path / "runs", + input_roots=[Path(sample_cif).parent], + profiles={"default": profile("--sleep")}, + timeout_s=2, + ) + + async def check(): + async with Client(server) as client: + bad = await client.call_tool( + "matkit_pores", + {"request": {"structure": {"path": "/etc/passwd"}}}, + ) + assert bad.is_error + result = await client.call_tool( + "matkit_pores", {"request": {"structure": {"path": sample_cif}}} + ) + assert result.structured_content["state"] == "interrupted" + assert not result.structured_content["accepted"] + + asyncio.run(check()) + + +def test_mcp_selected_catalog(sample_cif, tmp_path): + server = create_server( + run_root=tmp_path / "runs", + input_roots=[Path(sample_cif).parent], + tools=["matkit_capabilities"], + ) + + async def check(): + async with Client(server) as client: + result = await client.list_tools() + assert [t.name for t in result.tools] == ["matkit_capabilities"] + + asyncio.run(check()) + + +def test_mcp_cancellation_preserves_interrupted_run(sample_cif, tmp_path): + root = tmp_path / "runs" + server = create_server( + run_root=root, + input_roots=[Path(sample_cif).parent], + profiles={"default": profile("--sleep")}, + ) + + async def check(): + async with Client(server) as client: + task = asyncio.create_task( + client.call_tool( + "matkit_pores", + {"request": {"structure": {"path": sample_cif}}}, + ) + ) + for _ in range(200): + if list(root.glob("*/work/started")): + break + await asyncio.sleep(0.025) + else: + pytest.fail("worker did not start") + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + for _ in range(200): + runs = list(root.glob("*/result.json")) + if runs: + assert inspect_run(runs[0].parent).state == "interrupted" + break + await asyncio.sleep(0.025) + else: + pytest.fail("cancellation did not persist the interrupted run") + + asyncio.run(check()) diff --git a/tests/test_operation_cli.py b/tests/test_operation_cli.py new file mode 100644 index 0000000..8406b82 --- /dev/null +++ b/tests/test_operation_cli.py @@ -0,0 +1,130 @@ +"""Real CLI execution uses the same contracts as Python.""" + +import json +from pathlib import Path +import sys + +from click.testing import CliRunner + +from matkit.api import ExecutionConfig, PoreRequest, StructureRef, run +from matkit.cli import main + + +def inputs(tmp_path, sample_cif, *engine_args): + request = { + "operation": "pores", + "structure": {"path": sample_cif}, + "analyses": ["res", "sa"], + } + spec = tmp_path / "pores.json" + spec.write_text(json.dumps(request)) + fixture = Path(__file__).parent / "fixtures" / "fake_engine.py" + config = { + "executables": { + "zeopp": [sys.executable, str(fixture), "zeopp", *engine_args] + } + } + profile = tmp_path / "execution.json" + profile.write_text(json.dumps(config)) + return spec, profile, config + + +def test_cli_python_equivalent_results(sample_cif, tmp_path): + spec, profile, config = inputs(tmp_path, sample_cif) + result = CliRunner().invoke( + main, + [ + "pores", + "--spec", + str(spec), + "--execution", + str(profile), + "--outdir", + str(tmp_path / "cli"), + ], + ) + assert result.exit_code == 0, result.output + data = json.loads(result.stdout) + python = run( + PoreRequest( + structure=StructureRef(path=sample_cif), analyses=["res", "sa"] + ), + output_dir=tmp_path / "python", + execution=ExecutionConfig(**config), + ) + assert data["payload"] == python.payload.model_dump(mode="json") + assert data["checks"] == [c.model_dump(mode="json") for c in python.checks] + + +def test_cli_preparation_then_execution_and_inspection(sample_cif, tmp_path): + spec, profile, _ = inputs(tmp_path, sample_cif) + runner = CliRunner() + root = tmp_path / "run" + prepared = runner.invoke( + main, ["prepare", "--spec", str(spec), "--outdir", str(root)] + ) + assert prepared.exit_code == 0, prepared.output + assert json.loads(prepared.stdout)["state"] == "prepared" + executed = runner.invoke( + main, ["execute", str(root), "--execution", str(profile)] + ) + assert executed.exit_code == 0, executed.output + inspected = runner.invoke(main, ["inspect", str(root)]) + assert inspected.exit_code == 0 + assert json.loads(executed.stdout) == json.loads(inspected.stdout) + rerun = runner.invoke( + main, ["execute", str(root), "--execution", str(profile)] + ) + assert rerun.exit_code == 2 + + +def test_cli_failure_is_structured_and_nonzero(sample_cif, tmp_path): + spec, profile, _ = inputs(tmp_path, sample_cif, "--fail") + result = CliRunner().invoke( + main, + [ + "pores", + "--spec", + str(spec), + "--execution", + str(profile), + "--outdir", + str(tmp_path / "run"), + ], + ) + assert result.exit_code == 1, result.output + data = json.loads(result.stdout) + assert data["state"] == "failed" + assert "code 7" in data["failure"]["message"] + + +def test_cli_invalid_scientific_request_exits_two(sample_cif, tmp_path): + spec, _, _ = inputs(tmp_path, sample_cif) + data = json.loads(spec.read_text()) + data["probe_radius"] = -1 + spec.write_text(json.dumps(data)) + result = CliRunner().invoke( + main, ["pores", "--spec", str(spec), "--outdir", str(tmp_path / "run")] + ) + assert result.exit_code == 2 + assert not (tmp_path / "run").exists() + + +def test_cli_batch_executes_in_worker(sample_cif, tmp_path): + spec, profile, _ = inputs(tmp_path, sample_cif) + request = json.loads(spec.read_text()) + spec.write_text(json.dumps([request, request])) + result = CliRunner().invoke( + main, + [ + "batch", + "--spec", + str(spec), + "--execution", + str(profile), + "--outdir", + str(tmp_path / "batch"), + ], + ) + assert result.exit_code == 0, result.output + assert len(json.loads(result.stdout)["items"]) == 2 diff --git a/tests/test_packaged_resources.py b/tests/test_packaged_resources.py new file mode 100644 index 0000000..70a20a8 --- /dev/null +++ b/tests/test_packaged_resources.py @@ -0,0 +1,45 @@ +"""Run these checks from an installed wheel as well as editable checkouts.""" + +from importlib import resources +import os +from pathlib import Path +import subprocess +import sys + +import pytest + + +@pytest.mark.parametrize( + "package,path", + [ + ("graspa", "files/template/simulation.input"), + ("graspa", "files/template_mixture_isotherm/template_mixture/SO2.def"), + ("pygraspa", "files/template_mixture_isotherm/template_mixture/H2.def"), + ("graspa_sycl", "files/template/simulation.input"), + ("raspa2", "files/template/simulation.input"), + ("raspa3", "files/template/CO2.def"), + ("zeopp", "files/UFF.rad"), + ], +) +def test_bundled_resources(package, path): + assert resources.files(f"matkit.{package}").joinpath(path).read_bytes() + + +def test_core_import_keeps_optional_runtimes_unloaded(): + subprocess.run( + [ + sys.executable, + "-c", + "import sys; import matkit.api; " + "assert not {'torch', 'mace', 'rootstock', 'nvalchemi', 'mcp'} " + "& sys.modules.keys()", + ], + check=True, + ) + + +def test_installed_location(): + if os.environ.get("MATKIT_WHEEL_TEST") == "1": + import matkit + + assert "site-packages" in Path(matkit.__file__).parts diff --git a/tests/test_unified_smoke.py b/tests/test_unified_smoke.py new file mode 100644 index 0000000..72d746c --- /dev/null +++ b/tests/test_unified_smoke.py @@ -0,0 +1,45 @@ +"""Check execution evidence recording with a synthetic engine.""" + +import importlib.util +import json +from pathlib import Path +import sys + + +def test_evidence_recorder_continues_after_failure(sample_cif, tmp_path): + script = Path(__file__).parents[1] / "examples" / "unified_smoke.py" + spec = importlib.util.spec_from_file_location("unified_smoke", script) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + good = tmp_path / "good.json" + good.write_text( + json.dumps({"operation": "pores", "structure": {"path": sample_cif}}) + ) + bad = tmp_path / "bad.json" + bad.write_text("{}") + profile = tmp_path / "execution.json" + engine = Path(__file__).parent / "fixtures" / "fake_engine.py" + profile.write_text( + json.dumps( + {"executables": {"zeopp": [sys.executable, str(engine), "zeopp"]}} + ) + ) + root = tmp_path / "evidence" + assert ( + module.main( + [ + "--spec", + str(bad), + "--spec", + str(good), + "--execution", + str(profile), + "--outdir", + str(root), + ] + ) + == 1 + ) + report = json.loads((root / "execution_report.json").read_text()) + assert [case["accepted"] for case in report["cases"]] == [False, True] + assert len(report["cases"][1]["spec_sha256"]) == 64