diff --git a/dev_scripts/install_mpich.sh b/dev_scripts/install_mpich.sh index de865217..41d14bfe 100755 --- a/dev_scripts/install_mpich.sh +++ b/dev_scripts/install_mpich.sh @@ -11,7 +11,7 @@ CC=${CC:-gcc} FC=${FC:-gfortran} if [[ -z ${1-} ]]; then - echo "ERROR: Provide prefix path where install MPICH as first parameter" + echo "ERROR: Provide prefix path where to install MPICH as first parameter" exit 1 fi # Path as an optional first parameter diff --git a/interfaces/MACE/README.md b/interfaces/MACE/README.md index b42d768a..736ac4f3 100644 --- a/interfaces/MACE/README.md +++ b/interfaces/MACE/README.md @@ -4,7 +4,7 @@ This directory contains the Python-based server for the MACE (Machine Learning A ## Requirements -- Python >= 3.8 +- Python >= 3.9 - PyTorch >= 1.12 - mace-torch - ase @@ -15,7 +15,7 @@ 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 +```bash # 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 @@ -32,7 +32,7 @@ source .venv/bin/activate # Activates the environment 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 +```bash $ make test TEST="MACE MACE_ERROR MACE_ERROR2" Running tests in directories: MACE MACE_ERROR MACE_ERROR2 @@ -55,6 +55,8 @@ before running production calculations! 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`. +MACE model is required to use this interface. If you don't have one, you can check out MACE [foundational models](https://mace-docs.readthedocs.io/en/latest/guide/foundation_models.html). + Rough steps: 1. Specify `pot='_mace_'` in the ABIN input file, everything else is the same. diff --git a/interfaces/MACE/mace_server.py b/interfaces/MACE/mace_server.py index f642081d..86b6fd3c 100755 --- a/interfaces/MACE/mace_server.py +++ b/interfaces/MACE/mace_server.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 +# ruff: noqa: BLE001 # /// script -# requires-python = ">=3.8" +# requires-python = ">=3.9" # dependencies = [ # "ase>=3.18.0", # "mace-torch>=0.3.10", @@ -18,16 +19,79 @@ The server writes its MPI port to 'mace_port.txt' for ABIN to read. """ -# ruff: file-ignore[blind-except] import argparse import functools +import logging import sys +import warnings from pathlib import Path from time import perf_counter -from traceback import print_tb -LOG_NAME = "MaceMPIServer" +logger = logging.getLogger(__name__) + + +def setup_logger(debug=True): + """Configure standard library logging to output to sys.stdout.""" + level = logging.DEBUG if debug else logging.INFO + logging.basicConfig( + level=level, + format="%(asctime)s | %(levelname)-8s | %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + stream=sys.stdout, + force=True, + ) + logging.captureWarnings(True) + warnings.formatwarning = lambda msg, cat, fname, lineno, line=None: ( + f"{fname}:{lineno}: {cat.__name__}: {str(msg).strip()}" + ) + + +# Configure logger at module import time (debug=True by default) +setup_logger(debug=True) + + +def log_environment_info(config): + """Log debug information about Python, PyTorch, CUDA, dependencies, and environment.""" + import platform + + import ase + import mpi4py + import numpy as np + import torch + + logger.debug("=== MACE Server Environment ===") + logger.debug(f"Python Executable : {sys.executable}") + logger.debug(f"Python Version : {sys.version.replace('\n', ' ')}") + logger.debug(f"Platform / OS : {platform.platform()}") + logger.debug(f"Working Directory : {Path.cwd()}") + logger.debug(f"Model Path : {config.model_path}") + logger.debug(f"Configured Device : {config.device}") + + # Dependency versions + logger.debug(f"PyTorch Version : {torch.__version__}") + logger.debug(f"ASE Version : {ase.__version__}") + logger.debug(f"NumPy Version : {np.__version__}") + logger.debug(f"mpi4py Version : {mpi4py.__version__}") + + try: + import mace + + logger.debug(f"MACE Version : {getattr(mace, '__version__', 'unknown')}") + except ImportError: + logger.debug("MACE Version : Not installed") + + # CUDA & Hardware details + cuda_avail = torch.cuda.is_available() + logger.debug(f"CUDA Available : {cuda_avail}") + if cuda_avail: + logger.debug(f"CUDA PyTorch Build: {torch.version.cuda}") + logger.debug(f"GPU Device Count : {torch.cuda.device_count()}") + logger.debug(f"GPU Device Name : {torch.cuda.get_device_name(0)}") + mem_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3) + logger.debug(f"GPU Total Memory : {mem_gb:.2f} GB") + logger.debug("======================================================") + # MPI Tags (must match Fortran module mod_mace_mpi) MACE_TAG_EXIT = 666 @@ -37,15 +101,6 @@ 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) @@ -62,10 +117,19 @@ def parse_cmd(): required=True, help="Device for model inference", ) + parser.add_argument( + "--debug", + action=argparse.BooleanOptionalAction, + default=True, + help="Enable or disable debug logging output", + ) config = parser.parse_args() + if not config.debug: + setup_logger(debug=False) 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") + logger.error(f"File '{config.model_path}' not found") + sys.exit(1) return config @@ -78,7 +142,7 @@ class MaceModel: def __init__(self, config): from mace.calculators import MACECalculator - log("initializing MACE model") + logger.info("initializing MACE model") # Set ASE calculator self.calculator = MACECalculator( @@ -130,7 +194,7 @@ class HarmonicModel: def __init__(self, config): self.model_path = config.model_path - log("Using Harmonic Mock Model") + logger.info("Using Harmonic Mock Model") def evaluate(self, _atom_types, coords_bohr): """ @@ -160,24 +224,26 @@ def connect_to_abin(): # Open MPI port and write to file for ABIN to read port_name = MPI.Open_port() - log(f"MPI port opened: {port_name}") + logger.info(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}") + logger.info(f"Port written to {MACE_PORT_FILE}") # Accept connection from ABIN - log("Waiting for ABIN to connect...") + logger.info("Waiting for ABIN to connect...") abin_comm = MPI.COMM_WORLD.Accept(port_name) - log("Connection from ABIN accepted!") + logger.info("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) + logger.error( + f"Unexpected {exception_type.__name__}: {exception}", + exc_info=(exception_type, exception, traceback), + ) # Restore original exception handling to prevent endless loop # in case of uncaught excpetion during shutdown sys.excepthook = sys.__excepthook__ @@ -189,32 +255,33 @@ def main(config): import numpy as np from mpi4py import MPI + log_environment_info(config) port_name, abin_comm = connect_to_abin() def shutdown_communication(): """Gracefully shutdown communication with ABIN""" - log("Shutting down communication with ABIN...") + logger.info("Shutting down communication with ABIN...") try: abin_comm.Disconnect() except Exception as e: - log(e) + logger.error(f"Error disconnecting ABIN communicator: {e}") else: - log("ABIN communicator disconnected") + logger.info("ABIN communicator disconnected") try: MPI.Close_port(port_name) except Exception as e: - log(e) + logger.error(f"Error closing port {port_name}: {e}") else: - log(f"Port {port_name} closed") + logger.info(f"Port {port_name} closed") def error_shutdown(): - log("Sending ERROR tag to ABIN") + logger.warning("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) + logger.error(f"Error sending ERROR tag to ABIN: {e}") shutdown_communication() @@ -227,16 +294,16 @@ def check_incoming_msg(): if (tag := status.Get_tag()) in (MACE_TAG_EXIT, MACE_TAG_ERROR): if tag == MACE_TAG_EXIT: - log("Received graceful exit signal from ABIN") + logger.info("Received graceful exit signal from ABIN") exit_code = 0 else: - log("Received ERROR signal from ABIN. Stopping server") + logger.warning("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) + logger.error(f"Error receiving tag payload: {e}") shutdown_communication() sys.exit(exit_code) @@ -249,7 +316,7 @@ def 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}") + logger.info(f"Received number of atoms: {natom}") # Receive atom types check_incoming_msg() @@ -262,7 +329,7 @@ def check_incoming_msg(): 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}") + logger.info(f"Received atom types: {atom_types}") assert len(atom_types) == natom # Load MACE model @@ -272,7 +339,7 @@ def check_incoming_msg(): else: mace_model = MaceModel(config) - log("MACE model ready. Entering main loop.") + logger.info("MACE model ready. Entering main loop.") # Main loop: receive coordinates, compute, send results eval_count = 0 @@ -293,8 +360,8 @@ def check_incoming_msg(): 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") + logger.info(f"Step {eval_count} done in {time_ms:.3f} miliseconds") + logger.info(f"Energy = {energy:.15f} Hartree") # Send energy (1 double, in Hartree) energy_buf = np.array([energy], dtype=np.float64) @@ -312,7 +379,7 @@ def check_incoming_msg(): 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") + logger.info(f"Communication overhead = {loop_ms - time_ms:.3f} ms") eval_count += 1 diff --git a/utils/run.mace_mpi_abin.sh b/utils/run.mace_mpi_abin.sh index ff4f8604..50564d18 100755 --- a/utils/run.mace_mpi_abin.sh +++ b/utils/run.mace_mpi_abin.sh @@ -100,7 +100,7 @@ function cleanup { } function wait_for_portfile { - # Wait 10s for the MACE server to write the port file + # Wait 20s for the MACE server to write the port file MAX_WAIT=20 i=0 while [[ ! -f mace_port.txt ]]; do