Merge pull request #97 from pyiron/zmq

Communicated via zmq

Author: jan-janssen

.ci_support/environment-mpich.yml

@@ -0,0 +1,13 @@
+channels:
+  - conda-forge
+dependencies:
+  - coveralls
+  - coverage
+  - codacy-coverage
+  - lammps >=2022.06.23
+  - mpi4py =3.1.4
+  - mpich
+  - numpy
+  - distributed
+  - dask-jobqueue
+  - pyzmq =25.0.0

.ci_support/environment.yml renamed to .ci_support/environment-openmpi.yml

@@ -10,3 +10,4 @@ dependencies:
   - numpy
   - distributed 
   - dask-jobqueue
+  - pyzmq =25.0.0

.github/workflows/UpdateDependabotPR.yml

@@ -20,7 +20,8 @@ jobs:
           package=$(echo "${{ github.event.pull_request.title }}" | awk '{print $2}')
           from=$(echo "${{ github.event.pull_request.title }}" | awk '{print $4}')
           to=$(echo "${{ github.event.pull_request.title }}" | awk '{print $6}')
-          sed -i "/${package}/s/${from}/${to}/g" .ci_support/environment.yml
+          sed -i "/${package}/s/${from}/${to}/g" .ci_support/environment-openmpi.yml
+          sed -i "/${package}/s/${from}/${to}/g" .ci_support/environment-mpich.yml
       - name: UpdateDependabotPR commit
         run: |
           git config --local user.email "41898282+github-actions[bot]@users.noreply.github.com"

.github/workflows/coverage.yml

@@ -22,7 +22,7 @@ jobs:
         channels: conda-forge
         channel-priority: strict
         use-mamba: true
-        environment-file: .ci_support/environment.yml
+        environment-file: .ci_support/environment-openmpi.yml
     - name: Setup
       shell: bash -l {0}
       run: |

.github/workflows/pypicheck.yml

@@ -18,7 +18,7 @@ jobs:
         channels: conda-forge
         channel-priority: strict
         use-mamba: true
-        environment-file: .ci_support/environment.yml
+        environment-file: .ci_support/environment-openmpi.yml
     - name: Setup
       shell: bash -l {0}
       run: |

.github/workflows/unittests.yml renamed to .github/workflows/unittests-mpich.yml

@@ -1,7 +1,7 @@
 # This workflow will install Python dependencies, run tests and lint with a variety of Python versions
 # For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions
 
-name: Python package
+name: mpich
 
 on:
   push:
@@ -26,7 +26,7 @@ jobs:
         channels: conda-forge
         channel-priority: strict
         use-mamba: true
-        environment-file: .ci_support/environment.yml
+        environment-file: .ci_support/environment-mpich.yml
     - name: Setup
       shell: bash -l {0}
       run: |

.github/workflows/unittests-openmpi.yml

@@ -0,0 +1,36 @@
+# This workflow will install Python dependencies, run tests and lint with a variety of Python versions
+# For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions
+
+name: OpenMPI
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ['3.8', '3.9', '3.10', '3.11']
+
+    steps:
+    - uses: actions/checkout@v2
+    - uses: conda-incubator/setup-miniconda@v2
+      with:
+        python-version: ${{ matrix.python-version }}
+        miniforge-variant: Mambaforge
+        channels: conda-forge
+        channel-priority: strict
+        use-mamba: true
+        environment-file: .ci_support/environment-openmpi.yml
+    - name: Setup
+      shell: bash -l {0}
+      run: |
+        pip install --no-deps .
+    - name: Test
+      shell: bash -l {0}
+      run: coverage run --omit pylammpsmpi/_version.py -m unittest discover tests

pylammpsmpi/lammps_wrapper.py

