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, and the per-rank (count, accumulator.value()) payloads fold with comm.reduce – mpi4py’s lowercase, pickle-based object-mode reduce, which performs a genuine O(log W) binary reduction tree rather than a gather-to-root: each internal tree node combines two already-deserialized payloads via a fresh accumulator’s combine(), so no single rank ever folds more than O(log W) payloads. This is the same technique the Spark transport uses via RDD.treeReduce and the local executor uses via its in-process combine tree (see mixle.inference.spark_executor and mixle.inference.heterogeneous_executor). The root then 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.

This is mixle’s one MPI EM transport: an EncodedDataHandle that plugs directly into optimize, inheriting its convergence loop and logging. (A second, standalone comm.reduce-based transport lived at mixle.inference.mpi_executor pre-0.8.0; it was folded in here – its O(log W) tree-reduce technique is what _fold_and_share/_fold_value_and_share now use, in place of the old gather-to-root loop – and removed as a redundant second entry point.)

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 (MPI.Comm | None)

  • root (int)