From ac1b072f6e94b6fc71396724da7a977c4a73df36 Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@lanl.gov>
Date: Mon, 22 May 2023 07:40:58 -0600
Subject: [PATCH] switch to cloud pickle

---
 .ci_support/environment-mpich.yml   | 1 +
 .ci_support/environment-openmpi.yml | 1 +
 pylammpsmpi/mpi/lmpmpi.py           | 6 +++---
 pylammpsmpi/utils/lammps.py         | 6 +++---
 setup.py                            | 2 +-
 5 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/.ci_support/environment-mpich.yml b/.ci_support/environment-mpich.yml
index 8211da8..e558ea2 100644
--- a/.ci_support/environment-mpich.yml
+++ b/.ci_support/environment-mpich.yml
@@ -1,6 +1,7 @@
 channels:
   - conda-forge
 dependencies:
+  - cloudpickle =2.2.1
   - coveralls
   - coverage
   - codacy-coverage
diff --git a/.ci_support/environment-openmpi.yml b/.ci_support/environment-openmpi.yml
index 7899186..95e15de 100644
--- a/.ci_support/environment-openmpi.yml
+++ b/.ci_support/environment-openmpi.yml
@@ -1,6 +1,7 @@
 channels:
   - conda-forge
 dependencies:
+  - cloudpickle =2.2.1
   - coveralls
   - coverage
   - codacy-coverage
diff --git a/pylammpsmpi/mpi/lmpmpi.py b/pylammpsmpi/mpi/lmpmpi.py
index b933d35..1bef0a1 100644
--- a/pylammpsmpi/mpi/lmpmpi.py
+++ b/pylammpsmpi/mpi/lmpmpi.py
@@ -5,7 +5,7 @@
 from ctypes import c_double, c_int
 from mpi4py import MPI
 import numpy as np
-import pickle
+import cloudpickle
 import sys
 import zmq
 from lammps import lammps
@@ -492,7 +492,7 @@ def _gather_data_from_all_processors(data):
         socket.connect("tcp://localhost:" + port_selected)
     while True:
         if MPI.COMM_WORLD.rank == 0:
-            input_dict = pickle.loads(socket.recv())
+            input_dict = cloudpickle.loads(socket.recv())
             # with open('process.txt', 'a') as file:
             #     print('Input:', input_dict, file=file)
         else:
@@ -505,4 +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)
-            socket.send(pickle.dumps(output))
+            socket.send(cloudpickle.dumps(output))
diff --git a/pylammpsmpi/utils/lammps.py b/pylammpsmpi/utils/lammps.py
index 437cdf4..779d2eb 100644
--- a/pylammpsmpi/utils/lammps.py
+++ b/pylammpsmpi/utils/lammps.py
@@ -3,7 +3,7 @@
 # Distributed under the terms of "New BSD License", see the LICENSE file.
 
 import os
-import pickle
+import cloudpickle
 import subprocess
 import zmq
 
@@ -75,7 +75,7 @@ def _send(self, command, data=None):
         -------
         None
         """
-        self._socket.send(pickle.dumps({"c": command, "d": data}))
+        self._socket.send(cloudpickle.dumps({"c": command, "d": data}))
 
     def _receive(self):
         """
@@ -90,7 +90,7 @@ def _receive(self):
         data : string
             data from the command
         """
-        output = pickle.loads(self._socket.recv())
+        output = cloudpickle.loads(self._socket.recv())
         return output
 
     @property
diff --git a/setup.py b/setup.py
index 5d2548e..678bd59 100644
--- a/setup.py
+++ b/setup.py
@@ -26,7 +26,7 @@
     keywords='lammps, mpi4py',
     packages=find_packages(exclude=["*tests*"]),
     install_requires=[
-        'mpi4py==3.1.4', "pyzmq==25.0.0"
+         "cloudpickle==2.2.1", "mpi4py==3.1.4", "pyzmq==25.0.0"
     ],
     cmdclass=versioneer.get_cmdclass(),
 )