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"

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):

setup.py

@@ -26,7 +26,7 @@
     keywords='lammps, mpi4py',
     packages=find_packages(exclude=["*tests*"]),
     install_requires=[
-        'mpi4py==3.1.4'
+        'mpi4py==3.1.4', "pyzmq==25.0.0"
     ],
     cmdclass=versioneer.get_cmdclass(),
 )