Use executorlib directly (#371)

* Use executorlib directly

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix ruff

* error handling

* use contextlib

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix tests

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fixes

* extend tests

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* get_result() function

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* clean up

* increase coverage

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* remove old interface

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add backend test

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix test

* more tests

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* use more recent lammps version for old tests

* relax mpi4py

* downgrade lammps

* downgrade once more

* working version

* fix mpi4py

* add test for installed packages

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add test for command

* fix test

* test get_file

* extend test for extract_atom

* provide executor as additional argument

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix import

* use more recent version of executorlib

* remove broadcast

* support for external executor

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: jan-janssen

.ci_support/environment-old.yml

@@ -5,7 +5,7 @@ dependencies:
 - openmpi
 - numpy =1.23.5
 - mpi4py =4.0.1
-- executorlib =1.2.0
+- executorlib =1.3.0
 - ase =3.23.0
 - scipy =1.9.3
 - hatchling

pylammpsmpi/__init__.py

@@ -1,9 +1,9 @@
 import pylammpsmpi._version
 from pylammpsmpi.wrapper.base import LammpsBase
-from pylammpsmpi.wrapper.concurrent import LammpsConcurrent
+from pylammpsmpi.wrapper.concurrent import LammpsConcurrent, init_function
 from pylammpsmpi.wrapper.extended import LammpsLibrary
 
-__all__ = ["LammpsLibrary", "LammpsConcurrent", "LammpsBase"]
+__all__ = ["LammpsLibrary", "LammpsConcurrent", "LammpsBase", "init_function"]
 __version__ = pylammpsmpi._version.__version__
 
 

pylammpsmpi/mpi/lmpmpi.py

@@ -1,17 +1,9 @@
 # Copyright (c) Max-Planck-Institut für Eisenforschung GmbH - Computational Materials Design (CM) Department
 # Distributed under the terms of "New BSD License", see the LICENSE file.
 
-import sys
 from ctypes import c_double, c_int
 
 import numpy as np
-from executorlib.api import (
-    interface_connect,
-    interface_receive,
-    interface_send,
-    interface_shutdown,
-)
-from lammps import lammps
 from mpi4py import MPI
 
 __author__ = "Sarath Menon, Jan Janssen"
@@ -80,7 +72,10 @@ def convert_data(val, type, length, width):
             data=job.numpy.extract_compute(*filtered_args)
         )
         length = job.get_natoms()
-        return convert_data(val=val, type=type, length=length, width=width)
+        if MPI.COMM_WORLD.rank == 0:
+            return convert_data(val=val, type=type, length=length, width=width)
+        else:
+            return val
     else:  # Todo
         raise ValueError("Local style is currently not supported")
 
@@ -156,7 +151,10 @@ def extract_variable(job, funct_args):
         data = _gather_data_from_all_processors(
             data=job.numpy.extract_variable(*funct_args)
         )
-        return np.array(data)
+        if MPI.COMM_WORLD.rank == 0:
+            return np.array(data)
+        else:
+            return np.array([])
     else:
         # if type is 1 - reformat file
         try:
@@ -428,50 +426,7 @@ def select_cmd(argument):
 
 def _gather_data_from_all_processors(data):
     data_gather = MPI.COMM_WORLD.gather(data, root=0)
-    return [v for vl in data_gather for v in vl]
-
-
-def _run_lammps_mpi(argument_lst):
-    index_selected = argument_lst.index("--zmqport")
-    port_selected = argument_lst[index_selected + 1]
-    if "--host" in argument_lst:
-        index_selected = argument_lst.index("--host")
-        host = argument_lst[index_selected + 1]
-    else:
-        host = "localhost"
-    argument_red_lst = argument_lst[:index_selected]
     if MPI.COMM_WORLD.rank == 0:
-        context, socket = interface_connect(host=host, port=port_selected)
+        return [v for vl in data_gather for v in vl]
     else:
-        context, socket = None, None
-    # Lammps executable
-    args = ["-screen", "none"]
-    if len(argument_red_lst) > 1:
-        args.extend(argument_red_lst[1:])
-    job = lammps(cmdargs=args)
-    while True:
-        if MPI.COMM_WORLD.rank == 0:
-            input_dict = interface_receive(socket=socket)
-        else:
-            input_dict = None
-        input_dict = MPI.COMM_WORLD.bcast(input_dict, root=0)
-        if "shutdown" in input_dict and input_dict["shutdown"]:
-            job.close()
-            if MPI.COMM_WORLD.rank == 0:
-                interface_send(socket=socket, result_dict={"result": True})
-                interface_shutdown(socket=socket, context=context)
-            break
-        try:
-            output = select_cmd(input_dict["command"])(
-                job=job, funct_args=input_dict["args"]
-            )
-        except Exception as error:
-            if MPI.COMM_WORLD.rank == 0:
-                interface_send(socket=socket, result_dict={"error": error})
-        else:
-            if MPI.COMM_WORLD.rank == 0 and output is not None:
-                interface_send(socket=socket, result_dict={"result": output})
-
-
-if __name__ == "__main__":
-    _run_lammps_mpi(argument_lst=sys.argv)
+        return []

