Pickling objects for parallel calculations

[rozyczko]

If we want to make core parallel (not threaded! parallel = running in multiple processes), we need a way to send/receive basic objects between processes.
This communication requires properly serialized objects.

Currently, we serialize BUMPS FitProblem in order to send it for parallel processing in BUMPS.
This can be generalized to standard EasyList/ModelBase based objects.

Serialization in parallel BUMPS sampling

How the model is sent to worker processes when sampling with n_workers > 1 (Sampler.sample() / Sampler.extend()), as implemented in src/easyscience/fitting/samplers/parallel.py.

Summary

Parallel sampling does not use EasyScience's own serialization (to_dict / from_dict / SerializerBase), and it does not serialize only
BUMPS objects either. Exactly one object is pickled, the BUMPS FitProblem.
Pickling it pulls in every EasyScience object it references. Those objects are pickled as raw object state by a custom reducer.

What actually gets pickled

DreamSampler.run builds the problem with build_curve_problem and sends it to BumpsPoolMapper, which pickles it once. Following references from there:

FitProblem                    (BUMPS)
 └─ Curve                     (BUMPS; also holds BUMPS Parameter objects)
     └─ EvalCounter           (easyscience wrapper around the fit function)
         └─ _fit_function     (closure from EngineBase._generate_fit_function)
             ├─ self -> DreamSampler engine
             │    ├─ _object -> the user's model (ModelBase / EasyList / ...)
             │    └─ _cached_pars -> the easyscience Parameters being sampled
             └─ func -> reshaping closures -> the user's fit function (often the model itself)

These objects are pickled by cloudpickle, using the _BumpsProblemPickler subclass created by _problem_pickler_class().

How NewBase objects are handled

The pickler's reducer_override gets every NewBase instance. That covers ModelBase, EasyList, DescriptorBase and Parameter.

• Writing. _reduce_object_state takes the object's __dict__ minus _global_object, and passes each value through _serialize_worker_value. That function turns scipp scalars into plain {value, variance, unit} dicts and drops weak references. Nested dicts, lists, tuples and sets are treated recursively.

• Reading. _new_object calls cls.__new__(cls) without running __init__. This is deliberate: __init__ would register the object in the
  worker's global map and generate a new unique_name, or clash with an existing one. _set_object_state then restores the attributes and attaches the worker's own global_object.

EasyList has no special path. It is a NewBase, so its __dict__ (_data, _protected_types, names) is pickled as above, and each item in _data goes through the same reducer. Its to_dict / _convert_to_dict code
is not involved.

Classes and functions. Anything importable by name, such as ModelBase or Parameter, is pickled as a reference: module plus name. Classes and functions defined in __main__ or a notebook, and closures, are pickled by value by cloudpickle. That is why a custom model class derived from those will work.

Other special cases.

• Weak references become None.
• Text streams become the worker's own stdout / stderr. This covers the writer of the asteval interpreter behind dependent parameters.
• Everything else gets cloudpickled

How the workers deal with the pickled problem

_dump_problem pauses the cyclic garbage collector while it pickles. A collection in the middle of the dump can run finalizers for objects left over from an earlier run, for example the global map's clean-up of freed parameters.
Those finalizers can resize a dictionary the pickler is iterating, which fails with RuntimeError: dictionary changed size during iteration.

The bytes are written to a private temporary file, and only its path is passed to the pool initializer (_init_bumps_worker). Each worker reads and unpickles the file once; the file is removed when the pool shuts down.
Passing the bytes as an initializer argument instead would send them through the spawn pipe. The content usually does not fit in that pipe's buffer making this method scale VERY badly.

Known limitation: array-valued scipp data

_serialize_worker_value assumes every scipp Variable is a scalar. A model containing a DescriptorArray (array-valued) makes the dump fail with:

DimensionError: The 'value' property cannot be used with non-scalar Variables.

Scipp Variables cannot be pickled natively either (TypeError: cannot pickle 'scipp._scipp.core.Variable' object), so they do
need converting. Until this is fixed, any model containing a DescriptorArray cannot use n_workers > 1 and fails with the FitError above. The fix is to store dims, values, variances, unit and dtype, and rebuild with sc.array / sc.scalar.

Why not use to_dict / from_dict?

The dict serializer rebuilds objects through their constructors (_arg_spec). It cannot represent what a worker needs: the user's fit function, the closures from _generate_fit_function and the reshaping step, and the BUMPS problem around them. Capturing the exact runtime state with pickle avoids that. The trade-off is that it relies on internals (__dict__ layouts) rather than on the public save/load format.