mixle.utils.parallel.mpi module

mpi4py backend for distributed estimation (SPMD).

Every MPI rank runs the same script; rank 0 plays the Spark driver. Each rank holds a shard of the raw data, encodes it locally once, and keeps the encoded chunks resident (process persistence is free under MPI). Per EM iteration the model is broadcast from the root for bitwise consistency, every rank accumulates sufficient statistics over its shard, the per-rank (count, accumulator.value()) payloads are gathered to the root, the root folds them with combine(), applies key_merge/key_replace once globally, runs the M-step, and broadcasts the new model - so every rank returns an identical distribution and the surrounding optimize loop stays in lockstep across ranks with no further coordination.

Usage (run with mpiexec -n 4 python script.py):

from mixle.utils.parallel.mpi import MPIEncodedData, mpi_out
from mixle.inference.estimation import optimize

data = load_data()                       # every rank loads (or root_only=True)
enc = MPIEncodedData(data, estimator=est)
model = optimize(None, est, enc_data=enc, max_its=50, out=mpi_out())

mpi_out() returns sys.stdout on the root and a throwaway buffer on other ranks so iteration logging is printed once. Log-density sums are allreduce-d, so the convergence test in optimize agrees on all ranks.

mpi4py is an optional dependency (pip install mixle[mpi]); importing this module without it raises ImportError.

class MPIEncodedData(data, estimator=None, encoder=None, sub_chunks=1, comm=None, root=0, root_only=False)[source]

Bases: EncodedDataHandle

Encoded-data handle sharded across MPI ranks (SPMD).

Drop-in for the enc_data argument of optimize/best_of/ seq_estimate/seq_initialize/seq_log_density_sum; construct it and call those functions identically on every rank.

Parameters:

  • data (Sequence) – Raw observations. With root_only=False (default) every rank passes the SAME full dataset and keeps the round-robin shard data[rank::size]. With root_only=True only the root needs real data (other ranks may pass None); shards are scattered.

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

  • encoder (Optional[DataSequenceEncoder]) – Explicit encoder.

  • sub_chunks (int) – Encoded sub-chunks per rank (bounds peak memory of the vectorized update).

  • comm (Optional[MPI.Comm]) – Communicator (default COMM_WORLD).

  • root (int) – Driver rank for the combine/M-step.

  • root_only (bool)

pysp_seq_estimate(estimator, prev_estimate)[source]

One distributed EM step; every rank returns the identical model.

pysp_seq_initialize(estimator, rng, p)[source]

Distributed randomized initialization; identical model on all ranks.

Parameters:
pysp_seq_log_density_sum(estimate)[source]

Allreduced (count, log-density sum) - identical on every rank.

Return type:

tuple[float, float]

pysp_stream_accumulate(estimator, model)[source]

Globally folded batch sufficient statistics for streaming EM.

Return type:

tuple[float, Any]

mpi_out(comm=None, root=0)[source]

sys.stdout on the root rank, a discarded buffer elsewhere.

Parameters:

  • comm (mpi4py.MPI.Comm | None)

  • root (int)