pylammpsmpi/wrapper/ase.py

@@ -8,6 +8,7 @@
 from ase.atoms import Atoms
 from ase.calculators.lammps import Prism
 from ase.data import atomic_masses, atomic_numbers
+from executorlib import BaseExecutor
 from scipy import constants
 
 from pylammpsmpi.wrapper.base import LammpsBase
@@ -25,6 +26,7 @@ class LammpsASELibrary:
         log_file (str, optional): The log file path. Defaults to None.
         library (object, optional): The LAMMPS library object. Defaults to None.
         disable_log_file (bool, optional): Whether to disable the log file. Defaults to True.
+        executor: Executor to use for parallel execution (default: None)
     """
 
     def __init__(
@@ -36,6 +38,7 @@ def __init__(
         log_file: Optional[str] = None,
         library: Optional[object] = None,
         disable_log_file: bool = True,
+        executor: Optional[BaseExecutor] = None,
     ):
         self._logger = logger
         self._prism = None
@@ -60,7 +63,9 @@ def __init__(
                 )
         else:
             self._interactive_library = LammpsBase(
-                cores=self._cores, working_directory=working_directory
+                cores=self._cores,
+                working_directory=working_directory,
+                executor=executor,
             )
 
     def interactive_lib_command(self, command: str) -> None:

pylammpsmpi/wrapper/base.py

@@ -1,7 +1,8 @@
 # Copyright (c) Max-Planck-Institut für Eisenforschung GmbH - Computational Materials Design (CM) Department
 # Distributed under the terms of "New BSD License", see the LICENSE file.
 
-from typing import Union
+from concurrent.futures import Future
+from typing import Any, Union
 
 from pylammpsmpi.wrapper.concurrent import LammpsConcurrent
 
@@ -17,6 +18,13 @@
 __date__ = "Feb 28, 2020"
 
 
+def get_result(future: Future, cores: int) -> Any:
+    if cores > 1:
+        return future.result()[0]
+    else:
+        return future.result()
+
+
 class LammpsBase(LammpsConcurrent):
     @property
     def version(self) -> str:
@@ -27,7 +35,7 @@ def version(self) -> str:
         version: str
             version string of lammps
         """
-        return super().version.result()
+        return get_result(future=super().version, cores=self.cores)
 
     def file(self, inputfile: str) -> None:
         """
@@ -55,7 +63,7 @@ def extract_setting(self, *args) -> Union[int, float, str]:
         value: int, float, or str
             extracted setting value
         """
-        return super().extract_setting(*args).result()
+        return get_result(future=super().extract_setting(*args), cores=self.cores)
 
     def extract_global(self, name: str) -> Union[int, float, str]:
         """
@@ -69,7 +77,7 @@ def extract_global(self, name: str) -> Union[int, float, str]:
         value: int, float, or str
             extracted value of the global parameter
         """
-        return super().extract_global(name=name).result()
+        return get_result(future=super().extract_global(name=name), cores=self.cores)
 
     def extract_box(self) -> list[Union[float, list[float], list[int]]]:
         """
@@ -84,7 +92,7 @@ def extract_box(self) -> list[Union[float, list[float], list[int]]]:
             the box is periodic in three dimensions, and box_change is a list of booleans
             indicating if the box dimensions have changed
         """
-        return super().extract_box().result()
+        return get_result(future=super().extract_box(), cores=self.cores)
 
     def extract_atom(self, name: str) -> Union[list[int], list[float]]:
         """
@@ -99,7 +107,7 @@ def extract_atom(self, name: str) -> Union[list[int], list[float]]:
             If the requested name has multiple dimensions, output
             will be a multi-dimensional list.
         """
-        return super().extract_atom(name=name).result()
+        return get_result(future=super().extract_atom(name=name), cores=self.cores)
 
     def extract_fix(self, *args) -> Union[int, float, list[Union[int, float]]]:
         """
@@ -113,7 +121,7 @@ def extract_fix(self, *args) -> Union[int, float, list[Union[int, float]]]:
         value: int, float, or list of int or float
             extracted fix value corresponding to the requested dimensions
         """
-        return super().extract_fix(*args).result()
+        return get_result(future=super().extract_fix(*args), cores=self.cores)
 
     def extract_variable(self, *args) -> Union[int, float, list[Union[int, float]]]:
         """
@@ -127,7 +135,7 @@ def extract_variable(self, *args) -> Union[int, float, list[Union[int, float]]]:
         data: int, float, or list of int or float
             value of the variable depending on the requested dimension
         """
-        return super().extract_variable(*args).result()
+        return get_result(future=super().extract_variable(*args), cores=self.cores)
 
     @property
     def natoms(self) -> int:
@@ -148,7 +156,7 @@ def get_natoms(self) -> int:
         natoms : int
             number of atoms
         """
-        return super().get_natoms().result()
+        return get_result(future=super().get_natoms(), cores=self.cores)
 
     def set_variable(self, *args) -> int:
         """
