diff --git a/.github/workflows/gfortran.yml b/.github/workflows/gfortran.yml index 40f2d342..057eb414 100644 --- a/.github/workflows/gfortran.yml +++ b/.github/workflows/gfortran.yml @@ -332,6 +332,15 @@ jobs: if [ $GCC_V -ge 10 ];then export FFLAGS="-fallow-argument-mismatch";export FCFLAGS="-fallow-argument-mismatch";fi && \ ./dev_scripts/install_mpich.sh ${HOME}/mpich ${MPICH_V} + - name: Set up uv + uses: astral-sh/setup-uv@v7 + with: + python-version: 3.12 + activate-environment: true + + - name: Install MACE + run: uv pip install --torch-backend=cpu -r interfaces/MACE/mace_server.py + - name: build ABIN run: | if [ $GCC_V -ge 10 ];then export FFLAGS="-fallow-argument-mismatch $FFLAGS";fi && \ @@ -341,6 +350,7 @@ jobs: LDLIBS: ${{ env.ABIN_LDLIBS }} - name: test ABIN run: make test + - name: Codecov upload uses: codecov/codecov-action@v7 with: @@ -374,6 +384,15 @@ jobs: if: steps.openmpi-cache.outputs.cache-hit != 'true' run: ./dev_scripts/install_openmpi.sh ${HOME}/openmpi ${OPENMPI_V} ${OPENMPI_PATCH} + - name: Set up uv + uses: astral-sh/setup-uv@v7 + with: + python-version: 3.12 + activate-environment: true + + - name: Install MACE + run: uv pip install --torch-backend=cpu -r interfaces/MACE/mace_server.py + - name: build ABIN run: ./configure --mpi "${HOME}/openmpi/${OPENMPI_V}/install" && make env: diff --git a/README.md b/README.md index 816f7a33..7c740947 100644 --- a/README.md +++ b/README.md @@ -7,9 +7,9 @@ ABIN is a program for performing ab initio molecular dynamics. It is a general purpose program that was initially designed to model nuclear quantum effects (NQE). NQE can be most rigirously captured with path integral MD (PIMD), but also within the Quantum Thermostat based on General Langevin Equation framework developed by Michele Cerriotti. -ABIN can also simulate non-adiabatic events using Surface-hoping algorithm, using either the classical fewest-switches algorithm (FSSH) or simpler Landau-Zener approach which does not require non-adiabatic couplings. The LZ approach can also capture singlet-triplet transitions. +ABIN can also simulate non-adiabatic events using Surface-hoping algorithm, using either the classical fewest-switches algorithm (FSSH) or simpler Landau-Zener (LZ) approach which does not require non-adiabatic couplings. The LZ approach can also capture singlet-triplet transitions. -The basic philosophy of ABIN program is simple — +The basic philosophy of ABIN program is simple — while the program itself handles the propagation of the system according to the equations of motion, the forces and energies are taken from an external electronic structure program such as ORCA or TeraChem. The call to the chosen external program is handled via a simple shell script interface. @@ -72,11 +72,41 @@ To install the libraries, you can use the install scripts in `dev_scripts/`. We use these in our Continuous Integration testing suite on Github using the Ubuntu 18.04 image. The optional libraries are: - - [MPICH](https://www.mpich.org/): An MPI implementation used for Replica Exchange MD and MPI interface with TeraChem. + - [MPICH](https://www.mpich.org/): An MPI implementation used for Replica Exchange MD and MPI interface with TeraChem and MACE. - If you just need REMD you can also use other MPI libraries such as OpenMPI or IntelMPI. - [FFTW](http://www.fftw.org/): Fast Fourier Transform library used for normal mode transformation in Path Integral MD. - [PLUMED](https://www.plumed.org/): A collection of very useful tools for free energy calculations (MetaDynamics, Umbrella Sampling etc). - [TCPB-CPP](https://github.com/mtzgroup/tcpb-cpp): [EXPERIMENTAL] TCPB interface to TeraChem + - [MACE](https://github.com/ACEsuit/mace): Machine Learning Atomic Cluster Expansion potential. + - The interface to MACE requires an MPICH installation (see below) and MACE python environment. + See `interfaces/MACE/README.md` for details + + +### Installing with MPICH + +1. Install MPICH (adjust MPICH_PATH as needed) + +```bash +MPICH_PATH=/home/$USER/software/mpich +./dev_scripts/install_mpich.sh $MPICH_PATH +``` + +2. Compile ABIN and run tests + +(Substitute the `` for the actual version that was installed) +```bash +MPICH_PATH=$MPICH_PATH//install +./configure --mpi $MPICH_PATH +make clean && make +make test +``` + +Before running ABIN, you might need to adjust your PATH and LD_LIBRARY_PATH + +```bash +export LD_LIBRARY_PATH=$MPICH_PATH/lib:$LD_LIBRARY_PATH +export PATH=$MPICH_PATH/bin:$PATH +``` ## Structure of the repository diff --git a/configure b/configure index 102bb76e..8461f791 100755 --- a/configure +++ b/configure @@ -123,7 +123,7 @@ done # Set to default if FFLAGS is not defined yet, see # https://stackoverflow.com/questions/11362250/in-bash-how-do-i-test-if-a-variable-is-defined-in-u-mode -: ${FFLAGS:="-O2 -fopenmp -fimplicit-none $warning_flags"} +: ${FFLAGS:="-O2 -g -fopenmp -fimplicit-none $warning_flags"} if [[ -z ${FC-} ]];then if [[ $MPI = "TRUE" ]];then diff --git a/dev_scripts/install_mpich.sh b/dev_scripts/install_mpich.sh index 00862cab..de865217 100755 --- a/dev_scripts/install_mpich.sh +++ b/dev_scripts/install_mpich.sh @@ -7,19 +7,26 @@ # Exit script immediately upon error set -euo pipefail +CC=${CC:-gcc} +FC=${FC:-gfortran} + +if [[ -z ${1-} ]]; then + echo "ERROR: Provide prefix path where install MPICH as first parameter" + exit 1 +fi # Path as an optional first parameter -MPICH_DIR="${1-$HOME/mpich}" +MPICH_DIR="$(realpath "$1")" # We take current stable version as default # (as of 06 Nov 2020). -MPICH_VERSION="${2-"4.0.2"}" +MPICH_VERSION="${2-"4.3.2"}" + +# Detect number of CPUs +NCPUS=$(nproc) TAR_FILE="mpich-${MPICH_VERSION}.tar.gz" DOWNLOAD_URL="https://www.mpich.org/static/downloads/${MPICH_VERSION}/${TAR_FILE}" INSTALL_DIR="$MPICH_DIR/$MPICH_VERSION/install" -# Github Actions machines have two CPUs, per: -# https://docs.github.com/en/free-pro-team@latest/actions/reference/specifications-for-github-hosted-runners#supported-runners-and-hardware-resources -NCPUS=2 if [[ -d "$INSTALL_DIR" ]];then echo "Found existing MPICH installation in $INSTALL_DIR" @@ -56,9 +63,10 @@ cd "$MPICH_DIR/$MPICH_VERSION/src" && tar -xzf "../pkg/${TAR_FILE}" && cd "mpich # export FCFLAGS=-fallow-argument-mismatch ./configure FC=gfortran CC=gcc \ --enable-fortran=all \ - --with-pm=hydra --with-device=ch3:nemesis \ + --with-pm=hydra \ + --with-device=ch3:nemesis \ --with-namepublisher=pmi \ - --enable-static --disable-shared \ + --enable-static \ --prefix="${INSTALL_DIR}" 2>&1 |\ tee configure.log make -j $NCPUS 2>&1 | tee make.log diff --git a/interfaces/MACE/README.md b/interfaces/MACE/README.md new file mode 100644 index 00000000..b42d768a --- /dev/null +++ b/interfaces/MACE/README.md @@ -0,0 +1,64 @@ +# MACE MPI Interface for ABIN + +This directory contains the Python-based server for the MACE (Machine Learning Atomic Cluster Expansion) potential. + +## Requirements + +- Python >= 3.8 +- PyTorch >= 1.12 +- mace-torch +- ase +- mpi4py +- numpy + +The exact dependencies are specified as inline metadata in `mace_server.py`. +We highly recommend installing the dependencies in a fresh virtual environment +using the [uv package manager](https://github.com/astral-sh/uv): + +```console +# Install uv first, https://github.com/astral-sh/uv#installation +uv venv # Creates a new virtual environment in .venv/ folder +uv pip install -r interfaces/MACE/mace_server.py --torch-backend=auto +source .venv/bin/activate # Activates the environment +``` + +> [!IMPORTANT] +> The tricky part is to install the correct binary version of PyTorch, as it needs to match your CUDA version. +> The `--torch-backend=auto` should autodetect your environment and install the correct version. +> In HPE clusters, make sure you load your CUDA environment before running the installation. +> Please read the [uv PyTorch documentation](https://docs.astral.sh/uv/guides/integration/pytorch/#automatic-backend-selection) for more information. +> If your machine doesn't have a GPU accelerator, the installation will automatically pick up the CPU-only build. + +ABIN itself must be compiled using the MPICH compiler, see top-level [README.md](../../README.md#installing-with-mpich) for instructions. + +After installation, run the MACE tests to make sure the basic communication works. +```console +$ make test TEST="MACE MACE_ERROR MACE_ERROR2" +Running tests in directories: +MACE MACE_ERROR MACE_ERROR2 +MACE PASSED +======================= +MACE_ERROR PASSED +======================= +MACE_ERROR2 PASSED +======================= + +3 tests PASSED. +``` + +These tests use a fake potential, not the actual MACE implementation. +Make sure you test the energy conservation using a short NVE dynamics +before running production calculations! + +## Usage + +The MACE server must be launched alongside ABIN using mpirun. +It is recommended to use the provided launch script `utils/run.mace_mpi_abin.sh`. + +Rough steps: + +1. Specify `pot='_mace_'` in the ABIN input file, everything else is the same. +2. Copy `interfaces/MACE/mace_server.py` and `utils/run.mace_mpi_abin.sh` into your working directory. +3. In the launch script, specify the ABIN input files, and paths to MPICH and Python installation. +4. In the launch script, specify path to your MACE model file, and whether to evaluate the model on CPU or CUDA GPU. +5. Profit! diff --git a/interfaces/MACE/mace_server.py b/interfaces/MACE/mace_server.py new file mode 100755 index 00000000..f642081d --- /dev/null +++ b/interfaces/MACE/mace_server.py @@ -0,0 +1,322 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.8" +# dependencies = [ +# "ase>=3.18.0", +# "mace-torch>=0.3.10", +# "mpi4py>=4.1.2", +# "numpy>=1.26.0", +# ] +# /// +""" +MACE MPI SERVER + +This server communicates with ABIN via MPI (using mpi4py). + +Usage: + mpirun -n 1 python mace_server.py + + The server writes its MPI port to 'mace_port.txt' for ABIN to read. +""" +# ruff: file-ignore[blind-except] + +import argparse +import functools +import sys +from pathlib import Path +from time import perf_counter +from traceback import print_tb + +LOG_NAME = "MaceMPIServer" + +# MPI Tags (must match Fortran module mod_mace_mpi) +MACE_TAG_EXIT = 666 +MACE_TAG_DATA = 2 +MACE_TAG_ERROR = 13 + +MACE_PORT_FILE = "mace_port.txt" + + +# TODO: Use logging module +# TODO: Create log_exception helper +def log(message, should_print=True): + msg_formatted = f"[{LOG_NAME}]: {message!s}" + if should_print: + print(msg_formatted, flush=True) + return msg_formatted + + +def parse_cmd(): + desc = "MACE MPI server for ground state MD with ABIN" + parser = argparse.ArgumentParser(description=desc) + parser.add_argument( + "--model-path", + type=str, + required=True, + help="Path to MACE model file", + ) + parser.add_argument( + "--device", + type=str, + choices=("cpu", "cuda"), + required=True, + help="Device for model inference", + ) + config = parser.parse_args() + model = config.model_path + if not model.startswith("__MOCK_") and not Path(model).is_file(): + sys.exit(f"ERROR: file '{config.model_path}' not found") + return config + + +class MaceModel: + """ + Manages the MACE ML model for evaluating atomic configurations. + Configuration is received from ABIN via MPI. + """ + + def __init__(self, config): + from mace.calculators import MACECalculator + + log("initializing MACE model") + + # Set ASE calculator + self.calculator = MACECalculator( + model_paths=config.model_path, + device=config.device, + ) + + def evaluate(self, atom_types, coords_bohr): + """ + Evaluate energy and forces for a single configuration. + + Parameters: + atom_types: list of atomic symbols (e.g. ['H', 'O', 'H']) + coords_bohr: numpy array of shape (natom, 3) in Bohr + + Returns: + energy_hartree: energy in Hartree + forces_hartree_bohr: forces in Hartree/Bohr, shape (natom, 3) + """ + import ase + + # Unit conversions + bohr_to_ang = 0.529177249 + ev_to_hartree = 1.0 / 27.211399 + ev_per_ang_to_hartree_per_bohr = ev_to_hartree * bohr_to_ang + + # Convert coordinates from Bohr to Angstrom + coords_ang = coords_bohr.copy() * bohr_to_ang + + # Create ASE atoms object + pbc = (False, False, False) + # cell_size = 100.0 # Angstroms + # cell = ((cell_size, 0, 0), (0, cell_size, 0), (0, 0, cell_size)) + atoms = ase.Atoms(symbols=atom_types, positions=coords_ang, pbc=pbc) + atoms.calc = self.calculator + + energy_hartree = atoms.get_potential_energy() * ev_to_hartree + forces_hartree_bohr = atoms.get_forces() * ev_per_ang_to_hartree_per_bohr + + return energy_hartree, forces_hartree_bohr + + +class HarmonicModel: + """ + This is just a mock for testing, using harmonic potential. + """ + + k = 0.01 # force constant in Hartree/Bohr^2 + + def __init__(self, config): + self.model_path = config.model_path + log("Using Harmonic Mock Model") + + def evaluate(self, _atom_types, coords_bohr): + """ + Compute mock energy and forces using simple harmonic potential + E = 0.5 * k * sum(r^2) where r is displacement from origin + """ + import numpy as np + + if config.model_path == "__MOCK_ERROR__": + raise RuntimeError("Simulating error condition") + + coords_t = coords_bohr.T # (natom, 3) + + # Simple harmonic energy around the center of mass + com = np.mean(coords_t, axis=0) + displ = coords_t - com + energy = 0.5 * self.k * np.sum(displ**2) + + # Forces = -gradient = -k * displacement + forces = -self.k * displ # (natom, 3) + return energy, forces + + +def connect_to_abin(): + """Establish initial connection to ABIN""" + from mpi4py import MPI + + # Open MPI port and write to file for ABIN to read + port_name = MPI.Open_port() + log(f"MPI port opened: {port_name}") + + with open(MACE_PORT_FILE, "w", encoding="utf-8") as f: + f.write(port_name) + log(f"Port written to {MACE_PORT_FILE}") + + # Accept connection from ABIN + log("Waiting for ABIN to connect...") + abin_comm = MPI.COMM_WORLD.Accept(port_name) + log("Connection from ABIN accepted!") + return port_name, abin_comm + + +# https://docs.python.org/3/library/sys.html#sys.excepthook +def exception_handler(shutdown_callback, exception_type, exception, traceback): + """Try to gracefully shutdown communication with ABIN upon uncaught exceptions""" + print(f"Unexpected {exception_type.__name__}: {exception}") + print_tb(traceback) + # Restore original exception handling to prevent endless loop + # in case of uncaught excpetion during shutdown + sys.excepthook = sys.__excepthook__ + shutdown_callback() + sys.exit(1) + + +def main(config): + import numpy as np + from mpi4py import MPI + + port_name, abin_comm = connect_to_abin() + + def shutdown_communication(): + """Gracefully shutdown communication with ABIN""" + log("Shutting down communication with ABIN...") + try: + abin_comm.Disconnect() + except Exception as e: + log(e) + else: + log("ABIN communicator disconnected") + + try: + MPI.Close_port(port_name) + except Exception as e: + log(e) + else: + log(f"Port {port_name} closed") + + def error_shutdown(): + log("Sending ERROR tag to ABIN") + # This is best effort only, since ABIN might be dead already + try: + abin_comm.Send([MPI.BOTTOM, MPI.INT], dest=0, tag=MACE_TAG_ERROR) + except Exception as e: + log(e) + + shutdown_communication() + + sys.exit(1) + + def check_incoming_msg(): + """Check whether ABIN is sending ERROR or EXIT message""" + status = MPI.Status() + abin_comm.Probe(source=0, tag=MPI.ANY_TAG, status=status) + + if (tag := status.Get_tag()) in (MACE_TAG_EXIT, MACE_TAG_ERROR): + if tag == MACE_TAG_EXIT: + log("Received graceful exit signal from ABIN") + exit_code = 0 + else: + log("Received ERROR signal from ABIN. Stopping server") + exit_code = 1 + + try: + abin_comm.Recv([MPI.BOTTOM, MPI.INT], source=0, tag=tag) + except Exception as e: + log(e) + + shutdown_communication() + sys.exit(exit_code) + + # Call error_shutdown upon any unhandled exception + sys.excepthook = functools.partial(exception_handler, error_shutdown) + + # Receive number of atoms + check_incoming_msg() + natom_buf = np.empty(1, dtype=np.intc) + abin_comm.Recv([natom_buf, MPI.INT], source=0, tag=MACE_TAG_DATA) + natom = int(natom_buf[0]) + log(f"Received number of atoms: {natom}") + + # Receive atom types + check_incoming_msg() + byte_buf = bytearray(natom * 2) + mpi_status = MPI.Status() + abin_comm.Recv([byte_buf, MPI.CHAR], source=0, tag=MACE_TAG_DATA, status=mpi_status) + assert mpi_status.Get_count() == natom * 2 + # TODO: Check that size of received data! + atom_types_str = byte_buf.decode("ascii") + atom_types = [ + atom_types_str[i : i + 2].strip() for i in range(0, len(atom_types_str), 2) + ] + log(f"Received atom types: {atom_types}") + assert len(atom_types) == natom + + # Load MACE model + if config.model_path.startswith("__MOCK_"): + # This is for testing purposes only + mace_model = HarmonicModel(config) + else: + mace_model = MaceModel(config) + + log("MACE model ready. Entering main loop.") + + # Main loop: receive coordinates, compute, send results + eval_count = 0 + while True: + start_loop = perf_counter() + check_incoming_msg() + + # Receive coordinates (3*natom doubles, in Bohr) + coords = np.empty((natom, 3), dtype=np.float64) + abin_comm.Recv( + [coords, MPI.DOUBLE], source=0, tag=MACE_TAG_DATA, status=mpi_status + ) + assert mpi_status.Get_elements(MPI.DOUBLE) == natom * 3 + + start = perf_counter() + + energy, forces = mace_model.evaluate(atom_types, coords) + + end = perf_counter() + time_ms = (end - start) * 1000 + log(f"Step {eval_count} done in {time_ms:.3f} miliseconds") + log(f"Energy = {energy:.15f} Hartree") + + # Send energy (1 double, in Hartree) + energy_buf = np.array([energy], dtype=np.float64) + abin_comm.Send([energy_buf, MPI.DOUBLE], dest=0, tag=MACE_TAG_DATA) + + # Send forces (3*natom doubles, in Hartree/Bohr) + # Transpose back to (3, natom) to match Fortran column-major layout + if forces.dtype != np.float64: + forces_send = forces.T.astype(np.float64) + else: + forces_send = forces.T.copy() + abin_comm.Send([forces_send, MPI.DOUBLE], dest=0, tag=MACE_TAG_DATA) + + end_loop = perf_counter() + loop_ms = (end_loop - start_loop) * 1000 + # Note: Communication overhead includes the ABIN propagation time, + # which should however be negligible. + log(f"Communication overhead = {loop_ms - time_ms:.3f} ms") + + eval_count += 1 + + +if __name__ == "__main__": + config = parse_cmd() + main(config) diff --git a/prek.toml b/prek.toml index 39dbf541..e8d183fe 100644 --- a/prek.toml +++ b/prek.toml @@ -25,6 +25,23 @@ hooks = [ } ] +[[repos]] +# ruff: A Python linter, https://astral.sh/ruff +repo = "https://github.com/astral-sh/ruff-pre-commit" +rev = "v0.15.20" +hooks = [ + { + id = "ruff-check", + # The --preview flag enables a much larger set of rules by default + args = ["--fix", "--preview"], + files = "interfaces/MACE/mace_server.py", + }, + { + id = "ruff-format", + files = "interfaces/MACE/mace_server.py", + } +] + # Verify Github workflows have correct syntax [[repos]] repo = "https://github.com/python-jsonschema/check-jsonschema" diff --git a/src/Makefile b/src/Makefile index e2baf631..88aa6f06 100644 --- a/src/Makefile +++ b/src/Makefile @@ -1,7 +1,7 @@ F_OBJS := constants.o fortran_interfaces.o error.o modules.o mpi_wrapper.o files.o utils.o io.o random.o arrays.o qmmm.o fftw_interface.o \ shake.o nosehoover.o gle.o transform.o potentials.o force_spline.o estimators.o ekin.o vinit.o plumed.o \ remd.o force_bound.o water.o h2o_schwenke.o h2o_cvrqd.o force_h2o.o force_cp2k.o sh_integ.o surfacehop.o landau_zener.o potentials_sh.o\ - force_mm.o tera_mpi_api.o force_abin.o force_tcpb.o force_tera.o force_terash.o en_restraint.o analyze_ext_template.o geom_analysis.o analysis.o \ + force_mm.o tera_mpi_api.o mace_mpi_api.o force_abin.o force_tcpb.o force_tera.o force_terash.o force_mace.o en_restraint.o analyze_ext_template.o geom_analysis.o analysis.o \ minimizer.o mdstep.o forces.o cmdline.o init.o C_OBJS := water_interface.o diff --git a/src/error.F90 b/src/error.F90 index 4289e8c8..b9029bae 100644 --- a/src/error.F90 +++ b/src/error.F90 @@ -15,8 +15,7 @@ ! During initialization the testing framework can override the ! default behavior with a call to set_error_method(). This logic ! can comfortably live in your test code, and thus does not -! introduce any undesirable dependencies (just a bit of -! obscurity). +! introduce any undesirable dependencies (just a bit of obscurity). module mod_error implicit none private @@ -27,11 +26,12 @@ module mod_error public :: set_error_method abstract interface - subroutine error(filename, line_number, message) + subroutine error(filename, line_number, message, error_code) implicit none character(len=*), intent(in) :: filename integer, intent(in) :: line_number character(len=*), intent(in) :: message + integer, intent(in) :: error_code end subroutine error end interface @@ -48,12 +48,20 @@ end subroutine set_error_method ! filename and line_number parameters should be passed using the preprocessor ! defined constants __FILE__ and __LINE__ - subroutine fatal_error(filename, line_number, message) + subroutine fatal_error(filename, line_number, message, error_code) character(len=*), intent(in) :: filename integer, intent(in) :: line_number character(len=*), intent(in) :: message + integer, intent(in), optional :: error_code + integer :: err_code - call error_method(filename, line_number, message=message) + ! Default error code + err_code = 1 + if (present(error_code)) then + err_code = error_code + end if + + call error_method(filename, line_number, message=message, error_code=err_code) end subroutine fatal_error subroutine not_compiled_with(feature) @@ -63,20 +71,19 @@ subroutine not_compiled_with(feature) call fatal_error(__FILE__, __LINE__, error_msg) end subroutine not_compiled_with - subroutine print_error_and_stop(filename, line, message) + subroutine print_error_and_stop(filename, line, message, error_code) use mod_interfaces, only: finish character(*), intent(in) :: filename integer, intent(in) :: line character(*), intent(in) :: message - ! In case of an error, ABIN will return this to the shell. - integer, parameter :: ERROR_CODE = 1 + integer, intent(in) :: error_code call print_error(filename, line, message) ! Try to finalize various modules gracefully. - call finish(ERROR_CODE) + call finish(error_code) - stop ERROR_CODE + stop error_code end subroutine print_error_and_stop ! We print the error both to stderr and to file 'ERROR'. diff --git a/src/force_mace.F90 b/src/force_mace.F90 new file mode 100644 index 00000000..a890ea2f --- /dev/null +++ b/src/force_mace.F90 @@ -0,0 +1,121 @@ +module mod_force_mace +! ---------------------------------------------------------------- +! Interface for MACE (Machine Learning Atomic Cluster Expansion) +! Perform MPI communications with the MACE Python server. +! +! Modeled after mod_force_tera for TeraChem. +! ---------------------------------------------------------------- + use mod_const, only: DP + use mod_error, only: fatal_error + use mod_files, only: stdout, stderr + use mod_mace_mpi +#ifdef USE_MPI + use mpi + use mod_mpi, only: handle_mpi_error, check_recv_count +#endif + implicit none + private + public :: force_mace + save + +contains + + subroutine force_mace(x, y, z, fx, fy, fz, eclas, walkmax) + use mod_qmmm, only: natqm + real(DP), intent(in) :: x(:, :), y(:, :), z(:, :) + real(DP), intent(inout) :: fx(:, :), fy(:, :), fz(:, :) + real(DP), intent(inout) :: eclas + integer, intent(in) :: walkmax + integer :: iw + integer :: mace_comm + logical :: success + + success = .false. + + do iw = 1, walkmax + +#ifdef USE_MPI + mace_comm = get_mace_communicator() + ! Send coordinates in Bohr + call send_mace_coordinates(x, y, z, natqm, iw, mace_comm) + + success = receive_mace_data(fx, fy, fz, eclas, iw, walkmax, mace_comm) +#endif + if (.not. success) then + call fatal_error(__FILE__, __LINE__, 'MACE evaluation failed', error_code=MACE_TAG_ERROR) + end if + + end do + + end subroutine force_mace + +#ifdef USE_MPI + + logical function receive_mace_data(fx, fy, fz, eclas, iw, walkmax, comm) result(success) + use mod_general, only: idebug + use mod_qmmm, only: natqm + real(DP), intent(inout) :: fx(:, :), fy(:, :), fz(:, :) + real(DP), intent(inout) :: eclas + integer, intent(in) :: iw, walkmax + integer, intent(in) :: comm + real(DP) :: energy + real(DP) :: forces(3, size(fx, 1)) + integer :: status(MPI_STATUS_SIZE) + integer :: ierr, iat + + success = .true. + ! Receive energy (in Hartree, already converted by server) + if (idebug > 1) then + write (stdout, '(a)') 'MACE: Waiting to receive energy...' + call flush (stdout) + end if + if (.not. check_incoming_msg()) then + success = .false. + return + end if + call MPI_Recv(energy, 1, MPI_DOUBLE_PRECISION, MPI_ANY_SOURCE, & + MPI_ANY_TAG, comm, status, ierr) + call handle_mpi_error(ierr) + call check_recv_count(status, 1, MPI_DOUBLE_PRECISION) + + if (idebug > 1) then + print '(A,ES15.6)', 'MACE: Received energy [Hartree]:', energy + call flush (stdout) + end if + + ! Receive forces (in Hartree/Bohr, already converted by server) + if (idebug > 1) then + print '(a)', 'MACE: Waiting to receive forces...' + end if + if (.not. check_incoming_msg()) then + success = .false. + return + end if + call MPI_Recv(forces, 3 * natqm, MPI_DOUBLE_PRECISION, & + MPI_ANY_SOURCE, MPI_ANY_TAG, comm, status, ierr) + call handle_mpi_error(ierr) + call check_recv_count(status, 3 * natqm, MPI_DOUBLE_PRECISION) + + if (idebug > 1) then + print '(A)', 'MACE: Received forces [Hartree/Bohr]:' + do iat = 1, natqm + write (stdout, *) 'Atom ', iat, ': ', forces(:, iat) + end do + call flush (stdout) + end if + + ! Forces are received as forces (not gradients), so no sign flip needed + do iat = 1, natqm + fx(iat, iw) = forces(1, iat) + fy(iat, iw) = forces(2, iat) + fz(iat, iw) = forces(3, iat) + end do + + eclas = eclas + energy / walkmax + + end function receive_mace_data + +! USE_MPI +#endif + +end module mod_force_mace diff --git a/src/forces.F90 b/src/forces.F90 index 21006d0f..67893504 100644 --- a/src/forces.F90 +++ b/src/forces.F90 @@ -140,6 +140,7 @@ subroutine force_wrapper(x, y, z, fx, fy, fz, e_pot, chpot, walkmax) use mod_force_tcpb, only: force_tcpb use mod_force_tera, only: force_tera use mod_terampi_sh, only: force_terash + use mod_force_mace, only: force_mace implicit none ! allow(external-procedure) ! fortitude linter external :: force_water @@ -183,6 +184,8 @@ subroutine force_wrapper(x, y, z, fx, fy, fz, e_pot, chpot, walkmax) else call force_tera(x, y, z, fx, fy, fz, eclas, walkmax) end if + case ("_mace_") + call force_mace(x, y, z, fx, fy, fz, eclas, walkmax) case DEFAULT call force_abin(x, y, z, fx, fy, fz, eclas, chpot, walkmax) end select diff --git a/src/init.F90 b/src/init.F90 index 16c76c63..c2d30b4e 100644 --- a/src/init.F90 +++ b/src/init.F90 @@ -66,6 +66,7 @@ subroutine init(dt) use mod_force_tcpb, only: initialize_tcpb use mod_terampi use mod_terampi_sh + use mod_mace_mpi, only: initialize_mace_interface, initialize_mace_server use mod_mdstep, only: initialize_integrator, nabin, nstep_ref real(DP), intent(out) :: dt ! Input parameters for analytical potentials @@ -216,6 +217,11 @@ subroutine init(dt) call initialize_terachem_interface(trim(tc_server_name)) end if + ! Connect to MACE server as early as possible + if (pot == '_mace_') then + call initialize_mace_interface() + end if + if (mdtype /= '') then mdtype = tolower(mdtype) select case (mdtype) @@ -405,6 +411,10 @@ subroutine init(dt) end if end if + if (pot == '_mace_') then + call initialize_mace_server() + end if + if (pot == '_tcpb_' .or. restrain_pot == '_tcpb_' .or. pot_ref == '_tcpb_') then call initialize_tcpb(natqm, atnames, tcpb_port, tcpb_host, tcpb_input_file) end if @@ -1256,6 +1266,7 @@ subroutine finish(error_code) use mod_plumed, only: iplumed, finalize_plumed use mod_terampi, only: finalize_terachem use mod_terampi_sh, only: finalize_terash + use mod_mace_mpi, only: finalize_mace use mod_splined_grid, only: finalize_spline use mod_force_mm, only: finalize_mm use mod_force_tcpb, only: finalize_tcpb @@ -1270,6 +1281,10 @@ subroutine finish(error_code) call finalize_terachem(error_code) end if + if (pot == '_mace_') then + call finalize_mace(error_code) + end if + if (pot == '_tcpb_') then call finalize_tcpb() end if diff --git a/src/mace_mpi_api.F90 b/src/mace_mpi_api.F90 new file mode 100644 index 00000000..25288f36 --- /dev/null +++ b/src/mace_mpi_api.F90 @@ -0,0 +1,252 @@ +! MPI interface with MACE Python server +! Modeled after mod_terampi for TeraChem +module mod_mace_mpi + use mod_const, only: DP + use mod_error, only: fatal_error + use mod_files, only: stdout, stderr +#ifdef USE_MPI + use mpi + use mod_mpi, only: handle_mpi_error, get_mpi_error_string +#endif + implicit none + private + + ! MPI tags for MACE protocol + integer, parameter :: MACE_TAG_EXIT = 666 + integer, parameter :: MACE_TAG_DATA = 2 + integer, parameter :: MACE_TAG_ERROR = 13 + + ! Port file for MACE MPI connection + character(len=*), parameter :: MACE_PORT_FILE_NAME = 'mace_port.txt' + +#ifdef USE_MPI + integer :: mace_comm = MPI_COMM_NULL +#endif + + public :: MACE_TAG_ERROR +#ifdef USE_MPI + public :: get_mace_communicator + public :: check_incoming_msg + public :: send_mace_atom_types, send_mace_coordinates +#endif + public :: initialize_mace_interface, initialize_mace_server, finalize_mace + save + +contains + +#ifdef USE_MPI + + subroutine initialize_mace_interface() + integer :: ierr + + mace_comm = MPI_COMM_NULL + + ! Set MPI error handler to allow retries + call MPI_Comm_set_errhandler(MPI_COMM_WORLD, MPI_ERRORS_RETURN, ierr) + call handle_mpi_error(ierr) + + call connect_mace_server() + end subroutine initialize_mace_interface + + subroutine connect_mace_server() + character(len=MPI_MAX_PORT_NAME) :: port_name + integer :: ierr, newcomm + + call read_mpi_port_from_file(MACE_PORT_FILE_NAME, port_name) + + write (stdout, '(2A)') 'Found MACE port: ', trim(port_name) + write (stdout, '(A)') 'Establishing connection to MACE server...' + call flush (stdout) + + call MPI_Comm_connect(trim(port_name), MPI_INFO_NULL, 0, MPI_COMM_SELF, newcomm, ierr) + call handle_mpi_error(ierr) + write (stdout, '(A)') 'Connection to MACE server established!' + + mace_comm = newcomm + end subroutine connect_mace_server + + ! Read MPI port from a file. + subroutine read_mpi_port_from_file(portfile, port_name) + character(len=*), intent(in) :: portfile + character(len=MPI_MAX_PORT_NAME), intent(out) :: port_name + integer :: iunit + + write (stdout, '(A)') 'Reading MPI port name from file '//portfile + port_name = '' + + open (newunit=iunit, file=portfile, action="read", status="old") + read (iunit, '(A)') port_name + close (iunit, status='delete') + end subroutine read_mpi_port_from_file + + integer function get_mace_communicator() result(comm) + comm = mace_comm + end function get_mace_communicator + + logical function check_incoming_msg() result(success) + ! Before calling MPI_Recv, check we're not receiving error tag + ! from the MACE server + integer :: mpi_status(MPI_STATUS_SIZE) + integer :: tag, ierr + character(len=100) :: errormsg + + success = .false. + + call MPI_Probe(MPI_ANY_SOURCE, MPI_ANY_TAG, mace_comm, mpi_status, ierr) + tag = mpi_status(MPI_TAG) + + if (tag == MACE_TAG_ERROR) then + ! Consume the error message from MACE, the caller of this function should + ! stop the program. + write (stdout, '(a)') 'MACE server sent an ERROR tag' + call flush (stdout) + call MPI_Recv(MPI_BOTTOM, 0, MPI_INTEGER, MPI_ANY_SOURCE, MACE_TAG_ERROR, & + mace_comm, mpi_status, ierr) + + else if (tag /= MACE_TAG_DATA) then + + write (errormsg, '(a,i0)') 'Got invalid MPI tag from MACE server: ', tag + call fatal_error(__FILE__, __LINE__, errormsg) + + else + success = .true. + end if + end function check_incoming_msg + + subroutine initialize_mace_server() + use mod_qmmm, only: natqm + use mod_system, only: names + + call send_natom(natqm, mace_comm) + call send_mace_atom_types(names, natqm, mace_comm) + end subroutine initialize_mace_server + + subroutine finalize_mace(error_code) + integer, intent(in) :: error_code + integer :: ierr, mpi_tag + + if (mace_comm == MPI_COMM_NULL) return + + ! Set error handler to return so we can handle errors gracefully + call MPI_Comm_set_errhandler(MPI_COMM_WORLD, MPI_ERRORS_RETURN, ierr) + + ! if the error code is MACE_TAG_ERROR, then + ! MACE server crashed and notified us so no need to + ! send it another message. + if (error_code /= MACE_TAG_ERROR) then + if (error_code == 0) then + mpi_tag = MACE_TAG_EXIT + else + mpi_tag = MACE_TAG_ERROR + end if + + write (stdout, '(A)') 'Sending shutdown message to MACE server' + + call MPI_Send(MPI_BOTTOM, 0, MPI_INTEGER, 0, mpi_tag, mace_comm, ierr) + + if (ierr /= MPI_SUCCESS) then + write (stderr, '(A)') 'MPI ERROR during shutdown of MACE server' + write (stderr, '(A)') 'Verify manually that the MACE server was terminated.' + write (stderr, *) get_mpi_error_string(ierr) + end if + end if + + write (stdout, '(A)') 'Freeing MACE communicator' + call MPI_Comm_free(mace_comm, ierr) + if (ierr /= MPI_SUCCESS) then + write (stderr, *) get_mpi_error_string(ierr) + end if + write (stdout, '(A)') 'MACE communicator freed' + end subroutine finalize_mace + + subroutine send_natom(num_atom, comm) + use mod_general, only: idebug + integer, intent(in) :: num_atom + integer, intent(in) :: comm + integer :: ierr + + if (idebug > 1) then + write (stdout, '(A, I0)') 'MACE: Sending number of atoms = ', num_atom + call flush (stdout) + end if + call MPI_Send(num_atom, 1, MPI_INTEGER, 0, MACE_TAG_DATA, comm, ierr) + call handle_mpi_error(ierr) + end subroutine send_natom + + subroutine send_mace_atom_types(at_names, num_atom, comm) + use mod_general, only: idebug + character(len=2), intent(in) :: at_names(:) + integer, intent(in) :: num_atom + integer, intent(in) :: comm + character(len=2*num_atom) :: buffer + integer :: ierr, offset, iat + + buffer = '' + offset = 1 + do iat = 1, num_atom + write (buffer(offset:offset + 1), '(A2)') at_names(iat) + offset = offset + 2 + end do + + if (idebug > 1) then + write (stdout, '(A)') 'MACE: Sending atom types: ' + write (stdout, '(A)') trim(buffer) + call flush (stdout) + end if + + call MPI_Send(buffer, num_atom * 2, MPI_CHARACTER, 0, MACE_TAG_DATA, comm, ierr) + call handle_mpi_error(ierr) + end subroutine send_mace_atom_types + + subroutine send_mace_coordinates(x, y, z, num_atom, iw, comm) + use mod_general, only: idebug + real(DP), intent(in) :: x(:, :), y(:, :), z(:, :) + integer, intent(in) :: num_atom + integer, intent(in) :: iw + integer, intent(in) :: comm + real(DP), allocatable :: coords(:, :) + integer :: ierr, iat + + allocate (coords(3, num_atom)) + + ! Send coordinates in Bohr (server converts to Angstrom) + do iat = 1, num_atom + coords(1, iat) = x(iat, iw) + coords(2, iat) = y(iat, iw) + coords(3, iat) = z(iat, iw) + end do + + if (idebug > 1) then + write (stdout, '(A)') 'MACE: Sending coordinates [Bohr]: ' + do iat = 1, num_atom + write (stdout, *) 'Atom ', iat, ': ', coords(:, iat) + call flush (stdout) + end do + end if + call MPI_Send(coords, num_atom * 3, MPI_DOUBLE_PRECISION, 0, MACE_TAG_DATA, comm, ierr) + call handle_mpi_error(ierr) + end subroutine send_mace_coordinates + +#else + + subroutine initialize_mace_interface() + use mod_error, only: not_compiled_with + call not_compiled_with('MPI') + end subroutine initialize_mace_interface + + subroutine initialize_mace_server() + use mod_error, only: not_compiled_with + call not_compiled_with('MPI') + end subroutine initialize_mace_server + + ! This must be a no-op, since it is called from finish() + subroutine finalize_mace(error_code) + integer, intent(in) :: error_code + integer :: i + i = error_code + end subroutine finalize_mace + +! USE_MPI +#endif + +end module mod_mace_mpi diff --git a/src/mpi_wrapper.F90 b/src/mpi_wrapper.F90 index 0c58aa53..c3df419d 100644 --- a/src/mpi_wrapper.F90 +++ b/src/mpi_wrapper.F90 @@ -61,6 +61,7 @@ subroutine finalize_mpi(error_code) integer :: ierr call MPI_Initialized(initialized, ierr) + call handle_mpi_error(ierr) if (.not. initialized) then return end if @@ -75,8 +76,10 @@ subroutine finalize_mpi(error_code) ! so MPI_Abort is safer. if (error_code /= 0 .and. get_mpi_size() > 1) then call MPI_Abort(MPI_COMM_WORLD, error_code, ierr) + call handle_mpi_error(ierr) else call MPI_Finalize(ierr) + call handle_mpi_error(ierr) end if end if end subroutine finalize_mpi @@ -173,6 +176,7 @@ subroutine initialize_mpi(pot, pot_ref, nteraservers) i = nteraservers if (iremd == 1 .or. & & pot == '_tera_' .or. pot_ref == '_tera_' .or. & + & pot == '_mace_' .or. & & pot == '_cp2k_' .or. pot_ref == '__cp2k__') then call not_compiled_with('MPI') end if diff --git a/tests/.gitignore b/tests/.gitignore index c52d490b..140385a9 100644 --- a/tests/.gitignore +++ b/tests/.gitignore @@ -28,6 +28,8 @@ engrad.dat.??? # PLUMED backup files bck.* +ompi_uri.txt + # Files from TERAPI* tests tc_server? tc_server diff --git a/tests/MACE/ERROR.ref b/tests/MACE/ERROR.ref new file mode 100644 index 00000000..e69de29b diff --git a/tests/MACE/energies.dat.ref b/tests/MACE/energies.dat.ref new file mode 100644 index 00000000..bdba150c --- /dev/null +++ b/tests/MACE/energies.dat.ref @@ -0,0 +1,3 @@ + # Time[fs] E-potential E-kinetic E-Total E-Total-Avg + 0.48 0.1900249814E-01 0.4107749350E-04 0.1904357563E-01 0.1904357563E-01 + 0.97 0.1887955609E-01 0.1639526646E-03 0.1904350876E-01 0.1904354219E-01 diff --git a/tests/MACE/forces.xyz.ref b/tests/MACE/forces.xyz.ref new file mode 100644 index 00000000..700ba5d3 --- /dev/null +++ b/tests/MACE/forces.xyz.ref @@ -0,0 +1,10 @@ + 3 +net force: -0.51642E-02 -0.51642E-02 0.10328E-01 torque force: -0.82312E-02 -0.53959E-02 0.13627E-01 +O 0.7386054962E-03 0.7386054962E-03 -0.1477210992E-02 +H 0.1811684277E-02 -0.1247763221E-01 0.1066594794E-01 +H -0.7714526717E-02 0.6574789773E-02 0.1139736944E-02 + 3 +net force: -0.51451E-02 -0.51451E-02 0.10290E-01 torque force: -0.82038E-02 -0.53788E-02 0.13583E-01 +O 0.7384535391E-03 0.7384535391E-03 -0.1476907078E-02 +H 0.1805766333E-02 -0.1243687350E-01 0.1063110716E-01 +H -0.7689326886E-02 0.6553312942E-02 0.1136013944E-02 diff --git a/tests/MACE/input.in b/tests/MACE/input.in new file mode 100644 index 00000000..0c063a9d --- /dev/null +++ b/tests/MACE/input.in @@ -0,0 +1,19 @@ +&general +pot='_mace_' +mdtype='md', +nstep=2, +dt=20., +irandom=13131313, +idebug=2 + +nwrite=1, +nwritef=1, +nwritev=1, +nwritex=1, +nrest=1, +/ + +&nhcopt +inose=0, +temp=0.0d0 +/ diff --git a/tests/MACE/mini.xyz b/tests/MACE/mini.xyz new file mode 100644 index 00000000..37dcf4f2 --- /dev/null +++ b/tests/MACE/mini.xyz @@ -0,0 +1,5 @@ +3 +Water molecule +O 0.00000000 0.00000000 0.11726400 +H 0.00000000 0.75698200 -0.46905800 +H 0.00000000 -0.75698200 -0.46905800 diff --git a/tests/MACE/movie.xyz.ref b/tests/MACE/movie.xyz.ref new file mode 100644 index 00000000..79c63dcc --- /dev/null +++ b/tests/MACE/movie.xyz.ref @@ -0,0 +1,10 @@ + 3 +Units: angs.; Time step: 1 ; Sim. Time [au]: 20.00 +O 0.26805290E-05 0.26805290E-05 0.11725864E+00 +H 0.10446384E-03 0.75626253E+00 -0.46844299E+00 +H -0.44482865E-03 -0.75660289E+00 -0.46899228E+00 + 3 +Units: angs.; Time step: 2 ; Sim. Time [au]: 40.00 +O 0.10721748E-04 0.10721748E-04 0.11724256E+00 +H 0.41762794E-03 0.75410567E+00 -0.46659929E+00 +H -0.17783462E-02 -0.75546638E+00 -0.46879527E+00 diff --git a/tests/MACE/temper.dat.ref b/tests/MACE/temper.dat.ref new file mode 100644 index 00000000..e8d85f11 --- /dev/null +++ b/tests/MACE/temper.dat.ref @@ -0,0 +1,3 @@ + # Time[fs] Temperature T-Average Conserved_quantity_of_thermostat + 0.48 2.88 2.88 + 0.97 11.50 7.19 diff --git a/tests/MACE/test.sh b/tests/MACE/test.sh new file mode 100755 index 00000000..d83c8429 --- /dev/null +++ b/tests/MACE/test.sh @@ -0,0 +1,140 @@ +#!/bin/bash +set -uo pipefail + +ABINEXE=$1 + +ABINOUT=abin.out +ABININ=input.in +ABINGEOM=mini.xyz + +MACE_SERVER=../../interfaces/MACE/mace_server.py +MACE_OUT=mace_server.out + +ABIN_CMD="$ABINEXE -i $ABININ -x $ABINGEOM" +MACE_CMD="python3 $MACE_SERVER --device cpu --model-path __MOCK_HARMONIC__" +PORT_FILE=mace_port.txt + +# If $1 = "clean"; exit early. +if [[ "${1-}" = "clean" ]]; then + rm -f $MACE_OUT $ABINOUT ./*.dat ./*.diff + rm -f restart.xyz velocities.xyz forces.xyz movie.xyz restart.xyz.old + rm -f $PORT_FILE ERROR ompi_uri.txt + exit 0 +fi + +# Skip the test if the python environment does not contain necessary libraries +if ! python3 -c "import mpi4py, numpy, ase" 2>/dev/null; then + echo "MACE python environment not activated" + exit 3 +fi + +# Determine MPI paths +if [[ -z ${MPI_PATH-} ]]; then + MPIRUN=mpirun +else + MPIRUN=$MPI_PATH/bin/mpirun + export PATH=$MPI_PATH/bin:$PATH + export LD_LIBRARY_PATH=$MPI_PATH/lib:${LD_LIBRARY_PATH-} +fi + +# Detect OpenMPI vs MPICH +IS_OPENMPI=false +if $MPIRUN --version 2>&1 | grep -q "Open MPI"; then + IS_OPENMPI=true +fi + +ompi_server_pid="" +MPIRUN_EXTRA_ARGS="" + +if [[ "$IS_OPENMPI" = "true" ]]; then + # OpenMPI requires ompi-server for MPI_Comm_connect/accept + OMPI_SERVER=${MPI_PATH-}/bin/ompi-server + if [[ -z ${MPI_PATH-} ]]; then + OMPI_SERVER=ompi-server + fi + + if ! which $OMPI_SERVER &> /dev/null; then + echo "ERROR: Skipping MACE test: ompi-server not found (required for OpenMPI)" >> ERROR + exit 1 + fi + + OMPI_URI_FILE="$PWD/ompi_uri.txt" + $OMPI_SERVER --no-daemonize -r "$OMPI_URI_FILE" & + ompi_server_pid=$! + sleep 1 + + if [[ ! -f "$OMPI_URI_FILE" ]]; then + echo "ERROR: ompi-server did not create URI file" >> ERROR + kill $ompi_server_pid 2> /dev/null + exit 1 + fi + + MPIRUN_EXTRA_ARGS="--ompi-server file:$OMPI_URI_FILE" +fi + +MPIRUN_CMD="$MPIRUN -n 1 $MPIRUN_EXTRA_ARGS" + +# Cleanup function to stop the background processes +function cleanup { + if [[ -n ${mace_pid-} ]] && kill -0 $mace_pid >& /dev/null; then + echo "ERROR: MACE server $mace_pid is still running!" >> ERROR + kill ${mace_pid-} &> /dev/null || true + fi + + if [[ -n ${abinpid-} ]] && kill -0 $abinpid >& /dev/null; then + echo "ERROR: ABIN process $abinpid is still running!" >> ERROR + kill ${abinpid-} &> /dev/null || true + fi + + if [[ -n "${ompi_server_pid-}" ]]; then + kill $ompi_server_pid &> /dev/null || true + fi +} + +function wait_for_portfile { + # Wait 10s for the MACE server to write the port file + MAX_WAIT=20 + i=0 + while [[ ! -f $PORT_FILE ]]; do + if [[ $i -gt $MAX_WAIT ]] || ! kill -0 $mace_pid >& /dev/null; then + echo "ERROR: MACE server did not write port file 'mace_port.txt'" | tee ERROR + set -x + cat $MACE_OUTPUT + exit 1 + fi + sleep 0.5 + let i++ + done +} + +# Automatically call the cleanup function when the script +# exits or is interrupted by a signal +trap cleanup INT ABRT TERM EXIT + +# Launch mock MACE server +$MPIRUN_CMD $MACE_CMD &> $MACE_OUT & +mace_pid=$! + +wait_for_portfile + +# Launch ABIN +$MPIRUN_CMD $ABIN_CMD &> $ABINOUT & +abinpid=$! + +# Give both processes 10 seconds to finish +MAX_ITER=20 +iter=0 +# Note about 'kill -0' https://unix.stackexchange.com/questions/169898/what-does-kill-0-do +while ( (kill -0 $abinpid >& /dev/null) || (kill -0 $mace_pid >& /dev/null) ); do + if [[ $iter -gt $MAX_ITER ]]; then + echo "Test did not finish in time" >> ERROR + break + fi + sleep 0.5 + let iter++ +done + +# Any errors in this script should be echoed to file ERROR +# so that the overall test fails when comparing to empty ERROR.ref file +# Here we create an empty one in case no errors actually occured, as is expected. +touch ERROR diff --git a/tests/MACE/velocities.xyz.ref b/tests/MACE/velocities.xyz.ref new file mode 100644 index 00000000..b904af36 --- /dev/null +++ b/tests/MACE/velocities.xyz.ref @@ -0,0 +1,10 @@ +3 +Units: a.u.; Time step: 1 +O 0.5065291961E-06 0.5065291961E-06 -0.1013058392E-05 +H 0.1973006086E-04 -0.1358870561E-03 0.1161569952E-03 +H -0.8401468378E-04 0.7160243317E-04 0.1241225060E-04 +3 +Units: a.u.; Time step: 2 +O 0.1012988920E-05 0.1012988920E-05 -0.2025977840E-05 +H 0.3941717126E-04 -0.2714782991E-03 0.2320611278E-03 +H -0.1678464756E-03 0.1430489947E-03 0.2479748093E-04 diff --git a/tests/MACE_ERROR/ERROR.ref b/tests/MACE_ERROR/ERROR.ref new file mode 100644 index 00000000..072500c7 --- /dev/null +++ b/tests/MACE_ERROR/ERROR.ref @@ -0,0 +1 @@ +ERROR in force_mace.F90: MACE evaluation failed diff --git a/tests/MACE_ERROR/input.in b/tests/MACE_ERROR/input.in new file mode 100644 index 00000000..642d402b --- /dev/null +++ b/tests/MACE_ERROR/input.in @@ -0,0 +1,13 @@ +! Test that ABIN handles ERROR tag sent by the MACE server +&general +pot='_mace_' +mdtype='md', +nstep=2, +dt=20., +idebug=2 +/ + +&nhcopt +inose=0, +temp=0.0d0 +/ diff --git a/tests/MACE_ERROR/mini.xyz b/tests/MACE_ERROR/mini.xyz new file mode 100644 index 00000000..37dcf4f2 --- /dev/null +++ b/tests/MACE_ERROR/mini.xyz @@ -0,0 +1,5 @@ +3 +Water molecule +O 0.00000000 0.00000000 0.11726400 +H 0.00000000 0.75698200 -0.46905800 +H 0.00000000 -0.75698200 -0.46905800 diff --git a/tests/MACE_ERROR/test.sh b/tests/MACE_ERROR/test.sh new file mode 100755 index 00000000..c565730f --- /dev/null +++ b/tests/MACE_ERROR/test.sh @@ -0,0 +1,135 @@ +#!/bin/bash +set -uo pipefail + +ABINEXE=$1 + +ABINOUT=abin.out +ABININ=input.in +ABINGEOM=mini.xyz + +MACE_SERVER=../../interfaces/MACE/mace_server.py +MACE_OUT=mace_server.out + +ABIN_CMD="$ABINEXE -i $ABININ -x $ABINGEOM" +MACE_CMD="python3 $MACE_SERVER --device cpu --model-path __MOCK_ERROR__" +PORT_FILE=mace_port.txt + +# If $1 = "clean"; exit early. +if [[ "${1-}" = "clean" ]]; then + rm -f $MACE_OUT $ABINOUT ./*.dat ./*.diff + rm -f restart.xyz velocities.xyz forces.xyz movie.xyz restart.xyz.old + rm -f $PORT_FILE ERROR ompi_uri.txt + exit 0 +fi + +# Skip the test if the python environment does not contain necessary libraries +if ! python3 -c "import mpi4py, numpy, ase" 2>/dev/null; then + echo "MACE python environment not activated" + exit 3 +fi + +# Determine MPI paths +if [[ -z ${MPI_PATH-} ]]; then + MPIRUN=mpirun +else + MPIRUN=$MPI_PATH/bin/mpirun + export PATH=$MPI_PATH/bin:$PATH + export LD_LIBRARY_PATH=$MPI_PATH/lib:${LD_LIBRARY_PATH-} +fi + +# Detect OpenMPI vs MPICH +IS_OPENMPI=false +if $MPIRUN --version 2>&1 | grep -q "Open MPI"; then + IS_OPENMPI=true +fi + +ompi_server_pid="" +MPIRUN_EXTRA_ARGS="" + +if [[ "$IS_OPENMPI" = "true" ]]; then + # OpenMPI requires ompi-server for MPI_Comm_connect/accept + OMPI_SERVER=${MPI_PATH-}/bin/ompi-server + if [[ -z ${MPI_PATH-} ]]; then + OMPI_SERVER=ompi-server + fi + + if ! which $OMPI_SERVER &> /dev/null; then + echo "ERROR: Skipping MACE test: ompi-server not found (required for OpenMPI)" >> ERROR + exit 1 + fi + + OMPI_URI_FILE="$PWD/ompi_uri.txt" + $OMPI_SERVER --no-daemonize -r "$OMPI_URI_FILE" & + ompi_server_pid=$! + sleep 1 + + if [[ ! -f "$OMPI_URI_FILE" ]]; then + echo "ERROR: ompi-server did not create URI file" >> ERROR + kill $ompi_server_pid 2> /dev/null + exit 1 + fi + + MPIRUN_EXTRA_ARGS="--ompi-server file:$OMPI_URI_FILE" +fi + +MPIRUN_CMD="$MPIRUN -n 1 $MPIRUN_EXTRA_ARGS" + +# Cleanup function to stop the background processes +function cleanup { + if [[ -n ${mace_pid-} ]] && kill -0 $mace_pid >& /dev/null; then + echo "ERROR: MACE server $mace_pid is still running!" >> ERROR + kill ${mace_pid-} &> /dev/null || true + fi + + if [[ -n ${abinpid-} ]] && kill -0 $abinpid >& /dev/null; then + echo "ERROR: ABIN process $abinpid is still running!" >> ERROR + kill ${abinpid-} &> /dev/null || true + fi + + if [[ -n "${ompi_server_pid-}" ]]; then + kill $ompi_server_pid &> /dev/null || true + fi +} + +function wait_for_portfile { + # Wait 10s for the MACE server to write the port file + MAX_WAIT=20 + i=0 + while [[ ! -f $PORT_FILE ]]; do + if [[ $i -gt $MAX_WAIT ]] || ! kill -0 $mace_pid >& /dev/null; then + echo "ERROR: MACE server did not write port file 'mace_port.txt'" | tee ERROR + set -x + cat $MACE_OUTPUT + exit 1 + fi + sleep 0.5 + let i++ + done +} + +# Automatically call the cleanup function when the script +# exits or is interrupted by a signal +trap cleanup INT ABRT TERM EXIT + +# Launch mock MACE server +$MPIRUN_CMD $MACE_CMD &> $MACE_OUT & +mace_pid=$! + +wait_for_portfile + +# Launch ABIN +$MPIRUN_CMD $ABIN_CMD &> $ABINOUT & +abinpid=$! + +# Give both processes 10 seconds to finish +MAX_ITER=20 +iter=0 +# Note about 'kill -0' https://unix.stackexchange.com/questions/169898/what-does-kill-0-do +while ( (kill -0 $abinpid >& /dev/null) || (kill -0 $mace_pid >& /dev/null) ); do + if [[ $iter -gt $MAX_ITER ]]; then + echo "Test did not finish in time" >> ERROR + break + fi + sleep 0.5 + let iter++ +done diff --git a/tests/MACE_ERROR2/ERROR.ref b/tests/MACE_ERROR2/ERROR.ref new file mode 100644 index 00000000..7dc4ce8d --- /dev/null +++ b/tests/MACE_ERROR2/ERROR.ref @@ -0,0 +1 @@ +ERROR in init.F90: invalid thermostat (therm) in input.in diff --git a/tests/MACE_ERROR2/input.in b/tests/MACE_ERROR2/input.in new file mode 100644 index 00000000..1ebbafa1 --- /dev/null +++ b/tests/MACE_ERROR2/input.in @@ -0,0 +1,13 @@ +! Test that when ABIN encounters a fatal error, +! it sends an error message to mace_server and mace_server shuts down gracefully. +! Set invalid thermostat to trigger ABIN error +&thermostat +therm='invalid' +/ +&general +pot='_mace_' +mdtype='md', +nstep=1, +dt=20., +idebug=2 +/ diff --git a/tests/MACE_ERROR2/mini.xyz b/tests/MACE_ERROR2/mini.xyz new file mode 100644 index 00000000..37dcf4f2 --- /dev/null +++ b/tests/MACE_ERROR2/mini.xyz @@ -0,0 +1,5 @@ +3 +Water molecule +O 0.00000000 0.00000000 0.11726400 +H 0.00000000 0.75698200 -0.46905800 +H 0.00000000 -0.75698200 -0.46905800 diff --git a/tests/MACE_ERROR2/test.sh b/tests/MACE_ERROR2/test.sh new file mode 100755 index 00000000..c565730f --- /dev/null +++ b/tests/MACE_ERROR2/test.sh @@ -0,0 +1,135 @@ +#!/bin/bash +set -uo pipefail + +ABINEXE=$1 + +ABINOUT=abin.out +ABININ=input.in +ABINGEOM=mini.xyz + +MACE_SERVER=../../interfaces/MACE/mace_server.py +MACE_OUT=mace_server.out + +ABIN_CMD="$ABINEXE -i $ABININ -x $ABINGEOM" +MACE_CMD="python3 $MACE_SERVER --device cpu --model-path __MOCK_ERROR__" +PORT_FILE=mace_port.txt + +# If $1 = "clean"; exit early. +if [[ "${1-}" = "clean" ]]; then + rm -f $MACE_OUT $ABINOUT ./*.dat ./*.diff + rm -f restart.xyz velocities.xyz forces.xyz movie.xyz restart.xyz.old + rm -f $PORT_FILE ERROR ompi_uri.txt + exit 0 +fi + +# Skip the test if the python environment does not contain necessary libraries +if ! python3 -c "import mpi4py, numpy, ase" 2>/dev/null; then + echo "MACE python environment not activated" + exit 3 +fi + +# Determine MPI paths +if [[ -z ${MPI_PATH-} ]]; then + MPIRUN=mpirun +else + MPIRUN=$MPI_PATH/bin/mpirun + export PATH=$MPI_PATH/bin:$PATH + export LD_LIBRARY_PATH=$MPI_PATH/lib:${LD_LIBRARY_PATH-} +fi + +# Detect OpenMPI vs MPICH +IS_OPENMPI=false +if $MPIRUN --version 2>&1 | grep -q "Open MPI"; then + IS_OPENMPI=true +fi + +ompi_server_pid="" +MPIRUN_EXTRA_ARGS="" + +if [[ "$IS_OPENMPI" = "true" ]]; then + # OpenMPI requires ompi-server for MPI_Comm_connect/accept + OMPI_SERVER=${MPI_PATH-}/bin/ompi-server + if [[ -z ${MPI_PATH-} ]]; then + OMPI_SERVER=ompi-server + fi + + if ! which $OMPI_SERVER &> /dev/null; then + echo "ERROR: Skipping MACE test: ompi-server not found (required for OpenMPI)" >> ERROR + exit 1 + fi + + OMPI_URI_FILE="$PWD/ompi_uri.txt" + $OMPI_SERVER --no-daemonize -r "$OMPI_URI_FILE" & + ompi_server_pid=$! + sleep 1 + + if [[ ! -f "$OMPI_URI_FILE" ]]; then + echo "ERROR: ompi-server did not create URI file" >> ERROR + kill $ompi_server_pid 2> /dev/null + exit 1 + fi + + MPIRUN_EXTRA_ARGS="--ompi-server file:$OMPI_URI_FILE" +fi + +MPIRUN_CMD="$MPIRUN -n 1 $MPIRUN_EXTRA_ARGS" + +# Cleanup function to stop the background processes +function cleanup { + if [[ -n ${mace_pid-} ]] && kill -0 $mace_pid >& /dev/null; then + echo "ERROR: MACE server $mace_pid is still running!" >> ERROR + kill ${mace_pid-} &> /dev/null || true + fi + + if [[ -n ${abinpid-} ]] && kill -0 $abinpid >& /dev/null; then + echo "ERROR: ABIN process $abinpid is still running!" >> ERROR + kill ${abinpid-} &> /dev/null || true + fi + + if [[ -n "${ompi_server_pid-}" ]]; then + kill $ompi_server_pid &> /dev/null || true + fi +} + +function wait_for_portfile { + # Wait 10s for the MACE server to write the port file + MAX_WAIT=20 + i=0 + while [[ ! -f $PORT_FILE ]]; do + if [[ $i -gt $MAX_WAIT ]] || ! kill -0 $mace_pid >& /dev/null; then + echo "ERROR: MACE server did not write port file 'mace_port.txt'" | tee ERROR + set -x + cat $MACE_OUTPUT + exit 1 + fi + sleep 0.5 + let i++ + done +} + +# Automatically call the cleanup function when the script +# exits or is interrupted by a signal +trap cleanup INT ABRT TERM EXIT + +# Launch mock MACE server +$MPIRUN_CMD $MACE_CMD &> $MACE_OUT & +mace_pid=$! + +wait_for_portfile + +# Launch ABIN +$MPIRUN_CMD $ABIN_CMD &> $ABINOUT & +abinpid=$! + +# Give both processes 10 seconds to finish +MAX_ITER=20 +iter=0 +# Note about 'kill -0' https://unix.stackexchange.com/questions/169898/what-does-kill-0-do +while ( (kill -0 $abinpid >& /dev/null) || (kill -0 $mace_pid >& /dev/null) ); do + if [[ $iter -gt $MAX_ITER ]]; then + echo "Test did not finish in time" >> ERROR + break + fi + sleep 0.5 + let iter++ +done diff --git a/tests/WITHOUT_MACE/ERROR.ref b/tests/WITHOUT_MACE/ERROR.ref new file mode 100644 index 00000000..e4443bdb --- /dev/null +++ b/tests/WITHOUT_MACE/ERROR.ref @@ -0,0 +1 @@ +ERROR in error.F90: ABIN was not compiled with MPI diff --git a/tests/WITHOUT_MACE/input.in b/tests/WITHOUT_MACE/input.in new file mode 100644 index 00000000..fccf833b --- /dev/null +++ b/tests/WITHOUT_MACE/input.in @@ -0,0 +1,5 @@ +&general +pot='_mace_' +nstep=1, +dt=40., +/ diff --git a/tests/WITHOUT_MACE/mini.xyz b/tests/WITHOUT_MACE/mini.xyz new file mode 100644 index 00000000..099f6aa0 --- /dev/null +++ b/tests/WITHOUT_MACE/mini.xyz @@ -0,0 +1,4 @@ + 2 + + H 0.000000 0.000000 0.000000 + H 0.000000 0.000000 3.000000 diff --git a/tests/numdiff.py b/tests/numdiff.py index 9e50e326..be75e24b 100755 --- a/tests/numdiff.py +++ b/tests/numdiff.py @@ -110,9 +110,9 @@ def parse_diff(fname, absolute_tolerance): if len(diff1) != len(diff2): print("Number of columns differ!") print("Expected:") - print(split[1]) + print(split[1].strip()) print("Test:") - print(split[0]) + print(split[0].strip()) exit(1) compare_lines(diff1, diff2, absolute_tolerance) @@ -127,6 +127,6 @@ def parse_diff(fname, absolute_tolerance): absolute_tolerance = read_custom_threshold(THR_FNAME) inpfile = read_cmd() - print("Comparing numerical differences in file " + inpfile) + print("\nComparing numerical differences in file " + inpfile) parse_diff(inpfile, absolute_tolerance) exit(0) diff --git a/tests/test.sh b/tests/test.sh index bb803760..180d65e3 100755 --- a/tests/test.sh +++ b/tests/test.sh @@ -72,6 +72,14 @@ function diff_files { error_code=0 diff -q $test_file $ref_file > /dev/null || error_code=$? if [[ $error_code -ne 0 ]];then + + if [[ $test_file = "ERROR" ]]; then + echo "Files $ref_file and $test_file differ!" + diff --color=always $ref_file $test_file | tee $test_file.diff + return_status=1 + continue + fi + # The reference file is different, but maybe it's just numerical noise? error_code=0 diff -y -W 500 $test_file $ref_file | grep -e '|' -e '<' -e '>' > $test_file.diff @@ -157,9 +165,17 @@ if [[ $TESTS = "all" ]];then folders[index]=TERAPI-SH-S0 let index++ folders[index]=TERAPI-LZ + let index++ + folders[index]=MACE + let index++ + folders[index]=MACE_ERROR + let index++ + folders[index]=MACE_ERROR2 else let index=${#folders[@]}+1 folders[index]=WITHOUT_MPI + let index++ + folders[index]=WITHOUT_MACE fi if [[ $CP2K = "TRUE" ]];then @@ -214,6 +230,8 @@ echo "Running tests in directories:" echo ${folders[@]} errors=0 +skipped=0 +passed=0 for dir in ${folders[@]} do @@ -221,7 +239,7 @@ do echo "Directory $dir not found. Exiting prematurely." exit 1 fi - echo "Entering directory $dir" + echo -en "$dir\t" cd $dir # Always clean the test directory before runnning the test. @@ -249,6 +267,13 @@ do # TODO: Figure out a different solution #./test.sh $ABINEXE 2> /dev/null ./test.sh $ABINEXE + # exit code 3 indicates skipped test + if [[ $? -eq 3 ]]; then + echo -e "\033[0;33mSKIPPED\033[0m" + let skipped++ + cd $TESTDIR + continue + fi else if [[ -f "velocities.in" ]];then @@ -271,6 +296,7 @@ do if diff_files; then echo -e "\033[0;32mPASSED\033[0m" + let passed++ else let errors++ echo -e "$dir \033[0;31mFAILED\033[0m" @@ -283,10 +309,15 @@ do done echo " " +if [[ $ACTION = "makeref" ]]; then + exit 0 +fi +echo -e "\033[0;32m$passed tests PASSED.\033[0m" +if [[ ${skipped} -ne 0 ]]; then + echo -e "\033[0;33m$skipped tests SKIPPED.\033[0m" +fi if [[ ${errors} -ne 0 ]];then echo -e "$errors tests \033[0;31mFAILED\033[0m." exit 1 -else - echo -e "\033[0;32mAll tests PASSED.\033[0m" fi diff --git a/unit_tests/throw_with_pfunit.F90 b/unit_tests/throw_with_pfunit.F90 index a0376d76..f15b6eec 100644 --- a/unit_tests/throw_with_pfunit.F90 +++ b/unit_tests/throw_with_pfunit.F90 @@ -11,12 +11,13 @@ module throw_with_pfunit_mod ! This routine replaces src/error.F90:fatal_error() ! so that we can catch exceptions in the unit tests. - subroutine throw(file_name, line_number, message) + subroutine throw(file_name, line_number, message, error_code) use funit, only: SourceLocation use funit, only: pFUnit_throw => throw character(len=*), intent(in) :: file_name integer, intent(in) :: line_number character(len=*), intent(in) :: message + integer, intent(in) :: error_code call pFUnit_throw(message, SourceLocation(file_name, line_number)) diff --git a/utils/run.mace_mpi_abin.sh b/utils/run.mace_mpi_abin.sh new file mode 100755 index 00000000..ff4f8604 --- /dev/null +++ b/utils/run.mace_mpi_abin.sh @@ -0,0 +1,146 @@ +#!/bin/bash +# Launch script for ABIN + MACE MPI interface. +# +# This script launches both the Python MACE server and ABIN, +# connecting them via MPI. +# +# Prerequisites: +# - ABIN compiled with MPI=TRUE +# - Python environment with: mpi4py, mace-torch, torch, ase, numpy +# - MPICH (not OpenMPI) +# +# Usage: +# ./run.mace_mpi_abin.sh + +set -uo pipefail + +# ABIN SETUP +ABIN_IN=input.in +GEOM_IN=mini.xyz +VELOC_IN= +ABIN_OUTPUT=abin.out +ABINEXE="./abin" # Path to ABIN binary +# Path to MPICH installation. Leave commented if you're using system-wide installation +# or if you're loading it via e.g. `load module mpich` +# This must match the installation that you used for ABIN compilation! +# MPI_PATH=/home/${USER}/software/mpich/4.3.2/install + +# MACE SETUP +MODEL_PATH=/path/to/mace.model +MACE_DEVICE=cpu # 'cpu' or 'cuda' + +# If you're on a system with multiple GPUs, select which GPU to use +# with CUDA_VISIBLE_DEVICES +#export CUDA_VISIBLE_DEVICES=0 + +# Path to Python executable with MACE dependencies (mpi4py, mace-torch, torch, ase, numpy). +# Typically, you'll provide an absolute path to your conda or virtual environment +MACE_PYTHON="../.venv/bin/python" + +# Path to MACE server script (copy from interfaces/MACE) +MACE_SERVER=./mace_server.py +MACE_OUTPUT=mace_server.out + + +##### END OF USER INPUT ##### +# + +# Convert to absolute path if not provided as such +MACE_PYTHON=$(command -v ${MACE_PYTHON:-python3}) +if [[ -z ${MACE_PYTHON} ]]; then + echo "ERROR: Could not find python executable!" + exit 1 +fi +if [[ -z ${MPI_PATH-} ]]; then + MPIRUN=$(command -v mpirun) + if [[ -z $MPIRUN ]]; then + echo "Command mpirun not found. Please specify MPI_PATH variable" + exit 1 + fi +else + MPIRUN="$MPI_PATH/bin/mpirun" + export LD_LIBRARY_PATH="${MPI_PATH}/lib:${LD_LIBRARY_PATH-}" +fi + +function files_exist() { + local error="" + for file in "$@"; do + if [[ ! -f $file ]];then + echo "ERROR: Cannot find file $file" >&2 + error=1 + fi + done + if [[ -n ${error-} ]];then + exit 1 + fi +} + +function validate_inputs() { + files_exist "$ABIN_IN" "$GEOM_IN" "$MACE_SERVER" "$ABINEXE" "$MACE_PYTHON" "$MPIRUN" + + # Check pot='_mace_' in ABIN input + test=$(grep -E -o -e "^[^!]*pot[[:space:]]*=[[:space:]]*['\"]_mace_[\"']" $ABIN_IN || true) + if [[ -z $test ]];then + echo "ERROR: You did not specify pot='_mace_' in $ABIN_IN." >&2 + exit 1 + fi +} + +# Cleanup function to stop the background processes +function cleanup { + if [[ -n ${mace_pid-} ]] && kill -0 $mace_pid >& /dev/null; then + echo "WARNING: MACE server $mace_pid is still running!" + kill ${mace_pid} + fi + + if [[ -n ${abin_pid-} ]] && kill -0 $abin_pid >& /dev/null; then + echo "WARNING: ABIN process $abin_pid is still running!" + kill ${abin_pid} + fi +} + +function wait_for_portfile { + # Wait 10s for the MACE server to write the port file + MAX_WAIT=20 + i=0 + while [[ ! -f mace_port.txt ]]; do + if [[ $i -gt $MAX_WAIT ]] || ! kill -0 $mace_pid >& /dev/null; then + echo "ERROR: MACE server did not write port file 'mace_port.txt'" + echo "cat $MACE_OUTPUT" + cat $MACE_OUTPUT + exit 1 + fi + sleep 0.5 + (( i++ )) + done +} + +# Validate input files exist +validate_inputs + +# Automatically call the cleanup function when the script +# exits or is interrupted by a signal +trap cleanup INT ABRT TERM EXIT + +# LAUNCH MACE SERVER +$MPIRUN $MACE_PYTHON $MACE_SERVER --device "$MACE_DEVICE" --model-path "$MODEL_PATH" > "$MACE_OUTPUT" 2>&1 & +mace_pid=$! +echo "Launched MACE server (PID: ${mace_pid})" + +wait_for_portfile + +# LAUNCH ABIN +ABIN_CMD="$ABINEXE -i $ABIN_IN -x $GEOM_IN" +if [[ -n $VELOC_IN ]];then + ABIN_CMD="$ABIN_CMD -v $VELOC_IN" +fi +$MPIRUN $ABIN_CMD &> $ABIN_OUTPUT & +abin_pid=$! +echo "Launched ABIN (PID: ${abin_pid})" +echo "(Monitor $ABIN_OUTPUT and $MACE_OUTPUT for progress)" + +# Note about 'kill -0' https://unix.stackexchange.com/questions/169898/what-does-kill-0-do +while ( (kill -0 $abin_pid >& /dev/null) && (kill -0 $mace_pid >& /dev/null) ); do sleep 1; done +sleep 2 # grace time for program termination + +echo "Simulation finished."