I am stating the issues I faced with pickling and MPI with the PSO module here for posterity. I was getting the following pickling error (identified after 2c16876).

  File "/u/home/a/ajshajib/mybin/lenstronomy/lenstronomy/Workflow/fitting_sequence.py", line 80, in fit_sequence
    kwargs_result, chain, param = self.pso(**kwargs)
  File "/u/home/a/ajshajib/mybin/lenstronomy/lenstronomy/Workflow/fitting_sequence.py", line 238, in pso
    threadCount=threadCount, mpi=self._mpi, print_key=print_key)
  File "/u/home/a/ajshajib/mybin/lenstronomy/lenstronomy/Sampling/sampler.py", line 62, in pso
    result, [chi2_list, pos_list, vel_list] = pso.optimize(n_iterations)
  File "/u/home/a/ajshajib/mybin/lenstronomy/lenstronomy/Sampling/Samplers/pso.py", line 198, in optimize
    for _ in self.sample(max_iter, c1, c2, p, m, n):
  File "/u/home/a/ajshajib/mybin/lenstronomy/lenstronomy/Sampling/Samplers/pso.py", line 133, in sample
    self._get_fitness(self.swarm)
  File "/u/home/a/ajshajib/mybin/lenstronomy/lenstronomy/Sampling/Samplers/pso.py", line 223, in _get_fitness
    ln_probability = list(map_func(self.func, position))
  File "/u/home/a/ajshajib/mybin/schwimmbad/schwimmbad/mpi.py", line 185, in map
    self.comm.send(task, dest=worker, tag=taskid)
  File "mpi4py/MPI/Comm.pyx", line 1156, in mpi4py.MPI.Comm.send
  File "mpi4py/MPI/msgpickle.pxi", line 173, in mpi4py.MPI.PyMPI_send
  File "mpi4py/MPI/msgpickle.pxi", line 104, in mpi4py.MPI.Pickle.dump
  File "mpi4py/MPI/msgpickle.pxi", line 91, in mpi4py.MPI.Pickle.cdumps
TypeError: can't pickle module objects

I tried the following things that didn't work.

• putting the likelihood function within a wrapper class (idea borrowed from emcee)
• explicitly setting __getstate__ and __setstate__ methods in the classes
• change all class variable types to Python primitive types only

So, finally I resorted to using dill for pickling. Luckily, I added the dill functionality to schwimmbad sometime ago (adrn/schwimmbad@6bcc321). However, in my previous experience, I found dill to slow the computation down, expectedly because it's a pure-Python serialization, although I am not sure how much quantitatively. If it's a small hit, that should be fine. However, I am out of ideas to try for now (although, if emcee can work in MPI with default pickling, why can't the PSO module?) I am fine to use dill for now until someone (or me) finds the motivation to make things faster and to look for other ways to get MPI+PSO work with default pickling.

I ran a short test run on hoffman2, and the 4-core run took about 1.7 times shorter than a 2-core run, so the MPI-distribution is working.

Hi Anowar, thanks for looking into this in so much detail! I am happy to merge it for now and we can leave an issue open for the future. One way that might speed things up (but I don't know how to implement) is that the class instances do only need to be pickled/dill'ed once and only the functional arguments are then passed each time. All the likelihood computations can use the same instance of the likelihood class and the map function only needs to tell the arguments.

Yes, please merge if you find no issues in the code. Your suggestion sounds effective, but I'm also not sure exactly how to implement it. An issue to track it would be nice.