[Python] Pathos multiprocessing map not completing on Linux machine, works on Mac M2

I am using the pathos multiprocessing to parallelize a gradient calculation that is embarrassingly parallel, using finite difference.

Below is a high level example of how it is set up,

` ...

class Routine:
def __init__(self):


...


def create_grad(self):

from pathos.pools import ProcessPool

def numerical_grad_p(betas, h=1e-5):

pool = ProcessPool(nodes=6)

def compute_partial_derivative(i):
t_betas = betas.copy()
t_betas += h * t_betas
return (self.neg_log_likelihood(t_betas, d=True) - self.neg_log_likelihood_i) / (h * t_betas)

indices = range(len(betas))
grad = np.array(pool.map(compute_partial_derivative, indices))

return grad

self.grad_func = numerical_grad_p


def HMC(self):
...
gradient_func = self.create_grad()
grad = gradient_func(betas)
...
return samples


`

The code runs on my MacBook successfully but does when transferring to my Linux Ubuntu desktop the compute partial gradient function never finishes.

I have tried passing the full self object into the function, as well as defining the neg_log_likelihood function as a variable then passing that into the function.

The solver used is non-picklable and I haven't been able to find success using things like JAX mapping or the default multiprocessing package. It is a custom modeling software, where the HMC runs through parameters (betas) to optimize.

I am new to parallel computing and have tried things like changing to spawn instead of fork but there is a very real chance I made a mistake doing so. Any help is appreciated.

Continue reading...