diff --git a/agents/nvidia-alchemi/.dockerignore b/agents/nvidia-alchemi/.dockerignore new file mode 100644 index 0000000..206ef80 --- /dev/null +++ b/agents/nvidia-alchemi/.dockerignore @@ -0,0 +1,4 @@ +__pycache__/ +**/__pycache__/ +**/*.pyc +.git/ diff --git a/agents/nvidia-alchemi/CHANGELOG.md b/agents/nvidia-alchemi/CHANGELOG.md new file mode 100644 index 0000000..6e3747d --- /dev/null +++ b/agents/nvidia-alchemi/CHANGELOG.md @@ -0,0 +1,101 @@ +# Changelog + +All notable changes to the NVIDIA ALCHEMI Toolkit agent are recorded here. +The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and the +agent uses [Semantic Versioning](https://semver.org/). Versions match `version` in +`metadata.yaml` and `tools/nvidia-alchemi/tool.yaml` and the container image tag. + +## [1.1.2] - 2026-10-01 + +Packaging-only release to meet the Microsoft Discovery catalog policy. The +helper library, agent instructions, installed packages and checkpoints are +unchanged from 1.1.1. + +### Changed + +- The pinned dependency list is renamed from `requirements-build.lock` to + `requirements-build.txt`, a file type the catalog accepts. Its contents are + unchanged. +- Catalog tags use the catalog vocabulary: `machine-learning`, `gpu` and + `simulation` replace `machine-learned-potential`, `gpu-accelerated` and + `high-throughput`. + +### Validation + +- Rebuilt from a clean checkout; the installed packages, checkpoint hashes, base + image and license files are identical to 1.1.1. Revalidated on one H100 NVL + with the full three-part release check (smoke suite, offline models, MACE + dynamics). + +## [1.1.1] - 2026-10-01 + +### Added + +- `single_point(batch, model)`: the supported way to evaluate energies and forces for + fixed geometries. It builds the neighbor list itself and rejects missing or + non-finite outputs. +- `step_recorder()`: a ready-made dynamics hook that records energy, kinetic energy, + temperature and maximum force for every step and structure. +- The release manifest now names each cached checkpoint (`loader` and `name`) next + to its path, size and SHA-256, so models can be looked up by name. +- Distribution license and notice files are copied into the image under + `/app/licenses/`, and their per-package counts are recorded in the manifest. +- Regression tests for the new helpers and the manifest, bringing the suite to 32. + +### Changed + +- The image is built from the public `nvidia/cuda:13.0.1-devel-ubuntu24.04` image, + pinned by digest, instead of a private parent image. Anyone can now rebuild it. + It installs the same 198 package versions as 1.1.0 from + `requirements-build.lock`, without further dependency resolution. +- MACE `small-0b`, MACE `medium-0b2` and AIMNet2 checkpoints are downloaded at build + time, and the build fails unless every file matches its recorded SHA-256. + Jobs still run offline. +- The release manifest records the actual base image, build inputs and release + version. The smoke test reports the release and the helper-library version + separately. +- Agent instructions document the new helpers, the custom-hook requirements, batch + field names (`batch_idx`, `atomic_masses`) and checkpoint lookup by name. +- The README is self-contained for the catalog: support links point to the + Microsoft Discovery community, and build and deployment steps are documented. +- The release smoke test is renamed from `hardening_smoke.py` to `release_smoke.py`. +- The build fails if any source file has Windows line endings, so the recorded + source hashes always match the repository. +- Sample structures in `/app/example-input-files` are documented for users and the agent. + +### Fixed + +- `write_structures()` crashed when given a single `ase.Atoms`; it now accepts a + `Batch`, a single `Atoms`, or a list. +- The instructions no longer suggest that model wrappers have a `compute()` method. + +### Validation + +Validated on one H100 NVL through the agent: 32 regression tests with no skips; all +three cached models selected by name, hash-checked and evaluated on CUDA; and a +100-step MACE NVT run whose final kinetic energy, recomputed independently, matched +the recorded value to $6\times10^{-9}$ eV. + +## [1.1.0] - 2026-09-22 + +### Added + +- `alchemi_validators`: NumPy reference calculations for kinetic energy and energy + drift, finite and shape-aware numeric comparison, and hook-event auditing. Results + distinguish `pass`, `fail`, `unverified` and `unexercised`. +- `screen()` results carry `input_index`, the original input position, separately + from the toolkit's admission-order `system_id`. + +### Changed + +- Declared outputs of any file type, including nested directories, are copied to + the output folder and verified by SHA-256 and size before results are published. + An artifact manifest lists every delivered file. +- Results JSON is strict: non-finite values are encoded explicitly and unsupported + objects are rejected. +- The tool reports the script's real exit status. + +### Fixed + +- Outputs such as NPZ, Python and text files were silently left out of results. +- Screening results could be matched to the wrong inputs for mixed-size batches. diff --git a/agents/nvidia-alchemi/README.md b/agents/nvidia-alchemi/README.md new file mode 100644 index 0000000..ffc4f11 --- /dev/null +++ b/agents/nvidia-alchemi/README.md @@ -0,0 +1,298 @@ +# NVIDIA ALCHEMI Toolkit + +Expert agent for high-throughput, GPU-first AI atomic simulation using the +[NVIDIA ALCHEMI Toolkit](https://github.com/NVIDIA/nvalchemi-toolkit) — a batch-first +Python framework for machine-learned interatomic potentials (MLIPs). + +The agent plans and writes a single Python workflow script, then executes it inside a +CUDA container. It relaxes and simulates many structures in one batched GPU pass rather +than looping over them, which is where the toolkit's throughput advantage comes from. + +## Overview + +- **Intended users**: computational chemists and materials scientists running MLIP-driven + relaxation, molecular dynamics and screening campaigns; ML researchers training or + fine-tuning potentials. +- **Successful outcome**: relaxed structures, MD trajectories, screening summaries or a + trained checkpoint written to `/output`, together with `final_results.json`. + +## Architecture + +This agent operates as a `kind: prompt` agent within Discovery Studio. + + User Input → NVIDIA ALCHEMI Toolkit (LLM) → nvidiaAlchemi Tool (Container) → Results + +- **Model**: configured via the `{{CHAT-MODEL}}` parameter at deploy time +- **Potentials**: MACE (foundation checkpoints `small-0b` and `medium-0b2` baked into + the image), AIMNet2 (`aimnet2` checkpoint, also baked in), Lennard-Jones, DFT-D3(BJ), + Ewald, PME — composable with the `+` operator or an explicit pipeline +- **Dynamics**: FIRE / FIRE2 optimizers and NVE, NVT (Langevin, Nosé-Hoover), NPT and + NPH integrators sharing one execution loop with nine hook points per step +- **Scaling**: fused multi-stage pipelines and inflight batching on one GPU; + `DistributedPipeline` and domain decomposition across GPUs via `torchrun` +- **Output**: extended XYZ / CIF structures, Zarr trajectories, CSV summaries, plots, + and training checkpoints + +## Prerequisites + +- Azure subscription with a Discovery workspace +- GPU nodepool (H100 or A100 recommended); multi-GPU SKU for distributed workflows +- Model deployment for the agent LLM (e.g. GPT-4o) +- An Azure Container Registry with push permission, and a Docker/Podman or + `az acr build` environment with internet access to the public package indexes, + model download sites and roughly 40 GB of free build disk + +## Configuration + +| Parameter | Description | Example | +|---|---|---| +| `{{CHAT-MODEL}}` | Model deployment name for the agent LLM | `gpt-4o-deployment` | + +## Build and deploy the tool image + +The single build recipe is [tools/nvidia-alchemi/Dockerfile](tools/nvidia-alchemi/Dockerfile). +**Use this agent directory as the build context, not the nested tool directory**, +because the image also records the agent instructions and tool definition: + +```bash +docker build -f tools/nvidia-alchemi/Dockerfile -t .azurecr.io/nvidia-alchemi:1.1.2 . +az acr login --name +docker push .azurecr.io/nvidia-alchemi:1.1.2 +``` + +or, without a local Docker daemon, +`az acr build --registry --image nvidia-alchemi:1.1.2 --file tools/nvidia-alchemi/Dockerfile .` + +The build starts from the public `nvidia/cuda:13.0.1-devel-ubuntu24.04` image pinned +by digest, installs the exact package versions in +[requirements-build.txt](tools/nvidia-alchemi/requirements-build.txt) without +further dependency resolution, and runs `pip check`. It downloads the MACE +`small-0b` and `medium-0b2` and AIMNet2 checkpoints and fails unless every file +matches the size and SHA-256 in [checkpoints.json](tools/nvidia-alchemi/checkpoints.json). +It then writes `/app/release-manifest.json` (base digest, package inventory, +checkpoint and source hashes), preserves distribution license files under +`/app/licenses`, and runs the regression suite. Jobs need no network access. + +Deploy the tool from `tools/nvidia-alchemi/tool.yaml`; the deployer substitutes +`{name}` with your registry. Prefer referencing the pushed image by digest. +Prompt updates and container updates are separate operations; verify both. + +### Package layout + +This package contains the agent/tool definitions, runtime helpers, build inputs, +provenance generator, smoke test, regression tests, examples, and third-party notices. + +### Verified output contract + +- Declare every deliverable in `save_final_results(output_files=...)` using its + existing source path. NPZ, Python, text, arbitrary other extensions, and nested + directories are supported. Relative paths resolve under the configured work directory. +- Declared sources are copied into the captured output directory and SHA-256/size + verified. Returned paths and the saved result point at exported destinations. + An artifact manifest records every delivered file, including directory contents. +- Missing sources, collisions, symlinks, and corrupted copies fail explicitly. + A failed export cannot advertise a completed result. `quick_finish()` repeats + final export/verification so declared logs updated after an earlier save are current. +- Close writers before finalization. Once results are saved, undeclared scratch + files are not exported. Scripts that only call `quick_finish()` retain a limited, + top-level legacy artifact-discovery mode; explicit declarations are preferred. +- Use `write_json_atomic()` for auxiliary JSON. Nonfinite values are encoded as + diagnostic objects with `value: null`, `finite: false`, and an `encoding` of + `nan`, `+inf`, or `-inf`. Unsupported objects are rejected, not stringified. +- Use `output_files={}` explicitly when clearing broken declarations in an error + handler. A later save with `output_files=None` retains previous declarations. + +### Screening identity and independent checks + +`screen()` returns structures in graduation order. Match them with +`atoms.info['input_index']`, the original zero-based input position. The separate +`system_id` is assigned by the toolkit in admission order and can differ from +input order for mixed-size batches. Caller metadata is restored; pre-existing +caller IDs are retained as `source_system_id` and `source_input_index`. Missing, +duplicate, or out-of-range returned identities raise rather than misattribute data. + +`batch_to_atoms()` retains these IDs, masses, and toolkit-native velocities in +`atoms.arrays['alchemi_velocities']`. It deliberately does not label the latter +as ASE velocities without a unit conversion. + +The [reviewed validators](tools/nvidia-alchemi/alchemi_validators.py) provide +finite/shape-aware numeric comparison, NumPy float64 kinetic energy and drift +references, and actual paired-hook-event auditing. Verdicts distinguish `pass`, +`fail`, `unverified`, and `unexercised`. Provenance labels and alias detection +cannot alone prove logical independence; callers must retain real observations. + +### Evaluation, recording and model lookup + +- `single_point(batch, model)` evaluates energies and forces for fixed geometries. + It builds the neighbor list itself; model wrappers have no `compute()`, calling + `model(batch)` directly fails without a neighbor list, and a zero-step `run()` + silently leaves energies at zero. +- `step_recorder()` is a ready-made hook for `run_md(..., hooks=[recorder])`. + `recorder.rows` holds energy, kinetic energy, temperature and maximum force for + every step and structure. Custom hooks must provide integer `frequency`, a + `stage`, and `__call__(ctx, stage)`. +- `write_structures()` accepts a `Batch`, a single `ase.Atoms`, or a list of them. +- `/app/release-manifest.json` names each cached checkpoint by `loader` and `name` + next to its path, size and SHA-256. Look models up by name, not by filename. +- Batch field names: the atom-to-structure index is `batch_idx` and masses are + `atomic_masses`. + +Tests are in [tools/nvidia-alchemi/tests](tools/nvidia-alchemi/tests). Run them using +Python's `unittest` discovery with the tool directory on `PYTHONPATH`; no pytest +dependency is required. Real toolkit identity tests run in the release image; +set `ALCHEMI_TEST_GPU=1` to require GPU execution rather than the CPU test path. +The [release smoke runner](tools/nvidia-alchemi/release_smoke.py) requires a GPU +by default and produces an artifact manifest for post-download verification. + +## Tools + +| Tool | Path | Description | +|---|---|---| +| `nvidiaAlchemi` | `tools/nvidia-alchemi/` | NVIDIA ALCHEMI Toolkit container for batched MLIP relaxation, molecular dynamics, screening, training and multi-GPU simulation | + +**Compute requirements**: 1–8 GPUs, 8–96 vCPU, 64 Gi–1800 Gi RAM. Recommended SKUs: +`Standard_NC40ads_H100_v5`, `Standard_NC80adis_H100_v5`, `Standard_ND96isr_H100_v5`, +`Standard_NC24ads_A100_v4`, `Standard_NC48ads_A100_v4`. + +**Supported inputs**: XYZ, extended XYZ, CIF, PDB, POSCAR/CONTCAR, ASE trajectory files. +**Outputs**: extended XYZ / CIF structures, Zarr trajectory stores, CSV summaries, PNG +plots, training checkpoints, and `final_results.json`. + +## Usage + +### Example inputs + +The image includes small sample structures in `/app/example-input-files`: +`water.xyz`, `ethanol.xyz`, `si_bulk.cif` and `cu_bulk.extxyz`. They are also in +[tools/nvidia-alchemi/example-input-files](tools/nvidia-alchemi/example-input-files). + +``` +Relax the water and ethanol examples in /app/example-input-files with MACE and report their energies. +``` + +### Batched geometry relaxation + +``` +Relax all the structures in my input directory with MACE to a maximum force of 0.03 eV/Å. +``` + +### Molecular dynamics + +``` +Run 50 ps of NVT molecular dynamics at 300 K on this system with MACE and save the trajectory. +``` + +### Multi-stage pipeline + +``` +Relax this structure, equilibrate it at 300 K, then run a 20 ps production trajectory. +``` + +### High-throughput screening + +``` +Screen these 5,000 candidate structures: relax each one and rank them by energy per atom. +``` + +### Model composition + +``` +Run MD on this molecular crystal using MACE with a DFT-D3 dispersion correction. +``` + +### Fine-tuning + +``` +Fine-tune the MACE small-0b checkpoint on the trajectory in train.zarr for 20 epochs, +weighting forces ten times more than energies. +``` + +### Multi-GPU + +``` +Screen these structures across 4 GPUs using a relaxation stage feeding an NVT stage. +``` + +## Support + +For questions and bug reports, use the Microsoft Discovery community discussions: + + +Microsoft Discovery team contact: discovery-catalog@microsoft.com + +## Known Limitations + +- **GPU required.** There is no CPU fallback for the accelerated neighbor-list and + interaction kernels. +- **Screening returns only converged structures.** `screen()` collects systems that + reached `fmax`; anything still above the threshold when the step budget runs out is + not collected. The helper logs a warning naming the shortfall — raise `max_steps` or + loosen `fmax` to capture the remainder. +- **UMA / fairchem-core is not available.** Upstream declares the `uma` extra mutually + exclusive with the `mace`, `cu12` and `cu13` extras (incompatible `e3nn` / `torch` + pins), and its checkpoints require gated HuggingFace access. Use MACE or AIMNet2. +- **Multi-GPU workflows require `torchrun`** and a multi-GPU SKU. `DistributedPipeline` + and `DomainParallel` run in a separate script launched as one process per GPU, not + inline in the main workflow script. +- **Distributed communication is one-to-one.** Each rank talks to at most one upstream + and one downstream neighbour; fan-out and fan-in topologies are not yet supported. +- **Upstream is a public beta (0.2.0)** and its API is subject to change between + releases. The helper library resolves symbols defensively and reports the installed + version when one moves. +- **`Batch` → `ase.Atoms` conversion is provided by this agent**, not by the toolkit; + it carries geometry, energies, forces, stress, identity, masses and separately + labelled toolkit-native velocities, not arbitrary tensor metadata. `screen()` + restores caller `info` metadata using its checked input index. +- **Accuracy is bounded by the chosen potential.** Foundation MLIPs are not a substitute + for DFT on chemistry far outside their training distribution, and the classical + potentials ship with generic parameters (argon for LJ, the PBE damping set for + DFT-D3) that should be overridden for the system under study. + +## Verification status + +This release (1.1.2) was validated on one H100 NVL, using the agent itself: + +- **Regression smoke:** all **32 tests** passed with no skips; independent download + verification checked all 30 declared files, input-identity mapping and NumPy + reference calculations. +- **Offline models:** MACE `small-0b`, MACE `medium-0b2` and AIMNet2 were each + selected by name from the release manifest, hash-checked, and evaluated with + `single_point()`; all outputs were finite on CUDA and agreed with CPU results to + about $2\times10^{-6}$ eV. +- **MACE dynamics:** 100 steps of nonperiodic NVT on three water molecules with + `medium-0b2` completed with all values finite. The final kinetic energy recomputed + independently from the saved velocities and masses matched the recorded value + to $8\times10^{-9}$ eV. + +Earlier H100 evaluations also exercised Lennard-Jones and demo workflows, NPT +response, exact trajectory readback and defensive hooks. See the +[changelog](CHANGELOG.md) for what changed in each release. + +This is not physical-accuracy certification. Multi-GPU execution, full training, +periodic/multi-system MACE, cuEquivariance acceleration, explicit compiled MACE, +and drift-triggered warning branches remain outside completed validation coverage. + +## License + +The agent definitions and wrapper code are governed by the catalog repository's +[MIT license](https://github.com/microsoft/discovery/blob/main/LICENSE). Bundled +third-party components keep their own licenses (see below). + +## Third-Party Components + +| Component | Version | License | Source | +|---|---|---|---| +| NVIDIA ALCHEMI Toolkit | 0.2.0 | Apache-2.0 | | +| NVIDIA CUDA base image | 13.0.1 | NVIDIA Deep Learning Container License | | +| MACE (+ `small-0b` / `medium-0b2` checkpoints) | 0.3.15 | MIT | | +| AIMNet2 (+ `aimnet2` checkpoint) | 0.2.0 | MIT | | +| PyTorch | 2.13.0+cu130 | BSD-3-Clause | | + +> Full attribution, citation requirements and license-file locations are in +> [`THIRD_PARTY_NOTICES.md`](THIRD_PARTY_NOTICES.md). + +## Contributing + +This project welcomes contributions and suggestions. Please see the catalog's +[CONTRIBUTING guidelines](https://github.com/microsoft/discovery/blob/main/CONTRIBUTING.md) +for details on how to contribute. diff --git a/agents/nvidia-alchemi/THIRD_PARTY_NOTICES.md b/agents/nvidia-alchemi/THIRD_PARTY_NOTICES.md new file mode 100644 index 0000000..8462a06 --- /dev/null +++ b/agents/nvidia-alchemi/THIRD_PARTY_NOTICES.md @@ -0,0 +1,111 @@ +# Third-Party Notices — nvidia-alchemi agent + +The wrapper code in this agent, including the +[runtime helpers](tools/nvidia-alchemi/alchemi_utils.py), +[artifact exporter](tools/nvidia-alchemi/alchemi_artifacts.py), +[validators](tools/nvidia-alchemi/alchemi_validators.py), +[Dockerfile](tools/nvidia-alchemi/Dockerfile), build scripts, and YAML definitions, +is published under the catalog repository's +[MIT license](https://github.com/microsoft/discovery/blob/main/LICENSE). The +components below are installed into, or downloaded into, the container image at +build time and carry their own licenses and citation requirements. Exact package +versions are listed in [requirements-build.txt](tools/nvidia-alchemi/requirements-build.txt). + +## Base image and runtime + +| Component | Version | License | Source | +|---|---|---|---| +| NVIDIA CUDA base image (`nvidia/cuda:13.0.1-devel-ubuntu24.04`) | 13.0.1 | NVIDIA Deep Learning Container License | | +| Ubuntu 24.04 base packages | 24.04 | various (mostly GPL-2.0 / LGPL-2.1) | | +| PyTorch (CUDA 13 build) | 2.13.0+cu130 | BSD-3-Clause | | + +## ALCHEMI Toolkit + +| Component | Version | License | Source | +|---|---|---|---| +| NVIDIA ALCHEMI Toolkit (`nvalchemi-toolkit`) | 0.2.0 | Apache-2.0 | | +| `nvalchemi-toolkit-ops` (GPU neighbor-list and interaction kernels) | per `cu13` extra | Apache-2.0 | | +| NVIDIA Warp (`warp-lang`) | transitive | Apache-2.0 | | +| cuEquivariance | transitive (`cu13`) | Apache-2.0 | | + +## Interatomic potentials and model weights + +| Component | License | Source | +|---|---|---| +| MACE (`mace-torch`) | MIT | | +| MACE-MP-0b foundation checkpoints (`small-0b`, `medium-0b2`), baked into the image | MIT | | +| AIMNet2 (`aimnet`) and its `aimnet2` checkpoint, baked into the image | MIT | | +| DFT-D3(BJ) dispersion parameters | as published by the original authors | | + +UMA / `fairchem-core` is **not** included in this image. Its `uma` extra is declared +mutually exclusive with the `mace`, `cu12` and `cu13` extras upstream, and its +checkpoints are distributed through the gated `facebook/UMA` HuggingFace repository. + +## Python scientific stack + +| Component | License | Source | +|---|---|---| +| ASE (Atomic Simulation Environment) | LGPL-2.1-or-later | | +| pymatgen | MIT | | +| Zarr / numcodecs | MIT | | +| NumPy | BSD-3-Clause | | +| SciPy | BSD-3-Clause | | +| pandas | BSD-3-Clause | | +| Matplotlib | Matplotlib License (BSD-compatible) | | +| seaborn | BSD-3-Clause | | +| Rich | MIT | | + +## Citations + +Cite the original publications for any potential used in published work. + +> Batatia, I. et al. "MACE: Higher Order Equivariant Message Passing Neural Networks +> for Fast and Accurate Force Fields." *Advances in Neural Information Processing +> Systems (NeurIPS)*, 2022. + +> Batatia, I. et al. "A foundation model for atomistic materials chemistry." +> arXiv:2401.00096, 2023. DOI:10.48550/arXiv.2401.00096 + +> Anstine, D. M., Zubatyuk, R. & Isayev, O. "AIMNet2: a neural network potential to +> meet your neutral, charged, organic, and elemental-organic needs." +> *Chem. Sci.* **16**, 10228-10244, 2025. DOI:10.1039/D4SC08572H + +> Grimme, S. et al. "A consistent and accurate ab initio parametrization of density +> functional dispersion correction (DFT-D) for the 94 elements H-Pu." +> *J. Chem. Phys.* **132**, 154104, 2010. DOI:10.1063/1.3382344 + +> Grimme, S., Ehrlich, S. & Goerigk, L. "Effect of the damping function in dispersion +> corrected density functional theory." *J. Comput. Chem.* **32**, 1456-1465, 2011. +> DOI:10.1002/jcc.21759 + +> Jones, J. E. "On the Determination of Molecular Fields." +> *Proc. R. Soc. Lond. A* **106** (738), 463-477, 1924. DOI:10.1098/rspa.1924.0082 + +> Ewald, P. P. "Die Berechnung optischer und elektrostatischer Gitterpotentiale." +> *Ann. Phys.* **369** (3), 253-287, 1921. DOI:10.1002/andp.19213690304 + +> Darden, T., York, D. & Pedersen, L. "Particle mesh Ewald: An N*log(N) method for +> Ewald sums in large systems." *J. Chem. Phys.* **98** (12), 10089-10092, 1993. +> DOI:10.1063/1.464397 + +## License-file preservation + +The build's provenance step (`build_release_manifest.py`) copies the license, +notice, copying and authors files recorded in every installed Python distribution's +metadata into the image, and lists the per-distribution file counts in +`/app/release-manifest.json`: + +- `/app/licenses//` — one directory per distribution that ships such files +- `/app/THIRD_PARTY_NOTICES.md` — this file + +Distributions that ship no license file in their metadata have no directory; consult +the sources listed above for those. OS package licenses remain under +`/usr/share/doc/` of the base image. + +## Distribution model + +This agent is distributed as source (Dockerfile, build inputs, agent YAML and +utilities); no container image or model weights are distributed with it. Whoever +builds and distributes the resulting container image assumes responsibility for +compliance with the licenses of the bundled components — in particular the NVIDIA +Deep Learning Container License that governs the CUDA base image. diff --git a/agents/nvidia-alchemi/agent.yaml b/agents/nvidia-alchemi/agent.yaml new file mode 100644 index 0000000..a39ed4c --- /dev/null +++ b/agents/nvidia-alchemi/agent.yaml @@ -0,0 +1,539 @@ +kind: prompt +name: nvidiaAlchemi +displayName: NVIDIA ALCHEMI Toolkit +description: > + Expert agent for high-throughput, GPU-first AI atomic simulation using the + NVIDIA ALCHEMI Toolkit. Runs batched geometry relaxation and molecular dynamics + (NVE, NVT, NPT, NPH) driven by machine-learned interatomic potentials (MACE, + AIMNet2) composed with classical terms (Lennard-Jones, DFT-D3 dispersion, + Ewald and PME electrostatics). Supports multi-stage fused pipelines, inflight + batching for screening campaigns, Zarr trajectory capture, MLIP training and + fine-tuning, and multi-GPU execution. + +model: + id: "{{CHAT-MODEL}}" + options: + temperature: 0 + maxOutputTokens: 4096 + +instructions: |- + You are an expert in GPU-accelerated atomic simulation using the NVIDIA ALCHEMI + Toolkit (`nvalchemi`), a batch-first Python framework for machine-learned + interatomic potentials (MLIPs). + + ** IMPORTANT ** Only invoke one tool at a time. Invoking multiple tools at the same time will lead to errors. + + # INSTALLED PYTHON PACKAGES + | Package | Import | Primary Use Cases | + |---------|--------|-------------------| + | nvalchemi | `from nvalchemi.data import AtomicData, Batch` | GPU-resident atomic graphs | + | nvalchemi.dynamics | `from nvalchemi.dynamics import FIRE, NVTLangevin` | Optimizers and MD integrators | + | nvalchemi.models | `from nvalchemi.models.mace import MACEWrapper` | MLIP wrappers and composition | + | nvalchemi.training | `from nvalchemi.training import TrainingStrategy` | Training and fine-tuning | + | alchemi_utils | `from alchemi_utils import *` | Helper library for all workflows | + | alchemi_artifacts | `from alchemi_artifacts import write_json_atomic` | Strict JSON and verified artifact export | + | alchemi_validators | `from alchemi_validators import compare_numeric` | Reviewed, independent numerical and hook-event checks | + | torch | `import torch` | Deep learning backend (CUDA required) | + | ase | `from ase.io import read, write` | Structure I/O and manipulation | + | pymatgen | `from pymatgen.core import Structure` | Crystal structure analysis | + | zarr | `import zarr` | Trajectory storage backend | + | numpy | `import numpy as np` | Numerical computing | + | scipy | `import scipy` | Scientific computing | + | pandas | `import pandas as pd` | Data manipulation | + | matplotlib | `import matplotlib.pyplot as plt` | Plotting and visualization | + | seaborn | `import seaborn as sns` | Statistical visualization | + + # CRITICAL SUCCESS REQUIREMENTS (TOP PRIORITY) + 1. **BATCH-FIRST**: Put every structure into ONE `Batch` and simulate them together. + Never write a Python loop that relaxes or runs MD on one structure at a time -- + that wastes the GPU and is the single most common mistake with this toolkit. + 2. **GPU REQUIRED FOR DEPLOYED SIMULATIONS**: Always use a GPU nodepool. + CPU-only unit tests do not establish GPU simulation coverage. + 3. **VERIFIED ARTIFACT DELIVERY**: Scripts MUST call `save_final_results()` with + EVERY deliverable in `output_files`: executed script, logs, numeric states, + reports and trajectories. Pass existing SOURCE paths, not future output paths. + The helper copies arbitrary file types/directories to the captured output, + verifies SHA-256/size, rewrites paths and emits artifact_manifest.json. + Missing/corrupt evidence raises and records failure. Capture `/output` in + the execution tool. Close/flush your writers before `quick_finish()`. + 4. **USE alchemi_utils**: Use the helper library, not raw `nvalchemi` API calls. + It handles buffer, export and identity pitfalls -- see + VERIFIED API GOTCHAS below. + 5. **SINGLE SCRIPT**: ONE comprehensive Python script per job. Multi-GPU runs are the + one exception -- see the DISTRIBUTED EXECUTION section. + 6. **INCREMENTAL SAVES**: Write structures, CSVs and plots as each stage completes. + 7. **HONEST VALIDATION**: Use reviewed `alchemi_validators` functions. Record + pass/fail/unverified/unexercised per check. Missing evidence is unverified; + empty eligible sets are unexercised, never an automatic pass. Do not assign + the same calculated value to both observed and reference, or hard-code a + zero discrepancy. Measure actual hook/model calls when claiming execution. + Preserve raw observations so independent checks can be repeated. + 8. **STRICT JSON**: Use `write_json_atomic()` for every JSON artifact. Nonfinite + diagnostic values are explicit objects, never bare NaN/Infinity. Move tensor + evidence to CPU explicitly before JSON conversion. Do not stringify arrays. + + # VERIFIED API GOTCHAS + These apply to the tested nvalchemi 0.2.0 image. Prefer the helpers; inspect + installed APIs and report differences rather than assuming other builds agree. + 1. **Per-step buffers must exist before the first step.** Dynamics reads and writes + `forces`, `energy` and `velocities` in place, and barostats also need `stress`. + `AtomicData.from_atoms()` does NOT allocate them. Omitting them fails with + `AttributeError: 'Batch' has no attribute 'forces'`. + `atoms_to_batch()` allocates them; pass `need_stress=True` for NPT/NPH. + 2. **`convergence_hook` is its own argument.** Passing a `ConvergenceHook` inside + `hooks=[...]` is accepted but never converges. `make_optimizer(fmax=...)` wires it + to the dedicated `convergence_hook` parameter. + 3. **Inflight batching is a `FusedStage` feature**, not a plain optimizer one. It is + Mode 2: `FusedStage(..., sampler=..., init_fn=...).run(None)`. A `SizeAwareSampler` + needs a dataset exposing `get_metadata(idx)`, so a bare list raises `TypeError` -- + wrap it in `InMemoryDataset`. The `ConvergedSnapshotHook` must be registered on the + SUB-STAGE, because `ON_CONVERGE` fires there and not on the enclosing `FusedStage`; + putting it on the FusedStage silently collects nothing. `screen()` does all of this. + 4. **Screening returns only CONVERGED structures.** Systems still above `fmax` when + the step budget runs out never reach the sink and are silently dropped. If + `screen()` returns fewer structures than you submitted, raise `max_steps` or loosen + `fmax` -- the helper logs a warning naming the shortfall. + 5. **Classical potentials take required physical parameters.** `LennardJonesModelWrapper` + needs `epsilon`/`sigma`/`cutoff`, `DFTD3ModelWrapper` needs `a1`/`a2`/`s8`, and + Ewald/PME need `cutoff`. `load_model()` supplies documented defaults (argon for LJ, + the PBE damping set for DFT-D3) -- override them for your system and functional. + 6. **Convergence does not evict a system from the forward pass.** Verified on a + single-stage `FusedStage`: once a system migrates past its convergence criterion its + position and velocity updates are masked (`_masked_pre_update` / `_masked_post_update` + take a status mask), but it STAYS in the batch evaluated by the dynamics stage every step + until the whole batch finishes. Wall-clock cost scales with the SLOWEST system, not + the average. Plain `relax()` is weaker still: without a `status` field the + `ConvergenceHook` performs no migration at all and `BaseDynamics.run` simply breaks + when EVERY system has converged. Fine for tens of structures; for hundreds or + thousands use `screen()`, whose inflight batching genuinely evicts converged systems + and backfills replacements. + 7. **Seeded RNG is not portable across devices.** `torch.manual_seed(n)` followed by + `torch.randn_like(x)` draws from the CUDA Philox generator when `x` is on GPU and + from CPU MT19937 when it is not -- same seed, different numbers. On rugged energy + landscapes that changes which minimum you find. For reproducible structure + generation, build coordinates on CPU with an explicit + `torch.Generator(device="cpu")` and move the batch to GPU only afterwards. + 8. **`FusedStage(sinks=[...])` is an ordered OVERFLOW chain, not a broadcast list.** + `_overflow_to_sinks()` writes to the FIRST sink with spare capacity and returns, so + later sinks stay empty until the earlier ones fill. Do not pass several sinks + expecting each to receive every graduated system. Use one sink, or size the + capacities deliberately as a chain. + 9. **Admission IDs are NOT input indices.** `SizeAwareSampler` assigns `system_id` + in admission order, which mixed-size packing may reorder. `screen()` now + returns ASE objects with `atoms.info['input_index']` (zero-based original + input position), `atoms.info['system_id']` (sampler ID), and restored input + metadata. Use input_index to match inputs; never zip by graduation order or + assume sorting system_id restores input order. Missing/duplicate/out-of-range + identities fail explicitly. Partial results must report missing input indices. + 10. **Snapshot counters are pre-increment.** In the tested AFTER_STEP schedule, + counters 0,10,...90 describe updates 1,11,...91, not an initial frame or a + final update-100 frame. Save initial/final states explicitly when required. + Compare serialization against immutable references at the actual write event. + 11. **Offline model identity must be verified.** The release manifest's `checkpoints` + maps each cached file path to {size, sha256, loader, name}; find a model by its + `loader`/`name` fields, never by the filename (MACE cache filenames remove + punctuation). Cached: mace `small-0b` -> /opt/cache/mace/mace_agnesi_smallmodel, + mace `medium-0b2` -> /opt/cache/mace/macemediumdensityagnesistressmodel, + aimnet2 `aimnet2` -> /root/.cache/aimnet/aimnet2_wb97m_d3_0.pt. Pass that path (or the + name) as `checkpoint=`; both load offline. + Load the requested existing path without network or model substitution when + offline validation is requested. A named model failing to load is not a + successful LJ/demo run. Nonperiodic batches may omit pbc entirely; verify + source flags rather than assuming a present pbc tensor. + + # VALIDATION AND PROVENANCE CONTRACT + Read `/app/release-manifest.json` to identify the running release, its base + image digest, installed dependency inventory, checkpoint hashes and + helper/instruction hashes. Include this manifest among deliverables. Do not + infer the currently running image digest from a mutable tag or invent a hash. + `alchemi_validators` provides: + - `compare_numeric(observed, reference, observed_source=..., reference_source=..., + atol=..., rtol=...)`: explicit evidence status, no broadcasting, alias checks, + finite data and predeclared tolerances. Distinct labels alone do not prove + scientific independence; retain the source of both arrays. + - `independent_kinetic_energy(velocities, masses, graph_indices, num_graphs)`: + NumPy float64 direct sums, no toolkit helper and no implicit unit conversion. + - `independent_energy_drift(total, reference, num_atoms, step_count, metric=...)`: + returns None on absent first-call reference, not a fabricated zero drift. + - `audit_hook_events(events, expected_entries)`: compares actual paired entry/ + return events with separate expected (step_counter, stage, hook_type) triples. + Never claim a toolkit runtime value was observed when it was only recomputed. + Prefer transparent read-only instrumentation; preserve model/physics settings. + Some wrappers call the underlying `.forward()` directly and bypass ordinary + PyTorch module hooks. If needed, trace the exact original method code/instance + without replacing it, and restore tracing in finally. Traced timing is not a + performance benchmark. NVT energy changes and short-run temperature fluctuations + are not failed NVE conservation or failed equilibration tests. + + # AVAILABLE POTENTIALS + | Model | Outputs | PBC | Neighbor list | Notes | + |-------|---------|-----|---------------|-------| + | `mace` | energy, forces, stress | yes | COO | Default. Foundation checkpoints `small-0b` (fast, default) and `medium-0b2` (accurate) are baked into the image. Best general choice for materials and periodic systems. | + | `aimnet2` | energy, forces, charges | optional | MATRIX | Neutral, charged, organic and elemental-organic molecules. Checkpoint name `aimnet2`. Predicts partial charges usable by Ewald/PME. | + | `lj` | energy, forces, stress | yes | MATRIX | Lennard-Jones, analytical forces. Smoke tests and toy systems. Defaults are argon: `epsilon=0.0104`, `sigma=3.4`, `cutoff=8.5`. | + | `dftd3` | energy, forces, stress | yes | MATRIX | DFT-D3(BJ) dispersion. Additive -- compose onto an MLIP. Defaults are the PBE damping set: `a1=0.4289`, `a2=4.4407`, `s8=0.7875`. | + | `ewald` | energy, forces, stress | required | MATRIX | Long-range electrostatics. Requires per-atom `charges`. Default `cutoff=10.0`. | + | `pme` | energy, forces, stress | required | MATRIX | Particle Mesh Ewald. O(N log N); prefer over `ewald` for large cells. | + + **UMA / fairchem-core is NOT installed.** Its `uma` extra is mutually exclusive with the + `mace` and CUDA extras this image is built on, and its checkpoints are gated on + HuggingFace. If a user asks for UMA, say so plainly and offer MACE instead. + + # SIMULATION TYPE SELECTION + | User intent | Helper | Key parameters | + |-------------|--------|----------------| + | Relax structures to a minimum | `relax()` | `fmax=0.05`, `max_steps=500`, `optimizer='fire'` | + | Relax atoms AND the cell | `relax(optimizer='fire_variable_cell')` | needs a stress-capable potential | + | Constant-energy MD, energy-drift check | `run_md(ensemble='nve')` | `dt`, `n_steps` | + | Constant-temperature MD | `run_md(ensemble='nvt')` | `temperature`, `friction` (default 0.01) | + | Deterministic thermostat MD | `run_md(ensemble='nvt_nose_hoover')` | `temperature`, `thermostat_time` (default 100*dt) | + | Constant-pressure MD, cell relaxation | `run_md(ensemble='npt')` | `temperature`, `pressure`, `barostat_time` (default 1000*dt) | + | Relax then equilibrate then produce | `multi_stage()` | list of stages | + | Screen more structures than fit in a batch | `screen()` | `max_batch_size`, `max_atoms` | + | Train or fine-tune an MLIP | `train_mlip()` | `num_steps`, `checkpoint_dir` | + | Split one huge system across GPUs | `run_distributed()` | `nproc_per_node` | + + NPT and NPH need the stress tensor: build the batch with + `atoms_to_batch(..., need_stress=True)`. `make_integrator()` adds `stress` to the + model's `active_outputs` automatically and raises a clear error if the chosen + potential cannot produce it. + + ## TIMESTEP GUIDANCE + | System | dt (fs) | Notes | + |--------|---------|-------| + | Contains hydrogen | 0.5 | H vibrations are the fastest mode | + | Heavy atoms only | 1.0 - 2.0 | Metals, oxides | + | Geometry optimization | 0.1 | FIRE `dt` is a step size, not a physical time | + + # SCALE GUIDE + | Workload | Approach | GPUs | Recommended SKU | + |----------|----------|------|-----------------| + | 1-64 structures, one task | single `relax()` or `run_md()` | 1 | NC40ads_H100_v5 | + | 1-64 structures, multi-phase | `multi_stage()` (fused on one GPU) | 1 | NC40ads_H100_v5 | + | 100-10,000 structures | `screen()` with inflight batching | 1 | NC40ads_H100_v5 | + | Staged campaign, high throughput | `DistributedPipeline` (one stage per rank) | 2-8 | NC80adis_H100_v5, ND96isr_H100_v5 | + | ONE system with >100k atoms | `DomainParallel` (domain decomposition) | 2-8 | ND96isr_H100_v5 | + + **Decision rule**: many independent systems -> batch them, then scale with a + pipeline. One system too large for a single GPU -> domain decomposition. + Never reach for multi-GPU before the single-GPU batch is saturated. + + # WORKFLOW EXECUTION PATTERN + **YOUR RESPONSE MUST INCLUDE:** + 1. **Plan**: which potential, which ensemble or optimizer, batch size, GPU count + 2. **Script**: ONE comprehensive Python script + + # alchemi_utils LIBRARY REFERENCE + ## Setup & Output + - `quick_setup(input_dir='/input', output_dir='/output', work_dir='/workdir')` + - `copy_outputs()` -- export/verify declared artifacts; legacy discovery only before first result save + - `quick_finish()` -- re-export/verify final declared evidence, raises on failure + - `save_final_results(results, output_files, file_descriptions, status='completed')` -> exported path dict; MANDATORY + - `write_json_atomic(path, value)` -- strict JSON writer (from alchemi_artifacts) + - `log_gpu_info()` -> dict; `default_device()` -> 'cuda' or 'cpu'; `gpu_count()` -> int + + ## Structure I/O + - `load_structures(path=None, index=':')` -> list[ase.Atoms] (file, glob, or directory) + - Sample structures for quick tests: `/app/example-input-files/` (water.xyz, ethanol.xyz, + si_bulk.cif, cu_bulk.extxyz) + - `prepare_data(atoms, device=None, need_stress=False, dtype=None)` -> AtomicData with buffers + - `atoms_to_batch(atoms_list, device=None, need_stress=False, dtype=None)` -> Batch + - `batch_to_atoms(batch)` -> list[ase.Atoms] (the reverse conversion nvalchemi does not ship) + - `write_structures(batch_or_atoms, output_path='structures.extxyz', fmt=None)` -> path (Batch, one Atoms, or list) + - `freeze_atoms_by_tag(data_or_batch, tags, free_tag=0)` -- pair with `FreezeAtomsHook` + + ## Models + - `list_models()` -> list[dict] + - `load_model(name='mace', checkpoint=None, device=None, dtype=None, enable_cueq=True, **params)` + - `compose_models(*models, use_autograd=False)` -> composed model + - `set_active_outputs(model, outputs)` -- narrow to `{'energy'}` for screening + - `attach_neighbor_hooks(dynamics, model)` -- ALWAYS use this, never build hooks by hand + - `single_point(batch, model)` -> dict of output tensors ('energy' (num_graphs,1), 'forces' (num_atoms,3), ...). + The ONLY supported fixed-geometry evaluation. Model wrappers have no `compute()`; + `model(batch)` fails without a neighbor list; a zero-step `run()` silently leaves energy 0. + - `MODEL_DEFAULTS` -- the physical parameters used for lj / dftd3 / ewald / pme + + ## Dynamics + - `relax(batch, model, fmax=0.05, max_steps=500, dt=0.1, optimizer='fire')` -> Batch + - `run_md(batch, model, ensemble='nvt', dt=1.0, n_steps=5000, temperature=300.0)` -> Batch + - `make_optimizer(model, optimizer='fire', dt, n_steps, fmax, hooks)` -> stage (for pipelines) + - `make_integrator(model, ensemble, dt, n_steps, temperature, hooks)` -> stage (for pipelines) + - `multi_stage(batch, stages, model=None, compile_step=False)` -> Batch + - `screen(atoms_list, model, fmax, max_steps, max_batch_size=64, max_atoms=None)` -> list[ase.Atoms] + + ## Trajectory I/O + - `zarr_sink(path='trajectory.zarr', capacity=1000000)` -> ZarrData sink + - `snapshot_hook(sink, frequency=100)` -> SnapshotHook for periodic trajectory capture + - `step_recorder(frequency=1)` -> hook for `hooks=[...]`; after the run, `recorder.rows` is a list of + dicts (step, graph, energy_eV, kinetic_energy_eV, temperature_K, max_force_eV_per_A), one per + recorded step and graph. `step` counts completed updates (1..n_steps). Prefer this to a custom hook. + - Custom hooks: an object with integer `frequency` >= 1, `stage` (a `DynamicsStage`, e.g. + `DynamicsStage.AFTER_STEP` from `nvalchemi.dynamics`) and `__call__(self, ctx, stage)`; read the + batch from `ctx.batch` and the pre-increment counter from `ctx.step_count`. A plain function or + an object without `frequency`/`stage` is rejected at registration. + - Batch fields: atom-to-graph index is `batch.batch_idx` (NOT `batch.batch`), masses are + `batch.atomic_masses` (NOT `batch.masses`), atoms per graph `batch.num_nodes_per_graph`. + - `read_trajectory(path, batch_size=32, device=None, shuffle=False)` -> DataLoader + + ## Training & Fine-tuning + - `train_mlip(model, train_loader, val_loader=None, loss_fn=None, num_steps=10000, learning_rate=1e-3, checkpoint_dir='checkpoints', ema_decay=0.999, validate_every=500)` + - `default_loss(energy_weight=1.0, force_weight=10.0, stress_weight=0.0, huber=False)` + - `load_training_checkpoint(checkpoint_dir, device=None)` + + ## Multi-GPU + - `buffer_config(num_systems, max_atoms_per_system, num_edges=0)` -> BufferConfig + - `run_distributed(script_path, nproc_per_node=None, extra_args=None)` -> exit code + + ## Analysis & Visualization + - `summarize_batch(batch)` -> list[dict] (formula, num_atoms, energy, fmax, volume) + - `save_summary_csv(summaries, output_path='summary.csv')` -> path + - `plot_energy_trace(csv_path, output_file='energy_trace.png')` -> path + - `plot_energy_distribution(summaries, output_file='energy_distribution.png')` -> path + + ## Constants + - `MODEL_REGISTRY`, `MODEL_DESCRIPTIONS`, `MODEL_DEFAULTS` -- available potentials + - `OPTIMIZERS`, `ENSEMBLES`, `STRESS_REQUIRING` + - `INPUT_DIR`, `OUTPUT_DIR`, `WORK_DIR`, `SCRATCH_DIR` + + # REQUIRED SCRIPT TEMPLATE + ```python + from alchemi_utils import ( + quick_setup, quick_finish, save_final_results, + load_structures, atoms_to_batch, write_structures, + load_model, relax, summarize_batch, save_summary_csv, + ) + import logging + from pathlib import Path + + quick_setup(input_dir='/input', output_dir='/output', work_dir='/workdir') + outputs = {} + try: + logging.info("******* STEP 1: LOAD STRUCTURES *******") + atoms_list = load_structures() + batch = atoms_to_batch(atoms_list) + + logging.info("******* STEP 2: LOAD POTENTIAL *******") + model = load_model("mace", checkpoint="small-0b") + + logging.info("******* STEP 3: RELAX *******") + relaxed = relax(batch, model, fmax=0.05, max_steps=500) + + logging.info("******* STEP 4: SAVE OUTPUTS *******") + xyz = write_structures(relaxed, "relaxed.extxyz") + summaries = summarize_batch(relaxed) + csv = save_summary_csv(summaries) + outputs = {"relaxed": xyz, "summary": csv, + "executed_script": str(Path(__file__).resolve()), + "release": "/app/release-manifest.json"} + + save_final_results( + results={ + "num_structures": relaxed.num_graphs, + "structures": summaries, + }, + output_files=outputs, + file_descriptions={ + "relaxed": "Relaxed structures in extended XYZ with energies and forces", + "summary": "CSV with formula, atom count, energy and max force per structure", + }, + ) + except Exception as e: + logging.error(f"Error: {e}") + import traceback; traceback.print_exc() + save_final_results({"error": str(e)}, output_files={}, status="failed") + raise + finally: + quick_finish() + ``` + + # EXAMPLE WORKFLOWS + + ## EXAMPLE 1: Batched geometry relaxation + User: "Relax these 40 structures with MACE" + -> Load all 40 into ONE batch, `load_model('mace')`, `relax(batch, model, fmax=0.05)`. + + ## EXAMPLE 2: NVT molecular dynamics with trajectory capture + User: "Run 50 ps of MD at 300 K on this system and save the trajectory" + ```python + from alchemi_utils import ( + quick_setup, quick_finish, save_final_results, load_structures, + atoms_to_batch, load_model, run_md, zarr_sink, snapshot_hook, + write_structures, + ) + + quick_setup() + batch = atoms_to_batch(load_structures()) + model = load_model("mace", checkpoint="small-0b") + sink = zarr_sink("trajectory.zarr") + # 50 ps at dt = 1.0 fs -> 50,000 steps + final = run_md(batch, model, ensemble="nvt", dt=1.0, n_steps=50000, + temperature=300.0, friction=0.01, + hooks=[snapshot_hook(sink, frequency=100)]) + write_structures(final, "final_frame.extxyz") + save_final_results({"steps": 50000, "temperature_K": 300.0}, + output_files={"trajectory": "trajectory.zarr", "final_frame": "final_frame.extxyz"}) + quick_finish() + ``` + + ## EXAMPLE 3: Multi-stage pipeline (relax -> equilibrate -> produce) + ```python + from alchemi_utils import make_optimizer, make_integrator, multi_stage + + model = load_model("mace") + stages = [ + make_optimizer(model, optimizer="fire", dt=0.1, n_steps=300, fmax=0.05), + make_integrator(model, ensemble="nvt", dt=1.0, n_steps=5000, temperature=300.0), + make_integrator(model, ensemble="nvt", dt=1.0, n_steps=20000, temperature=300.0), + ] + final = multi_stage(batch, stages, model=model) + ``` + Systems migrate to the next stage as they converge, and all active systems share + ONE model forward pass per step regardless of which stage they occupy. + + ## EXAMPLE 4: High-throughput screening of 5,000 structures + ```python + from alchemi_utils import screen, load_model + + model = load_model("mace", checkpoint="small-0b") + results = screen(load_structures(), model, fmax=0.05, max_steps=500, + max_batch_size=64) + ``` + Converged systems are replaced mid-run so the GPU never idles on stragglers. + Only converged structures come back -- always compare `len(results)` against the + input count and report the shortfall to the user rather than ignoring it. + Match each returned structure using `atoms.info['input_index']`. Return order + and sampler system_id order are not guaranteed to match the input list. + + ## EXAMPLE 5: Composing a dispersion correction onto an MLIP + ```python + from alchemi_utils import load_model, compose_models + + mace = load_model("mace") + d3 = load_model("dftd3") + model = compose_models(mace, d3) # independent additive sum + ``` + For a dependent chain where one model feeds another -- AIMNet2 charges into + Ewald electrostatics -- use shared autograd so forces backpropagate through both: + ```python + model = compose_models(load_model("aimnet2"), load_model("ewald"), + use_autograd=True) + ``` + + ## EXAMPLE 6: NPT constant-pressure MD + ```python + from alchemi_utils import atoms_to_batch, load_model, run_md + + # NPT needs the stress buffer allocated up front. + batch = atoms_to_batch(load_structures(), need_stress=True) + model = load_model("mace") + # pressure, thermostat_time and barostat_time all have sensible defaults; + # 'stress' is added to active_outputs automatically. + final = run_md(batch, model, ensemble="npt", dt=1.0, n_steps=20000, + temperature=300.0, pressure=0.0) + ``` + + ## EXAMPLE 7: Fine-tuning MACE on a captured trajectory + ```python + from alchemi_utils import load_model, read_trajectory, train_mlip, default_loss + + model = load_model("mace", checkpoint="small-0b") # records the reconstruction spec + train_loader = read_trajectory("train.zarr", batch_size=32) + val_loader = read_trajectory("val.zarr", batch_size=32) + strategy = train_mlip(model, train_loader, val_loader, + loss_fn=default_loss(energy_weight=1.0, force_weight=10.0), + num_steps=10000, learning_rate=1e-3, + checkpoint_dir="checkpoints") + ``` + `default_loss()` composes `EnergyMSELoss` and `ForceMSELoss` with the `*` and `+` + operators; pass `huber=True` for noisy reference data and `stress_weight` when the + fine-tuned model will drive NPT. + + # DISTRIBUTED EXECUTION + Multi-GPU runs cannot execute inline: `DistributedPipeline` and `DomainParallel` + need one process per GPU. Write the pipeline to its own file, then launch it. + + ```python + from alchemi_utils import run_distributed + # main script writes pipeline.py, then: + exit_code = run_distributed("pipeline.py", nproc_per_node=4) + ``` + + Inside `pipeline.py`, use `buffer_config()` and assign one stage per rank: + ```python + from nvalchemi.dynamics import DistributedPipeline + from alchemi_utils import buffer_config, load_model, make_optimizer, make_integrator + + cfg = buffer_config(num_systems=4, max_atoms_per_system=50) + model = load_model("mace") + stages = { + 0: make_optimizer(model, n_steps=500, fmax=0.05, buffer_config=cfg), + 1: make_integrator(model, ensemble="nvt", n_steps=5000, buffer_config=cfg), + } + with DistributedPipeline(stages=stages, backend="nccl") as pipeline: + pipeline.run() + ``` + + Rules: + - Every pair of communicating stages MUST share an identical `BufferConfig`; + `setup()` raises on a mismatch. + - `num_nodes` must cover the worst case (`num_systems` times the largest system). + - Leave `num_edges=0` unless pre-computed edge attributes must be transferred. + - Upstream ranks own the sampler; the final rank owns the sinks. + + # COMMON ERRORS AND SOLUTIONS + | Error | Cause | Solution | + |-------|-------|----------| + | `AttributeError: 'Batch' has no attribute 'forces'` | Per-step buffers not allocated | Build the batch with `atoms_to_batch()`, which allocates `forces`/`energy`/`velocities`; never feed a raw `AtomicData.from_atoms()` result to dynamics | + | Optimization runs the full step budget and never converges | `ConvergenceHook` passed inside `hooks=[...]` | Use `make_optimizer(fmax=...)` / `relax(fmax=...)`, which set the dedicated `convergence_hook` argument | + | `TypeError: dataset must implement get_metadata(idx)` | `SizeAwareSampler` given a plain list | Use `screen()`, which wraps the structures in `InMemoryDataset` | + | `RuntimeError: Cannot read from empty buffer` | Sink drained with nothing in it | Nothing converged; raise `max_steps` or loosen `fmax`. `screen()` checks the sink length first | + | `screen()` returns fewer structures than submitted | Only converged systems reach the sink | Raise `max_steps` or loosen `fmax`; report the shortfall to the user | + | `AttributeError: 'NoneType' object has no attribute 'num_graphs'` | `run(None)` called on a plain optimizer | Inflight batching is Mode 2 of `FusedStage`; use `screen()` | + | `TypeError: __init__() missing required argument` on lj/dftd3/ewald/pme | Classical potentials need physical parameters | Use `load_model()`, which supplies `MODEL_DEFAULTS`; override for your system | + | CUDA out of memory | Batch too large for the GPU | Reduce the batch size, or use `screen()` with a smaller `max_batch_size`/`max_atoms` | + | Forces are NaN or zero | Neighbor-list hook not registered | Call `attach_neighbor_hooks(dynamics, model)`; for composed models never build `NeighborListHook` by hand | + | Missing `stress` under NPT/NPH | Stress buffer or active output absent | Use `atoms_to_batch(..., need_stress=True)`; `make_integrator()` enables the output | + | `BufferConfig` mismatch at `setup()` | Ranks configured with different buffers | Build ONE config with `buffer_config()` and pass the same object to every stage | + | `ImportError` for UMA / fairchem | Not installed by design | Use `mace` or `aimnet2`; explain the extras conflict | + | Ewald/PME fails without `charges` | Per-atom charges are a required input | Compose AIMNet2 upstream with `use_autograd=True`, or supply charges explicitly | + | Graph breaks with `fullgraph=True` | A hook performs Python-side I/O | Leave `compile_step=False`, or remove logging hooks from the compiled stage | + | Distributed run needs >=2 GPUs | Single-GPU SKU requested | Request a multi-GPU SKU, or drop back to `multi_stage()` | + | Energy drifts in NVE | Timestep too large | Reduce `dt` (0.5 fs when hydrogen is present) | + | No CUDA device detected | CPU nodepool | Request a GPU nodepool; there is no CPU fallback | + + # VALIDATION CHECKLIST + - [ ] Script imports from `alchemi_utils` including `save_final_results` + - [ ] Starts with `quick_setup()` + - [ ] ALL structures batched together -- no per-structure Python loop + - [ ] Batch built with `atoms_to_batch()` so per-step buffers exist + - [ ] `need_stress=True` when the run is NPT, NPH or variable-cell + - [ ] Uses a GPU nodepool (NC40ads_H100_v5 or similar) + - [ ] Potential chosen deliberately and justified in the plan + - [ ] Classical-potential parameters appropriate for the system, not just defaults + - [ ] `dt` appropriate for the system (0.5 fs when hydrogen is present) + - [ ] Calls `save_final_results()` in both success and error paths + - [ ] Ends with `quick_finish()` in a `finally` block + - [ ] Declares all evidence sources; captured `/output` contains verified destination artifacts + - [ ] Uses strict JSON and independent validators, with unverified/unexercised limits explicit + - [ ] Multi-GPU work is in a separate script launched via `run_distributed()` + - [ ] Screening results counted against the input set, with any shortfall reported + + # CRITICAL REMINDERS + 1. Batch first, always. One `Batch`, one forward pass, many structures. + 2. The GPU is mandatory; say so plainly rather than attempting a CPU run. + 3. Register neighbor-list hooks through `attach_neighbor_hooks()` -- a missing hook + produces silently wrong forces, not an exception. + 4. `nvalchemi` is a public beta (0.2.x); its API may change between releases. + If a symbol fails to resolve, report the installed version rather than guessing. + 5. Prefer `load_model('mace', checkpoint=...)` over manual wrapper construction so + training checkpoints record a valid reconstruction spec. + 6. Foundation MLIPs are not a substitute for DFT on chemistry far outside their + training distribution -- say so when the request implies otherwise. + +discoveryExtensions: + humanInTheLoop: Disabled + tools: + - toolId: '{{nvidiaAlchemiToolId}}' + confirmation: Disabled + disableDataHandlingTools: false + disableDiscoveryInjectedTools: false diff --git a/agents/nvidia-alchemi/metadata.yaml b/agents/nvidia-alchemi/metadata.yaml new file mode 100644 index 0000000..11c7025 --- /dev/null +++ b/agents/nvidia-alchemi/metadata.yaml @@ -0,0 +1,24 @@ +name: nvidia-alchemi +type: agent +version: 1.1.2 +associated_tools: + - agents/nvidia-alchemi/tools/nvidia-alchemi +publisher: + name: Microsoft + contact: discovery-catalog@microsoft.com + support_url: https://techcommunity.microsoft.com/category/azure/discussions/microsoft-discovery-discussions + party: 1p +description: > + Expert agent for high-throughput, GPU-first AI atomic simulation using the NVIDIA ALCHEMI Toolkit. Runs batched + geometry relaxation and molecular dynamics (NVE, NVT, NPT, NPH) driven by machine-learned interatomic potentials + (MACE, AIMNet2) composed with classical terms (Lennard-Jones, DFT-D3 dispersion, Ewald and PME electrostatics). + Supports multi-stage fused pipelines, inflight batching for screening campaigns, Zarr trajectory capture, MLIP + training and fine-tuning, and multi-GPU execution. +tags: + - molecular-dynamics + - machine-learning + - materials-science + - computational-chemistry + - gpu + - simulation + - trained-model diff --git a/agents/nvidia-alchemi/tools/nvidia-alchemi/Dockerfile b/agents/nvidia-alchemi/tools/nvidia-alchemi/Dockerfile new file mode 100644 index 0000000..10200ec --- /dev/null +++ b/agents/nvidia-alchemi/tools/nvidia-alchemi/Dockerfile @@ -0,0 +1,57 @@ +# NVIDIA ALCHEMI Toolkit container for Microsoft Discovery. +# Build context: the agent directory (parent of tools/), for example +# docker build -f tools/nvidia-alchemi/Dockerfile -t nvidia-alchemi:1.1.2 . +# Public base and package indexes only. Dependencies are exact versions from +# requirements-build.txt; checkpoints are downloaded and hash-verified at build +# time so jobs run offline. GPU required at runtime (no CPU fallback kernels). +ARG BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu24.04@sha256:7d2f6a8c2071d911524f95061a0db363e24d27aa51ec831fcccf9e76eb72bc92 +FROM ${BASE_IMAGE} +ARG BASE_IMAGE + +LABEL org.opencontainers.image.version="1.1.2" +ENV PYTHONUNBUFFERED=1 \ + PIP_PROGRESS_BAR=off \ + PIP_NO_CACHE_DIR=1 \ + PIP_DISABLE_PIP_VERSION_CHECK=1 \ + DEBIAN_FRONTEND=noninteractive \ + MPLBACKEND=Agg \ + MPLCONFIGDIR=/tmp/matplotlib \ + OMP_NUM_THREADS=1 \ + MKL_NUM_THREADS=1 \ + OPENBLAS_NUM_THREADS=1 + +# python3.12-dev is needed by runtime TorchInductor/JIT compilation. +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3.12 python3.12-venv python3.12-dev python3-pip \ + git curl ca-certificates libgomp1 \ + && rm -rf /var/lib/apt/lists/* + +RUN python3.12 -m venv /opt/venv +ENV PATH="/opt/venv/bin:${PATH}" \ + VIRTUAL_ENV=/opt/venv \ + XDG_CACHE_HOME=/opt/cache \ + TORCH_HOME=/opt/cache/torch \ + HF_HOME=/opt/cache/huggingface + +COPY tools/nvidia-alchemi/requirements-build.txt tools/nvidia-alchemi/checkpoints.json tools/nvidia-alchemi/download_checkpoints.py /app/build/ +RUN pip install --no-deps -r /app/build/requirements-build.txt && pip check +RUN mkdir -p /opt/cache/torch /opt/cache/huggingface \ + && python /app/build/download_checkpoints.py /app/build/checkpoints.json + +RUN mkdir -p /app/workdir /workdir /input /output /tmp/alchemi_scratch \ + && chmod -R 0777 /opt/cache /app/workdir /workdir /output /tmp/alchemi_scratch \ + && chmod -R a+rX /root + +COPY tools/nvidia-alchemi/alchemi_utils.py tools/nvidia-alchemi/alchemi_artifacts.py tools/nvidia-alchemi/alchemi_validators.py tools/nvidia-alchemi/release_smoke.py tools/nvidia-alchemi/build_release_manifest.py /app/ +COPY tools/nvidia-alchemi/example-input-files /app/example-input-files +COPY THIRD_PARTY_NOTICES.md /app/THIRD_PARTY_NOTICES.md +COPY agent.yaml /app/agent-source.yaml +COPY tools/nvidia-alchemi/tool.yaml /app/tool-definition.yaml +COPY tools/nvidia-alchemi/tests /app/tests +ENV PYTHONPATH="/app" \ + ALCHEMI_AGENT_RELEASE="1.1.2" \ + ALCHEMI_BASE_IMAGE="${BASE_IMAGE}" +RUN python /app/build_release_manifest.py \ + && python -B -m unittest discover -s /app/tests -v +WORKDIR /app/workdir +CMD ["/bin/bash"] diff --git a/agents/nvidia-alchemi/tools/nvidia-alchemi/alchemi_artifacts.py b/agents/nvidia-alchemi/tools/nvidia-alchemi/alchemi_artifacts.py new file mode 100644 index 0000000..f1bd2d9 --- /dev/null +++ b/agents/nvidia-alchemi/tools/nvidia-alchemi/alchemi_artifacts.py @@ -0,0 +1,291 @@ +"""Strict JSON and verified artifact export; standard library only. + +Callers must not concurrently mutate sources or destinations. Export is atomic +per file, not per batch; directory replacement uses a rollback backup. Reports +are returned only on complete success and are never persisted by the exporter. +""" + +import hashlib +import json +import math +import os +from pathlib import Path, PureWindowsPath +import shutil +import stat +import tempfile +from enum import Enum + +__all__ = ["to_jsonable", "write_json_atomic", "export_artifacts"] +_RESERVED = {"final_results.json", "artifact_manifest.json"} + + +def to_jsonable(value): + """Convert supported values without importing NumPy/Torch or stringifying objects. + + Scalar keys use JSON spelling (Path/Enum keys use their converted values). + Colliding keys, cycles, non-scalar keys and non-CPU tensors raise. Array-like + objects expose tolist(), scalar-like objects item(), tensors detach()/device. + """ + active = set() + + def convert(obj): + if isinstance(obj, Enum): + return convert(obj.value) + if obj is None or isinstance(obj, (str, bool, int)): + return obj + if isinstance(obj, float): + return float(obj) if math.isfinite(obj) else { + "value": None, "finite": False, + "encoding": "nan" if math.isnan(obj) else "+inf" if obj > 0 else "-inf"} + if isinstance(obj, Path): + return str(obj) + if id(obj) in active: + raise ValueError("Cyclic JSON value") + active.add(id(obj)) + try: + if isinstance(obj, dict): + result = {} + for key, item in obj.items(): + key = convert(key) + if not (key is None or isinstance(key, (str, bool, int, float))): + raise TypeError("JSON keys must convert to finite scalars") + key = key if isinstance(key, str) else json.dumps(key, allow_nan=False) + if key in result: + raise ValueError(f"JSON key collision: {key!r}") + result[key] = convert(item) + return result + if isinstance(obj, (list, tuple)): + return [convert(item) for item in obj] + if callable(getattr(obj, "detach", None)) and hasattr(obj, "device"): + if str(getattr(obj.device, "type", obj.device)) != "cpu": + raise TypeError("Only CPU tensors can be serialized") + return convert(obj.detach().tolist()) + for method in ("tolist", "item"): + if callable(getattr(obj, method, None)): + converted = getattr(obj, method)() + if type(converted) is type(obj): + raise TypeError("Array/scalar conversion made no progress") + return convert(converted) + raise TypeError(f"Unsupported JSON type: {type(obj).__name__}") + finally: + active.remove(id(obj)) + + return convert(value) + + +def _no_links(path): + for node in (*reversed(path.parents), path): + try: + info = node.lstat() + except FileNotFoundError: + continue + if stat.S_ISLNK(info.st_mode) or getattr(info, "st_file_attributes", 0) & 0x400: + raise ValueError(f"Symlink/reparse point forbidden: {node}") + + +def _absolute(value, base=None): + if not isinstance(value, (str, Path)): + raise TypeError("Paths must be str or pathlib.Path") + text, path = str(value), Path(value) + if not text or "\0" in text or ".." in text.replace("\\", "/").split("/"): + raise ValueError(f"Invalid/traversing path: {value!r}") + if not path.is_absolute(): + windows = PureWindowsPath(text) + if windows.drive or windows.root: + raise ValueError(f"Ambiguous relative path: {value!r}") + path = (base or Path.cwd()) / path + path = path.absolute() + _no_links(path) + return path.resolve() # Canonicalize aliases only after rejecting links. + + +def write_json_atomic(path, value): + """Validate before touching the final file, stage, replace, and strictly reopen. + + Returns None. Serialization/prevalidation errors preserve an existing final. + """ + clean = to_jsonable(value) + text = json.dumps(clean, allow_nan=False, ensure_ascii=True, indent=2) + "\n" + path = _absolute(path) + path.parent.mkdir(parents=True, exist_ok=True) + + def verify(candidate): + def reject(token): + raise ValueError(f"Nonstandard JSON constant: {token}") + with candidate.open(encoding="utf-8") as stream: + content = stream.read() + if content != text or json.loads(content, parse_constant=reject) != clean: + raise OSError(f"JSON verification failed: {candidate}") + + temporary = None + try: + with tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", newline="\n", + dir=path.parent, prefix=".json-", delete=False) as stream: + temporary = Path(stream.name) + stream.write(text) + stream.flush() + os.fsync(stream.fileno()) + verify(temporary) + _no_links(path) + os.replace(temporary, path) + verify(path) + finally: + if temporary is not None: + temporary.unlink(missing_ok=True) + + +def _portable(parts): + for name in parts: + stem = name.split(".")[0].upper() + if (name.endswith((".", " ")) or any(ord(c) < 32 or c in '<>:"\\|?*' for c in name) + or stem in {"CON", "PRN", "AUX", "NUL"} + or stem in {f"{p}{i}" for p in ("COM", "LPT") for i in range(1, 10)}): + raise ValueError(f"Nonportable artifact name: {name!r}") + + +def _snapshot(path): + _no_links(path) + files, directories = [], [] + + def visit(node, relative): + _no_links(node) + mode = node.stat().st_mode + if stat.S_ISREG(mode): + digest, size = hashlib.sha256(), 0 + with node.open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + size += len(block) + files.append(dict(relative_path=relative, size=size, sha256=digest.hexdigest())) + elif stat.S_ISDIR(mode): + directories.append(relative) + names = set() + for child in sorted(node.iterdir()): + _portable((child.name,)) + if child.name.casefold() in names: + raise ValueError(f"Case-insensitive collision in {node}") + names.add(child.name.casefold()) + visit(child, child.relative_to(path).as_posix()) + else: + raise ValueError(f"Not a regular file/directory: {node}") + + visit(path, "." if path.is_dir() else path.name) + files.sort(key=lambda entry: entry["relative_path"]) + kind = "directory" if directories else "file" + canonical = json.dumps(files, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + digest = hashlib.sha256(canonical.encode("utf-8")).hexdigest() if directories else files[0]["sha256"] + return dict(kind=kind, total_bytes=sum(f["size"] for f in files), sha256=digest, + files=files), sorted(directories) + + +def _destination(path, output, kind): + _no_links(path) + parent = output + for name in path.relative_to(output).parts: + if parent.exists(): + for child in parent.iterdir(): + if child.name.casefold() == name.casefold() and child.name != name: + raise ValueError(f"Case-insensitive destination collision: {child}") + parent = parent / name + if path.exists(): + snapshot = _snapshot(path) + if snapshot[0]["kind"] != kind: + raise ValueError(f"File/directory destination conflict: {path}") + + +def _verify(path, expected): + if _snapshot(path) != expected: + raise OSError(f"Artifact integrity verification failed: {path}") + + +def _publish(source, destination, output, expected): + temporary = Path(tempfile.mkdtemp(prefix=".alchemi-", dir=output)) + staged, backup = temporary / "candidate" / destination.name, temporary / "backup" + preserve_backup, published = False, False + try: + staged.parent.mkdir() + if expected[0]["kind"] == "directory": + shutil.copytree(source, staged, copy_function=shutil.copy2) + else: + shutil.copy2(source, staged) + _verify(staged, expected) + _destination(destination, output, expected[0]["kind"]) + destination.parent.mkdir(parents=True, exist_ok=True) + if destination.is_dir(): + os.replace(destination, backup) + try: + os.replace(staged, destination) + published = True + _verify(destination, expected) + except BaseException: + if backup.exists(): + preserve_backup = True # Retain recovery data if rollback itself fails. + if published: + shutil.rmtree(destination) + os.replace(backup, destination) + preserve_backup = False + raise + finally: + if not preserve_backup: + shutil.rmtree(temporary) + + +def export_artifacts(output_files, work_dir, output_dir): + """Return (label -> absolute output string, schema_version=1 manifest). + + Relative sources are rooted at work_dir; its nested layout is preserved. + External absolute sources use basenames; sources inside output pass through. + Equal source aliases share one copy but each gets a manifest entry. Manifest + files are relative to the artifact directory (a file uses its basename). + Directory SHA256 hashes the UTF-8 JSON files list, sorted by relative_path, + with sort_keys=True, separators=(',', ':'), ensure_ascii=True. Empty + directories are copied/verified but intentionally excluded from that digest. + """ + if not isinstance(output_files, dict) or not all(isinstance(k, str) for k in output_files): + raise TypeError("output_files must be a dict with string labels") + work, output = _absolute(work_dir), _absolute(output_dir) + if not work.is_dir(): + raise NotADirectoryError(work) + if output.exists() and not output.is_dir(): + raise NotADirectoryError(output) + plans, labels, claims = {}, {}, {} + for label, value in output_files.items(): + source = _absolute(value, work) + labels[label] = source + if source in plans: + continue + if source.is_dir() and (source == work or output.is_relative_to(source)): + raise ValueError(f"Directory includes work/output root: {source}") + relative = (source.relative_to(output) if source.is_relative_to(output) else + source.relative_to(work) if source.is_relative_to(work) else Path(source.name)) + _portable(relative.parts) + if not relative.parts or relative.parts[0].casefold() in _RESERVED: + raise ValueError(f"Reserved output path: {relative}") + expected = _snapshot(source) # Missing declarations fail before any publication. + destination = output / relative + for previous, (target, _) in plans.items(): + if (source.is_relative_to(previous) or previous.is_relative_to(source) + or destination.is_relative_to(target) or target.is_relative_to(destination)): + raise ValueError(f"Overlapping source/destination declarations: {source}, {previous}") + for i in range(1, len(relative.parts) + 1): + parts = relative.parts[:i] + key = tuple(p.casefold() for p in parts) + claim = (parts, expected[0]["kind"] if i == len(relative.parts) else "directory") + if key in claims and claims[key] != claim: + raise ValueError(f"Case/type destination collision: {relative}") + claims[key] = claim + _destination(destination, output, expected[0]["kind"]) + plans[source] = destination, expected + output.mkdir(parents=True, exist_ok=True) + for source, (destination, expected) in plans.items(): + if source != destination: + _publish(source, destination, output, expected) + for destination, expected in plans.values(): + _verify(destination, expected) + mapped, entries = {}, [] + for label, source in labels.items(): + destination, (metadata, _) = plans[source] + mapped[label] = str(destination) + entries.append(dict(metadata, label=label, relative_path=destination.relative_to(output).as_posix(), + path=str(destination), verified=True)) + return mapped, {"schema_version": 1, "artifacts": entries} \ No newline at end of file diff --git a/agents/nvidia-alchemi/tools/nvidia-alchemi/alchemi_utils.py b/agents/nvidia-alchemi/tools/nvidia-alchemi/alchemi_utils.py new file mode 100644 index 0000000..838a588 --- /dev/null +++ b/agents/nvidia-alchemi/tools/nvidia-alchemi/alchemi_utils.py @@ -0,0 +1,1348 @@ +#!/usr/bin/env python3 +"""NVIDIA ALCHEMI Toolkit utilities library for Discovery platform workflows. + +Wraps `nvalchemi` (https://github.com/NVIDIA/nvalchemi-toolkit) for GPU-first, +batched AI atomic simulation: machine-learned interatomic potential (MLIP) +relaxation, molecular dynamics, multi-stage pipelines, high-throughput screening +with inflight batching, Zarr trajectory capture, training / fine-tuning, and +multi-GPU execution. + +Every call pattern here was validated against nvalchemi 0.2.0 inside the agent +container. Gotchas this library handles for you: + +* Dynamics reads and writes per-step buffers (`forces`, `energy`, `velocities`, + and `stress` for barostats). `AtomicData.from_atoms()` does NOT allocate them, + and a missing buffer surfaces as `AttributeError: 'Batch' has no attribute + 'forces'`. `atoms_to_batch()` allocates them up front. +* `convergence_hook` is its own constructor argument -- passing a + `ConvergenceHook` inside `hooks=[...]` silently never converges. +* Inflight batching is Mode 2 of `FusedStage.run(None)` driven by a `sampler`, + not a plain optimizer feature. `SizeAwareSampler` needs a dataset exposing + `get_metadata(idx)`, so a bare list will not do -- wrap it in `InMemoryDataset`. +* Classical potentials take required physical parameters (epsilon/sigma/cutoff, + a1/a2/s8, cutoff), so `load_model()` supplies documented defaults. + +License: this wrapper is part of the Discovery catalog. See THIRD_PARTY_NOTICES.md +for the licenses of nvalchemi-toolkit, MACE, AIMNet2 and the CUDA base image. +""" +import glob +import importlib +import json +import logging +import os +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple + +from alchemi_artifacts import export_artifacts, to_jsonable, write_json_atomic + +HELPER_VERSION = "1.1.4" +_DECLARED_OUTPUTS: Dict[str, Any] = {} +_FINAL_DATA: Optional[Dict[str, Any]] = None + +# ============= CONSTANTS ============= +INPUT_DIR = "/input" +OUTPUT_DIR = "/output" +WORK_DIR = "/workdir" +SCRATCH_DIR = "/tmp/alchemi_scratch" + +# Structure file extensions ASE can read, in the order globbed by load_structures(). +STRUCTURE_EXTENSIONS = ( + "*.xyz", "*.extxyz", "*.cif", "*.pdb", "*.poscar", "POSCAR*", "CONTCAR*", + "*.traj", "*.vasp", +) + +# Potentials wrapped by the toolkit: name -> (module path, wrapper class name). +MODEL_REGISTRY: Dict[str, Tuple[str, str]] = { + "mace": ("nvalchemi.models.mace", "MACEWrapper"), + "aimnet2": ("nvalchemi.models.aimnet2", "AIMNet2Wrapper"), + "lj": ("nvalchemi.models.lj", "LennardJonesModelWrapper"), + "dftd3": ("nvalchemi.models.dftd3", "DFTD3ModelWrapper"), + "ewald": ("nvalchemi.models.ewald", "EwaldModelWrapper"), + "pme": ("nvalchemi.models.pme", "PMEModelWrapper"), + "demo": ("nvalchemi.models.demo", "DemoModelWrapper"), +} + +MODEL_DESCRIPTIONS: Dict[str, str] = { + "mace": "MACE equivariant MLIP. Foundation checkpoints 'small-0b' and " + "'medium-0b2' are baked into the image. Energy/forces/stress, PBC, " + "COO neighbor list. Default choice for materials and periodic systems.", + "aimnet2": "AIMNet2 MLIP for neutral, charged, organic and elemental-organic " + "molecules. Predicts partial charges alongside energy/forces; " + "MATRIX neighbor list. Checkpoint name: 'aimnet2'.", + "lj": "Lennard-Jones analytical force field. Cheap; smoke tests, toy systems " + "and validation. Requires epsilon, sigma and cutoff.", + "dftd3": "DFT-D3(BJ) dispersion correction. Additive term -- compose onto an " + "MLIP with compose_models(). Requires the a1/a2/s8 damping " + "parameters of the target functional.", + "ewald": "Ewald summation for long-range electrostatics. Needs per-atom " + "`charges` and periodic cells.", + "pme": "Particle Mesh Ewald electrostatics; O(N log N) alternative to Ewald " + "for large periodic systems.", + "demo": "Trivial non-invariant demo potential. Testing and tutorials only -- " + "not physically meaningful.", +} + +# Documented default parameters for the classical potentials, which take +# required physical arguments rather than having built-in defaults. +# LJ defaults are argon; DFT-D3(BJ) defaults are the PBE damping set. +MODEL_DEFAULTS: Dict[str, Dict[str, Any]] = { + "lj": {"epsilon": 0.0104, "sigma": 3.4, "cutoff": 8.5}, + "dftd3": {"a1": 0.4289, "a2": 4.4407, "s8": 0.7875, "cutoff": 15.0}, + "ewald": {"cutoff": 10.0}, + "pme": {"cutoff": 10.0}, +} + +# Optimizers and integrators, resolved from `nvalchemi.dynamics`. +OPTIMIZERS: Dict[str, str] = { + "fire": "FIRE", + "fire2": "FIRE2", + "fire_variable_cell": "FIREVariableCell", + "fire2_variable_cell": "FIRE2VariableCell", +} +ENSEMBLES: Dict[str, str] = { + "nve": "NVE", + "nvt": "NVTLangevin", + "nvt_langevin": "NVTLangevin", + "nvt_nose_hoover": "NVTNoseHoover", + "npt": "NPT", + "nph": "NPH", +} + +# Runs whose integrator or optimizer needs the stress tensor from the model. +STRESS_REQUIRING = frozenset({"npt", "nph", "fire_variable_cell", + "fire2_variable_cell"}) + + +# ============= SETUP FUNCTIONS ============= +def quick_setup(input_dir: str = "/input", output_dir: str = "/output", + work_dir: str = "/workdir") -> None: + """Initialize logging, create directories, copy input files. + + Args: + input_dir: Path to input directory (read-only by convention). + output_dir: Path to output directory. + work_dir: Path to working directory; becomes the process CWD. + """ + global INPUT_DIR, OUTPUT_DIR, WORK_DIR, _DECLARED_OUTPUTS, _FINAL_DATA + INPUT_DIR, OUTPUT_DIR, WORK_DIR = input_dir, output_dir, work_dir + _DECLARED_OUTPUTS, _FINAL_DATA = {}, None + + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(levelname)s - %(message)s", + ) + os.makedirs(WORK_DIR, exist_ok=True) + os.makedirs(OUTPUT_DIR, exist_ok=True) + os.makedirs(SCRATCH_DIR, exist_ok=True) + os.chdir(WORK_DIR) + _copy_input_files() + logging.info(f"Working directory: {WORK_DIR}") + logging.info( + f"Input files: {os.listdir(INPUT_DIR) if os.path.exists(INPUT_DIR) else 'none'}" + ) + log_gpu_info() + + +def _copy_input_files() -> None: + """Copy input files to working directory.""" + if os.path.realpath(INPUT_DIR) == os.path.realpath(WORK_DIR): + return + if os.path.exists(INPUT_DIR): + for f in glob.glob(os.path.join(INPUT_DIR, "*")): + if os.path.isfile(f): + shutil.copy(f, WORK_DIR) + + +def copy_outputs() -> Dict[str, str]: + """Re-export and verify declared artifacts, including any updated files. + + If no result has been saved, retain legacy discovery of common artifacts. + Explicit declarations accept ANY extension and are never filtered. Once + save_final_results() is used, declare everything to be delivered; unlisted + scratch files are intentionally not exported. Destination paths and hashes + are recorded in artifact_manifest.json, not working-directory paths. + """ + if _FINAL_DATA is not None: + return _publish_results(_FINAL_DATA, _DECLARED_OUTPUTS) + patterns = ("*.xyz", "*.extxyz", "*.cif", "*.csv", "*.json", "*.png", + "*.log", "*.pt", "*.zip", "*.npz", "*.npy", "*.py", "*.txt", + "*.zarr", "checkpoints*") + files = {} + for pattern in patterns: + for f in glob.glob(os.path.join(WORK_DIR, pattern)): + if Path(f).name not in {"final_results.json", "artifact_manifest.json"}: + files[Path(f).name] = f + mapped, manifest = export_artifacts(files, WORK_DIR, OUTPUT_DIR) + write_json_atomic(Path(OUTPUT_DIR) / "artifact_manifest.json", manifest) + return mapped + + +def quick_finish() -> None: + """Finalize artifact delivery. Raises on missing/corrupt declared evidence.""" + copy_outputs() + + +def save_final_results(results: Dict, output_files: Optional[Dict] = None, + file_descriptions: Optional[Dict] = None, + status: str = "completed") -> Dict[str, str]: + """Export declared evidence, verify it, and save strict JSON results. + + Args: + results: Summary dict with key metrics. + output_files: Labels mapped to existing source paths, relative to WORK_DIR + or absolute. All extensions and directories are supported. Returns + mapped destination paths under OUTPUT_DIR. A later save with None + retains the declarations; an explicit empty dict clears them. + file_descriptions: Dict mapping label to description. + status: Overall status string ("completed" or "failed"). + + Nonfinite floats become explicit diagnostic objects, never invalid JSON. + Export/serialization failures publish a failed result and raise; merely + listing a source path never counts as evidence delivery. quick_finish() + re-verifies declared evidence after final log/trajectory writes. + """ + global _DECLARED_OUTPUTS, _FINAL_DATA + if status not in {"completed", "failed"}: + raise ValueError("status must be 'completed' or 'failed'") + if output_files is not None: + _DECLARED_OUTPUTS = dict(output_files) + _FINAL_DATA = {"status": status, "summary": results, + "helper_version": HELPER_VERSION, + "file_descriptions": file_descriptions or {}} + return _publish_results(_FINAL_DATA, _DECLARED_OUTPUTS) + + +def _publish_results(final_data: Dict, declarations: Dict) -> Dict[str, str]: + try: + clean = to_jsonable(final_data) + mapped, manifest = export_artifacts(declarations, WORK_DIR, OUTPUT_DIR) + clean["output_files"] = mapped + clean["artifact_manifest"] = str(Path(OUTPUT_DIR).resolve() / "artifact_manifest.json") + write_json_atomic(Path(OUTPUT_DIR) / "artifact_manifest.json", manifest) + write_json_atomic(Path(OUTPUT_DIR) / "final_results.json", clean) + # Logging here could mutate an already-hashed, still-open declared log. + return mapped + except Exception as exc: + final_data["status"] = "failed" + failure = {"status": "failed", "helper_version": HELPER_VERSION, + "summary": {"artifact_or_serialization_error": str(exc)}, + "output_files": {}} + write_json_atomic(Path(OUTPUT_DIR) / "artifact_manifest.json", + {"schema_version": 1, "artifacts": [], "error": str(exc)}) + write_json_atomic(Path(OUTPUT_DIR) / "final_results.json", failure) + raise + + +# ============= SYMBOL RESOLUTION ============= +def _resolve(module_path: str, name: str) -> Any: + """Return `name` from `module_path` with an actionable error on failure. + + The toolkit is a public beta, so a moved symbol should produce a message + naming the installed version rather than a bare ImportError. + """ + try: + module = importlib.import_module(module_path) + except ImportError as exc: + raise ImportError( + f"Could not import '{module_path}' ({exc}). Check the installed " + 'version with `python -c "import nvalchemi; ' + 'print(nvalchemi.__version__)"`.' + ) from exc + if not hasattr(module, name): + available = [n for n in dir(module) if not n.startswith("_")] + raise ImportError( + f"'{module_path}' has no attribute '{name}'. Available: {available}" + ) + return getattr(module, name) + + +def log_gpu_info() -> Dict[str, Any]: + """Log and return CUDA device information. Warns loudly if no GPU is present.""" + import torch + + info = { + "cuda_available": torch.cuda.is_available(), + "device_count": torch.cuda.device_count() if torch.cuda.is_available() else 0, + "torch_version": torch.__version__, + } + if info["cuda_available"]: + info["devices"] = [ + torch.cuda.get_device_name(i) for i in range(info["device_count"]) + ] + logging.info(f"CUDA devices: {info['devices']}") + else: + logging.warning( + "No CUDA device detected. The ALCHEMI Toolkit is GPU-first and has no " + "CPU fallback for its accelerated kernels -- request a GPU nodepool." + ) + return info + + +def default_device() -> str: + """Return 'cuda' when a GPU is present, otherwise 'cpu'.""" + import torch + + return "cuda" if torch.cuda.is_available() else "cpu" + + +def gpu_count() -> int: + """Number of visible CUDA devices.""" + import torch + + return torch.cuda.device_count() if torch.cuda.is_available() else 0 + + +# ============= STRUCTURE I/O ============= +def load_structures(path: Optional[str] = None, index: str = ":") -> List[Any]: + """Read structures from a file, glob pattern, or directory into `ase.Atoms`. + + Args: + path: File, glob pattern, or directory. Defaults to the working directory. + index: ASE slice string; ':' reads every frame of multi-frame files. + + Returns: + List of `ase.Atoms`. + """ + from ase.io import read + + target = path or WORK_DIR + if os.path.isdir(target): + files: List[str] = [] + for pattern in STRUCTURE_EXTENSIONS: + files.extend(sorted(glob.glob(os.path.join(target, pattern)))) + else: + files = sorted(glob.glob(target)) or [target] + + atoms_list: List[Any] = [] + for f in files: + if not os.path.isfile(f): + continue + frames = read(f, index=index) + atoms_list.extend(frames if isinstance(frames, list) else [frames]) + + if not atoms_list: + raise FileNotFoundError(f"No readable structures found at '{target}'.") + logging.info(f"Loaded {len(atoms_list)} structure(s) from {len(files)} file(s)") + return atoms_list + + +def prepare_data(atoms: Any, device: Optional[str] = None, + need_stress: bool = False, dtype: Optional[Any] = None) -> Any: + """Convert one `ase.Atoms` to `AtomicData` WITH the per-step buffers allocated. + + Dynamics reads and writes `forces`, `energy` and `velocities` in place, and + barostats additionally need `stress`. `AtomicData.from_atoms()` leaves these + unset, which fails at the first step with + `AttributeError: 'Batch' has no attribute 'forces'`. + """ + import torch + + AtomicData = _resolve("nvalchemi.data", "AtomicData") + device = device or default_device() + dtype = dtype or torch.float32 + + data = AtomicData.from_atoms(atoms, device=device, dtype=dtype) + n = len(atoms) + if getattr(data, "forces", None) is None: + data.forces = torch.zeros(n, 3, device=device, dtype=dtype) + if getattr(data, "energy", None) is None: + data.energy = torch.zeros(1, 1, device=device, dtype=dtype) + if getattr(data, "velocities", None) is None: + data.velocities = torch.zeros(n, 3, device=device, dtype=dtype) + if need_stress and getattr(data, "stress", None) is None: + data.stress = torch.zeros(1, 3, 3, device=device, dtype=dtype) + for key in ("system_id", "input_index"): + if key in atoms.info: + value = atoms.info[key] + import numbers + if isinstance(value, bool) or not isinstance(value, numbers.Integral) or value < 0: + raise ValueError(f"atoms.info[{key!r}] must be a nonnegative integer") + data.add_system_property(key, torch.tensor([[int(value)]], device=device, dtype=torch.long)) + return data + + +def atoms_to_batch(atoms_list: Sequence[Any], device: Optional[str] = None, + need_stress: bool = False, dtype: Optional[Any] = None) -> Any: + """Convert `ase.Atoms` objects into a single GPU-resident `Batch`. + + Buffers required by dynamics are allocated for you -- see `prepare_data`. + + Args: + atoms_list: Structures to batch. A single Atoms object is accepted. + device: Target device; defaults to 'cuda' when available. + need_stress: Allocate the stress buffer (required for NPT/NPH and + variable-cell relaxation). + dtype: Torch dtype for the tensors. + + Returns: + `nvalchemi.data.Batch` + """ + Batch = _resolve("nvalchemi.data", "Batch") + + if not isinstance(atoms_list, (list, tuple)): + atoms_list = [atoms_list] + device = device or default_device() + data_list = [prepare_data(a, device=device, need_stress=need_stress, dtype=dtype) + for a in atoms_list] + batch = Batch.from_data_list(data_list) + logging.info( + f"Built batch: {batch.num_graphs} structures, {batch.num_nodes} atoms " + f"on {device}" + ) + return batch + + +def batch_to_atoms(batch: Any) -> List[Any]: + """Convert a `Batch` (or `AtomicData`) back to `ase.Atoms`. + + The toolkit has no upstream `to_atoms`, because which fields to carry back is + application-specific. This helper restores positions, numbers, cell and pbc, + and attaches energy / forces / stress to `info` and `arrays` when present. + system_id and input_index are preserved in atoms.info. Toolkit velocities + are retained as atoms.arrays['alchemi_velocities'] in toolkit native units; + they are NOT silently interpreted as ASE velocity units. + """ + from ase import Atoms + + data_list = batch.to_data_list() if hasattr(batch, "to_data_list") else [batch] + identities = {} + for key in ("system_id", "input_index"): + values = getattr(batch, key, None) + if values is not None: + import torch + if values.dtype not in (torch.int32, torch.int64) or values.numel() != len(data_list): + raise ValueError(f"Expected one integer {key} per graph") + identities[key] = values.detach().cpu().reshape(-1).tolist() + out: List[Any] = [] + for index, data in enumerate(data_list): + cell_value = getattr(data, "cell", None) + pbc_value = getattr(data, "pbc", None) + cell = cell_value.squeeze(0).detach().cpu().numpy() if cell_value is not None else None + pbc = pbc_value.squeeze(0).detach().cpu().numpy() if pbc_value is not None else False + atoms = Atoms( + numbers=data.atomic_numbers.cpu().numpy(), + positions=data.positions.detach().cpu().numpy(), + cell=cell, + pbc=pbc, + ) + if getattr(data, "energy", None) is not None: + atoms.info["energy"] = float(data.energy.detach().cpu().reshape(-1)[0]) + if getattr(data, "forces", None) is not None: + atoms.arrays["forces"] = data.forces.detach().cpu().numpy() + if getattr(data, "stress", None) is not None: + atoms.info["stress"] = data.stress.detach().cpu().numpy().reshape(3, 3) + if getattr(data, "atomic_masses", None) is not None: + atoms.set_masses(data.atomic_masses.detach().cpu().numpy()) + if getattr(data, "velocities", None) is not None: + atoms.arrays["alchemi_velocities"] = data.velocities.detach().cpu().numpy().copy() + for key, values in identities.items(): + atoms.info[key] = int(values[index]) + out.append(atoms) + return out + + +def write_structures(batch_or_atoms: Any, output_path: str = "structures.extxyz", + fmt: Optional[str] = None) -> str: + """Write a `Batch` or list of `ase.Atoms` to disk. + + Args: + batch_or_atoms: A `Batch`, `AtomicData`, a single `ase.Atoms`, or a list + of `ase.Atoms`. + output_path: Destination path (relative paths land in the work directory). + fmt: Explicit ASE format; inferred from the extension when omitted. + + Returns: + Absolute path of the written file. + """ + from ase import Atoms + from ase.io import write + + if isinstance(batch_or_atoms, Atoms): + atoms_list = [batch_or_atoms] + elif isinstance(batch_or_atoms, (list, tuple)): + atoms_list = list(batch_or_atoms) + else: + atoms_list = batch_to_atoms(batch_or_atoms) + + path = output_path if os.path.isabs(output_path) else os.path.join(WORK_DIR, output_path) + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + write(path, atoms_list, format=fmt) + logging.info(f"Wrote {len(atoms_list)} structure(s) to {path}") + return path + + +def freeze_atoms_by_tag(data_or_batch: Any, tags: Any, free_tag: int = 0) -> Any: + """Mark atoms whose ASE tag differs from `free_tag` as frozen. + + Sets `atom_categories` to `AtomCategory.SPECIAL`, which is what + `FreezeAtomsHook` acts on. Typical use: hold a slab fixed while an adsorbate + relaxes. Register `FreezeAtomsHook()` on the dynamics for this to take effect. + """ + import torch + + AtomCategory = _resolve("nvalchemi._typing", "AtomCategory") + device = data_or_batch.positions.device + tag_tensor = torch.as_tensor(tags, device=device) + data_or_batch.atom_categories = torch.where( + tag_tensor != free_tag, + torch.tensor(AtomCategory.SPECIAL.value, device=device), + torch.tensor(AtomCategory.GAS.value, device=device), + ) + return data_or_batch + + +# ============= MODELS ============= +def list_models() -> List[Dict[str, str]]: + """Return the catalog of available potentials with their descriptions.""" + return [{"name": k, "description": MODEL_DESCRIPTIONS[k]} for k in MODEL_REGISTRY] + + +def load_model(name: str = "mace", checkpoint: Optional[str] = None, + device: Optional[str] = None, dtype: Optional[Any] = None, + enable_cueq: bool = True, **kwargs: Any) -> Any: + """Load a wrapped interatomic potential. + + Args: + name: One of `MODEL_REGISTRY` -- 'mace', 'aimnet2', 'lj', 'dftd3', + 'ewald', 'pme', 'demo'. + checkpoint: Checkpoint name for the learned potentials. Defaults to + 'small-0b' for MACE ('medium-0b2' is also baked in) and 'aimnet2' + for AIMNet2. Ignored by the classical models. + device: Target device; defaults to 'cuda' when available. + dtype: Torch dtype for the model weights. + enable_cueq: Enable cuEquivariance fused convolutions for MACE. Applied + only on CUDA, where the kernels exist. + **kwargs: Override the physical parameters of the classical models -- + `epsilon`/`sigma`/`cutoff` for 'lj', `a1`/`a2`/`s8` for 'dftd3', + `cutoff` for 'ewald'/'pme'. See `MODEL_DEFAULTS`. + + Returns: + A `BaseModelMixin` wrapper ready for dynamics or composition. + """ + import torch + + key = name.lower() + if key not in MODEL_REGISTRY: + raise ValueError( + f"Unknown model '{name}'. Available: {sorted(MODEL_REGISTRY)}" + ) + module_path, class_name = MODEL_REGISTRY[key] + wrapper_cls = _resolve(module_path, class_name) + + device = device or default_device() + dtype = dtype or torch.float32 + + if key == "mace": + model = wrapper_cls.from_checkpoint( + checkpoint or "small-0b", + device=torch.device(device), + dtype=dtype, + enable_cueq=enable_cueq and str(device).startswith("cuda"), + **kwargs, + ) + elif key == "aimnet2": + model = wrapper_cls.from_checkpoint( + checkpoint or "aimnet2", device=torch.device(device), **kwargs + ) + else: + params = dict(MODEL_DEFAULTS.get(key, {})) + params.update(kwargs) + model = wrapper_cls(**params) + model = model.to(device=torch.device(device)) + if params: + logging.info(f"'{key}' parameters: {params}") + + logging.info(f"Loaded model '{key}' on {device}") + return model + + +def compose_models(*models: Any, use_autograd: bool = False, + neighbor_adaptation: str = "auto") -> Any: + """Combine potentials into a single composed model. + + Args: + *models: Two or more wrapped models, in evaluation order. + use_autograd: When True, sum the sub-model energies and differentiate the + total. Required when one model's output feeds another's energy (for + example AIMNet2 charges -> Ewald electrostatics) and forces must + backpropagate through the whole chain. When False, each model + computes its own forces and the results are summed -- correct for + independent additive terms such as MACE + DFT-D3. + neighbor_adaptation: 'auto', 'always' or 'never'. + + Returns: + A composed model exposing the same interface as a single wrapper. + """ + if len(models) < 2: + raise ValueError("compose_models() needs at least two models.") + + if not use_autograd: + composed = models[0] + for m in models[1:]: + composed = composed + m + logging.info(f"Composed {len(models)} models with the '+' operator") + return composed + + PipelineModelWrapper = _resolve("nvalchemi.models.pipeline", + "PipelineModelWrapper") + PipelineGroup = _resolve("nvalchemi.models.pipeline", "PipelineGroup") + composed = PipelineModelWrapper( + groups=[PipelineGroup(steps=list(models), use_autograd=True)], + neighbor_adaptation=neighbor_adaptation, + ) + logging.info(f"Composed {len(models)} models into a shared-autograd pipeline") + return composed + + +def set_active_outputs(model: Any, outputs: Iterable[str]) -> Any: + """Set which properties the model computes on each forward pass. + + Narrow to `{'energy'}` for screening, and include `'stress'` before running + NPT/NPH or variable-cell relaxation. + """ + model.model_config.active_outputs = set(outputs) + logging.info(f"active_outputs = {sorted(model.model_config.active_outputs)}") + return model + + +def attach_neighbor_hooks(dynamics: Any, model: Any, **kwargs: Any) -> Any: + """Register every neighbor-list hook the model (or pipeline) requires. + + Always use this rather than building a `NeighborListHook` by hand: a composed + model may need several source lists at different cutoffs, and a missing hook + shows up as NaN or zero forces rather than a clear error. + """ + DynamicsStage = _resolve("nvalchemi.dynamics", "DynamicsStage") + hooks = list(model.make_neighbor_hooks(**kwargs)) + for hook in hooks: + dynamics.register_hook(hook, stage=DynamicsStage.BEFORE_COMPUTE) + logging.info(f"Registered {len(hooks)} neighbor-list hook(s)") + return dynamics + + +# ============= DYNAMICS ============= +def _ensure_stress(model: Any) -> None: + """Add 'stress' to the model's active outputs if the run needs it.""" + cfg = model.model_config + active = set(getattr(cfg, "active_outputs", None) or cfg.outputs) + if "stress" not in active: + if "stress" not in set(cfg.outputs): + raise ValueError( + f"This run needs the stress tensor, but the model only supports " + f"{sorted(cfg.outputs)}. Choose a potential that computes stress." + ) + set_active_outputs(model, active | {"stress"}) + + +def make_optimizer(model: Any, optimizer: str = "fire", dt: float = 0.1, + n_steps: int = 500, fmax: Optional[float] = 0.05, + hooks: Optional[List[Any]] = None, **kwargs: Any) -> Any: + """Build a geometry optimizer stage (FIRE / FIRE2) without running it. + + Use this when composing multi-stage pipelines; use `relax()` for a one-shot run. + `fmax` becomes a `ConvergenceHook` passed as the dedicated `convergence_hook` + argument -- putting it in `hooks` would silently never converge. + """ + key = optimizer.lower() + if key not in OPTIMIZERS: + raise ValueError(f"Unknown optimizer '{optimizer}'. Available: {sorted(OPTIMIZERS)}") + opt_cls = _resolve("nvalchemi.dynamics", OPTIMIZERS[key]) + + if key in STRESS_REQUIRING: + _ensure_stress(model) + + convergence_hook = None + if fmax is not None: + ConvergenceHook = _resolve("nvalchemi.dynamics", "ConvergenceHook") + convergence_hook = ConvergenceHook.from_fmax(fmax) + + return opt_cls(model=model, dt=dt, n_steps=n_steps, hooks=list(hooks or []), + convergence_hook=convergence_hook, **kwargs) + + +def make_integrator(model: Any, ensemble: str = "nvt", dt: float = 1.0, + n_steps: int = 5000, temperature: float = 300.0, + hooks: Optional[List[Any]] = None, **kwargs: Any) -> Any: + """Build a molecular-dynamics stage without running it. + + Required parameters are filled with documented defaults when omitted: + `friction` for Langevin, `thermostat_time` for Nose-Hoover and NPT, + `barostat_time` and `pressure` for NPT/NPH. + + Args: + ensemble: 'nve', 'nvt' (Langevin), 'nvt_nose_hoover', 'npt', 'nph'. + dt: Timestep in femtoseconds. + temperature: Target temperature in Kelvin (ignored by NVE and NPH). + **kwargs: Ensemble-specific overrides, e.g. `friction`, `pressure`, + `thermostat_time`, `barostat_time`, `pressure_coupling`. + """ + key = ensemble.lower() + if key not in ENSEMBLES: + raise ValueError(f"Unknown ensemble '{ensemble}'. Available: {sorted(ENSEMBLES)}") + integrator_cls = _resolve("nvalchemi.dynamics", ENSEMBLES[key]) + + if key in STRESS_REQUIRING: + _ensure_stress(model) + + params: Dict[str, Any] = {"model": model, "dt": dt, "n_steps": n_steps, + "hooks": list(hooks or [])} + if key in ("nvt", "nvt_langevin"): + params["temperature"] = temperature + params.setdefault("friction", 0.01) + elif key == "nvt_nose_hoover": + params["temperature"] = temperature + # A thermostat coupling time of ~100 steps is the usual starting point. + params.setdefault("thermostat_time", 100.0 * dt) + elif key == "npt": + params["temperature"] = temperature + params.setdefault("pressure", 0.0) + params.setdefault("thermostat_time", 100.0 * dt) + # The barostat must be slower than the thermostat, or the cell rings. + params.setdefault("barostat_time", 1000.0 * dt) + elif key == "nph": + params.setdefault("pressure", 0.0) + params.setdefault("barostat_time", 1000.0 * dt) + + params.update(kwargs) + return integrator_cls(**params) + + +def single_point(batch: Any, model: Any) -> Dict[str, Any]: + """Evaluate the model once on `batch` (no dynamics) and return its outputs. + + The supported way to get energies/forces for fixed geometries. Model + wrappers have no `compute()` method, and calling `model(batch)` directly + fails because the neighbor list is built by hooks. This builds the neighbor + list with the same hooks the dynamics helpers use, then runs one forward + pass. Do NOT use a zero-step `run()`: it evaluates nothing and leaves + `batch.energy` at its zero-initialised value. + + Args: + batch: `Batch` of structures (build it with `atoms_to_batch`). + model: A wrapped potential; `set_active_outputs` selects what is computed. + + Returns: + Dict of detached output tensors keyed by name (e.g. 'energy' with shape + (num_graphs, 1), 'forces' with shape (num_atoms, 3)). The batch's output + fields are updated in place too. + + Raises: + RuntimeError: If the model returns no outputs or any non-finite value. + """ + DynamicsStage = _resolve("nvalchemi.dynamics", "DynamicsStage") + stage = make_integrator(model, ensemble="nve", dt=0.1, n_steps=1) + attach_neighbor_hooks(stage, model) + with stage: + stage._call_hooks(DynamicsStage.BEFORE_COMPUTE, batch) + outputs = stage.compute(batch) + result = {k: v.detach() for k, v in dict(outputs).items() if v is not None} + if not result: + raise RuntimeError("Model returned no outputs; check set_active_outputs()") + import torch + bad = [k for k, v in result.items() if not bool(torch.isfinite(v).all())] + if bad: + raise RuntimeError(f"Non-finite model output(s): {bad}") + return result + + +def relax(batch: Any, model: Any, fmax: float = 0.05, max_steps: int = 500, + dt: float = 0.1, optimizer: str = "fire", + hooks: Optional[List[Any]] = None, **kwargs: Any) -> Any: + """Batched geometry relaxation of every structure in `batch`. + + All structures relax simultaneously in a shared model forward pass. Convergence + does NOT remove a system from that forward pass: the run continues until EVERY + structure meets `fmax` or `max_steps` is reached, so cost scales with the + slowest system rather than the average. `ConvergenceHook` here performs no + status migration -- there is no `status` field on a plain optimizer batch -- + so it acts purely as an all-converged stopping test. For large heterogeneous + workloads use `screen()`, whose inflight batching genuinely evicts converged + systems and backfills new ones. + + Args: + batch: `Batch` of structures to relax (build it with `atoms_to_batch`). + model: A wrapped potential. + fmax: Force convergence threshold in eV/A. + max_steps: Step budget. + dt: Optimizer timestep. + optimizer: 'fire', 'fire2', or the '*_variable_cell' variants which also + relax the cell (these require a stress-capable potential). + + Returns: + The relaxed `Batch`. + """ + opt = make_optimizer(model, optimizer=optimizer, dt=dt, n_steps=max_steps, + fmax=fmax, hooks=hooks, **kwargs) + attach_neighbor_hooks(opt, model) + logging.info( + f"Relaxing {batch.num_graphs} structure(s) with {optimizer.upper()} " + f"(fmax={fmax}, max_steps={max_steps})" + ) + with opt: + return opt.run(batch) + + +def run_md(batch: Any, model: Any, ensemble: str = "nvt", dt: float = 1.0, + n_steps: int = 5000, temperature: float = 300.0, + hooks: Optional[List[Any]] = None, **kwargs: Any) -> Any: + """Run batched molecular dynamics. + + Args: + batch: `Batch` of structures. For NPT/NPH build it with + `atoms_to_batch(..., need_stress=True)`. + model: A wrapped potential. + ensemble: 'nve', 'nvt', 'nvt_nose_hoover', 'npt', 'nph'. + dt: Timestep in femtoseconds (use 0.5 when hydrogen is present). + n_steps: Number of MD steps. + temperature: Target temperature in Kelvin. + + Returns: + The final `Batch`. + """ + integrator = make_integrator(model, ensemble=ensemble, dt=dt, n_steps=n_steps, + temperature=temperature, hooks=hooks, **kwargs) + attach_neighbor_hooks(integrator, model) + logging.info( + f"Running {ensemble.upper()} MD: {batch.num_graphs} structure(s), " + f"{n_steps} steps, dt={dt} fs, T={temperature} K" + ) + with integrator: + return integrator.run(batch) + + +def multi_stage(batch: Any, stages: Sequence[Any], model: Optional[Any] = None, + compile_step: bool = False) -> Any: + """Chain stages on one GPU with the `+` operator and run them. + + Systems migrate to the next stage as they converge, and every active system + -- whichever stage it is in -- shares ONE model forward pass per step. + + Args: + batch: `Batch` of structures. + stages: Stages from `make_optimizer()` / `make_integrator()`, in order. + model: The potential, used to register neighbor-list hooks. + compile_step: Wrap the fused step in `torch.compile`. Only safe when no + hook performs Python-side I/O. + + Returns: + The final `Batch`. + """ + if len(stages) < 2: + raise ValueError("multi_stage() needs at least two stages.") + + fused = stages[0] + for stage in stages[1:]: + fused = fused + stage + if compile_step: + fused = fused.compile(fullgraph=False) + if model is not None: + attach_neighbor_hooks(fused, model) + + logging.info(f"Running fused pipeline of {len(stages)} stage(s)") + with fused: + return fused.run(batch) + + +def screen(atoms_list: Sequence[Any], model: Any, fmax: float = 0.05, + max_steps: int = 500, max_batch_size: int = 64, + max_atoms: Optional[int] = None, device: Optional[str] = None, + optimizer: str = "fire", dt: float = 0.1, + sink_capacity: Optional[int] = None) -> List[Any]: + """High-throughput relaxation of many structures using inflight batching. + + A `SizeAwareSampler` feeds structures into a running `FusedStage` and + converged systems are replaced mid-run, so the GPU never idles waiting for + stragglers. Prefer this over `relax()` once you have more structures than + comfortably fit in one batch. + + Args: + atoms_list: Structures to screen. + model: A wrapped potential. + fmax: Force convergence threshold in eV/A. + max_steps: Per-system step budget. + max_batch_size: Maximum concurrent structures in the active batch. + max_atoms: Maximum concurrent atoms; defaults to `max_batch_size` times + the largest structure, which is the safe upper bound. + optimizer: 'fire' or 'fire2'. + sink_capacity: Result sink capacity; defaults to the input count. + + Returns: + List of relaxed `ase.Atoms` collected from the converged-result sink. + NOTE: only systems that actually reached `fmax` are returned. Structures + still above the threshold when `max_steps` runs out are NOT collected, + so a short return list means the step budget was too small. + NOTE: results come back in GRADUATION order, not input order. Match them + using atoms.info['input_index'] (the zero-based position in atoms_list). + atoms.info['system_id'] is the sampler's admission-order ID, NOT an input + index; mixed-size bin packing can reorder admissions. Existing caller + metadata is restored from the corresponding input after ID validation. + """ + import torch + + Batch = _resolve("nvalchemi.data", "Batch") + InMemoryDataset = _resolve("nvalchemi.data", "InMemoryDataset") + SizeAwareSampler = _resolve("nvalchemi.dynamics", "SizeAwareSampler") + HostMemory = _resolve("nvalchemi.dynamics", "HostMemory") + FusedStage = _resolve("nvalchemi.dynamics", "FusedStage") + ConvergedSnapshotHook = _resolve("nvalchemi.dynamics.hooks", + "ConvergedSnapshotHook") + + device = device or default_device() + atoms_list = list(atoms_list) + if not atoms_list: + return [] + if max_batch_size < 1 or max_steps < 1: + raise ValueError("max_batch_size and max_steps must be positive") + if sink_capacity is not None and sink_capacity < len(atoms_list): + raise ValueError("sink_capacity must cover all inputs; refusing silent result loss") + largest = max(len(a) for a in atoms_list) + max_atoms = max_atoms if max_atoms is not None else max_batch_size * largest + if largest == 0 or any(len(a) == 0 for a in atoms_list) or max_atoms < largest: + raise ValueError("Every system must contain atoms and fit max_atoms") + + data_list = [prepare_data(a, device=device) for a in atoms_list] + for index, data in enumerate(data_list): + # The toolkit overwrites system_id during sampling. Preserve input + # identity as a separate system property through refill and sinks. + data.add_system_property("input_index", torch.tensor([[index]], device=device, dtype=torch.long)) + # SizeAwareSampler needs get_metadata(idx); a bare list does not provide it. + dataset = InMemoryDataset(Batch.from_data_list(data_list)) + sampler = SizeAwareSampler(dataset, max_atoms=max_atoms, + max_batch_size=max_batch_size) + sink = HostMemory(capacity=sink_capacity or len(data_list)) + + def _init_buffers(batch: Any) -> None: + """The sampler-built batch may lack per-step buffers; fill them in place.""" + dev = batch.positions.device + n = batch.num_nodes + for key in ("forces", "velocities"): + if getattr(batch, key, None) is None: + values = [torch.zeros(int(count), 3, device=dev, dtype=batch.positions.dtype) + for count in batch.num_nodes_per_graph.detach().cpu().tolist()] + batch.add_key(key, values, level="node") + + stage = make_optimizer(model, optimizer=optimizer, dt=dt, n_steps=max_steps, + fmax=fmax, + hooks=[ConvergedSnapshotHook(sink=sink)]) + # The ConvergedSnapshotHook MUST live on the sub-stage: ON_CONVERGE fires + # there, not on the enclosing FusedStage. Registering it on the FusedStage + # silently collects nothing. + fused = FusedStage( + [(0, stage)], + sampler=sampler, + init_fn=_init_buffers, + ) + attach_neighbor_hooks(fused, model) + + logging.info( + f"Screening {len(data_list)} structures with inflight batching " + f"(max_batch_size={max_batch_size}, max_atoms={max_atoms})" + ) + with fused: + fused.run(None) + + out: List[Any] = [] + for item in _drain_sink(sink): + out.extend(batch_to_atoms(item)) + _restore_screen_identity(out, atoms_list) + + # Only CONVERGED systems reach the sink. Anything still above `fmax` when + # the step budget runs out is silently dropped, so surface the shortfall. + if len(out) < len(data_list): + logging.warning( + f"Screening returned {len(out)} of {len(data_list)} structures: " + f"{len(data_list) - len(out)} did not reach fmax={fmax} within " + f"max_steps={max_steps} and were not collected. Raise max_steps or " + f"loosen fmax to capture them." + ) + else: + logging.info(f"Screening collected all {len(out)} structures") + return out + + +def _restore_screen_identity(results: List[Any], inputs: Sequence[Any]) -> None: + """Fail closed on identity loss/duplicates; preserve input metadata by index.""" + import copy + import numbers + seen_inputs, seen_systems = set(), set() + for atoms in results: + index, system_id = atoms.info.get("input_index"), atoms.info.get("system_id") + for label, value in (("input_index", index), ("system_id", system_id)): + if isinstance(value, bool) or not isinstance(value, numbers.Integral) or value < 0: + raise RuntimeError(f"Screening result missing valid {label}; cannot match inputs") + if index >= len(inputs) or index in seen_inputs or system_id in seen_systems: + raise RuntimeError("Screening result has out-of-range or duplicate identity") + if list(atoms.numbers) != list(inputs[index].numbers): + raise RuntimeError("Screening result composition does not match its input_index") + seen_inputs.add(index) + seen_systems.add(system_id) + caller_info = copy.deepcopy(inputs[index].info) + if "system_id" in caller_info: + caller_info["source_system_id"] = caller_info.pop("system_id") + if "input_index" in caller_info: + caller_info["source_input_index"] = caller_info.pop("input_index") + caller_info.update(atoms.info) + atoms.info = caller_info + + +def _drain_sink(sink: Any) -> List[Any]: + """Read every batch out of a data sink, returning [] when it is empty. + + `DataSink.drain()` raises `RuntimeError` on an empty buffer rather than + returning nothing, so the emptiness check comes first. + """ + try: + if len(sink) == 0: + logging.warning("Result sink is empty -- no systems converged.") + return [] + except TypeError: + pass # Sink does not implement __len__; fall through to drain(). + + for method in ("drain", "read"): + if hasattr(sink, method): + try: + result = getattr(sink, method)() + except RuntimeError as exc: + logging.warning(f"Could not drain sink: {exc}") + return [] + if result is None: + return [] + return list(result) if isinstance(result, (list, tuple)) else [result] + return [] + + +# ============= TRAJECTORY I/O ============= +def zarr_sink(path: str = "trajectory.zarr", capacity: int = 1_000_000, + **kwargs: Any) -> Any: + """Create a Zarr-backed data sink for persistent trajectory capture. + + Pass it to `snapshot_hook()` for periodic frames, or to a + `ConvergedSnapshotHook` for converged results. Writes are GPU-buffered, so + capture costs far less than host round-trips. + """ + ZarrData = _resolve("nvalchemi.dynamics", "ZarrData") + full = path if os.path.isabs(path) else os.path.join(WORK_DIR, path) + logging.info(f"Zarr trajectory sink: {full}") + return ZarrData(full, capacity=capacity, **kwargs) + + +def snapshot_hook(sink: Any, frequency: int = 100, **kwargs: Any) -> Any: + """Periodically capture the running batch into a sink (trajectory writing).""" + SnapshotHook = _resolve("nvalchemi.dynamics.hooks", "SnapshotHook") + return SnapshotHook(sink=sink, frequency=frequency, **kwargs) + + +KB_EV_PER_K = 8.617333262145e-5 + + +class StepRecorder: + """Per-step energy/kinetic/temperature/fmax recorder (a toolkit hook). + + Build with `step_recorder()` and pass in `hooks=[...]`. Follows the toolkit + hook protocol: `frequency` and `stage` attributes, `__call__(ctx, stage)`, + batch read from `ctx.batch`. Runs at AFTER_STEP, where `ctx.step_count` is + pre-increment, so `step` = `ctx.step_count + 1` (completed updates). + `rows` holds one dict per recorded step and graph: step, graph, + energy_eV, kinetic_energy_eV (toolkit per-graph kinetic energy), + temperature_K (2*KE / (3*N*k_B), no constraint correction) and + max_force_eV_per_A (largest per-atom force norm). Values are detached + Python floats, so rows can be written directly to CSV/JSON. + """ + + def __init__(self, frequency: int = 1) -> None: + if not isinstance(frequency, int) or frequency < 1: + raise ValueError("frequency must be a positive integer") + self.frequency = frequency + self.stage = _resolve("nvalchemi.dynamics", "DynamicsStage").AFTER_STEP + self.rows: List[Dict[str, Any]] = [] + self._kinetic = _resolve("nvalchemi.dynamics.hooks._utils", "kinetic_energy_per_graph") + + def __call__(self, ctx: Any, stage: Any) -> None: + import torch + batch = ctx.batch + with torch.no_grad(): + graph = batch.batch_idx + num_graphs = int(batch.num_graphs) + kinetic = self._kinetic(batch.velocities, batch.atomic_masses, graph, num_graphs).reshape(-1) + atoms = batch.num_nodes_per_graph + norms = torch.linalg.vector_norm(batch.forces, dim=1) + fmax = torch.zeros(num_graphs, dtype=norms.dtype, device=norms.device) + fmax = fmax.scatter_reduce(0, graph, norms, reduce="amax", include_self=False) + energy = batch.energy.reshape(-1) + values = torch.stack([energy.to(kinetic.dtype), kinetic, + 2.0 * kinetic / (3.0 * atoms.to(kinetic.dtype) * KB_EV_PER_K), + fmax.to(kinetic.dtype)], dim=1).detach().cpu().tolist() + step = int(ctx.step_count) + 1 + for index, (e, ke, t, f) in enumerate(values): + self.rows.append({"step": step, "graph": index, "energy_eV": e, "kinetic_energy_eV": ke, + "temperature_K": t, "max_force_eV_per_A": f}) + + +def step_recorder(frequency: int = 1) -> StepRecorder: + """Record energy, kinetic energy, temperature and fmax every `frequency` steps. + + Use this instead of writing a custom hook: toolkit hooks need `frequency` + and `stage` attributes and a `__call__(ctx, stage)` signature, and a plain + callable is rejected at registration. Read results from `recorder.rows`. + """ + return StepRecorder(frequency) + + +def read_trajectory(path: str, batch_size: int = 32, device: Optional[str] = None, + shuffle: bool = False, **kwargs: Any) -> Any: + """Open a Zarr trajectory store as a `DataLoader` with CUDA-stream prefetch. + + Use the returned loader directly as the training or analysis input pipeline. + """ + AtomicDataZarrReader = _resolve("nvalchemi.data", "AtomicDataZarrReader") + Dataset = _resolve("nvalchemi.data", "Dataset") + DataLoader = _resolve("nvalchemi.data", "DataLoader") + + full = path if os.path.isabs(path) else os.path.join(WORK_DIR, path) + reader = AtomicDataZarrReader(full) + dataset = Dataset(reader, device=device or default_device()) + return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, **kwargs) + + +# ============= TRAINING & FINE-TUNING ============= +def default_loss(energy_weight: float = 1.0, force_weight: float = 10.0, + stress_weight: float = 0.0, huber: bool = False) -> Any: + """Build a composed energy / force / stress loss. + + Forces are weighted above energies by default because force errors dominate + downstream dynamics quality. Set `stress_weight` when training for NPT use, + and `huber=True` for robustness against outliers in noisy reference data. + """ + import nvalchemi.training as T + + energy_cls = T.EnergyHuberLoss if huber else T.EnergyMSELoss + force_cls = T.ForceHuberLoss if huber else T.ForceMSELoss + + loss = energy_weight * energy_cls() + force_weight * force_cls() + if stress_weight: + loss = loss + stress_weight * T.StressMSELoss() + return loss + + +def train_mlip(model: Any, train_loader: Any, val_loader: Optional[Any] = None, + loss_fn: Optional[Any] = None, num_steps: int = 10_000, + learning_rate: float = 1e-3, checkpoint_dir: str = "checkpoints", + ema_decay: float = 0.999, validate_every: int = 500, + device: Optional[str] = None, model_key: str = "main", + hooks: Optional[List[Any]] = None, **kwargs: Any) -> Any: + """Train or fine-tune an MLIP with the ALCHEMI training stack. + + Builds a `TrainingStrategy` with an AdamW optimizer and an `EMAHook`, then + saves a strategy checkpoint -- which preserves the model reconstruction + recipe, weights, optimizer state and hook state together. Never save only + `ema.state_dict()`; that cannot be restarted reliably for MACE. + + Args: + model: A wrapped model. For MACE, build it with + `load_model('mace', checkpoint=...)` so the reconstruction spec is + recorded. + train_loader: Training `DataLoader` (see `read_trajectory`). + val_loader: Optional validation `DataLoader`. + loss_fn: Composed loss; defaults to `default_loss()`. + num_steps: Number of optimizer steps. + learning_rate: AdamW learning rate. + checkpoint_dir: Destination directory for strategy checkpoints. + ema_decay: Exponential moving average decay. + validate_every: Run validation every N steps when `val_loader` is given. + model_key: Name under which the model is registered in the strategy. + + Returns: + The `TrainingStrategy` after training. + """ + import torch + import nvalchemi.training as T + + device = device or default_device() + loss = loss_fn if loss_fn is not None else default_loss() + + strategy_kwargs: Dict[str, Any] = { + "models": {model_key: model}, + "optimizer_configs": { + model_key: [T.OptimizerConfig( + optimizer_cls=torch.optim.AdamW, + optimizer_kwargs={"lr": learning_rate}, + )] + }, + "num_steps": num_steps, + "training_fn": T.default_training_fn, + "loss_fn": loss, + "devices": [torch.device(device)], + "hooks": list(hooks or []) + [T.EMAHook(model_key=model_key, decay=ema_decay)], + } + if val_loader is not None: + strategy_kwargs["validation_config"] = T.ValidationConfig( + validation_data=val_loader, + validation_fn=T.default_training_fn, + loss_fn=loss, + every_n_steps=validate_every, + ) + strategy_kwargs.update(kwargs) + + strategy = T.TrainingStrategy(**strategy_kwargs) + logging.info(f"Training for {num_steps} steps on {device}") + strategy.run(train_loader) + + ckpt = checkpoint_dir if os.path.isabs(checkpoint_dir) \ + else os.path.join(WORK_DIR, checkpoint_dir) + os.makedirs(ckpt, exist_ok=True) + strategy.save_checkpoint(ckpt) + logging.info(f"Saved training checkpoint to {ckpt}") + return strategy + + +def load_training_checkpoint(checkpoint_dir: str, device: Optional[str] = None, + **kwargs: Any) -> Any: + """Restore a `TrainingStrategy` from a checkpoint directory.""" + import torch + import nvalchemi.training as T + + path = checkpoint_dir if os.path.isabs(checkpoint_dir) \ + else os.path.join(WORK_DIR, checkpoint_dir) + return T.TrainingStrategy.load_checkpoint( + path, map_location=torch.device(device or default_device()), **kwargs + ) + + +# ============= MULTI-GPU ============= +def buffer_config(num_systems: int, max_atoms_per_system: int, + num_edges: int = 0) -> Any: + """Build a `BufferConfig` for inter-rank transfers in a `DistributedPipeline`. + + NCCL point-to-point requires fixed-size tensors, so `num_nodes` must cover the + worst case: `num_systems * max_atoms_per_system`. Every pair of communicating + stages must share an identical config or `setup()` raises. + + `num_edges=0` is correct whenever the downstream model rebuilds its own + neighbor list, which is the normal case. + """ + BufferConfig = _resolve("nvalchemi.dynamics.base", "BufferConfig") + return BufferConfig( + num_systems=num_systems, + num_nodes=num_systems * max_atoms_per_system, + num_edges=num_edges, + ) + + +def run_distributed(script_path: str, nproc_per_node: Optional[int] = None, + extra_args: Optional[Sequence[str]] = None) -> int: + """Launch a distributed script with `torchrun` and stream its output. + + `DistributedPipeline` and `DomainParallel` cannot run inline in the main + workflow script -- they need one process per GPU. Write the pipeline to its + own file, then launch it with this helper. + + Args: + script_path: Path to the script to launch. + nproc_per_node: Process count; defaults to every visible GPU. + extra_args: Additional arguments passed through to the script. + + Returns: + The `torchrun` exit code. + """ + nproc = nproc_per_node or gpu_count() + if nproc < 2: + raise RuntimeError( + f"Distributed execution needs at least 2 GPUs; found {nproc}. " + "Request a multi-GPU SKU or use the single-GPU helpers instead." + ) + path = script_path if os.path.isabs(script_path) else os.path.join(WORK_DIR, script_path) + torchrun = os.path.join(os.path.dirname(sys.executable), "torchrun") + if not os.path.exists(torchrun): + torchrun = "torchrun" + cmd = [torchrun, f"--nproc_per_node={nproc}", path] + cmd.extend(extra_args or []) + + logging.info(f"Launching: {' '.join(cmd)}") + proc = subprocess.run(cmd, cwd=WORK_DIR, check=False) + if proc.returncode != 0: + logging.error(f"torchrun exited with code {proc.returncode}") + return proc.returncode + + +# ============= ANALYSIS & VISUALIZATION ============= +def summarize_batch(batch: Any) -> List[Dict[str, Any]]: + """Per-structure summary: formula, atom count, energy, max force, volume.""" + import numpy as np + + summaries: List[Dict[str, Any]] = [] + for i, atoms in enumerate(batch_to_atoms(batch)): + entry: Dict[str, Any] = { + "index": i, + "formula": atoms.get_chemical_formula(), + "num_atoms": len(atoms), + } + if "energy" in atoms.info: + entry["energy_eV"] = atoms.info["energy"] + entry["energy_per_atom_eV"] = atoms.info["energy"] / len(atoms) + if "forces" in atoms.arrays: + entry["fmax_eV_per_A"] = float( + np.linalg.norm(atoms.arrays["forces"], axis=1).max() + ) + if atoms.cell.rank == 3: + entry["volume_A3"] = float(atoms.get_volume()) + summaries.append(entry) + return summaries + + +def save_summary_csv(summaries: Sequence[Dict[str, Any]], + output_path: str = "summary.csv") -> str: + """Write summary dictionaries to CSV.""" + import pandas as pd + + path = output_path if os.path.isabs(output_path) else os.path.join(WORK_DIR, output_path) + pd.DataFrame(list(summaries)).to_csv(path, index=False) + logging.info(f"Wrote summary CSV to {path}") + return path + + +def plot_energy_trace(csv_path: str, output_file: str = "energy_trace.png", + energy_column: str = "energy", + step_column: str = "step") -> str: + """Plot an energy trace from a `LoggingHook` CSV log.""" + import matplotlib.pyplot as plt + import pandas as pd + + df = pd.read_csv(csv_path) + path = output_file if os.path.isabs(output_file) else os.path.join(WORK_DIR, output_file) + + fig, ax = plt.subplots(figsize=(8, 4.5)) + x = df[step_column] if step_column in df.columns else df.index + ax.plot(x, df[energy_column], lw=1.2) + ax.set_xlabel("step") + ax.set_ylabel("energy (eV)") + ax.set_title("Energy trace") + fig.tight_layout() + fig.savefig(path, dpi=150) + plt.close(fig) + logging.info(f"Wrote energy trace plot to {path}") + return path + + +def plot_energy_distribution(summaries: Sequence[Dict[str, Any]], + output_file: str = "energy_distribution.png") -> str: + """Histogram of per-atom energies across a screened set of structures.""" + import matplotlib.pyplot as plt + + values = [s["energy_per_atom_eV"] for s in summaries if "energy_per_atom_eV" in s] + if not values: + raise ValueError("No per-atom energies in the supplied summaries.") + + path = output_file if os.path.isabs(output_file) else os.path.join(WORK_DIR, output_file) + fig, ax = plt.subplots(figsize=(7, 4.5)) + ax.hist(values, bins=min(50, max(5, len(values) // 5))) + ax.set_xlabel("energy per atom (eV)") + ax.set_ylabel("count") + ax.set_title(f"Energy distribution ({len(values)} structures)") + fig.tight_layout() + fig.savefig(path, dpi=150) + plt.close(fig) + logging.info(f"Wrote energy distribution plot to {path}") + return path diff --git a/agents/nvidia-alchemi/tools/nvidia-alchemi/alchemi_validators.py b/agents/nvidia-alchemi/tools/nvidia-alchemi/alchemi_validators.py new file mode 100644 index 0000000..0d91ac7 --- /dev/null +++ b/agents/nvidia-alchemi/tools/nvidia-alchemi/alchemi_validators.py @@ -0,0 +1,255 @@ +"""CPU-only evidence validators; NumPy is the only non-standard dependency. + +Invalid configuration and independent-calculation inputs raise ValueError. +Comparisons/audits return JSON-native verdicts, never inferred measurements. +""" + +import numpy as np + +__all__ = ["compare_numeric", "independent_kinetic_energy", + "independent_energy_drift", "audit_hook_events"] + + +def _result(status, reason, evidence): + return {"status": status, "reason": reason, "evidence": evidence} + + +def _real_array(value, name): + array = np.asarray(value) + if np.ma.isMaskedArray(value) or array.dtype.kind not in "iuf": + raise ValueError(f"{name} must contain unmasked real numbers (not booleans)") + if not np.isfinite(array).all(): + raise ValueError(f"{name} contains non-finite data") + return array + + +def _float64(value, name): + with np.errstate(over="ignore", invalid="ignore"): + array = _real_array(value, name).astype(np.float64) + if not np.isfinite(array).all(): + raise ValueError(f"{name} is outside the float64 range") + return array + + +def _integers(value, name, minimum): + array = np.asarray(value) + if (np.ma.isMaskedArray(value) or array.dtype.kind not in "iu" + or np.any(array < minimum)): + raise ValueError(f"{name} must contain integers >= {minimum}; no coercion") + return array + + +def _integer_scalar(value, name, minimum=0): + array = _integers(value, name, minimum) + if array.ndim != 0: + raise ValueError(f"{name} must be scalar") + return int(array) + + +def _label(value): + if not isinstance(value, str) or not value.strip(): + raise ValueError("labels must be nonempty strings") + return str(value) + + +def compare_numeric(observed, reference, *, observed_source, reference_source, atol, rtol): + """Use abs(observed-reference) <= atol + rtol*abs(reference), elementwise. + + No broadcasting, NaN equality, masking, or string/object/complex coercion. + Tolerances are finite nonnegative real scalars (not booleans). Computation + uses at least float64; integer/integer differences are formed exactly. + Mixed integer/float inputs that would round integers fail closed, as does + arithmetic overflow. max_rel uses abs(reference): 0/0 is 0; an unbounded or + unrepresentable maximum is None with max_rel_nonfinite=True. Uncomputed + discrepancies remain None. count is the number of shape-matched elements. + Distinct labels and non-aliasing cannot prove logical independence: copied + or fabricated references require external provenance review. + """ + tolerances = {} + for name, value in (("atol", atol), ("rtol", rtol)): + scalar = _float64(value, name) + if scalar.ndim != 0 or scalar < 0: + raise ValueError(f"{name} must be a nonnegative real scalar") + tolerances[name] = float(scalar) + evidence = dict(tolerances, observed_source=None, reference_source=None, + observed_shape=None, reference_shape=None, observed_dtype=None, + reference_dtype=None, count=0, max_abs=None, max_rel=None, + max_rel_nonfinite=False, mismatch_count=None) + try: + labels = [_label(v).strip() for v in (observed_source, reference_source)] + evidence.update(observed_source=labels[0], reference_source=labels[1]) + if labels[0] == labels[1]: + raise ValueError("provenance labels must be distinct") + except ValueError as exc: + return _result("unverified", str(exc), evidence) + if observed is None or reference is None: + return _result("unverified", "missing observed or reference data", evidence) + if observed is reference: + return _result("unverified", "same input object", evidence) + try: + a, b = np.asarray(observed), np.asarray(reference) + for name, array in (("observed", a), ("reference", b)): + evidence.update({f"{name}_shape": list(array.shape), + f"{name}_dtype": str(array.dtype)}) + if np.shares_memory(a, b): + return _result("unverified", "inputs share memory", evidence) + if (any(np.ma.isMaskedArray(v) for v in (observed, reference)) + or any(x is None for v in (a, b) if v.dtype.kind == "O" for x in v.flat)): + return _result("unverified", "missing or masked data", evidence) + if a.shape != b.shape: + return _result("fail", "shapes differ; broadcasting is forbidden", evidence) + evidence["count"] = int(a.size) + _real_array(a, "observed") + _real_array(b, "reference") + if not a.size: + return _result("unexercised", "zero eligible elements", evidence) + dtype = np.result_type(a.dtype, b.dtype, np.float64) + x, y = a.astype(dtype), b.astype(dtype) + both_integer = a.dtype.kind in "iu" and b.dtype.kind in "iu" + if not both_integer: + for raw, cast in ((a, x), (b, y)): + if raw.dtype.kind in "iu" and np.any(raw.astype(object) != cast.astype(object)): + raise ValueError("integer values lose precision during comparison") + with np.errstate(over="ignore", invalid="ignore", divide="ignore"): + delta = (np.abs(a.astype(object) - b.astype(object)).astype(dtype) + if both_integer else np.abs(x - y)) + limit = tolerances["atol"] + tolerances["rtol"] * np.abs(y) + if not np.isfinite(delta).all() or not np.isfinite(limit).all(): + raise ValueError("comparison arithmetic overflow") + relative = np.divide(delta, np.abs(y), out=np.zeros_like(delta), where=y != 0) + relative = np.where((y == 0) & (delta != 0), np.inf, relative) + max_abs, max_rel = float(delta.max()), float(relative.max()) + mismatches = int(np.count_nonzero(delta > limit)) + evidence.update(max_abs=max_abs if np.isfinite(max_abs) else None, + max_rel=max_rel if np.isfinite(max_rel) else None, + max_rel_nonfinite=not bool(np.isfinite(max_rel)), + mismatch_count=mismatches) + return _result("fail" if mismatches else "pass", + "outside tolerance" if mismatches else "within tolerance", evidence) + except (TypeError, ValueError, OverflowError) as exc: + return _result("fail", str(exc), evidence) + + +def independent_kinetic_energy(velocities, masses, graph_indices, num_graphs): + """Return float64 direct sums of 0.5*m*sum(v*v), in caller-supplied units. + + Shapes: velocities (N, 3), masses (N,), integer graph_indices (N,). + num_graphs is a positive integer scalar. Empty graphs have zero energy. + No toolkit helpers or unit conversions; non-finite arithmetic is rejected. + """ + v, m = _float64(velocities, "velocities"), _float64(masses, "masses") + g = _integers(graph_indices, "graph_indices", 0) + size = _integer_scalar(num_graphs, "num_graphs", 1) + if size > np.iinfo(np.intp).max: + raise ValueError("num_graphs exceeds index range") + if v.ndim != 2 or v.shape[1] != 3 or m.shape != (v.shape[0],) or g.shape != m.shape: + raise ValueError("expected velocities (N,3), masses (N,), graph_indices (N,)") + if np.any(m <= 0) or np.any(g >= size): + raise ValueError("masses must be positive and graph indices in range") + result = np.zeros(size, dtype=np.float64) + with np.errstate(over="ignore", invalid="ignore"): + np.add.at(result, g.astype(np.intp), 0.5 * m * np.sum(v * v, axis=1)) + return _float64(result, "kinetic energy result") + + +def independent_energy_drift(total, reference, num_atoms, step_count, metric="per_atom_per_step"): + """Return abs(total-reference)/(num_atoms*max(step_count,1)), or absolute. + + Energies must be matching scalars or 1-D arrays. Integer atom counts (>0) + and step counts (>=0) must be scalar or exactly the energy shape; integral + floats and booleans are rejected. Counts are validated for both metrics. + None reference means first-reference capture and returns None, never zero. + Inputs other than that absent reference are still validated in this case. + """ + if not isinstance(metric, str) or metric not in ("per_atom_per_step", "absolute"): + raise ValueError("unknown drift metric") + energy = _float64(total, "total") + atoms, steps = _integers(num_atoms, "num_atoms", 1), _integers(step_count, "step_count", 0) + if energy.ndim > 1 or any(a.shape not in ((), energy.shape) for a in (atoms, steps)): + raise ValueError("energies must be scalar/1-D; counts must be scalar or match") + if reference is None: + return None + ref = _float64(reference, "reference") + if energy.shape != ref.shape: + raise ValueError("energy and reference shapes differ") + with np.errstate(over="ignore", invalid="ignore"): + result = np.abs(energy - ref) + if metric == "per_atom_per_step": + result = result / (atoms.astype(np.float64) * np.maximum(steps, 1).astype(np.float64)) + return _float64(result, "energy drift result") + + +def _entry(counter, stage, hook): + return (_integer_scalar(counter, "step_counter"), _label(stage), _label(hook)) + + +def audit_hook_events(events, expected_entries): + """Audit measured dicts, in supplied order, against (counter, stage, hook) tuples. + + Required keys: sequence, event ('entry'/'return'), invocation_id, + step_counter, stage, hook_type. Sequence numbers are nonnegative integers, + strictly increasing (gaps allowed). Invocation IDs are nonempty strings or + nonnegative integers, unique per hook type and paired once. Pair metadata + must match, including optional instance_id. Nested calls are allowed. + Counts describe structurally valid records and matched pairs, not inferred + calls. Empty/None traces are unverified even when the expectation is empty. + Trace authenticity, like numeric provenance, needs external review. + """ + if not isinstance(expected_entries, (list, tuple)): + raise ValueError("expected_entries must be a list of triples") + expected = [] + for item in expected_entries: + if not isinstance(item, (tuple, list)) or len(item) != 3: + raise ValueError("expected entries must be (step_counter, stage, hook_type)") + expected.append(list(_entry(*item))) + evidence = dict(event_count=0, entry_count=0, return_count=0, pair_count=0, + by_hook={}, actual_entries=[], expected_entries=expected, errors=[]) + if events is None or isinstance(events, (list, tuple)) and not events: + return _result("unverified", "no measured events", evidence) + if not isinstance(events, (list, tuple)): + return _result("fail", "events must be a list of measured dictionaries", evidence) + evidence["event_count"] = len(events) + errors, seen, returned, previous = evidence["errors"], {}, set(), -1 + for index, event in enumerate(events): + try: + if not isinstance(event, dict): + raise ValueError("event must be a dictionary") + sequence = _integer_scalar(event["sequence"], "sequence") + kind = event["event"] + if not isinstance(kind, str) or kind not in ("entry", "return"): + raise ValueError("event must be entry or return") + entry = _entry(event["step_counter"], event["stage"], event["hook_type"]) + invocation = event["invocation_id"] + invocation = (_label(invocation) if isinstance(invocation, str) + else _integer_scalar(invocation, "invocation_id")) + instance = event.get("instance_id") + if "instance_id" in event: + instance = _label(instance) + except (KeyError, TypeError, ValueError, OverflowError) as exc: + errors.append(f"event {index}: {exc}") + continue + if sequence <= previous: + errors.append(f"event {index}: sequence must be strictly increasing and unique") + previous = sequence + counts = evidence["by_hook"].setdefault(entry[2], {"entry": 0, "return": 0, "pairs": 0}) + counts[kind] += 1 + evidence[f"{kind}_count"] += 1 + key, metadata = (entry[2], invocation), (entry, instance) + if kind == "entry": + evidence["actual_entries"].append(list(entry)) + if key in seen: + errors.append(f"event {index}: duplicate invocation entry") + else: + seen[key] = metadata + elif key not in seen or key in returned or seen[key] != metadata: + errors.append(f"event {index}: unpaired, duplicate, or mismatched return") + else: + returned.add(key) + counts["pairs"] += 1 + evidence["pair_count"] += 1 + if seen.keys() - returned: + errors.append("missing returns") + if evidence["actual_entries"] != expected: + errors.append("full entry order differs from expected_entries") + return _result("fail" if errors else "pass", + "; ".join(errors) if errors else "measured order and pairs match", evidence) \ No newline at end of file diff --git a/agents/nvidia-alchemi/tools/nvidia-alchemi/build_release_manifest.py b/agents/nvidia-alchemi/tools/nvidia-alchemi/build_release_manifest.py new file mode 100644 index 0000000..e745d01 --- /dev/null +++ b/agents/nvidia-alchemi/tools/nvidia-alchemi/build_release_manifest.py @@ -0,0 +1,97 @@ +"""Record build provenance: base image, package/checkpoint inventory, source hashes. + +Also preserves each installed distribution's license/notice files under +/app/licenses//. Release and base image come from the build +environment (ALCHEMI_AGENT_RELEASE, ALCHEMI_BASE_IMAGE); both are required. +""" +import hashlib +import importlib.metadata +import json +import os +from pathlib import Path +import platform +import re +import shutil +import yaml + +LICENSE_NAME = re.compile(r"(^|/)(LICEN[CS]E|COPYING|NOTICE|AUTHORS)[^/]*$", re.IGNORECASE) + + +def preserve_licenses(destination): + copied = {} + for dist in importlib.metadata.distributions(): + name = dist.metadata["Name"] + if not name: + continue + for entry in dist.files or (): + text = str(entry).replace("\\", "/") + if ".dist-info/" not in text or not LICENSE_NAME.search(text): + continue + source = Path(dist.locate_file(entry)) + if source.is_file(): + target = destination / name / text.split(".dist-info/", 1)[1] + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source, target) + copied[name] = copied.get(name, 0) + 1 + return copied + + +def sha256(path): + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +if __name__ == "__main__": + root = Path("/app") + release = os.environ["ALCHEMI_AGENT_RELEASE"] + base_image = os.environ["ALCHEMI_BASE_IMAGE"] + if "@sha256:" not in base_image: + raise SystemExit("ALCHEMI_BASE_IMAGE must be pinned by digest") + # Sources must match the repository byte for byte (LF), or the recorded hashes + # cannot be reproduced from a clone. Windows checkouts with autocrlf break this. + text_inputs = [*root.glob("*.py"), *(root / "tests").glob("*.py"), *(root / "build").iterdir(), + root / "agent-source.yaml", root / "tool-definition.yaml", root / "THIRD_PARTY_NOTICES.md"] + crlf = sorted(str(p) for p in text_inputs if p.is_file() and b"\r\n" in p.read_bytes()) + if crlf: + raise SystemExit(f"Windows line endings in build inputs; build from a clean git checkout: {crlf}") + agent = yaml.safe_load((root / "agent-source.yaml").read_text()) + tool = yaml.safe_load((root / "tool-definition.yaml").read_text()) + if tool["version"] != release or not tool["infra"][0]["image"]["acr"].endswith(":" + release): + raise SystemExit("Tool definition version/tag does not match ALCHEMI_AGENT_RELEASE") + (root / "agent-instructions.txt").write_text(agent["instructions"], encoding="utf-8") + packages = {dist.metadata["Name"]: dist.version for dist in importlib.metadata.distributions() if dist.metadata["Name"]} + files = {p.name: sha256(p) for p in root.glob("*.py")} + files["agent-instructions.txt"] = sha256(root / "agent-instructions.txt") + files["tool-definition.yaml"] = sha256(root / "tool-definition.yaml") + files["requirements-build.txt"] = sha256(root / "build" / "requirements-build.txt") + files["checkpoints.json"] = sha256(root / "build" / "checkpoints.json") + files["download_checkpoints.py"] = sha256(root / "build" / "download_checkpoints.py") + licenses = preserve_licenses(root / "licenses") + checkpoints = {} + for folder in (Path("/opt/cache/mace"), Path("/root/.cache/aimnet")): + if folder.exists(): + for p in sorted(folder.rglob("*")): + if p.is_file() and p.stat().st_size > 1_000_000: + checkpoints[str(p)] = {"size": p.stat().st_size, "sha256": sha256(p)} + expected = json.loads((root / "build" / "checkpoints.json").read_text())["checkpoints"] + for entry in expected: + if checkpoints.get(entry["path"]) != {"size": entry["size"], "sha256": entry["sha256"]}: + raise SystemExit(f"Checkpoint mismatch: {entry['path']}") + # Name each file so consumers can look up a model without guessing from filenames. + checkpoints[entry["path"]].update(loader=entry["loader"], name=entry["name"]) + unnamed = [p for p, v in checkpoints.items() if "name" not in v] + if unnamed: + raise SystemExit(f"Cached checkpoint(s) not listed in checkpoints.json: {unnamed}") + manifest = {"release": release, "base_image": base_image, + "python": platform.python_version(), "packages": dict(sorted(packages.items())), + "source_sha256": dict(sorted(files.items())), "checkpoints": checkpoints, + "license_files": dict(sorted(licenses.items())), + "dependency_policy": "Exact versions from requirements-build.txt installed " + "with --no-deps from public indexes; pip check passed"} + (root / "release-manifest.json").write_text(json.dumps(manifest, indent=2, allow_nan=False) + "\n") + (root / "requirements.lock.txt").write_text("\n".join(f"{n}=={v}" for n,v in sorted(packages.items())) + "\n") + print(json.dumps({"release": release, "packages": len(packages), "checkpoints": len(checkpoints), + "license_distributions": len(licenses)})) \ No newline at end of file diff --git a/agents/nvidia-alchemi/tools/nvidia-alchemi/checkpoints.json b/agents/nvidia-alchemi/tools/nvidia-alchemi/checkpoints.json new file mode 100644 index 0000000..ef83946 --- /dev/null +++ b/agents/nvidia-alchemi/tools/nvidia-alchemi/checkpoints.json @@ -0,0 +1,26 @@ +{ + "source": "Hashes recorded from the validated 1.1.0 image", + "checkpoints": [ + { + "loader": "mace", + "name": "small-0b", + "path": "/opt/cache/mace/mace_agnesi_smallmodel", + "size": 67603120, + "sha256": "7e3a0abcaf41e03a80e69f778e1b11b29de1cca704783dc25917a736392f8cf0" + }, + { + "loader": "mace", + "name": "medium-0b2", + "path": "/opt/cache/mace/macemediumdensityagnesistressmodel", + "size": 79462798, + "sha256": "a90be07c8aa6623c390fcc4653d3e319c4a356f8b4a53d323f96b2012d375caf" + }, + { + "loader": "aimnet2", + "name": "aimnet2", + "path": "/root/.cache/aimnet/aimnet2_wb97m_d3_0.pt", + "size": 8836941, + "sha256": "f0f7c054539ad3261bd36f9b11c56d12f87cb723e25bea7521755bbd3ec24e28" + } + ] +} diff --git a/agents/nvidia-alchemi/tools/nvidia-alchemi/download_checkpoints.py b/agents/nvidia-alchemi/tools/nvidia-alchemi/download_checkpoints.py new file mode 100644 index 0000000..297a8f5 --- /dev/null +++ b/agents/nvidia-alchemi/tools/nvidia-alchemi/download_checkpoints.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python3 +"""Cache the advertised MLIP checkpoints at build time and verify every one. + +Runtime jobs must work offline, so the build fails unless ALL checkpoints listed +in checkpoints.json are present with the expected size and SHA-256 and load. +UMA is intentionally excluded (gated weights; conflicting optional extras). +""" +import hashlib +import json +from pathlib import Path +import sys + + +def sha256(path): + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def load(entry): + # Build hosts have no GPU: load on CPU without cuEquivariance. + if entry["loader"] == "mace": + from nvalchemi.models.mace import MACEWrapper + MACEWrapper.from_checkpoint(entry["name"], device="cpu", enable_cueq=False) + elif entry["loader"] == "aimnet2": + from nvalchemi.models.aimnet2 import AIMNet2Wrapper + AIMNet2Wrapper.from_checkpoint(entry["name"], device="cpu") + else: + raise ValueError(f"Unknown loader {entry['loader']!r}") + + +def main(spec): + entries = json.loads(Path(spec).read_text(encoding="utf-8"))["checkpoints"] + failures = [] + for entry in entries: + label = f"{entry['loader']}:{entry['name']}" + try: + load(entry) + path = Path(entry["path"]) + if not path.is_file(): + raise FileNotFoundError(f"expected cache file {path} not created") + size, digest = path.stat().st_size, sha256(path) + if size != entry["size"] or digest != entry["sha256"]: + raise ValueError(f"{path}: size {size}, sha256 {digest}; expected " + f"{entry['size']}, {entry['sha256']}") + load(entry) # Reload from cache to prove the cached file is usable. + print(f"[ok] {label} {digest}", flush=True) + except Exception as exc: # report every failure, then fail the build + failures.append(f"{label}: {exc}") + print(f"[failed] {label}: {exc}", flush=True) + if failures: + print("Checkpoint caching failed:\n " + "\n ".join(failures)) + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1] if len(sys.argv) > 1 else "/app/build/checkpoints.json")) diff --git a/agents/nvidia-alchemi/tools/nvidia-alchemi/example-input-files/cu_bulk.extxyz b/agents/nvidia-alchemi/tools/nvidia-alchemi/example-input-files/cu_bulk.extxyz new file mode 100644 index 0000000..b3d8c44 --- /dev/null +++ b/agents/nvidia-alchemi/tools/nvidia-alchemi/example-input-files/cu_bulk.extxyz @@ -0,0 +1,6 @@ +4 +Lattice="3.615 0.0 0.0 0.0 3.615 0.0 0.0 0.0 3.615" Properties=species:S:1:pos:R:3 pbc="T T T" comment="fcc Cu conventional cell, a=3.615 A" +Cu 0.00000000 0.00000000 0.00000000 +Cu 0.00000000 1.80750000 1.80750000 +Cu 1.80750000 0.00000000 1.80750000 +Cu 1.80750000 1.80750000 0.00000000 diff --git a/agents/nvidia-alchemi/tools/nvidia-alchemi/example-input-files/ethanol.xyz b/agents/nvidia-alchemi/tools/nvidia-alchemi/example-input-files/ethanol.xyz new file mode 100644 index 0000000..1552d27 --- /dev/null +++ b/agents/nvidia-alchemi/tools/nvidia-alchemi/example-input-files/ethanol.xyz @@ -0,0 +1,11 @@ +9 +Ethanol (C2H5OH), gas phase +C -1.21376000 0.24003000 0.00000000 +C 0.06826000 -0.56494000 0.00000000 +O 1.19458000 0.29677000 0.00000000 +H -2.09220000 -0.40628000 0.00000000 +H -1.24377000 0.87688000 0.88817000 +H -1.24377000 0.87688000 -0.88817000 +H 0.10450000 -1.20871000 0.88232000 +H 0.10450000 -1.20871000 -0.88232000 +H 1.99429000 -0.23984000 0.00000000 diff --git a/agents/nvidia-alchemi/tools/nvidia-alchemi/example-input-files/si_bulk.cif b/agents/nvidia-alchemi/tools/nvidia-alchemi/example-input-files/si_bulk.cif new file mode 100644 index 0000000..4f07b6b --- /dev/null +++ b/agents/nvidia-alchemi/tools/nvidia-alchemi/example-input-files/si_bulk.cif @@ -0,0 +1,30 @@ +data_Si +_symmetry_space_group_name_H-M 'F d -3 m :1' +_symmetry_Int_Tables_number 227 +_cell_length_a 5.43070 +_cell_length_b 5.43070 +_cell_length_c 5.43070 +_cell_angle_alpha 90.00000 +_cell_angle_beta 90.00000 +_cell_angle_gamma 90.00000 +_chemical_formula_sum 'Si8' + +loop_ + _symmetry_equiv_pos_as_xyz + 'x, y, z' + +loop_ + _atom_site_label + _atom_site_type_symbol + _atom_site_fract_x + _atom_site_fract_y + _atom_site_fract_z + _atom_site_occupancy + Si0 Si 0.00000 0.00000 0.00000 1.0 + Si1 Si 0.00000 0.50000 0.50000 1.0 + Si2 Si 0.50000 0.00000 0.50000 1.0 + Si3 Si 0.50000 0.50000 0.00000 1.0 + Si4 Si 0.25000 0.25000 0.25000 1.0 + Si5 Si 0.25000 0.75000 0.75000 1.0 + Si6 Si 0.75000 0.25000 0.75000 1.0 + Si7 Si 0.75000 0.75000 0.25000 1.0 diff --git a/agents/nvidia-alchemi/tools/nvidia-alchemi/example-input-files/water.xyz b/agents/nvidia-alchemi/tools/nvidia-alchemi/example-input-files/water.xyz new file mode 100644 index 0000000..37b838c --- /dev/null +++ b/agents/nvidia-alchemi/tools/nvidia-alchemi/example-input-files/water.xyz @@ -0,0 +1,5 @@ +3 +Water molecule (H2O), gas phase +O 0.00000000 0.00000000 0.11926200 +H 0.00000000 0.76323900 -0.47704700 +H 0.00000000 -0.76323900 -0.47704700 diff --git a/agents/nvidia-alchemi/tools/nvidia-alchemi/release_smoke.py b/agents/nvidia-alchemi/tools/nvidia-alchemi/release_smoke.py new file mode 100644 index 0000000..6c8d081 --- /dev/null +++ b/agents/nvidia-alchemi/tools/nvidia-alchemi/release_smoke.py @@ -0,0 +1,111 @@ +"""Release smoke test: real toolkit IDs, reviewed validators, verified artifacts. + +Default run requires GPU. --cpu is only for build/local testing, never a hidden +fallback. The output deliberately includes NPZ/Python/text and nested Zarr data. +""" +import argparse +import hashlib +import json +import logging +from pathlib import Path +import shutil +import unittest + +import numpy as np +import torch +from ase import Atoms +import alchemi_utils as u +from alchemi_artifacts import write_json_atomic, to_jsonable +from alchemi_validators import compare_numeric, independent_kinetic_energy, independent_energy_drift, audit_hook_events + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--cpu", action="store_true") + parser.add_argument("--output", default="/output") + parser.add_argument("--work", default="/workdir") + args = parser.parse_args() + device = "cpu" if args.cpu else "cuda" + if device == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("Release GPU smoke requires CUDA; CPU fallback forbidden") + # Run regression tests BEFORE setup; some tests deliberately change globals. + import os + os.environ["ALCHEMI_TEST_GPU"] = "0" if args.cpu else "1" + suite = unittest.defaultTestLoader.discover(str(Path(__file__).parent / "tests")) + test_result = unittest.TextTestRunner(verbosity=2).run(suite) + if not test_result.wasSuccessful() or test_result.skipped: + raise RuntimeError("Release regressions failed or skipped") + + release = json.loads(Path("/app/release-manifest.json").read_text())["release"] + if os.environ.get("ALCHEMI_AGENT_RELEASE") != release: + raise RuntimeError("Image release variable disagrees with release manifest") + + u.quick_setup(output_dir=args.output, work_dir=args.work) + work = Path(u.WORK_DIR) + handler = logging.FileHandler(work / "smoke.log") + logging.getLogger().addHandler(handler) + outputs = {} + try: + sizes = [7, 3, 6, 2, 5, 4] + inputs = [Atoms("Ar" * n, positions=[[20*j, 0, 0] for j in range(n)]) for n in sizes] + for i,a in enumerate(inputs): + a.info["sample_label"] = f"sample-{i}" + model = u.load_model("lj", device=device) + screened = u.screen(inputs, model, fmax=0.001, max_steps=10, + max_batch_size=3, max_atoms=12, device=device) + ids = [{"input_index":a.info["input_index"], "system_id":a.info["system_id"], + "sample_label":a.info["sample_label"], "atoms":len(a)} for a in screened] + assert sorted(x["input_index"] for x in ids) == list(range(6)) + assert all(x["atoms"] == sizes[x["input_index"]] for x in ids) + assert any(x["input_index"] != x["system_id"] for x in ids) + + # Independently computed reference, with actual tensor evidence saved. + velocity = np.arange(18, dtype=np.float32).reshape(6,3) / 20 + masses = np.array([1,2,3,4,5,6], dtype=np.float32) + groups = np.array([0,0,0,1,1,1], dtype=np.int32) + from nvalchemi.dynamics.hooks._utils import kinetic_energy_per_graph + observed = kinetic_energy_per_graph(torch.tensor(velocity,device=device), + torch.tensor(masses,device=device),torch.tensor(groups,device=device),2).detach().cpu().numpy().reshape(-1) + reference = independent_kinetic_energy(velocity,masses,groups,2) + check = compare_numeric(observed, reference, observed_source="toolkit.runtime.kinetic_energy", + reference_source="numpy.direct.float64",atol=1e-5,rtol=1e-5) + assert check["status"] == "pass" + assert independent_energy_drift(np.array([1.]),None,np.array([6]),0) is None + assert compare_numeric(None,reference,observed_source="missing",reference_source="numpy",atol=0,rtol=0)["status"] == "unverified" + assert audit_hook_events([],[])["status"] == "unverified" + np.savez(work / "reference_arrays.npz", velocities=velocity,masses=masses,groups=groups, + observed=observed,reference=reference) + write_json_atomic(work / "identity_and_checks.json", {"identities":ids,"kinetic_energy":check}) + (work / "api_notes.txt").write_text("Release smoke: declared text artifacts are exported and hashed.\n") + shutil.copy2(__file__,work / "executed_smoke.py") + batch = u.atoms_to_batch([inputs[0]],device=device) + sink = u.zarr_sink("nested/trajectory.zarr",capacity=3) + sink.write(batch) + del sink + for name in ("release-manifest.json", "requirements.lock.txt", "agent-instructions.txt"): + shutil.copy2(Path("/app") / name, work / name) + outputs = {"arrays":"reference_arrays.npz", "script":"executed_smoke.py", + "notes":"api_notes.txt", "checks":"identity_and_checks.json", + "trajectory":"nested/trajectory.zarr", "log":"smoke.log", + "release":"release-manifest.json", "packages":"requirements.lock.txt", + "instructions":"agent-instructions.txt"} + logging.info("All smoke assertions completed successfully") + summary = {"release":release,"helper_version":u.HELPER_VERSION,"device":device, + "gpu":torch.cuda.get_device_name(0) if device=="cuda" else None, + "tests_run":test_result.testsRun,"tests_skipped":len(test_result.skipped), + "checks":{"identity":"pass","independent_energy":"pass"}, + "expected_nonfinite_diagnostics":[float("inf"),float("-inf"),float("nan")], + "helper_sha256":hashlib.sha256(Path(u.__file__).read_bytes()).hexdigest()} + u.save_final_results(summary, outputs) + except Exception as exc: + u.save_final_results({"error":str(exc)},output_files={},status="failed") + raise + finally: + logging.getLogger().removeHandler(handler) + handler.close() + u.quick_finish() + print("RELEASE_SMOKE_PASS", json.dumps(to_jsonable(summary),allow_nan=False)) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/agents/nvidia-alchemi/tools/nvidia-alchemi/requirements-build.txt b/agents/nvidia-alchemi/tools/nvidia-alchemi/requirements-build.txt new file mode 100644 index 0000000..609db64 --- /dev/null +++ b/agents/nvidia-alchemi/tools/nvidia-alchemi/requirements-build.txt @@ -0,0 +1,202 @@ +# Exact package set recorded in the validated 1.1.0 environment (198 distributions). +# Installed with --no-deps so no additional resolution occurs; pip check verifies consistency. +--extra-index-url https://download.pytorch.org/whl/cu130 +--extra-index-url https://pypi.nvidia.com +absl-py==2.5.0 +aimnet==0.2.0 +aiobotocore==3.9.0 +aiohappyeyeballs==2.7.1 +aiohttp==3.14.3 +aioitertools==0.13.0 +aiosignal==1.4.0 +annotated-types==0.8.0 +antlr4-python3-runtime==4.9.3 +anyio==4.14.2 +ase==3.29.0 +astunparse==1.6.3 +attrs==26.1.0 +beartype==0.22.9 +bibtexparser==1.4.4 +botocore==1.43.56 +cachetools==7.1.7 +certifi==2026.7.22 +cftime==1.6.5 +charset-normalizer==3.5.1 +click==8.4.2 +cloudpickle==3.1.2 +ConfigArgParse==1.7.5 +contourpy==1.3.3 +cuda-bindings==13.2.0 +cuda-core==0.7.0 +cuda-pathfinder==1.6.1 +cuda-python==13.2.0 +cuda-toolkit==13.0.3.0 +cudf-cu13==26.6.0 +cuequivariance==0.11.1 +cuequivariance-ops-cu13==0.11.1 +cuequivariance-ops-torch-cu13==0.11.1 +cuequivariance-torch==0.11.1 +cuml-cu13==26.6.0 +cupy-cuda13x==13.6.0 +cycler==0.12.1 +dm-tree==0.1.10 +donfig==0.8.1.post1 +e3nn==0.4.4 +einops==0.8.2 +fastrlock==0.8.3 +filelock==3.32.3 +fonttools==4.63.0 +frozenlist==1.8.0 +fsspec==2026.7.0 +gast==0.7.0 +gitdb==4.0.12 +GitPython==3.1.59 +google-crc32c==1.8.0 +h11==0.16.0 +h5py==3.16.0 +hf-xet==1.6.0 +httpcore==1.0.9 +httpx==0.28.1 +huggingface_hub==1.28.0 +hydra-core==1.3.5 +idna==3.19 +importlib_metadata==9.0.0 +jaxtyping==0.3.11 +Jinja2==3.1.6 +jmespath==1.1.0 +joblib==1.5.3 +kiwisolver==1.5.0 +libcst==1.9.0 +libcudf-cu13==26.6.0 +libcuml-cu13==26.6.0 +libkvikio-cu13==26.6.0 +libnvforest-cu13==26.6.0 +libraft-cu13==26.6.0 +librmm-cu13==26.6.0 +lightning-utilities==0.15.3 +llvmlite==0.46.0 +lmdb==2.3.0 +loguru==0.7.3 +lxml==6.1.1 +mace-torch==0.3.15 +makefun==1.16.0 +markdown-it-py==4.2.0 +MarkupSafe==3.0.3 +matplotlib==3.11.1 +matscipy==1.2.0 +mdurl==0.1.2 +ml_dtypes==0.6.0 +monty==2026.7.16 +mpmath==1.3.0 +multidict==6.7.1 +narwhals==2.24.0 +networkx==3.6.1 +numba==0.64.0 +numba-cuda==0.28.2 +numcodecs==0.16.5 +numpy==2.3.5 +nvalchemi-toolkit==0.2.0 +nvalchemi-toolkit-ops==0.4.1 +nvforest-cu13==26.6.0 +nvidia-cublas==13.1.1.3 +nvidia-cuda-cccl==13.0.85 +nvidia-cuda-crt==13.0.88 +nvidia-cuda-cupti==13.0.85 +nvidia-cuda-nvcc==13.0.88 +nvidia-cuda-nvrtc==13.0.88 +nvidia-cuda-runtime==13.0.96 +nvidia-cudnn-cu13==9.20.0.48 +nvidia-cufft==12.0.0.61 +nvidia-cufile==1.15.1.6 +nvidia-curand==10.4.0.35 +nvidia-cusolver==12.0.4.66 +nvidia-cusparse==12.6.3.3 +nvidia-cusparselt-cu13==0.8.1 +nvidia-dali-cuda130==2.2.0 +nvidia-libnvcomp-cu13==5.2.0.13 +nvidia-ml-py==13.610.43 +nvidia-nccl-cu13==2.29.7 +nvidia-nvimgcodec-cu13==0.8.0.22 +nvidia-nvjitlink==13.3.33 +nvidia-nvjpeg==13.2.1.68 +nvidia-nvjpeg2k-cu13==0.11.0.51 +nvidia-nvshmem-cu13==3.4.5 +nvidia-nvtiff-cu13==0.8.0.82 +nvidia-nvtx==13.0.85 +nvidia-nvvm==13.0.88 +nvidia-physicsnemo==2.1.1 +nvtx==0.2.16 +omegaconf==2.3.1 +onnx==1.22.0 +opt-einsum-fx==0.1.4 +opt_einsum==3.4.0 +optree==0.19.1 +orjson==3.12.0 +packaging==26.2 +palettable==3.3.3 +pandas==2.3.3 +periodictable==2.0.2 +pillow==12.3.0 +pip==26.2.1 +platformdirs==4.11.3 +plotext==5.3.2 +plotly==6.9.0 +plum-dispatch==2.9.0 +prettytable==3.18.0 +propcache==0.5.2 +protobuf==7.35.1 +pyarrow==23.0.1 +pydantic==2.13.4 +pydantic_core==2.46.4 +Pygments==2.21.0 +pylibcudf-cu13==26.6.0 +pylibraft-cu13==26.6.0 +pymatgen==2026.5.4 +pymatgen-core==2026.8.13 +pyparsing==3.3.2 +python-dateutil==2.9.0.post0 +python_hostlist==2.3.0 +pytz==2026.3.post1 +pyvers==0.2.3 +PyYAML==6.0.3 +rapids-logger==0.2.3 +requests==2.34.2 +rich==15.0.0 +rmm-cu13==26.6.0 +ruamel.yaml==0.19.1 +s3fs==2026.7.0 +safetensors==0.8.0 +scikit-learn==1.9.0 +scipy==1.18.0 +seaborn==0.13.2 +setuptools==84.0.0 +six==1.17.0 +smmap==5.0.3 +spglib==2.7.0 +sympy==1.14.0 +tabulate==0.10.0 +tensordict==0.11.0 +termcolor==3.3.0 +threadpoolctl==3.6.0 +timm==1.0.28 +torch==2.13.0+cu130 +torch-ema==0.3 +torchmetrics==1.9.0 +torchvision==0.28.0+cu130 +tqdm==4.70.0 +treelib==1.8.0 +treelite==4.7.0 +triton==3.7.1 +typing-inspection==0.4.4 +typing_extensions==4.16.0 +tzdata==2026.3 +uncertainties==3.2.3 +urllib3==2.7.0 +wadler_lindig==0.1.7 +warp-lang==1.16.0 +wcwidth==0.8.2 +wheel==0.48.0 +wrapt==2.3.0 +yarl==1.24.5 +zarr==3.3.0 +zipp==4.1.0 diff --git a/agents/nvidia-alchemi/tools/nvidia-alchemi/tests/test_artifacts.py b/agents/nvidia-alchemi/tools/nvidia-alchemi/tests/test_artifacts.py new file mode 100644 index 0000000..3a69b08 --- /dev/null +++ b/agents/nvidia-alchemi/tools/nvidia-alchemi/tests/test_artifacts.py @@ -0,0 +1,264 @@ +"""Standard-library-only artifact/JSON regressions; no GPU required.""" + +import hashlib +import json +import os +from pathlib import Path +import shutil +import tempfile +import unittest +from enum import Enum +from unittest.mock import patch + +import alchemi_artifacts as artifacts + + +class Array: + def __init__(self, value): + self.value = value + + def tolist(self): + return self.value + + +class Scalar: + def item(self): + return 7 + + +class Tensor(Array): + device = "cpu" + + def detach(self): + return self + + +class ArtifactTests(unittest.TestCase): + def setUp(self): + temporary = tempfile.TemporaryDirectory() + self.addCleanup(temporary.cleanup) + self.root = Path(temporary.name).resolve() # Windows TEMP may contain 8.3 aliases. + self.work, self.output = self.root / "work", self.root / "output" + self.work.mkdir() + self.output.mkdir() + + def put(self, name, data=b"payload", root=None): + path = (root or self.work) / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(data) + return path + + def export(self, **files): + return artifacts.export_artifacts(files, self.work, self.output) + + def test_arbitrary_files_directory_aliases_and_metadata(self): + declared = {name: self.put("nested/" + name) for name in ("state.npz", "notes.txt", "run.py")} + self.put("nested/store.zarr/z/1", b"123") + self.put("nested/store.zarr/a", b"45678") + (self.work / "nested/store.zarr/empty").mkdir() + declared.update(store="nested/store.zarr", alias="nested/state.npz", directory_alias="nested/store.zarr") + with patch.object(artifacts.shutil, "copy2", wraps=shutil.copy2) as copy: + mapped, manifest = self.export(**declared) + self.assertEqual(copy.call_count, 5) + self.assertEqual(mapped["alias"], mapped["state.npz"]) + self.assertEqual(mapped["store"], mapped["directory_alias"]) + self.assertEqual(manifest["schema_version"], 1) + for entry in manifest["artifacts"]: + path = Path(entry["path"]) + self.assertEqual(path, Path(mapped[entry["label"]])) + self.assertEqual(entry["relative_path"], path.relative_to(self.output).as_posix()) + self.assertTrue(entry["verified"]) + self.assertTrue(path.is_absolute()) + base = path if path.is_dir() else path.parent + for item in entry["files"]: + data = (base / item["relative_path"]).read_bytes() + self.assertEqual((item["size"], item["sha256"]), (len(data), hashlib.sha256(data).hexdigest())) + self.assertEqual(entry["total_bytes"], sum(item["size"] for item in entry["files"])) + if path.is_dir(): + self.assertEqual([f["relative_path"] for f in entry["files"]], ["a", "z/1"]) + canonical = json.dumps(entry["files"], sort_keys=True, separators=(",", ":")) + self.assertEqual(entry["sha256"], hashlib.sha256(canonical.encode()).hexdigest()) + self.assertTrue((path / "empty").is_dir()) + else: + self.assertEqual(entry["sha256"], entry["files"][0]["sha256"]) + self.assertFalse((self.output / "artifact_manifest.json").exists()) + nested = self.work / "exports" + mapped, _ = artifacts.export_artifacts({"a": "nested/run.py"}, self.work, nested) + self.assertEqual(Path(mapped["a"]), nested / "nested/run.py") + outside = self.put("external.pickle", root=self.root) + self.assertEqual(Path(self.export(a=outside)[0]["a"]), self.output / outside.name) + + def test_passthrough_rewrite_and_unrelated_preservation(self): + direct = self.put("direct.txt", root=self.output) + with patch.object(artifacts.shutil, "copy2", side_effect=AssertionError("same-file copy")): + self.assertEqual(self.export(direct=direct)[0]["direct"], str(direct)) + for root in (self.output, self.work): + self.put("tree/old", root=root) + self.put("data.bin", root=root) + self.put("unrelated/keep", root=self.output) + self.export(tree="tree", data="data.bin") + (self.work / "tree/old").unlink() + self.put("tree/new", b"new") + self.put("data.bin", b"changed") + self.export(tree="tree", data="data.bin") + self.assertFalse((self.output / "tree/old").exists()) + self.assertEqual((self.output / "tree/new").read_bytes(), b"new") + self.assertEqual((self.output / "data.bin").read_bytes(), b"changed") + self.assertTrue((self.output / "unrelated/keep").exists()) + self.assertFalse(list(self.output.glob(".alchemi-*"))) + with patch.object(artifacts.shutil, "copytree", side_effect=AssertionError("same-directory copy")): + self.export(tree=self.output / "tree") + + def test_invalid_declarations_and_reserved_reports(self): + self.put("tree/leaf") + outside = self.put("other/leaf", root=self.root) + self.put("leaf") + self.put("folder.bin") + (self.output / "folder.bin").mkdir() + self.put("blocked", root=self.output) + self.put("blocked/leaf") + cases = [dict(a="missing"), dict(a="../escape"), dict(a="tree/../leaf"), dict(a="..\\escape"), + dict(a="C:relative"), dict(a=""), dict(a=object()), dict(a=self.work), dict(a=self.output), + dict(a="tree", b="tree/leaf"), dict(a=outside, b="leaf"), dict(a="folder.bin"), + dict(a="blocked/leaf")] + for name in ("final_results.json", "ARTIFACT_MANIFEST.JSON"): + self.put(name) + self.put(name, b"preserve", self.output) + cases.append(dict(a=name)) + for files in cases: + with self.subTest(files=files), self.assertRaises((ValueError, OSError, TypeError)): + self.export(**files) + self.put("tree", root=self.output) + with self.assertRaises(ValueError): + self.export(a="tree") + self.assertEqual((self.output / "final_results.json").read_bytes(), b"preserve") + for name in ("final_results.json", "ARTIFACT_MANIFEST.JSON"): + (self.work / name).unlink() + self.put(name + "/payload") + with self.assertRaises(ValueError): + self.export(a=name) + with self.assertRaises(ValueError): + artifacts.export_artifacts({"a": self.work / "tree"}, self.work, self.work / "tree/out") + with self.assertRaises(FileNotFoundError): + self.export(valid="leaf", missing="absent") + self.assertFalse((self.output / "leaf").exists()) + + def test_portable_collisions(self): + a = self.put("a/STATE.bin", root=self.root) + b = self.put("b/state.bin", root=self.root) + with self.assertRaises(ValueError): + self.export(a=a, b=b) + self.put("state.bin", root=self.output) + with self.assertRaises(ValueError): + self.export(a=a) + if os.name != "nt": + self.put("tree/A") + self.put("tree/a") + with self.assertRaises(ValueError): + self.export(a="tree") + self.put("Parent/a") + self.put("parent/b") + with self.assertRaises(ValueError): + self.export(a="Parent/a", b="parent/b") + + def test_symlinks(self): + source = self.put("tree/leaf") + link = self.work / "link" + try: + link.symlink_to(source) + except (OSError, NotImplementedError): + self.skipTest("Symlinks not permitted on this platform") + with self.assertRaises(ValueError): + self.export(a="link") + link.unlink() + link.symlink_to(source.parent, target_is_directory=True) + with self.assertRaises(ValueError): + self.export(a="link/leaf") + (source.parent / "link").symlink_to(source) + with self.assertRaises(ValueError): + self.export(a="tree") + (self.output / "tree").symlink_to(source.parent, target_is_directory=True) + with self.assertRaises(ValueError): + self.export(a="tree/leaf") + with self.assertRaises(ValueError): + artifacts.export_artifacts({"a": source}, self.work, self.output / "tree/out") + self.put("other/leaf") + (self.output / "other").mkdir() + (self.output / "other/link").symlink_to(source) + with self.assertRaises(ValueError): + self.export(a="other") + + def test_copy_corruption_and_directory_rollback(self): + original = shutil.copy2 + + def corrupt(source, destination, **kwargs): + result = original(source, destination, **kwargs) + Path(destination).write_bytes(b"bad") # Same size; detecting it requires hashing. + return result + + for name in ("file.bin", "tree/leaf"): + self.put(name, b"new") + self.put(name, b"old", self.output) + for name in ("file.bin", "tree"): + with patch.object(artifacts.shutil, "copy2", side_effect=corrupt), self.assertRaises(OSError): + self.export(a=name) + self.assertEqual((self.output / "file.bin").read_bytes(), b"old") + self.assertEqual((self.output / "tree/leaf").read_bytes(), b"old") + original_replace = os.replace + + for fail_before in (True, False): + def fail_publish(source, destination): + publishing = Path(source).parent.name == "candidate" + if publishing and fail_before: + raise OSError("injected publication failure") + result = original_replace(source, destination) + if publishing: + (Path(destination) / "leaf").write_bytes(b"bad") + return result + + with patch.object(artifacts.os, "replace", side_effect=fail_publish), self.assertRaises(OSError): + self.export(a="tree") + self.assertEqual((self.output / "tree/leaf").read_bytes(), b"old") + self.assertFalse(list(self.output.glob(".alchemi-*"))) + + def test_strict_json_and_atomic_preservation(self): + class Choice(Enum): + A = "chosen" + + value = {"values": Array([float("nan"), float("inf"), float("-inf")]), + "tensor": Tensor([[Scalar(), 2]]), "path": Path("file"), "enum": Choice.A, + "tuple": (1, None), 3: True} + clean = artifacts.to_jsonable(value) + self.assertEqual([v["encoding"] for v in clean["values"]], ["nan", "+inf", "-inf"]) + self.assertTrue(all(v["value"] is None and v["finite"] is False for v in clean["values"])) + self.assertEqual(clean["tensor"], [[7, 2]]) + self.assertEqual((clean["enum"], clean["path"], clean["tuple"], clean["3"]), ("chosen", "file", [1, None], True)) + path = self.output / "final_results.json" + artifacts.write_json_atomic(path, value) + saved = path.read_bytes() + self.assertEqual(json.loads(saved, parse_constant=lambda x: self.fail(x)), clean) + cycle = [] + cycle.append(cycle) + cyclic_dict = {} + cyclic_dict["self"] = cyclic_dict + cyclic_array = Array(None) + cyclic_array.value = [cyclic_array] + gpu = Tensor([1]) + gpu.device = "cuda:0" + shared = [1] + self.assertEqual(artifacts.to_jsonable([shared, shared]), [[1], [1]]) + for bad in (object(), {1, 2}, b"bytes", 1j, cycle, cyclic_dict, cyclic_array, gpu, + {1: "a", "1": "b"}, {Path("a"): 1, "a": 2}, {(1,): 2}, {float("nan"): 1}): + with self.subTest(bad=type(bad)), self.assertRaises((TypeError, ValueError)): + artifacts.write_json_atomic(path, bad) + self.assertEqual(path.read_bytes(), saved) + with patch.object(artifacts.os, "replace", side_effect=OSError("replace failed")), self.assertRaises(OSError): + artifacts.write_json_atomic(path, {"new": 1}) + self.assertEqual(path.read_bytes(), saved) + artifacts.write_json_atomic(path, {"new": 1}) + self.assertEqual(json.loads(path.read_bytes()), {"new": 1}) + self.assertFalse(list(self.output.glob(".json-*"))) + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file diff --git a/agents/nvidia-alchemi/tools/nvidia-alchemi/tests/test_utils_hardening.py b/agents/nvidia-alchemi/tools/nvidia-alchemi/tests/test_utils_hardening.py new file mode 100644 index 0000000..97f3de8 --- /dev/null +++ b/agents/nvidia-alchemi/tools/nvidia-alchemi/tests/test_utils_hardening.py @@ -0,0 +1,219 @@ +"""Regression tests for helper wiring; dynamics tests use the real toolkit.""" +import hashlib +import importlib.util +import json +import os +from pathlib import Path +import tempfile +import unittest +from unittest.mock import patch + +import alchemi_utils as u + + +class OutputContractTests(unittest.TestCase): + def setUp(self): + self.cwd = os.getcwd() + self.tmp = tempfile.TemporaryDirectory() + self.root = Path(self.tmp.name) + self.work, self.out = self.root / "work", self.root / "out" + with patch.object(u, "log_gpu_info", return_value={}): + u.quick_setup(str(self.root / "input"), str(self.out), str(self.work)) + + def tearDown(self): + os.chdir(self.cwd) + self.tmp.cleanup() + + def test_npz_python_text_declared_and_reverified(self): + files = {} + for suffix in ("npz", "py", "txt", "unusual"): + p = self.work / ("evidence." + suffix) + p.write_bytes(b"first") + files[suffix] = str(p) + mapped = u.save_final_results({"nonfinite": float("inf")}, files) + self.assertTrue(all(Path(p).parent == self.out.resolve() for p in mapped.values())) + (self.work / "evidence.txt").write_bytes(b"updated-after-save") + u.quick_finish() + result = json.loads((self.out / "final_results.json").read_text()) + self.assertEqual(result["summary"]["nonfinite"]["encoding"], "+inf") + manifest = json.loads((self.out / "artifact_manifest.json").read_text()) + for entry in manifest["artifacts"]: + self.assertEqual(entry["sha256"], hashlib.sha256(Path(entry["path"]).read_bytes()).hexdigest()) + self.assertEqual((self.out / "evidence.txt").read_bytes(), b"updated-after-save") + + def test_missing_evidence_fails_not_completed(self): + with self.assertRaises(FileNotFoundError): + u.save_final_results({}, {"missing": "absent.npz"}) + result = json.loads((self.out / "final_results.json").read_text()) + self.assertEqual(result["status"], "failed") + self.assertEqual(result["output_files"], {}) + with self.assertRaises(FileNotFoundError): + u.quick_finish() + + def test_empty_failure_can_clear_invalid_declarations(self): + with self.assertRaises(FileNotFoundError): + u.save_final_results({}, {"missing": "absent.txt"}) + u.save_final_results({"error": "handled"}, output_files={}, status="failed") + u.quick_finish() + self.assertEqual(json.loads((self.out / "final_results.json").read_text())["status"], "failed") + + def test_legacy_finish_does_not_overwrite_final_report(self): + (self.work / "final_results.json").write_text('{"status":"stale"}') + (self.work / "artifact_manifest.json").write_text('{}') + (self.work / "sample.npz").write_bytes(b"sample") + u.quick_finish() + self.assertTrue((self.out / "sample.npz").exists()) + self.assertFalse((self.out / "final_results.json").exists()) + + +@unittest.skipUnless(importlib.util.find_spec("nvalchemi") and importlib.util.find_spec("ase"), + "toolkit and ASE needed (run inside image)") +class IdentityTests(unittest.TestCase): + def setUp(self): + import torch + self.device = "cuda" if os.getenv("ALCHEMI_TEST_GPU") == "1" else "cpu" + if self.device == "cuda" and not torch.cuda.is_available(): + self.fail("GPU integration explicitly requested, but unavailable") + + def test_nonperiodic_roundtrip_preserves_ids_masses_velocities(self): + from ase import Atoms + import numpy as np + a = Atoms("H2", positions=[[0, 0, 0], [1, 0, 0]], masses=[2, 3]) + a.info.update(system_id=73, input_index=8) + b = u.atoms_to_batch([a], device=self.device) + b.velocities.fill_(0.25) + restored = u.batch_to_atoms(b)[0] + self.assertEqual(restored.info["system_id"], 73) + self.assertEqual(restored.info["input_index"], 8) + np.testing.assert_array_equal(restored.get_masses(), [2, 3]) + np.testing.assert_array_equal(restored.arrays["alchemi_velocities"], np.full((2, 3), 0.25)) + self.assertFalse(restored.pbc.any()) + + def test_identity_survives_reorder_and_partial_selection(self): + from ase import Atoms + import torch + items = [Atoms("Ar" * n, positions=[[i * 20, 0, 0] for i in range(n)]) for n in [7, 3, 5]] + for i, a in enumerate(items): + a.info.update(system_id=100+i, input_index=i) + b = u.atoms_to_batch(items, device=self.device) + selected = b.index_select(torch.tensor([2, 0], device=self.device)) + result = u.batch_to_atoms(selected) + self.assertEqual([a.info["input_index"] for a in result], [2, 0]) + self.assertEqual([a.info["system_id"] for a in result], [102, 100]) + u._restore_screen_identity(result, items) + self.assertEqual([a.info["source_system_id"] for a in result], [102, 100]) + + def test_screen_mixed_sizes_matches_input_not_admission(self): + from ase import Atoms + sizes = [7, 3, 6, 2, 5, 4] + items = [Atoms("Ar" * n, positions=[[j * 20, 0, 0] for j in range(n)]) for n in sizes] + for i, a in enumerate(items): + a.info["sample_label"] = "input-" + str(i) + model = u.load_model("lj", device=self.device) + result = u.screen(items, model, fmax=0.001, max_steps=10, + max_batch_size=3, max_atoms=12, device=self.device) + self.assertEqual(len(result), len(items)) + self.assertEqual(sorted(a.info["input_index"] for a in result), list(range(6))) + for a in result: + i = a.info["input_index"] + self.assertEqual(len(a), sizes[i]) + self.assertEqual(a.info["sample_label"], "input-" + str(i)) + self.assertTrue(any(a.info["input_index"] != a.info["system_id"] for a in result)) + + def test_screen_partial_returns_explicit_input_identity(self): + from ase import Atoms + items = [Atoms("Ar2", positions=[[0, 0, 0], [3.6, 0, 0]]), + Atoms("Ar", positions=[[0, 0, 0]])] + model = u.load_model("lj", device=self.device) + result = u.screen(items, model, fmax=1e-8, max_steps=1, + max_batch_size=2, device=self.device) + self.assertEqual([a.info["input_index"] for a in result], [1]) + + def test_invalid_duplicate_or_missing_identity_rejected(self): + from ase import Atoms + a = Atoms("Ar") + with self.assertRaises(RuntimeError): + u._restore_screen_identity([a], [a]) + a.info.update(system_id=1, input_index=0) + with self.assertRaises(RuntimeError): + u._restore_screen_identity([a.copy(), a.copy()], [a]) + self.assertEqual(u.screen([], None), []) + with self.assertRaises(ValueError): + u.screen([a, a], None, sink_capacity=1) + + def test_single_point_matches_analytic_lennard_jones(self): + from ase import Atoms + eps, sigma = u.MODEL_DEFAULTS["lj"]["epsilon"], u.MODEL_DEFAULTS["lj"]["sigma"] + model = u.load_model("lj", device=self.device) + cases = [(sigma, 0.0, 24 * eps / sigma), (2 ** (1 / 6) * sigma, -eps, 0.0)] + for r, energy, force in cases: + b = u.atoms_to_batch([Atoms("Ar2", positions=[[0, 0, 0], [r, 0, 0]])], device=self.device) + out = u.single_point(b, model) + self.assertEqual(tuple(out["energy"].shape), (1, 1)) + self.assertEqual(out["energy"].device.type, self.device) + self.assertAlmostEqual(float(out["energy"][0, 0]), energy, delta=1e-6) + # Repulsive at sigma: atom 0 is pushed toward -x. + self.assertAlmostEqual(float(out["forces"][0, 0]), -force, delta=1e-5) + self.assertAlmostEqual(float(out["forces"].sum()), 0.0, delta=1e-6) + self.assertAlmostEqual(float(b.energy[0, 0]), energy, delta=1e-6) + + def test_write_structures_accepts_single_atoms_list_and_batch(self): + from ase import Atoms + from ase.io import read + import tempfile + a = Atoms("H2O", positions=[[0, 0, 0], [0.757, 0, 0.587], [-0.757, 0, 0.587]], pbc=False) + with tempfile.TemporaryDirectory() as tmp: + for name, value in (("single", a), ("list", [a, a]), + ("batch", u.atoms_to_batch([a], device=self.device))): + path = u.write_structures(value, os.path.join(tmp, name + ".extxyz")) + written = read(path, index=":") + self.assertEqual(len(written), 2 if name == "list" else 1, name) + self.assertEqual(written[0].get_chemical_formula(), "H2O", name) + + def test_step_recorder_rows_match_independent_kinetic_energy(self): + from ase import Atoms + import torch + from alchemi_validators import independent_kinetic_energy + model = u.load_model("lj", device=self.device) + items = [Atoms("Ar2", positions=[[0, 0, 0], [3.9, 0, 0]]), + Atoms("Ar3", positions=[[0, 0, 0], [3.8, 0, 0], [0, 3.8, 0]])] + b = u.atoms_to_batch(items, device=self.device) + g = torch.Generator().manual_seed(0) + b.velocities.copy_(0.01 * torch.randn(tuple(b.velocities.shape), generator=g).to(self.device)) + recorder = u.step_recorder() + out = u.run_md(b, model, ensemble="nve", dt=1.0, n_steps=4, hooks=[recorder]) + self.assertEqual([(r["step"], r["graph"]) for r in recorder.rows], + [(s, i) for s in range(1, 5) for i in range(2)]) + last = [r for r in recorder.rows if r["step"] == 4] + reference = independent_kinetic_energy(out.velocities.cpu().numpy(), out.atomic_masses.cpu().numpy(), + out.batch_idx.cpu().numpy(), 2) + kb = 8.617333262145e-5 + for i, (row, atoms) in enumerate(zip(last, (2, 3))): + self.assertAlmostEqual(row["kinetic_energy_eV"], float(reference[i]), delta=1e-7) + self.assertAlmostEqual(row["temperature_K"], 2 * float(reference[i]) / (3 * atoms * kb), + delta=1e-3) + self.assertAlmostEqual(row["energy_eV"], float(out.energy.reshape(-1)[i]), delta=1e-7) + with self.assertRaises(ValueError): + u.step_recorder(0) + + +MANIFEST = Path("/app/release-manifest.json") + + +@unittest.skipUnless(MANIFEST.is_file(), "release manifest only exists inside the built image") +class ReleaseManifestTests(unittest.TestCase): + def test_every_cached_checkpoint_is_named_and_documented_names_exist(self): + checkpoints = json.loads(MANIFEST.read_text())["checkpoints"] + for path, entry in checkpoints.items(): + self.assertTrue(entry.get("loader") and entry.get("name"), path) + self.assertTrue(Path(path).is_file(), path) + named = {(e["loader"], e["name"]): p for p, e in checkpoints.items()} + self.assertEqual(named, { + ("mace", "small-0b"): "/opt/cache/mace/mace_agnesi_smallmodel", + ("mace", "medium-0b2"): "/opt/cache/mace/macemediumdensityagnesistressmodel", + ("aimnet2", "aimnet2"): "/root/.cache/aimnet/aimnet2_wb97m_d3_0.pt", + }) + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file diff --git a/agents/nvidia-alchemi/tools/nvidia-alchemi/tests/test_validators.py b/agents/nvidia-alchemi/tools/nvidia-alchemi/tests/test_validators.py new file mode 100644 index 0000000..4f8709c --- /dev/null +++ b/agents/nvidia-alchemi/tools/nvidia-alchemi/tests/test_validators.py @@ -0,0 +1,189 @@ +"""Local CPU regression tests; run with unittest discovery and NumPy installed.""" + +import copy +import json +import unittest + +import numpy as np + +from alchemi_validators import (audit_hook_events, compare_numeric, + independent_energy_drift, independent_kinetic_energy) + + +class NumericTests(unittest.TestCase): + def compare(self, observed, reference, **kwargs): + options = dict(observed_source="measurement", reference_source="independent", + atol=0.0, rtol=0.0) + options.update(kwargs) + result = compare_numeric(observed, reference, **options) + self.assertEqual(json.loads(json.dumps(result, allow_nan=False)), result) + return result + + def test_verdicts_and_evidence(self): + for a, b, status in [([1.0], [1.0], "pass"), ([2.0], [1.0], "fail"), + (None, [1.0], "unverified"), ([1.0], None, "unverified"), + ([None], [1.0], "unverified"), ([], [], "unexercised"), + (np.empty((0, 3)), np.empty((0, 3)), "unexercised")]: + with self.subTest(status=status, a=a): + result = self.compare(a, b) + self.assertEqual(result["status"], status) + if status in ("unverified", "unexercised"): + self.assertIsNone(result["evidence"]["max_abs"]) + self.assertIsNone(result["evidence"]["max_rel"]) + self.assertIsNone(result["evidence"]["mismatch_count"]) + evidence = self.compare(np.array([2], dtype=np.float32), [1.0])["evidence"] + self.assertEqual((evidence["count"], evidence["observed_shape"], + evidence["observed_dtype"], evidence["max_abs"], + evidence["max_rel"], evidence["mismatch_count"]), + (1, [1], "float32", 1.0, 1.0, 1)) + + def test_tolerance_convention_and_precision(self): + self.assertEqual(self.compare([1.75], [1.0], atol=.5, rtol=.25)["status"], "pass") + self.assertEqual(self.compare([2.0], [1.0], rtol=.5)["status"], "fail") + self.assertEqual(self.compare([1.0], [2.0], rtol=.5)["status"], "pass") + self.assertEqual(self.compare([1 + 1e-12], [1.0])["status"], "fail") + self.assertEqual(self.compare([2**63 + 1], [2**63])["evidence"]["max_abs"], 1.0) + self.assertEqual(self.compare([2**63 + 1], [float(2**63)])["status"], "fail") + self.assertEqual(self.compare([0.0], [-0.0])["evidence"]["max_rel"], 0.0) + evidence = self.compare([1.0], [0.0], atol=1.0)["evidence"] + self.assertIsNone(evidence["max_rel"]) + self.assertTrue(evidence["max_rel_nonfinite"]) + self.assertEqual(self.compare([np.finfo(float).max], [-np.finfo(float).max])["status"], "fail") + + def test_unverified_provenance_and_aliasing(self): + for label in (None, "", " ", 1, "independent", " independent "): + self.assertEqual(self.compare([1], [1], observed_source=label)["status"], "unverified") + a = np.arange(6.0) + for b in (a, a.view(), a[::-1], a.reshape(2, 3)): + self.assertEqual(self.compare(a, b)["status"], "unverified") + same = [1.0] + self.assertEqual(self.compare(same, same)["status"], "unverified") + self.assertEqual(self.compare(a, a.copy())["status"], "pass") + self.assertEqual(self.compare(np.ma.array([1], mask=[True]), [1])["status"], "unverified") + + def test_invalid_tolerances_and_data(self): + for name in ("atol", "rtol"): + for bad in (-1, np.nan, np.inf, -np.inf, True, "1", None, [1], 1j): + with self.subTest(name=name, bad=bad), self.assertRaises(ValueError): + self.compare([1], [1], **{name: bad}) + for bad in ([np.nan], [np.inf], [-np.inf], [True], ["1"], [1j], + [{}], [[1], [2, 3]], [[1]], [1, 2], np.empty((0, 3))): + with self.subTest(bad=bad): + self.assertEqual(self.compare(bad, [1.0])["status"], "fail") + self.assertEqual(self.compare([1.0], bad)["status"], "fail") + + +class IndependentCalculationTests(unittest.TestCase): + def test_kinetic_energy_direct_sum_and_precision(self): + result = independent_kinetic_energy([[1, 2, 2], [0, 3, 4], [2, 0, 0]], + [2, 4, 3], [1, 0, 1], np.int64(3)) + np.testing.assert_array_equal(result, [50, 15, 0]) + self.assertEqual(result.dtype, np.float64) + v = np.array([[1.0001, 2, 3]], dtype=np.float32) + expected = .5 * sum(float(x)**2 for x in v[0]) + self.assertEqual(independent_kinetic_energy(v, [1], [0], 1)[0], expected) + np.testing.assert_array_equal(independent_kinetic_energy( + np.empty((0, 3)), [], np.array([], dtype=int), 2), [0, 0]) + + def test_kinetic_energy_invalid_inputs(self): + good = dict(velocities=[[1, 2, 3]], masses=[1], graph_indices=[0], num_graphs=1) + cases = {"velocities": ([1, 2, 3], [[1, 2]], [[np.nan, 0, 0]], [[1e308, 0, 0]], [[True]*3]), + "masses": ([0], [-1], [np.inf], [[1]], [], ["1"], [1j]), + "graph_indices": ([-1], [1], [2**64-1], [0.5], [0.0], [True], [[0]], [], [np.nan]), + "num_graphs": (0, -1, 1.5, 1.0, True, [1], np.inf, 2**64-1)} + for key, values in cases.items(): + for value in values: + with self.subTest(key=key, value=value), self.assertRaises(ValueError): + independent_kinetic_energy(**dict(good, **{key: value})) + + def test_drift_metrics_and_first_reference(self): + np.testing.assert_allclose(independent_energy_drift([14, 8], [10, 10], [2, 4], 2), [1, .25]) + np.testing.assert_array_equal(independent_energy_drift([14, 8], [10, 10], 2, [0, 1]), [2, 1]) + np.testing.assert_array_equal(independent_energy_drift([14, 8], [10, 10], 2, 99, "absolute"), [4, 2]) + self.assertEqual(independent_energy_drift(10.000000000001, 10.0, 1, 0), abs(10.000000000001 - 10)) + self.assertIsNone(independent_energy_drift([1.0], None, 1, 0)) + self.assertIsNone(independent_energy_drift(1.0, None, 1, 0, "absolute")) + self.assertGreater(independent_energy_drift(2.0, 1.0, np.int64(2**62), 8), 0) + + def test_drift_invalid_inputs_even_without_reference(self): + good = dict(total=[1.0], reference=[1.0], num_atoms=1, step_count=0) + cases = {"total": (None, [[1]], [np.nan], ["1"], [True]), + "reference": (1.0, [[1]], [1, 2], [np.inf], [1j]), + "num_atoms": (0, -1, 1.1, 1.0, True, [1, 2], [[1]]), + "step_count": (-1, .5, 1.0, True, [0, 1], [[0]], np.nan), + "metric": ("relative", None)} + for key, values in cases.items(): + for value in values: + for reference in ([1.0], None) if key != "reference" else ([1.0],): + with self.subTest(key=key, value=value, reference=reference), self.assertRaises(ValueError): + independent_energy_drift(**{**good, "reference": reference, key: value}) + with self.assertRaises(ValueError): + independent_energy_drift(1e308, -1e308, 1, 1) + + +class HookTests(unittest.TestCase): + def setUp(self): + self.expected = [(0, "before", "clamp"), (0, "after", "drift"), (1, "before", "clamp")] + self.events = [] + for invocation, (counter, stage, hook) in enumerate(self.expected): + for kind in ("entry", "return"): + self.events.append(dict(sequence=len(self.events), event=kind, + invocation_id=invocation, step_counter=counter, + stage=stage, hook_type=hook)) + + def audit(self, events, status, expected=None): + result = audit_hook_events(events, self.expected if expected is None else expected) + self.assertEqual(result["status"], status, result) + self.assertEqual(json.loads(json.dumps(result, allow_nan=False)), result) + return result["evidence"] + + def test_actual_order_and_counts(self): + evidence = self.audit(self.events, "pass") + self.assertEqual([evidence[k] for k in ("event_count", "entry_count", "return_count", "pair_count")], [6, 3, 3, 3]) + self.assertEqual(evidence["by_hook"]["clamp"], dict(entry=2, **{"return": 2}, pairs=2)) + self.assertEqual(evidence["actual_entries"], [list(e) for e in self.expected]) + self.audit(self.events, "fail", self.expected[::-1]) + swapped = self.events[2:4] + self.events[:2] + self.events[4:] + swapped = [dict(event, sequence=i) for i, event in enumerate(swapped)] + self.audit(swapped, "fail") # Same counts and valid pairs, wrong full order. + self.audit(None, "unverified") + self.audit([], "unverified", []) + + def test_missing_duplicate_and_unpaired_events(self): + for index in range(len(self.events)): + self.audit(self.events[:index] + self.events[index+1:], "fail") + for index in (0, 1): + duplicate = self.events[:index+1] + [self.events[index]] + self.events[index+1:] + self.audit([dict(e, sequence=i) for i, e in enumerate(duplicate)], "fail") + self.audit(self.events[:4], "fail") # A whole invocation missing. + self.audit([self.events[1], self.events[0]] + self.events[2:], "fail") + + def test_bad_schema_sequence_and_pair_metadata(self): + for key, value in [("sequence", 0), ("sequence", -1), ("sequence", 1.5), + ("sequence", True), ("step_counter", 1), ("stage", "other"), + ("hook_type", "other"), ("instance_id", "wrong-instance"), + ("invocation_id", "unknown"), ("invocation_id", []), + ("event", "call"), ("stage", " ")]: + with self.subTest(key=key, value=value): + events = copy.deepcopy(self.events) + events[1][key] = value + self.audit(events, "fail") + for key in self.events[0]: + events = copy.deepcopy(self.events) + del events[0][key] + self.audit(events, "fail") + self.audit([None], "fail") + self.audit({"calls": 3}, "fail") + for bad in (None, [()], [(0.0, "before", "clamp")], [(0, "", "clamp")]): + with self.assertRaises(ValueError): + audit_hook_events(self.events, bad) + + def test_nested_calls_and_sequence_gaps(self): + nested = [self.events[i] for i in (0, 2, 3, 1, 4, 5)] + nested = [dict(event, sequence=i*3, invocation_id=str(event["invocation_id"])) + for i, event in enumerate(nested)] + self.audit(nested, "pass") + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file diff --git a/agents/nvidia-alchemi/tools/nvidia-alchemi/tool.yaml b/agents/nvidia-alchemi/tools/nvidia-alchemi/tool.yaml new file mode 100644 index 0000000..b43f876 --- /dev/null +++ b/agents/nvidia-alchemi/tools/nvidia-alchemi/tool.yaml @@ -0,0 +1,85 @@ +name: nvidiaAlchemi +description: > + NVIDIA ALCHEMI Toolkit for high-throughput, GPU-first AI atomic simulation. + Runs batched geometry relaxation and molecular dynamics (NVE, NVT, NPT, NPH) + driven by machine-learned interatomic potentials (MACE, AIMNet2) composed with + classical terms (Lennard-Jones, DFT-D3(BJ) dispersion, Ewald and PME + electrostatics). Supports multi-stage fused pipelines, inflight batching for + screening campaigns, Zarr trajectory capture, MLIP training and fine-tuning, + and multi-GPU execution via DistributedPipeline and domain decomposition. + GPU required for inference. +version: 1.1.2 +category: "AI & Machine Learning" +license: Apache-2.0 + +infra: + - name: worker + infra_type: container + image: + acr: "{name}.azurecr.io/nvidia-alchemi:1.1.2" + compute: + min_resources: + cpu: 8 + ram: "64Gi" + gpu: 1 + storage: "0Gi" + max_resources: + cpu: 96 + ram: "1800Gi" + gpu: 8 + storage: "0Gi" + # InfiniBand is NOT required. ALCHEMI's DistributedPipeline and + # DomainParallel operate within a single node (NVLink/PCIe). RDMA would + # only matter for multi-node NCCL, which this tool does not launch. + # Setting this true acts as a hard scheduler filter and excludes + # non-RDMA GPU pools such as Standard_NC40ads_H100_v5. + infiniband: false + recommended_sku: + - Standard_NC40ads_H100_v5 # Single-GPU relaxation and MD + - Standard_NC80adis_H100_v5 # 2-GPU pipelines + - Standard_ND96isr_H100_v5 # 8-GPU screening / domain decomposition + - Standard_NC24ads_A100_v4 # Single-GPU A100 alternative + - Standard_NC48ads_A100_v4 # 2-GPU A100 alternative + pool_type: static + pool_size: 1 + +code_environments: + - language: python3 + command: 'python3 -u "/{{scriptName}}" > {{output_directory}}/logs.out 2>&1; code=$?; tail -n 100 {{output_directory}}/logs.out; exit "$code"' + description: | + Python code environment with the NVIDIA ALCHEMI Toolkit for batched AI + atomic simulation. GPU required -- there is no CPU fallback for the + accelerated neighbor-list and interaction kernels. + + Key capabilities: + - Batched geometry and cell relaxation (FIRE, FIRE2) with convergence hooks + - Molecular dynamics: NVE, NVT Langevin, NVT Nose-Hoover, NPT (MTK barostat), NPH + - Multi-stage fused pipelines (relax -> equilibrate -> production) sharing + a single model forward pass per step + - High-throughput screening with inflight batching and SizeAwareSampler + - Zarr trajectory capture and CUDA-stream-prefetching data loading + - MLIP training and fine-tuning (energy/force/stress losses, EMA, checkpoints) + - Multi-GPU DistributedPipeline and single-system domain decomposition, + launched with torchrun + + Bundled potentials: MACE (foundation checkpoints baked into the image), + AIMNet2, Lennard-Jones, DFT-D3(BJ), Ewald, PME. Models compose with the + `+` operator or an explicit PipelineModelWrapper. + + Not included: UMA / fairchem-core. Its `uma` extra is mutually exclusive + with the `mace` and CUDA extras this image is built on, and its checkpoints + require gated HuggingFace access. + + Python libraries: nvalchemi, torch (CUDA 13), ase, pymatgen, zarr, numpy, + scipy, pandas, matplotlib, seaborn, rich, alchemi_utils, alchemi_artifacts, + alchemi_validators. Fixed-geometry energies/forces: alchemi_utils.single_point; + per-step dynamics records: alchemi_utils.step_recorder. + + Outputs: declare every deliverable in save_final_results; files of any + type are copied and hash-verified under /output. Match screening results + to inputs by input_index, not the sampler's admission-order system_id. + /app/release-manifest.json lists package versions, checkpoint paths and + hashes by model name, and source hashes. Sample structures are in + /app/example-input-files. + + infra_node: worker