mixle.utils.parallel.multiprocessing module

Multiprocessing backend for distributed estimation.

Mirrors the Spark estimation path with a pool of persistent local worker processes: the driver builds one encoder and ships each worker a shard of the raw data once; every worker encodes its own shard and keeps the encoded chunks resident for the lifetime of the handle. Per EM iteration only the (small) pickled model crosses the process boundary outward and only the per-worker (count, accumulator.value()) sufficient-statistic payloads come back; the driver folds them with combine(), applies key_merge/key_replace once globally (parameter tying must happen after the full combine), and runs the M-step.

Usage - the handle plugs into the ordinary estimation entry points:

from mixle.utils.parallel.multiprocessing import MPEncodedData
from mixle.inference.estimation import optimize

with MPEncodedData(data, estimator=est, num_workers=8) as enc:
    model = optimize(None, est, enc_data=enc, max_its=50)

optimize/best_of need no changes: mixle.inference.seq_estimate, seq_initialize and mixle.stats.seq_log_density_sum recognize the handle and delegate to it. Validation data can stay locally encoded (a plain seq_encode result) alongside an MPEncodedData training handle.

Notes

  • Worker processes are started with the spawn method for cross-platform consistency; each worker imports mixle.stats on first use (fast when the numba cache is warm).

  • The M-step receives the true observation count (as the Spark path does), not None (as the in-process path does).

  • Runtime kernel/engine objects are not picklable; ship plain distributions/estimators (the handle does this for you).

class MPEncodedData(data, estimator=None, encoder=None, num_workers=None, sub_chunks=1)[source]

Bases: EncodedDataHandle

Encoded-data handle sharded across persistent local worker processes.

Drop-in for the enc_data argument of optimize/best_of/ seq_estimate/seq_initialize/seq_log_density_sum. The raw data is split round-robin into num_workers shards; each worker encodes its shard once and keeps it resident across EM iterations.

Parameters:

  • data (Sequence) – Raw observations (anything the model’s encoder accepts). 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 worker (bounds peak memory of the vectorized update inside each worker).

pysp_seq_estimate(estimator, prev_estimate)[source]

One distributed EM step: returns the re-estimated distribution.

pysp_seq_initialize(estimator, rng, p)[source]

Distributed randomized initialization (mirrors seq_initialize).

Parameters:
pysp_seq_log_density_sum(estimate)[source]

Total observation count and summed log density across all workers.

Return type:

tuple[float, float]

pysp_stream_accumulate(estimator, model)[source]

Return globally folded batch sufficient statistics for streaming EM.

Return type:

tuple[float, Any]

close()[source]

Shut the worker pool down. Idempotent.

Return type:

None