mixle.utils.parallel.resilient_em module

Resilient multiprocessing EM backend: retry + rank blacklisting + mid-fit checkpointing (K4).

Status audit this module leans on (see the roadmap): fits are deterministic given seed (#115) and sufficient statistics are ADDITIVE (combine() folds any partition of the data, in any grouping, into the same total). Those two facts together make worker-failure recovery exact, not approximate:

1. Checkpointing is trivial and exact. An accumulator’s value() payload IS the sufficient statistic, not opaque optimizer state, so serializing it mid-fold and restoring it later via from_value() reconstructs the identical accumulator. See checkpointed_fold().

2. Only the failed shard needs to be redone. If a worker dies mid-E-step, its surviving peers’ already-computed (count, accumulator.value()) payloads are trusted as-is; only the dead worker’s shard is recomputed – on a surviving worker, from the SAME raw shard bytes the driver still holds (the “elastic re-partition”).

3. Recovery is bit-identical, not just close. seq_update (the E-step) is a pure, deterministic function of (encoded data, weights, model) – no RNG is involved – so recomputing a shard on a different physical worker produces byte-identical floats to the original owner computing it. The one place determinism could quietly break is fold ORDER: floating-point summation is not associative, so this module always folds per-shard payloads back together in canonical shard-id order (matching what a failure-free run would have done), never in “whichever worker replied first” order.

4. Retry + rank blacklisting. A worker that dies is retried by respawning a fresh process for the same rank and re-registering its shard (a transient hiccup does not cost that rank its place). A rank that fails repeatedly (failures >= max_retries) is blacklisted for the rest of the fit: it is never respawned again and its shard is migrated permanently onto a surviving worker.

This is the EM-side sibling of F2 (see the roadmap’s checkpoint/resume line for the model-parallel path); this module is the mp-backend line for ordinary (non-model-parallel) distributed EM.

class ResilientMPEncodedData(data, estimator=None, encoder=None, num_workers=None, sub_chunks=1, max_retries=2)

Bases: EncodedDataHandle

MPEncodedData with retry, rank blacklisting, and exact chaos-tolerant recovery.

Drop-in for the enc_data argument of optimize/best_of/seq_estimate/ seq_initialize/seq_log_density_sum, exactly like MPEncodedData. Data is split round-robin into num_workers SHARDS (a fixed id space, 0..num_workers-1, that outlives any one worker process); each shard is initially resident on the worker of the same id, but the driver also keeps the shard’s raw (pre-encode) bytes so a shard can be recomputed elsewhere, or migrated permanently, if its worker dies.

Parameters:

• data (Sequence) – Raw observations. Must be an in-memory sequence.

• estimator (Optional[ParameterEstimator]) – Used to build the encoder when encoder is not given.

• encoder (Optional[DataSequenceEncoder]) – Explicit encoder; overrides estimator.

• num_workers (Optional[int]) – Worker process count (default: CPU count, capped at the number of observations).

• sub_chunks (int) – Encoded sub-chunks per shard (bounds peak memory of the vectorized update inside each worker); also carried along to ad hoc shard recovery so a recomputed shard’s encode/accumulate split – and therefore its floating-point summation order – matches what the shard’s original owner would have done.

• max_retries (int) – A rank is blacklisted once its cumulative failure count reaches this threshold; below it, a dead rank is respawned and keeps its place.

Testing hook:

arm_kill() registers a one-shot callback invoked, for every worker, right after that worker acknowledges it has started an update command and while it is still blocked waiting for the driver’s “go” – strictly before any accumulation happens – the deterministic rendezvous a chaos test uses to kill a real OS process mid-E-step with no timing race.

arm_kill(hook)

Register a one-shot hook fired for each worker right after its “started” ack, while that worker is still blocked waiting for the driver’s “go” (see _worker_main).

hook(worker_id, proc) may kill proc (e.g. proc.kill(); proc.join()) to simulate a real worker death mid-E-step, with no timing race: the worker cannot have started accumulating yet. It is consumed (cleared) the moment the next pysp_seq_estimate/pysp_stream_accumulate call begins, so it fires for exactly one round.

Parameters:

hook (Callable[[int, Any], None])

Return type:

None

pysp_seq_estimate(estimator, prev_estimate)

One distributed EM step, tolerant of a worker dying mid-accumulation.

Parameters:

• estimator (Any)

• prev_estimate (Any)

Return type:

Any

pysp_seq_initialize(estimator, rng, p)

Distributed randomized initialization; seeds are anchored to shard id, not worker identity, so a shard reassigned to a different worker still uses its own fixed seed.

Parameters:

• estimator (Any)

• rng (RandomState)

• p (float)

Return type:

Any

pysp_seq_log_density_sum(estimate)

Total observation count and summed log density across all live workers.

Parameters:

estimate (Any)

Return type:

tuple[float, float]

pysp_stream_accumulate(estimator, model)

Return globally folded batch sufficient statistics for streaming EM, chaos-tolerant.

Parameters:

• estimator (Any)

• model (Any)

Return type:

tuple[float, Any]

close()

Shut the worker pool down. Idempotent.

Return type:

None

checkpointed_fold(estimator, payloads, checkpoint_after=None)

Fold pickled (count, accumulator.value()) payloads into one sufficient statistic.

This is the additive fold every backend in this repo performs (see MPEncodedData._fold_stats / MPIEncodedData._fold_and_share), pulled out standalone so a checkpoint can be taken mid-fold: pass checkpoint_after=k to, immediately after combining payload index k, serialize the running accumulator via value(), DISCARD the in-memory accumulator object entirely, and rebuild a fresh one from that serialized value via from_value() before continuing. Because value()/from_value() is an exact round-trip of the accumulator’s own state (not lossy optimizer state), the returned (nobs, value) is identical whether or not a checkpoint was taken partway through – that identity is the mid-fit-checkpointing acceptance criterion for K4.

Parameters:

• estimator (Any)

• payloads (Sequence[bytes])

• checkpoint_after (int | None)

Return type:

tuple[float, Any]