@@ -162,7 +170,7 @@ def set_variable(self, *args) -> int:
         flag : int
             0 if successful, -1 otherwise
         """
-        return super().set_variable(*args).result()
+        return get_result(future=super().set_variable(*args), cores=self.cores)
 
     def reset_box(self, *args) -> None:
         """
@@ -270,27 +278,27 @@ def create_atoms(
     @property
     def has_exceptions(self) -> bool:
         """Return whether the LAMMPS shared library was compiled with C++ exceptions handling enabled"""
-        return super().has_exceptions.result()
+        return get_result(future=super().has_exceptions, cores=self.cores)
 
     @property
     def has_gzip_support(self) -> bool:
-        return super().has_gzip_support.result()
+        return get_result(future=super().has_gzip_support, cores=self.cores)
 
     @property
     def has_png_support(self) -> bool:
-        return super().has_png_support.result()
+        return get_result(future=super().has_png_support, cores=self.cores)
 
     @property
     def has_jpeg_support(self) -> bool:
-        return super().has_jpeg_support.result()
+        return get_result(future=super().has_jpeg_support, cores=self.cores)
 
     @property
     def has_ffmpeg_support(self) -> bool:
-        return super().has_ffmpeg_support.result()
+        return get_result(future=super().has_ffmpeg_support, cores=self.cores)
 
     @property
     def installed_packages(self) -> list[str]:
-        return super().installed_packages.result()
+        return get_result(future=super().installed_packages, cores=self.cores)
 
     def set_fix_external_callback(self, *args) -> None:
         _ = super().set_fix_external_callback(*args).result()
@@ -302,7 +310,7 @@ def get_neighlist(self, *args):
         :return: an instance of :class:`NeighList` wrapping access to neighbor list data
         :rtype:  NeighList
         """
-        return super().get_neighlist(*args).result()
+        return get_result(future=super().get_neighlist(*args), cores=self.cores)
 
     def find_pair_neighlist(self, *args) -> int:
         """Find neighbor list index of pair style neighbor list
@@ -324,7 +332,7 @@ def find_pair_neighlist(self, *args) -> int:
         :return: neighbor list index if found, otherwise -1
         :rtype:  int
         """
-        return super().find_pair_neighlist(*args).result()
+        return get_result(future=super().find_pair_neighlist(*args), cores=self.cores)
 
     def find_fix_neighlist(self, *args):
         """Find neighbor list index of fix neighbor list
@@ -335,7 +343,7 @@ def find_fix_neighlist(self, *args):
         :return: neighbor list index if found, otherwise -1
         :rtype:  int
         """
-        return super().find_fix_neighlist(*args).result()
+        return get_result(future=super().find_fix_neighlist(*args), cores=self.cores)
 
     def find_compute_neighlist(self, *args):
         """Find neighbor list index of compute neighbor list
@@ -346,7 +354,9 @@ def find_compute_neighlist(self, *args):
         :return: neighbor list index if found, otherwise -1
         :rtype:  int
         """
-        return super().find_compute_neighlist(*args).result()
+        return get_result(
+            future=super().find_compute_neighlist(*args), cores=self.cores
+        )
 
     def get_neighlist_size(self, *args):
         """Return the number of elements in neighbor list with the given index
@@ -355,10 +365,12 @@ def get_neighlist_size(self, *args):
         :return: number of elements in neighbor list with index idx
         :rtype:  int
         """
-        return super().get_neighlist_size(*args).result()
+        return get_result(future=super().get_neighlist_size(*args), cores=self.cores)
 
     def get_neighlist_element_neighbors(self, *args):
-        return super().get_neighlist_element_neighbors(*args).result()
+        return get_result(
+            future=super().get_neighlist_element_neighbors(*args), cores=self.cores
+        )
 
     def command(self, cmd):
         """
@@ -407,7 +419,9 @@ def gather_atoms(self, *args, concat=False, ids=None):
         --------
         extract_atoms
         """
-        return super().gather_atoms(*args, concat=concat, ids=ids).result()
+        return get_result(
+            future=super().gather_atoms(*args, concat=concat, ids=ids), cores=self.cores
+        )
 
     def scatter_atoms(self, *args, ids=None):
         """
@@ -433,7 +447,7 @@ def get_thermo(self, *args):
             value of the thermo keyword
 
         """
-        return super().get_thermo(*args).result()
+        return get_result(future=super().get_thermo(*args), cores=self.cores)
 
     # TODO
     def extract_compute(self, id, style, type, length=0, width=0):
@@ -468,8 +482,9 @@ def extract_compute(self, id, style, type, length=0, width=0):
             data computed by the fix depending on the chosen inputs
 
         """
-        return (
-            super()
-            .extract_compute(id=id, style=style, type=type, length=length, width=width)
-            .result()
+        return get_result(
+            future=super().extract_compute(
+                id=id, style=style, type=type, length=length, width=width
+            ),
+            cores=self.cores,
         )