From b36ed2e1fb0d1c6afa595f3d7c54bb4b9b1b8c70 Mon Sep 17 00:00:00 2001 From: FelyCZ Date: Wed, 5 Aug 2026 22:04:20 +0200 Subject: [PATCH 1/5] fix: better syntax highlighting --- interfaces/MACE/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/interfaces/MACE/README.md b/interfaces/MACE/README.md index b42d768a..847ad52f 100644 --- a/interfaces/MACE/README.md +++ b/interfaces/MACE/README.md @@ -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 From 63c7852c38831b5bcb8cfc8425801936cad73523 Mon Sep 17 00:00:00 2001 From: FelyCZ Date: Wed, 5 Aug 2026 22:17:04 +0200 Subject: [PATCH 2/5] fix: typo --- dev_scripts/install_mpich.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 From c50e4e8e8221e2b1f6ccadcd472d3db966f246f3 Mon Sep 17 00:00:00 2001 From: FelyCZ Date: Thu, 6 Aug 2026 01:29:29 +0200 Subject: [PATCH 3/5] feat: implement logging --- interfaces/MACE/README.md | 2 +- interfaces/MACE/mace_server.py | 137 ++++++++++++++++++++++++--------- utils/run.mace_mpi_abin.sh | 2 +- 3 files changed, 102 insertions(+), 39 deletions(-) diff --git a/interfaces/MACE/README.md b/interfaces/MACE/README.md index 847ad52f..9881493c 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 diff --git a/interfaces/MACE/mace_server.py b/interfaces/MACE/mace_server.py index f642081d..cf022770 100755 --- a/interfaces/MACE/mace_server.py +++ b/interfaces/MACE/mace_server.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 # /// script -# requires-python = ">=3.8" +# requires-python = ">=3.9" # dependencies = [ # "ase>=3.18.0", # "mace-torch>=0.3.10", @@ -22,12 +22,72 @@ 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" +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 + + logging.debug("=== MACE Server Environment ===") + logging.debug(f"Python Executable : {sys.executable}") + logging.debug(f"Python Version : {sys.version.replace('\n', ' ')}") + logging.debug(f"Platform / OS : {platform.platform()}") + logging.debug(f"Working Directory : {Path.cwd()}") + logging.debug(f"Model Path : {config.model_path}") + logging.debug(f"Configured Device : {config.device}") + + # Dependency versions + logging.debug(f"PyTorch Version : {torch.__version__}") + logging.debug(f"ASE Version : {ase.__version__}") + logging.debug(f"NumPy Version : {np.__version__}") + logging.debug(f"mpi4py Version : {mpi4py.__version__}") + + try: + import mace + + logging.debug(f"MACE Version : {getattr(mace, '__version__', 'unknown')}") + except ImportError: + logging.debug("MACE Version : Not installed") + + # CUDA & Hardware details + cuda_avail = torch.cuda.is_available() + logging.debug(f"CUDA Available : {cuda_avail}") + if cuda_avail: + logging.debug(f"CUDA PyTorch Build: {torch.version.cuda}") + logging.debug(f"GPU Device Count : {torch.cuda.device_count()}") + logging.debug(f"GPU Device Name : {torch.cuda.get_device_name(0)}") + mem_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3) + logging.debug(f"GPU Total Memory : {mem_gb:.2f} GB") + logging.debug("======================================================") + # MPI Tags (must match Fortran module mod_mace_mpi) MACE_TAG_EXIT = 666 @@ -37,15 +97,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 +113,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") + logging.error(f"File '{config.model_path}' not found") + sys.exit(1) return config @@ -78,7 +138,7 @@ class MaceModel: def __init__(self, config): from mace.calculators import MACECalculator - log("initializing MACE model") + logging.info("initializing MACE model") # Set ASE calculator self.calculator = MACECalculator( @@ -130,7 +190,7 @@ class HarmonicModel: def __init__(self, config): self.model_path = config.model_path - log("Using Harmonic Mock Model") + logging.info("Using Harmonic Mock Model") def evaluate(self, _atom_types, coords_bohr): """ @@ -160,24 +220,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}") + logging.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}") + logging.info(f"Port written to {MACE_PORT_FILE}") # Accept connection from ABIN - log("Waiting for ABIN to connect...") + logging.info("Waiting for ABIN to connect...") abin_comm = MPI.COMM_WORLD.Accept(port_name) - log("Connection from ABIN accepted!") + logging.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) + logging.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 +251,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...") + logging.info("Shutting down communication with ABIN...") try: abin_comm.Disconnect() except Exception as e: - log(e) + logging.error(f"Error disconnecting ABIN communicator: {e}") else: - log("ABIN communicator disconnected") + logging.info("ABIN communicator disconnected") try: MPI.Close_port(port_name) except Exception as e: - log(e) + logging.error(f"Error closing port {port_name}: {e}") else: - log(f"Port {port_name} closed") + logging.info(f"Port {port_name} closed") def error_shutdown(): - log("Sending ERROR tag to ABIN") + logging.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) + logging.error(f"Error sending ERROR tag to ABIN: {e}") shutdown_communication() @@ -227,16 +290,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") + logging.info("Received graceful exit signal from ABIN") exit_code = 0 else: - log("Received ERROR signal from ABIN. Stopping server") + logging.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) + logging.error(f"Error receiving tag payload: {e}") shutdown_communication() sys.exit(exit_code) @@ -249,7 +312,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}") + logging.info(f"Received number of atoms: {natom}") # Receive atom types check_incoming_msg() @@ -262,7 +325,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}") + logging.info(f"Received atom types: {atom_types}") assert len(atom_types) == natom # Load MACE model @@ -272,7 +335,7 @@ def check_incoming_msg(): else: mace_model = MaceModel(config) - log("MACE model ready. Entering main loop.") + logging.info("MACE model ready. Entering main loop.") # Main loop: receive coordinates, compute, send results eval_count = 0 @@ -293,8 +356,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") + logging.info(f"Step {eval_count} done in {time_ms:.3f} miliseconds") + logging.info(f"Energy = {energy:.15f} Hartree") # Send energy (1 double, in Hartree) energy_buf = np.array([energy], dtype=np.float64) @@ -312,7 +375,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") + logging.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 From a0b26cfaf8dd589e9b8b068c2c38f609805a00ee Mon Sep 17 00:00:00 2001 From: FelyCZ Date: Thu, 6 Aug 2026 01:48:16 +0200 Subject: [PATCH 4/5] docs: mention MACE foundational model --- interfaces/MACE/README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/interfaces/MACE/README.md b/interfaces/MACE/README.md index 9881493c..736ac4f3 100644 --- a/interfaces/MACE/README.md +++ b/interfaces/MACE/README.md @@ -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. From 2e76d97079c22ff91ddcdb0d6e6d96f95ccc1d4e Mon Sep 17 00:00:00 2001 From: FelyCZ Date: Thu, 6 Aug 2026 02:03:20 +0200 Subject: [PATCH 5/5] fix: reformat to pass ruff check --- interfaces/MACE/mace_server.py | 96 ++++++++++++++++++---------------- 1 file changed, 50 insertions(+), 46 deletions(-) diff --git a/interfaces/MACE/mace_server.py b/interfaces/MACE/mace_server.py index cf022770..86b6fd3c 100755 --- a/interfaces/MACE/mace_server.py +++ b/interfaces/MACE/mace_server.py @@ -1,4 +1,5 @@ #!/usr/bin/env python3 +# ruff: noqa: BLE001 # /// script # requires-python = ">=3.9" # dependencies = [ @@ -18,7 +19,6 @@ The server writes its MPI port to 'mace_port.txt' for ABIN to read. """ -# ruff: file-ignore[blind-except] import argparse import functools @@ -28,6 +28,9 @@ from pathlib import Path from time import perf_counter +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 @@ -39,8 +42,8 @@ def setup_logger(debug=True): force=True, ) logging.captureWarnings(True) - warnings.formatwarning = ( - lambda msg, cat, fname, lineno, line=None: f"{fname}:{lineno}: {cat.__name__}: {str(msg).strip()}" + warnings.formatwarning = lambda msg, cat, fname, lineno, line=None: ( + f"{fname}:{lineno}: {cat.__name__}: {str(msg).strip()}" ) @@ -51,42 +54,43 @@ def 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 - logging.debug("=== MACE Server Environment ===") - logging.debug(f"Python Executable : {sys.executable}") - logging.debug(f"Python Version : {sys.version.replace('\n', ' ')}") - logging.debug(f"Platform / OS : {platform.platform()}") - logging.debug(f"Working Directory : {Path.cwd()}") - logging.debug(f"Model Path : {config.model_path}") - logging.debug(f"Configured Device : {config.device}") + 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 - logging.debug(f"PyTorch Version : {torch.__version__}") - logging.debug(f"ASE Version : {ase.__version__}") - logging.debug(f"NumPy Version : {np.__version__}") - logging.debug(f"mpi4py Version : {mpi4py.__version__}") + 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 - logging.debug(f"MACE Version : {getattr(mace, '__version__', 'unknown')}") + logger.debug(f"MACE Version : {getattr(mace, '__version__', 'unknown')}") except ImportError: - logging.debug("MACE Version : Not installed") + logger.debug("MACE Version : Not installed") # CUDA & Hardware details cuda_avail = torch.cuda.is_available() - logging.debug(f"CUDA Available : {cuda_avail}") + logger.debug(f"CUDA Available : {cuda_avail}") if cuda_avail: - logging.debug(f"CUDA PyTorch Build: {torch.version.cuda}") - logging.debug(f"GPU Device Count : {torch.cuda.device_count()}") - logging.debug(f"GPU Device Name : {torch.cuda.get_device_name(0)}") + 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) - logging.debug(f"GPU Total Memory : {mem_gb:.2f} GB") - logging.debug("======================================================") + logger.debug(f"GPU Total Memory : {mem_gb:.2f} GB") + logger.debug("======================================================") # MPI Tags (must match Fortran module mod_mace_mpi) @@ -124,7 +128,7 @@ def parse_cmd(): setup_logger(debug=False) model = config.model_path if not model.startswith("__MOCK_") and not Path(model).is_file(): - logging.error(f"File '{config.model_path}' not found") + logger.error(f"File '{config.model_path}' not found") sys.exit(1) return config @@ -138,7 +142,7 @@ class MaceModel: def __init__(self, config): from mace.calculators import MACECalculator - logging.info("initializing MACE model") + logger.info("initializing MACE model") # Set ASE calculator self.calculator = MACECalculator( @@ -190,7 +194,7 @@ class HarmonicModel: def __init__(self, config): self.model_path = config.model_path - logging.info("Using Harmonic Mock Model") + logger.info("Using Harmonic Mock Model") def evaluate(self, _atom_types, coords_bohr): """ @@ -220,23 +224,23 @@ def connect_to_abin(): # Open MPI port and write to file for ABIN to read port_name = MPI.Open_port() - logging.info(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) - logging.info(f"Port written to {MACE_PORT_FILE}") + logger.info(f"Port written to {MACE_PORT_FILE}") # Accept connection from ABIN - logging.info("Waiting for ABIN to connect...") + logger.info("Waiting for ABIN to connect...") abin_comm = MPI.COMM_WORLD.Accept(port_name) - logging.info("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""" - logging.error( + logger.error( f"Unexpected {exception_type.__name__}: {exception}", exc_info=(exception_type, exception, traceback), ) @@ -256,28 +260,28 @@ def main(config): def shutdown_communication(): """Gracefully shutdown communication with ABIN""" - logging.info("Shutting down communication with ABIN...") + logger.info("Shutting down communication with ABIN...") try: abin_comm.Disconnect() except Exception as e: - logging.error(f"Error disconnecting ABIN communicator: {e}") + logger.error(f"Error disconnecting ABIN communicator: {e}") else: - logging.info("ABIN communicator disconnected") + logger.info("ABIN communicator disconnected") try: MPI.Close_port(port_name) except Exception as e: - logging.error(f"Error closing port {port_name}: {e}") + logger.error(f"Error closing port {port_name}: {e}") else: - logging.info(f"Port {port_name} closed") + logger.info(f"Port {port_name} closed") def error_shutdown(): - logging.warning("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: - logging.error(f"Error sending ERROR tag to ABIN: {e}") + logger.error(f"Error sending ERROR tag to ABIN: {e}") shutdown_communication() @@ -290,16 +294,16 @@ def check_incoming_msg(): if (tag := status.Get_tag()) in (MACE_TAG_EXIT, MACE_TAG_ERROR): if tag == MACE_TAG_EXIT: - logging.info("Received graceful exit signal from ABIN") + logger.info("Received graceful exit signal from ABIN") exit_code = 0 else: - logging.warning("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: - logging.error(f"Error receiving tag payload: {e}") + logger.error(f"Error receiving tag payload: {e}") shutdown_communication() sys.exit(exit_code) @@ -312,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]) - logging.info(f"Received number of atoms: {natom}") + logger.info(f"Received number of atoms: {natom}") # Receive atom types check_incoming_msg() @@ -325,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) ] - logging.info(f"Received atom types: {atom_types}") + logger.info(f"Received atom types: {atom_types}") assert len(atom_types) == natom # Load MACE model @@ -335,7 +339,7 @@ def check_incoming_msg(): else: mace_model = MaceModel(config) - logging.info("MACE model ready. Entering main loop.") + logger.info("MACE model ready. Entering main loop.") # Main loop: receive coordinates, compute, send results eval_count = 0 @@ -356,8 +360,8 @@ def check_incoming_msg(): end = perf_counter() time_ms = (end - start) * 1000 - logging.info(f"Step {eval_count} done in {time_ms:.3f} miliseconds") - logging.info(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) @@ -375,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. - logging.info(f"Communication overhead = {loop_ms - time_ms:.3f} ms") + logger.info(f"Communication overhead = {loop_ms - time_ms:.3f} ms") eval_count += 1