ScipyDistribution can't be pickled

[cliffckerr]

Something is happening with ss.ScipyDistribution to break pickling, which oddly only comes up in the context of parallelization:

import sciris as sc
import starsim as ss
import multiprocessing as mp

vals = [0.5, 0.7]

def run(x=0.5):
    dist = ss.ScipyDistribution(ss.bernoulli(p=x))
    return dist

# Works
o = run()

but

# Doesn't work
with mp.Pool(2) as pool:
    out = pool.map(run, vals)

gives

MaybeEncodingError: Error sending result: '[<starsim.distributions.ScipyDistribution object at 0x7f97c5c1f290>]'. Reason: 'AttributeError("Can't pickle local object 'ScipyDistribution.__init__.<locals>.starsim_gen'")'

while

# Doesn't work
out = sc.parallelize(run, vals)

gives

Exception in thread Thread-13 (_handle_results):
Traceback (most recent call last):
  File "/software/anaconda3/lib/python3.11/threading.py", line 1045, in _bootstrap_inner
    self.run()
  File "/software/anaconda3/lib/python3.11/threading.py", line 982, in run
    self._target(*self._args, **self._kwargs)
  File "/software/anaconda3/lib/python3.11/site-packages/multiprocess/pool.py", line 579, in _handle_results
    task = get()
           ^^^^^
  File "/software/anaconda3/lib/python3.11/site-packages/multiprocess/connection.py", line 253, in recv
    return _ForkingPickler.loads(buf.getbuffer())
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/software/anaconda3/lib/python3.11/site-packages/dill/_dill.py", line 301, in loads
    return load(file, ignore, **kwds)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/software/anaconda3/lib/python3.11/site-packages/dill/_dill.py", line 287, in load
    return Unpickler(file, ignore=ignore, **kwds).load()
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/software/anaconda3/lib/python3.11/site-packages/dill/_dill.py", line 442, in load
    obj = StockUnpickler.load(self)
          ^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: 'NoneType' object is not callable

Even more bizarrely, if you don't return the distribution, it's fine!

[kaa2102]

Pickeling works best on top-level functions. I'm not certain what level the functions are in the grand scheme of the design. Also, can the function be turned into a global (raises more issues) or another package like dill be leveraged to improve handling of serialization/deserialization?