diff --git a/alcf/polaris/mlip/benchmark/README.md b/alcf/polaris/mlip/benchmark/README.md new file mode 100644 index 0000000..e55688f --- /dev/null +++ b/alcf/polaris/mlip/benchmark/README.md @@ -0,0 +1,104 @@ +# MLIP throughput benchmarks on Polaris + +These scripts measure **performance** of MatKit's MLIP backends — how many +structures per second a backend sustains, how timing scales with structure size, +and how NVIDIA ALCHEMI native batching compares to sequential ASE MACE. + +This is distinct from `../smoke.py`, which is a correctness/integration test and +explicitly **not** a performance benchmark. Run the smoke test first to confirm +the backend works on the node, then run these to characterize throughput. + +All numbers are hardware-, checkpoint-, and version-specific. Every run writes a +`bench_meta.json` capturing the GPU model, package versions, and arguments — +keep it with any results you report. + +## Prerequisites + +Install the MLIP environment and pass the smoke test as described in +[`../README.md`](../README.md): + +```bash +export MATKIT_MLIP_ENV=/lus/eagle/projects///envs/matkit-mlip +bash alcf/polaris/mlip/install.sh +qsub -v MATKIT_MLIP_ENV="$MATKIT_MLIP_ENV" alcf/polaris/mlip/smoke.pbs +``` + +Edit the `#PBS -A PROJECT` line in the `.pbs` scripts before submitting. + +## Single-GPU sweeps + +```bash +qsub -v MATKIT_MLIP_ENV="$MATKIT_MLIP_ENV" \ + alcf/polaris/mlip/benchmark/bench.pbs +``` + +Runs two sweeps for both `nvalchemi-mace` and `ase-mace`: + +- **batch_size** — a fixed set of `--n-structures` copies, ALCHEMI batch size + swept over `--batch-sizes`. ASE MACE is sequential (batch size 1) and serves as + the baseline. Throughput should rise with batch size, then plateau when the GPU + saturates. Peak GPU memory is sampled from `nvidia-smi`. +- **structure_size** — one structure per supercell factor (`--size-factors`), + measuring per-structure time vs atom count. + +Override defaults at submit time: + +```bash +qsub -v MATKIT_MLIP_ENV="$MATKIT_MLIP_ENV",\ +MATKIT_BENCH_INPUT=/path/to/structure.cif,\ +MACE_CHECKPOINT=medium,\ +MATKIT_BENCH_DRIVER=opt,\ +MATKIT_BENCH_NSTRUCT=128 \ + alcf/polaris/mlip/benchmark/bench.pbs +``` + +Run `bench.py --help` for the full flag list (checkpoint, dtype, driver, steps, +batch sizes, size factors). + +## Multi-GPU node throughput + +Polaris nodes have 4x A100. The runner uses one GPU per process, so node-level +throughput is measured by launching one process per GPU over a round-robin shard +of the inputs: + +```bash +qsub -v MATKIT_MLIP_ENV="$MATKIT_MLIP_ENV" \ + alcf/polaris/mlip/benchmark/bench_multigpu.pbs +``` + +`node_result.json` reports aggregate `node_throughput_structs_per_s`. Compare it +to the single-GPU throughput at the same batch size — ideal scaling is ~4x. + +## Reading results + +`bench.py` writes `bench_results.jsonl`, one JSON object per configuration: + +| field | meaning | +|-------|---------| +| `sweep` | `batch_size` or `structure_size` | +| `backend` | `nvalchemi-mace` or `ase-mace` | +| `batch_size`, `n_structures`, `n_atoms` | configuration | +| `setup_time_s` | model load time (from the manifest) | +| `wall_time_s` | batch execution wall time (from the manifest) | +| `elapsed_s` | end-to-end `run_mlip_batch` time | +| `throughput_structs_per_s` | `succeeded / elapsed_s` | +| `calc_time_s_mean` / `_max` | per-item calculation time | +| `peak_gpu_mem_mib` | max used GPU memory sampled during the run | + +Quick summary with `jq`: + +```bash +jq -r 'select(.sweep=="batch_size") | + [.backend, .batch_size, .throughput_structs_per_s, .peak_gpu_mem_mib] | @tsv' \ + bench_results.jsonl +``` + +## Caveats + +- Let model weights finish downloading before trusting timings. The PBS scripts + do a throwaway warm-up run first; the ALCF HTTP proxy is exported for + compute-node downloads. +- `setup_time_s` (model load) is reported separately from per-item time; when + comparing backends, look at steady-state throughput, not the first item. +- These measure MatKit's execution path, not raw kernel performance, and do not + establish model parity or scientific accuracy. diff --git a/alcf/polaris/mlip/benchmark/bench.pbs b/alcf/polaris/mlip/benchmark/bench.pbs new file mode 100644 index 0000000..12832ed --- /dev/null +++ b/alcf/polaris/mlip/benchmark/bench.pbs @@ -0,0 +1,59 @@ +#!/bin/bash -l +#PBS -N matkit-mlip-bench +#PBS -l select=1:system=polaris +#PBS -l place=scatter +#PBS -l walltime=01:00:00 +#PBS -l filesystems=home:eagle +#PBS -q debug +#PBS -A PROJECT + +# Single-GPU throughput benchmark for MatKit MLIP backends. +# Submit from the MatKit checkout: +# qsub -v MATKIT_MLIP_ENV="$MATKIT_MLIP_ENV" alcf/polaris/mlip/benchmark/bench.pbs +# Override inputs/checkpoint/output as needed, e.g.: +# qsub -v MATKIT_MLIP_ENV=...,MATKIT_BENCH_INPUT=/path/to.cif,\ +# MACE_CHECKPOINT=medium alcf/polaris/mlip/benchmark/bench.pbs + +set -euo pipefail + +: "${MATKIT_MLIP_ENV:?Submit with -v MATKIT_MLIP_ENV=/path/to/env}" + +MATKIT_REPO="${MATKIT_REPO:-${PBS_O_WORKDIR}}" +INPUT_FILE="${MATKIT_BENCH_INPUT:-${MATKIT_REPO}/tests/data/test_structure.cif}" +OUTPUT_DIR="${MATKIT_BENCH_OUTPUT:-${MATKIT_REPO}/projects/mlip_bench_${PBS_JOBID}}" +MACE_CHECKPOINT="${MACE_CHECKPOINT:-medium}" +DRIVER="${MATKIT_BENCH_DRIVER:-energy}" + +module use /soft/modulefiles +module load conda/2025-09-25 +source "${MATKIT_MLIP_ENV}/bin/activate" + +export HTTP_PROXY="http://proxy.alcf.anl.gov:3128" +export HTTPS_PROXY="http://proxy.alcf.anl.gov:3128" +export http_proxy="${HTTP_PROXY}" +export https_proxy="${HTTPS_PROXY}" +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" +export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}" + +cd "${MATKIT_REPO}" + +# Pre-warm model weights so downloads do not pollute timed runs. A throwaway +# single-structure batch forces the checkpoint to resolve before the sweeps. +python -m matkit.cli mlip run \ + --backend nvalchemi-mace \ + --checkpoint "${MACE_CHECKPOINT}" \ + --device cuda \ + --dtype float32 \ + --driver energy \ + --input "${INPUT_FILE}" \ + --output "${OUTPUT_DIR}_warmup/result.json" || true + +python alcf/polaris/mlip/benchmark/bench.py \ + --input "${INPUT_FILE}" \ + --output-dir "${OUTPUT_DIR}" \ + --checkpoint "${MACE_CHECKPOINT}" \ + --driver "${DRIVER}" \ + --backends nvalchemi-mace ase-mace \ + --n-structures "${MATKIT_BENCH_NSTRUCT:-64}" \ + --batch-sizes ${MATKIT_BENCH_BATCH_SIZES:-1 2 4 8 16 32 64} \ + --size-factors ${MATKIT_BENCH_SIZE_FACTORS:-1 2 3} diff --git a/alcf/polaris/mlip/benchmark/bench.py b/alcf/polaris/mlip/benchmark/bench.py new file mode 100644 index 0000000..80c1d6e --- /dev/null +++ b/alcf/polaris/mlip/benchmark/bench.py @@ -0,0 +1,305 @@ +#!/usr/bin/env python3 +"""Throughput benchmark for MatKit MLIP backends on a single GPU. + +Unlike ``smoke.py`` (correctness/integration evidence), this script measures +performance: how many structures per second a backend sustains, how timing +scales with structure size, and how NVIDIA ALCHEMI native batching compares to +sequential ASE MACE. Results are hardware- and checkpoint-specific; record the +GPU model, MatKit commit, and package versions alongside any numbers. + +Each sweep runs ``run_mlip_batch`` in-process and reads timings from the +returned manifest. One backend per process keeps GPU runtime state isolated. +""" + +from __future__ import annotations + +import argparse +from importlib.metadata import PackageNotFoundError, version +import json +import subprocess +import sys +import time +from pathlib import Path + +from ase.build import make_supercell +from ase.io import read, write +import numpy as np + +from matkit.mlip import ( + ASEMACEConfig, + MLIPCalculationConfig, + NVAlchemiMACEConfig, + run_mlip_batch, +) + + +def _backend_config(name, checkpoint, dtype): + """Build a backend config with GPU defaults matched across backends.""" + if name == "ase-mace": + return ASEMACEConfig( + checkpoint=checkpoint or "medium", device="cuda", dtype=dtype + ) + if name == "nvalchemi-mace": + return NVAlchemiMACEConfig( + checkpoint=checkpoint or "medium", device="cuda", dtype=dtype + ) + raise ValueError(f"Unsupported benchmark backend: {name}") + + +def _gpu_memory_mib(): + """Best-effort peak used GPU memory via nvidia-smi; None if unavailable.""" + try: + result = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=memory.used", + "--format=csv,noheader,nounits", + ], + capture_output=True, + text=True, + timeout=30, + ) + except (OSError, subprocess.TimeoutExpired): + return None + if result.returncode != 0: + return None + values = [ + int(line) for line in result.stdout.split("\n") if line.strip().isdigit() + ] + return max(values) if values else None + + +def _gpu_name(): + """Best-effort GPU model name via nvidia-smi; None if unavailable.""" + try: + result = subprocess.run( + ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], + capture_output=True, + text=True, + timeout=30, + ) + except (OSError, subprocess.TimeoutExpired): + return None + if result.returncode != 0: + return None + return result.stdout.strip() or None + + +def _run_once(input_files, backend, calculation, output_dir, batch_size): + """Run one batch and summarize timing from the manifest and results.""" + started = time.perf_counter() + summary = run_mlip_batch( + input_files, + backend, + calculation=calculation, + output_dir=output_dir, + batch_size=batch_size, + ) + elapsed = time.perf_counter() - started + results = summary["results"] + calc_times = [ + r["calculation_time_s"] + for r in results + if r.get("success") and r.get("calculation_time_s") is not None + ] + succeeded = summary["succeeded"] + throughput = succeeded / elapsed if elapsed > 0 else 0.0 + return { + "status": summary["status"], + "total": summary["total"], + "succeeded": succeeded, + "failed": summary["failed"], + "setup_time_s": summary["setup_time_s"], + "wall_time_s": summary["wall_time_s"], + "elapsed_s": elapsed, + "throughput_structs_per_s": throughput, + "calc_time_s_mean": float(np.mean(calc_times)) if calc_times else None, + "calc_time_s_max": float(np.max(calc_times)) if calc_times else None, + "peak_gpu_mem_mib": _gpu_memory_mib(), + "manifest_file": summary["manifest_file"], + } + + +def _write_row(handle, row): + handle.write(json.dumps(row, allow_nan=False) + "\n") + handle.flush() + + +def _prepare_replicas(base_atoms, count, work_dir): + """Write ``count`` copies of a base structure; return their paths.""" + work_dir.mkdir(parents=True, exist_ok=True) + paths = [] + for index in range(count): + path = work_dir / f"struct_{index:04d}.extxyz" + write(path, base_atoms) + paths.append(str(path)) + return paths + + +def _prepare_sizes(base_atoms, factors, work_dir): + """Write supercells scaled by cubic diagonal factors; return (path, n).""" + work_dir.mkdir(parents=True, exist_ok=True) + entries = [] + for factor in factors: + supercell = make_supercell(base_atoms, np.diag([factor, factor, factor])) + path = work_dir / f"size_{len(supercell):06d}.extxyz" + write(path, supercell) + entries.append((str(path), len(supercell))) + return entries + + +def _sweep_batch_size(args, base_atoms, out, handle): + """Fixed structure set; vary ALCHEMI batch size (ASE ignores batch).""" + input_files = _prepare_replicas( + base_atoms, args.n_structures, out / "inputs_batch" + ) + for backend_name in args.backends: + backend = _backend_config(backend_name, args.checkpoint, args.dtype) + calculation = MLIPCalculationConfig(driver=args.driver, steps=args.steps) + sizes = args.batch_sizes if backend_name == "nvalchemi-mace" else [1] + for batch_size in sizes: + run_dir = out / "runs" / f"batch_{backend_name}_bs{batch_size}" + metrics = _run_once( + input_files, backend, calculation, run_dir, batch_size + ) + _write_row( + handle, + { + "sweep": "batch_size", + "backend": backend_name, + "driver": args.driver, + "checkpoint": args.checkpoint, + "dtype": args.dtype, + "batch_size": batch_size, + "n_structures": args.n_structures, + "n_atoms": len(base_atoms), + **metrics, + }, + ) + print( + f"[batch_size] {backend_name} bs={batch_size}: " + f"{metrics['throughput_structs_per_s']:.2f} struct/s" + ) + + +def _sweep_structure_size(args, base_atoms, out, handle): + """One structure per size; measure time vs atom count.""" + entries = _prepare_sizes(base_atoms, args.size_factors, out / "inputs_size") + for backend_name in args.backends: + backend = _backend_config(backend_name, args.checkpoint, args.dtype) + calculation = MLIPCalculationConfig(driver=args.driver, steps=args.steps) + for path, n_atoms in entries: + run_dir = out / "runs" / f"size_{backend_name}_{n_atoms}" + metrics = _run_once([path], backend, calculation, run_dir, 1) + _write_row( + handle, + { + "sweep": "structure_size", + "backend": backend_name, + "driver": args.driver, + "checkpoint": args.checkpoint, + "dtype": args.dtype, + "batch_size": 1, + "n_structures": 1, + "n_atoms": n_atoms, + **metrics, + }, + ) + print( + f"[structure_size] {backend_name} n_atoms={n_atoms}: " + f"{metrics['calc_time_s_mean']} s" + ) + + +def _package_versions(): + packages = {} + for name in ("matkit", "ase", "mace-torch", "nvalchemi-toolkit", "torch"): + try: + packages[name] = version(name) + except PackageNotFoundError: + packages[name] = None + return packages + + +def build_parser(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input", required=True, type=Path) + parser.add_argument("--output-dir", required=True, type=Path) + parser.add_argument( + "--sweeps", + nargs="+", + default=["batch_size", "structure_size"], + choices=["batch_size", "structure_size"], + ) + parser.add_argument( + "--backends", + nargs="+", + default=["nvalchemi-mace", "ase-mace"], + choices=["nvalchemi-mace", "ase-mace"], + help="ase-mace serves as the sequential baseline.", + ) + parser.add_argument("--checkpoint", default="medium") + parser.add_argument("--dtype", default="float32", choices=["float32", "float64"]) + parser.add_argument("--driver", default="energy", choices=["energy", "opt"]) + parser.add_argument("--steps", type=int, default=1000) + parser.add_argument( + "--n-structures", + type=int, + default=64, + help="Structure count for the batch-size sweep.", + ) + parser.add_argument( + "--batch-sizes", + nargs="+", + type=int, + default=[1, 2, 4, 8, 16, 32, 64], + ) + parser.add_argument( + "--size-factors", + nargs="+", + type=int, + default=[1, 2, 3], + help="Cubic supercell factors for the structure-size sweep.", + ) + return parser + + +def main(argv=None): + args = build_parser().parse_args(argv) + if args.n_structures < 1 or args.steps < 1: + raise SystemExit("--n-structures and --steps must be positive") + if any(size < 1 for size in args.batch_sizes): + raise SystemExit("--batch-sizes must be positive") + if any(factor < 1 for factor in args.size_factors): + raise SystemExit("--size-factors must be positive") + out = args.output_dir.resolve() + if out.exists(): + raise SystemExit("--output-dir must be a new directory") + out.mkdir(parents=True) + + base_atoms = read(args.input) + meta = { + "input_file": str(args.input.resolve()), + "base_n_atoms": len(base_atoms), + "packages": _package_versions(), + "gpu": _gpu_name(), + "args": { + key: (str(value) if isinstance(value, Path) else value) + for key, value in vars(args).items() + }, + } + (out / "bench_meta.json").write_text(json.dumps(meta, indent=2)) + + results_path = out / "bench_results.jsonl" + with results_path.open("w", encoding="utf-8") as handle: + if "batch_size" in args.sweeps: + _sweep_batch_size(args, base_atoms, out, handle) + if "structure_size" in args.sweeps: + _sweep_structure_size(args, base_atoms, out, handle) + + print(f"Benchmark results: {results_path}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/alcf/polaris/mlip/benchmark/bench_multigpu.pbs b/alcf/polaris/mlip/benchmark/bench_multigpu.pbs new file mode 100644 index 0000000..ac05773 --- /dev/null +++ b/alcf/polaris/mlip/benchmark/bench_multigpu.pbs @@ -0,0 +1,56 @@ +#!/bin/bash -l +#PBS -N matkit-mlip-bench-4gpu +#PBS -l select=1:system=polaris +#PBS -l place=scatter +#PBS -l walltime=01:00:00 +#PBS -l filesystems=home:eagle +#PBS -q debug +#PBS -A PROJECT + +# Node-level (4x A100) throughput benchmark. One process per GPU over a sharded +# input set. Submit from the MatKit checkout: +# qsub -v MATKIT_MLIP_ENV="$MATKIT_MLIP_ENV" \ +# alcf/polaris/mlip/benchmark/bench_multigpu.pbs + +set -euo pipefail + +: "${MATKIT_MLIP_ENV:?Submit with -v MATKIT_MLIP_ENV=/path/to/env}" + +MATKIT_REPO="${MATKIT_REPO:-${PBS_O_WORKDIR}}" +INPUT_FILE="${MATKIT_BENCH_INPUT:-${MATKIT_REPO}/tests/data/test_structure.cif}" +OUTPUT_DIR="${MATKIT_BENCH_OUTPUT:-${MATKIT_REPO}/projects/mlip_bench_4gpu_${PBS_JOBID}}" +MACE_CHECKPOINT="${MACE_CHECKPOINT:-medium}" + +module use /soft/modulefiles +module load conda/2025-09-25 +source "${MATKIT_MLIP_ENV}/bin/activate" + +export HTTP_PROXY="http://proxy.alcf.anl.gov:3128" +export HTTPS_PROXY="http://proxy.alcf.anl.gov:3128" +export http_proxy="${HTTP_PROXY}" +export https_proxy="${HTTPS_PROXY}" +export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}" +# Do NOT pin CUDA_VISIBLE_DEVICES here; bench_multigpu.py sets it per worker. +unset CUDA_VISIBLE_DEVICES || true + +cd "${MATKIT_REPO}" + +# Pre-warm weights on one GPU before the sharded run. +CUDA_VISIBLE_DEVICES=0 python -m matkit.cli mlip run \ + --backend nvalchemi-mace \ + --checkpoint "${MACE_CHECKPOINT}" \ + --device cuda \ + --dtype float32 \ + --driver energy \ + --input "${INPUT_FILE}" \ + --output "${OUTPUT_DIR}_warmup/result.json" || true + +python alcf/polaris/mlip/benchmark/bench_multigpu.py \ + --input "${INPUT_FILE}" \ + --output-dir "${OUTPUT_DIR}" \ + --n-gpus "${MATKIT_BENCH_NGPUS:-4}" \ + --backend nvalchemi-mace \ + --checkpoint "${MACE_CHECKPOINT}" \ + --driver "${MATKIT_BENCH_DRIVER:-energy}" \ + --n-structures "${MATKIT_BENCH_NSTRUCT:-256}" \ + --batch-size "${MATKIT_BENCH_BATCH_SIZE:-16}" diff --git a/alcf/polaris/mlip/benchmark/bench_multigpu.py b/alcf/polaris/mlip/benchmark/bench_multigpu.py new file mode 100644 index 0000000..ea2deee --- /dev/null +++ b/alcf/polaris/mlip/benchmark/bench_multigpu.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python3 +"""Node-level throughput: one MLIP process per GPU over a sharded input set. + +The MatKit runner uses a single GPU per process (``_execute_inputs`` loads one +model on ``config.device``). A Polaris node has 4x A100, so node throughput is +measured by launching one ``run_mlip_batch`` process per GPU, each pinned with +``CUDA_VISIBLE_DEVICES`` to a disjoint shard of the inputs, then aggregating the +per-process manifests. + +Each worker runs ``examples`` -style batching via a small inline driver so no +runner changes are needed. Structures are round-robin sharded across GPUs. +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import time +from pathlib import Path + +from ase.io import read, write + + +_WORKER = r""" +import json, sys +from pathlib import Path +from matkit.mlip import ( + ASEMACEConfig, NVAlchemiMACEConfig, MLIPCalculationConfig, run_mlip_batch, +) +cfg = json.loads(sys.argv[1]) +inputs = cfg["inputs"] +if cfg["backend"] == "nvalchemi-mace": + backend = NVAlchemiMACEConfig( + checkpoint=cfg["checkpoint"], device="cuda", dtype=cfg["dtype"]) +else: + backend = ASEMACEConfig( + checkpoint=cfg["checkpoint"], device="cuda", dtype=cfg["dtype"]) +calc = MLIPCalculationConfig(driver=cfg["driver"], steps=cfg["steps"]) +summary = run_mlip_batch( + inputs, backend, calculation=calc, + output_dir=cfg["output_dir"], batch_size=cfg["batch_size"]) +print(json.dumps({ + "succeeded": summary["succeeded"], + "failed": summary["failed"], + "wall_time_s": summary["wall_time_s"], + "manifest_file": summary["manifest_file"], +})) +""" + + +def _prepare_inputs(base_atoms, count, work_dir): + work_dir.mkdir(parents=True, exist_ok=True) + paths = [] + for index in range(count): + path = work_dir / f"struct_{index:04d}.extxyz" + write(path, base_atoms) + paths.append(str(path)) + return paths + + +def _shard(paths, n_gpus): + """Round-robin split so each GPU gets a comparable load.""" + shards = [[] for _ in range(n_gpus)] + for index, path in enumerate(paths): + shards[index % n_gpus].append(path) + return shards + + +def build_parser(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input", required=True, type=Path) + parser.add_argument("--output-dir", required=True, type=Path) + parser.add_argument("--n-gpus", type=int, default=4) + parser.add_argument( + "--backend", + default="nvalchemi-mace", + choices=["nvalchemi-mace", "ase-mace"], + ) + parser.add_argument("--checkpoint", default="medium") + parser.add_argument("--dtype", default="float32", choices=["float32", "float64"]) + parser.add_argument("--driver", default="energy", choices=["energy", "opt"]) + parser.add_argument("--steps", type=int, default=1000) + parser.add_argument("--n-structures", type=int, default=256) + parser.add_argument("--batch-size", type=int, default=16) + return parser + + +def main(argv=None): + args = build_parser().parse_args(argv) + if args.n_gpus < 1 or args.n_structures < 1: + raise SystemExit("--n-gpus and --n-structures must be positive") + out = args.output_dir.resolve() + if out.exists(): + raise SystemExit("--output-dir must be a new directory") + out.mkdir(parents=True) + + base_atoms = read(args.input) + paths = _prepare_inputs(base_atoms, args.n_structures, out / "inputs") + shards = _shard(paths, args.n_gpus) + + procs = [] + started = time.perf_counter() + for gpu_index, shard in enumerate(shards): + if not shard: + continue + cfg = { + "inputs": shard, + "backend": args.backend, + "checkpoint": args.checkpoint, + "dtype": args.dtype, + "driver": args.driver, + "steps": args.steps, + "batch_size": args.batch_size, + "output_dir": str(out / f"gpu_{gpu_index}"), + } + env = dict(os.environ, CUDA_VISIBLE_DEVICES=str(gpu_index)) + proc = subprocess.Popen( + [sys.executable, "-c", _WORKER, json.dumps(cfg)], + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + procs.append((gpu_index, proc)) + + workers = [] + total_succeeded = 0 + total_failed = 0 + for gpu_index, proc in procs: + stdout, stderr = proc.communicate() + record = {"gpu": gpu_index, "return_code": proc.returncode} + try: + record.update(json.loads(stdout.strip().splitlines()[-1])) + except (ValueError, IndexError): + record["error"] = stderr.strip()[-2000:] + total_succeeded += record.get("succeeded", 0) + total_failed += record.get("failed", 0) + workers.append(record) + (out / f"gpu_{gpu_index}_worker.log").write_text(stderr) + + wall = time.perf_counter() - started + node_result = { + "n_gpus": args.n_gpus, + "backend": args.backend, + "batch_size": args.batch_size, + "n_structures": args.n_structures, + "wall_time_s": wall, + "node_throughput_structs_per_s": ( + total_succeeded / wall if wall > 0 else 0.0 + ), + "total_succeeded": total_succeeded, + "total_failed": total_failed, + "workers": workers, + } + (out / "node_result.json").write_text(json.dumps(node_result, indent=2)) + print(json.dumps(node_result, indent=2)) + return 1 if total_failed else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/matkit/cli.py b/src/matkit/cli.py index 1c66f89..b3a363a 100644 --- a/src/matkit/cli.py +++ b/src/matkit/cli.py @@ -1044,7 +1044,7 @@ def _mlip_options(function): "--driver", default="energy", show_default=True, - type=click.Choice(["energy", "opt"]), + type=click.Choice(["energy", "opt", "md"]), ), click.option( "--optimizer", @@ -1116,6 +1116,41 @@ def _mlip_options(function): is_flag=True, help="Enable cuEquivariance in ALCHEMI MACE.", ), + click.option( + "--ensemble", + default="nvt", + show_default=True, + type=click.Choice(["nve", "nvt"]), + help="MD ensemble (--driver md).", + ), + click.option( + "--temperature", + default=300.0, + show_default=True, + type=float, + help="MD target temperature in Kelvin (--driver md).", + ), + click.option( + "--timestep", + default=1.0, + show_default=True, + type=float, + help="MD timestep in femtoseconds (--driver md).", + ), + click.option( + "--md-steps", + default=1000, + show_default=True, + type=int, + help="Number of MD steps (--driver md).", + ), + click.option( + "--friction", + default=0.01, + show_default=True, + type=float, + help="Langevin friction for NVT (--driver md).", + ), ] for decorator in reversed(decorators): function = decorator(function) @@ -1204,6 +1239,11 @@ def _build_mlip_configs(options): optimizer=options["optimizer"], fmax=options["fmax"], steps=options["steps"], + ensemble=options["ensemble"], + temperature=options["temperature"], + timestep=options["timestep"], + md_steps=options["md_steps"], + friction=options["friction"], ) if ( backend_name == "nvalchemi-mace" @@ -1211,6 +1251,10 @@ def _build_mlip_configs(options): and calculation.optimizer != "fire" ): raise ValueError("NVIDIA ALCHEMI supports only the FIRE optimizer") + if calculation.driver == "md" and backend_name != "nvalchemi-mace": + raise ValueError( + "MD driver is currently supported only by --backend nvalchemi-mace" + ) return backend, calculation diff --git a/src/matkit/mlip/config.py b/src/matkit/mlip/config.py index 9499594..6d6ac5e 100644 --- a/src/matkit/mlip/config.py +++ b/src/matkit/mlip/config.py @@ -11,6 +11,8 @@ _DTYPES = {"float32", "float64"} _OPTIMIZERS = {"bfgs", "lbfgs", "gpmin", "fire", "mdmin"} +_DRIVERS = {"energy", "opt", "md"} +_ENSEMBLES = {"nve", "nvt"} def _positive_number(name: str, value: Any) -> None: @@ -48,6 +50,11 @@ def _validate_explicit_options( "enable_cueq", "batch_size", "max_atoms", + "ensemble", + "temperature", + "timestep", + "md_steps", + "friction", }, } for owner, names in backend_options.items(): @@ -65,11 +72,19 @@ def _validate_explicit_options( raise ValueError("--dtype is controlled by the mace_anicc factory") if "dispersion" in provided and calculator_type != "mace_mp": raise ValueError("--dispersion requires --calculator-type mace_mp") - if driver == "energy": - opt_only = provided & {"optimizer", "fmax", "steps", "dt"} - if opt_only: - flag = sorted(opt_only)[0].replace("_", "-") + # dt is the ALCHEMI FIRE timestep; MD uses the separate --timestep instead. + md_only = {"ensemble", "temperature", "timestep", "md_steps", "friction"} + opt_only = {"optimizer", "fmax", "steps", "dt"} + if driver != "opt": + used = provided & opt_only + if used: + flag = sorted(used)[0].replace("_", "-") raise ValueError(f"--{flag} requires --driver opt") + if driver != "md": + used = provided & md_only + if used: + flag = sorted(used)[0].replace("_", "-") + raise ValueError(f"--{flag} requires --driver md") @dataclass(frozen=True) @@ -168,20 +183,37 @@ def to_dict(self) -> dict[str, Any]: @dataclass(frozen=True) class MLIPCalculationConfig: - """Calculation settings shared by all MLIP backends.""" + """Calculation settings shared by all MLIP backends. - driver: Literal["energy", "opt"] = "energy" + ``opt``-only fields (``optimizer``, ``fmax``) and ``md``-only fields + (``ensemble``, ``temperature``, ``timestep``, ``md_steps``, ``friction``) + are ignored by the other drivers. Batched MD is currently implemented on the + ``nvalchemi-mace`` backend only. + """ + + driver: Literal["energy", "opt", "md"] = "energy" optimizer: Literal["bfgs", "lbfgs", "gpmin", "fire", "mdmin"] = "fire" fmax: float = 0.01 steps: int = 1000 + ensemble: Literal["nve", "nvt"] = "nvt" + temperature: float = 300.0 + timestep: float = 1.0 + md_steps: int = 1000 + friction: float = 0.01 def __post_init__(self) -> None: - if self.driver not in {"energy", "opt"}: + if self.driver not in _DRIVERS: raise ValueError(f"Unsupported MLIP driver: {self.driver}") if self.optimizer not in _OPTIMIZERS: raise ValueError(f"Unsupported ASE optimizer: {self.optimizer}") + if self.ensemble not in _ENSEMBLES: + raise ValueError(f"Unsupported MD ensemble: {self.ensemble}") _positive_number("fmax", self.fmax) _positive_integer("steps", self.steps) + _positive_number("temperature", self.temperature) + _positive_number("timestep", self.timestep) + _positive_integer("md_steps", self.md_steps) + _positive_number("friction", self.friction) def to_dict(self) -> dict[str, Any]: return asdict(self) diff --git a/src/matkit/mlip/runner.py b/src/matkit/mlip/runner.py index 71cf0fa..69a8028 100644 --- a/src/matkit/mlip/runner.py +++ b/src/matkit/mlip/runner.py @@ -373,6 +373,12 @@ def _nvalchemi_result( stress = data.stress.detach().cpu().numpy() if stress.shape == (1, 3, 3): stress = stress[0] + if calculation.driver == "opt": + n_steps = None + elif calculation.driver == "md": + n_steps = calculation.md_steps + else: + n_steps = 0 return _success_result( input_file, backend, @@ -382,11 +388,56 @@ def _nvalchemi_result( forces, stress, converged, - None if calculation.driver == "opt" else 0, + n_steps, time.perf_counter() - started, ) +# Integration seam. The NVIDIA ALCHEMI MD integrator class names in +# nvalchemi.dynamics are NOT verifiable off-GPU (the package installs only with +# a CUDA extra). Confirm the real names on the target platform with +# python -c "import nvalchemi.dynamics as d; print(dir(d))" +# and update this mapping if they differ. Each class must accept the keyword +# arguments assembled in _nvalchemi_md_dynamics and expose the same +# context-manager + run(batch) protocol as BaseDynamics/FIRE. +_NVALCHEMI_MD_INTEGRATORS = { + "nve": "VelocityVerlet", + "nvt": "Langevin", +} + + +def _nvalchemi_md_dynamics( + dynamics_module, + model, + hooks, + backend: NVAlchemiMACEConfig, + calculation: MLIPCalculationConfig, +): + """Construct a batched MD integrator from nvalchemi.dynamics. + + Isolated so adapting to the confirmed ALCHEMI MD API is a localized change. + """ + ensemble = calculation.ensemble + class_name = _NVALCHEMI_MD_INTEGRATORS[ensemble] + integrator = getattr(dynamics_module, class_name, None) + if integrator is None: + raise RuntimeError( + f"nvalchemi.dynamics has no {class_name!r} integrator for the " + f"{ensemble!r} ensemble; confirm the class name on this platform " + "and update _NVALCHEMI_MD_INTEGRATORS." + ) + kwargs = { + "model": model, + "hooks": hooks, + "dt": calculation.timestep, + "n_steps": calculation.md_steps, + "temperature": calculation.temperature, + } + if ensemble == "nvt": + kwargs["friction"] = calculation.friction + return integrator(**kwargs) + + def _run_nvalchemi_chunk( model, entries: Sequence[tuple[int, str, Any]], @@ -394,6 +445,7 @@ def _run_nvalchemi_chunk( calculation: MLIPCalculationConfig, ) -> list[tuple[int, dict[str, Any]]]: try: + import nvalchemi.dynamics as nvalchemi_dynamics from nvalchemi.data import Batch from nvalchemi.dynamics import BaseDynamics, ConvergenceHook, FIRE except ImportError as exc: @@ -411,6 +463,10 @@ def _run_nvalchemi_chunk( convergence = None if calculation.driver == "energy": dynamics = BaseDynamics(model=model, hooks=hooks, n_steps=1) + elif calculation.driver == "md": + dynamics = _nvalchemi_md_dynamics( + nvalchemi_dynamics, model, hooks, backend, calculation + ) else: convergence = ConvergenceHook.from_fmax(calculation.fmax) dynamics = FIRE( @@ -533,6 +589,13 @@ def _validate_execution_request( and calculation.optimizer != "fire" ): raise ValueError("NVIDIA ALCHEMI supports only the FIRE optimizer") + if calculation.driver == "md" and not isinstance( + backend, NVAlchemiMACEConfig + ): + raise ValueError( + "MD driver is currently supported only by the " + "nvalchemi-mace backend" + ) def _execute_inputs( diff --git a/tests/test_cli.py b/tests/test_cli.py index 5754214..f3317d3 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -171,6 +171,21 @@ def test_mlip_batch_requires_one_input_source(self): "only the FIRE optimizer", ), ("ase-mace", ["--driver", "opt", "--fmax", "nan"], "finite"), + ( + "nvalchemi-mace", + ["--temperature", "300"], + "requires --driver md", + ), + ( + "nvalchemi-mace", + ["--driver", "energy", "--ensemble", "nve"], + "requires --driver md", + ), + ( + "ase-mace", + ["--driver", "md"], + "nvalchemi-mace", + ), ], ) def test_mlip_rejects_unsupported_explicit_options( diff --git a/tests/test_mlip_nvalchemi.py b/tests/test_mlip_nvalchemi.py index 97dce69..f6ad23d 100644 --- a/tests/test_mlip_nvalchemi.py +++ b/tests/test_mlip_nvalchemi.py @@ -126,7 +126,7 @@ def run(self, batch): data.energy = Tensor([[index + 1.25]]) data.forces = Tensor(np.zeros((data.num_nodes, 3))) data.stress = Tensor(np.eye(3)[None]) - if isinstance(self, FIRE): + if isinstance(self, (FIRE, VelocityVerlet, Langevin)): data.positions = Tensor(data.positions.numpy() + 0.1) if state.corrupt and index == 1: state.corrupt(data) @@ -135,6 +135,12 @@ def run(self, batch): class FIRE(BaseDynamics): pass + class VelocityVerlet(BaseDynamics): + pass + + class Langevin(BaseDynamics): + pass + class ConvergenceHook: @classmethod def from_fmax(cls, value): @@ -147,7 +153,13 @@ def evaluate(self, batch): modules["nvalchemi.data"].AtomicData = AtomicData modules["nvalchemi.data"].Batch = Batch modules["nvalchemi.models.mace"].MACEWrapper = MACEWrapper - for cls in (BaseDynamics, FIRE, ConvergenceHook): + for cls in ( + BaseDynamics, + FIRE, + VelocityVerlet, + Langevin, + ConvergenceHook, + ): setattr(modules["nvalchemi.dynamics"], cls.__name__, cls) state.model = Model() return state @@ -232,6 +244,48 @@ def test_dynamics_and_result_mapping(alchemi, driver): assert ("fmax", 0.02) in alchemi.events +@pytest.mark.parametrize( + "ensemble,integrator", + [("nvt", "Langevin"), ("nve", "VelocityVerlet")], +) +def test_md_driver_selects_integrator_and_maps_results( + alchemi, ensemble, integrator +): + inputs = entries() + backend = NVAlchemiMACEConfig("medium") + calculation = MLIPCalculationConfig( + driver="md", + ensemble=ensemble, + temperature=250.0, + timestep=2.0, + md_steps=5, + friction=0.05, + ) + outputs = runner._run_nvalchemi_chunk( + alchemi.model, inputs, backend, calculation + ) + assert [index for index, _ in outputs] == [3, 7] + settings = next( + event[1] for event in alchemi.events if event[0] == integrator + ) + assert settings["n_steps"] == 5 + assert settings["dt"] == 2.0 + assert settings["temperature"] == 250.0 + if ensemble == "nvt": + assert settings["friction"] == 0.05 + else: + assert "friction" not in settings + for (_, result), (_, _, original) in zip(outputs, inputs): + assert result["success"] + # MD integrators advance positions and report step count, not + # convergence. + assert result["converged"] is True + assert result["n_steps"] == 5 + assert np.allclose( + result["final_structure"]["positions"], original.positions + 0.1 + ) + + @pytest.mark.parametrize( "field,value", [ diff --git a/tests/test_mlip_validation.py b/tests/test_mlip_validation.py index 82a7d36..c4d68af 100644 --- a/tests/test_mlip_validation.py +++ b/tests/test_mlip_validation.py @@ -73,6 +73,36 @@ def test_integer_limits(copper, tmp_path, value): ) +def test_md_config_fields_validated(): + MLIPCalculationConfig(driver="md", ensemble="nve") + MLIPCalculationConfig(driver="md", ensemble="nvt") + with pytest.raises(ValueError, match="ensemble"): + MLIPCalculationConfig(ensemble="npt") + with pytest.raises(ValueError, match="temperature"): + MLIPCalculationConfig(temperature=0) + with pytest.raises(ValueError, match="timestep"): + MLIPCalculationConfig(timestep=float("nan")) + with pytest.raises(ValueError, match="md_steps"): + MLIPCalculationConfig(md_steps=0) + with pytest.raises(ValueError, match="friction"): + MLIPCalculationConfig(friction=-1) + + +def test_unknown_driver_rejected(): + with pytest.raises(ValueError, match="driver"): + MLIPCalculationConfig(driver="npt") + + +def test_md_driver_requires_nvalchemi_backend(copper, tmp_path): + with pytest.raises(ValueError, match="MD driver"): + run_mlip_batch( + [copper], + ASEMACEConfig(), + MLIPCalculationConfig(driver="md"), + output_dir=tmp_path / "md", + ) + + def test_rootstock_kwargs_must_be_serializable(): for value in (float("nan"), object()): with pytest.raises(ValueError, match="finite JSON"):