@@ -23,17 +23,25 @@ class LammpsLibrary:
     """
 
     def __init__(
-        self, cores=1, working_directory=".", client=None, mode="local", cmdargs=None
+        self,
+        cores=1,
+        oversubscribe=False,
+        working_directory=".",
+        client=None,
+        mode="local",
+        cmdargs=None,
     ):
         self.cores = cores
         self.working_directory = working_directory
+        self.oversubscribe = oversubscribe
         self.client = client
         self.mode = mode
 
         if self.mode == "dask":
             fut = self.client.submit(
                 LammpsBase,
                 cores=self.cores,
+                oversubscribe=self.oversubscribe,
                 working_directory=self.working_directory,
                 cmdargs=cmdargs,
                 actor=True,
@@ -46,6 +54,7 @@ def __init__(
         elif self.mode == "local":
             self.lmp = LammpsBase(
                 cores=self.cores,
+                oversubscribe=self.oversubscribe,
                 working_directory=self.working_directory,
                 cmdargs=cmdargs,
             )

pylammpsmpi/mpi/lmpmpi.py

@@ -7,6 +7,7 @@
 import numpy as np
 import pickle
 import sys
+import zmq
 from lammps import lammps
 
 __author__ = "Sarath Menon, Jan Janssen"
@@ -43,8 +44,8 @@
 
 # Lammps executable
 args = ["-screen", "none"]
-if len(sys.argv) > 1:
-    args.extend(sys.argv[1:])
+if len(sys.argv) > 3:
+    args.extend(sys.argv[3:])
 job = lammps(cmdargs=args)
 
 
@@ -483,9 +484,15 @@ def _gather_data_from_all_processors(data):
 
 
 if __name__ == "__main__":
+    if MPI.COMM_WORLD.rank == 0:
+        context = zmq.Context()
+        socket = context.socket(zmq.PAIR)
+        argument_lst = sys.argv
+        port_selected = argument_lst[argument_lst.index("--zmqport") + 1]
+        socket.connect("tcp://localhost:" + port_selected)
     while True:
         if MPI.COMM_WORLD.rank == 0:
-            input_dict = pickle.load(sys.stdin.buffer)
+            input_dict = pickle.loads(socket.recv())
             # with open('process.txt', 'a') as file:
             #     print('Input:', input_dict, file=file)
         else:
@@ -498,5 +505,4 @@ def _gather_data_from_all_processors(data):
         if MPI.COMM_WORLD.rank == 0 and output is not None:
             # with open('process.txt', 'a') as file:
             #     print('Output:', output, file=file)
-            pickle.dump(output, sys.stdout.buffer)
-            sys.stdout.flush()
+            socket.send(pickle.dumps(output))

pylammpsmpi/utils/lammps.py

@@ -5,6 +5,7 @@
 import os
 import pickle
 import subprocess
+import zmq
 
 
 __author__ = "Sarath Menon, Jan Janssen"
@@ -20,23 +21,33 @@
 
 
 class LammpsBase:
-    def __init__(self, cores=8, working_directory=".", cmdargs=None):
+    def __init__(
+        self, cores=8, oversubscribe=False, working_directory=".", cmdargs=None
+    ):
         self.cores = cores
         self.working_directory = working_directory
         self._process = None
+        self._oversubscribe = oversubscribe
         self._cmdargs = cmdargs
+        self._socket = None
 
     def start_process(self):
         executable = os.path.join(
             os.path.dirname(os.path.abspath(__file__)), "../mpi", "lmpmpi.py"
         )
-        cmds = [
-            "mpiexec",
-            "--oversubscribe",
+        context = zmq.Context()
+        self._socket = context.socket(zmq.PAIR)
+        port_selected = self._socket.bind_to_random_port("tcp://*")
+        cmds = ["mpiexec"]
+        if self._oversubscribe:
+            cmds += ["--oversubscribe"]
+        cmds += [
             "-n",
             str(self.cores),
             "python",
             executable,
+            "--zmqport",
+            str(port_selected),
         ]
         if self._cmdargs is not None:
             cmds.extend(self._cmdargs)
@@ -64,8 +75,7 @@ def _send(self, command, data=None):
         -------
         None
         """
-        pickle.dump({"c": command, "d": data}, self._process.stdin)
-        self._process.stdin.flush()
+        self._socket.send(pickle.dumps({"c": command, "d": data}))
 
     def _receive(self):
         """
@@ -80,7 +90,7 @@ def _receive(self):
         data : string
             data from the command
         """
-        output = pickle.load(self._process.stdout)
+        output = pickle.loads(self._socket.recv())
         return output
 
     @property
@@ -670,6 +680,7 @@ def close(self):
         except AttributeError:
             pass
         self._process = None
+        self._socket = None
 
     # TODO
     def __del__(self):