mixle.utils.parallel.model_parallel module

In-process component-parallel estimation over a fully resident model.

This runtime parallelizes independent model work across threads. It does not partition model storage or execute a planner’s device placements: every caller must be able to hold the full model and statistic state in one process.

• optimize(..., backend="component_parallel") uses a replicated in-process data handle and threads independent factors/components.

• optimize(ComponentParallelEstimator(est), backend=...) composes that threading with a data backend. Each backend process still owns the full model.

The fold (model_parallel_fold()) is recursive: it walks the whole model tree and threads the axes that a per-node compute-cost model says save the most wall-time (_parallel_ids, consistent with the structural planner – a narrow batch of heavy MVGaussians beats a wider batch of low-cost leaves), recursing serially below any threaded node so no two pools ever nest. Each recursive case reproduces the corresponding accumulator’s seq_update exactly:

• FACTOR (Composite/Record) – the per-factor accumulators are independent, so the per-factor seq_update calls are distributed (bit-identical).

• COMPONENT (mixtures) – the responsibility logsumexp couples the components, so the low-cost normalization runs centrally on the gathered score matrix while the expensive per-component scoring and accumulation are distributed – a bit-identical mirror of MixtureAccumulator.seq_update.

• atomic / unknown – the replicated base case acc.seq_update(enc, weights, model).

So the whole fold is bit-identical to the single-node path (the data-axis reduce across partitions is the usual additive combine, exact up to float reassociation like every data-parallel backend). Correct for every family; never worse than backend="local".

class ModelParallelEncodedData(data, *, estimator=None, model=None, encoder=None, num_workers=None, **_)

Bases: EncodedDataHandle

Compatibility name for a fully resident, component-threaded data handle.

Parameters:

• data (Any)

• estimator (Any | None)

• model (Any | None)

• encoder (Any | None)

• num_workers (int | None)

• _ (Any)

pysp_seq_log_density_sum(estimate)

Return the encoded-data size and total log likelihood under estimate.

Parameters:

estimate (Any)

Return type:

tuple[float, float]

pysp_seq_estimate(estimator, prev_estimate)

Run one model-parallel E/M update from prev_estimate.

Parameters:

• estimator (Any)

• prev_estimate (Any)

Return type:

Any

pysp_seq_initialize(estimator, rng, p)

Initialize a model by randomly selecting observations with probability p.

Parameters:

• estimator (Any)

• rng (RandomState)

• p (float)

Return type:

Any

pysp_stream_accumulate(estimator, model)

Accumulate model-parallel sufficient statistics for streaming backends.

Parameters:

• estimator (Any)

• model (Any)

Return type:

tuple[float, Any]

class ModelParallelEstimator(inner, num_workers=None)

Bases: ParameterEstimator

Compatibility wrapper for fully resident component-parallel EM.

optimize(ModelParallelEstimator(est), backend="spark"|"mpi"|"mp"|"local") shards the data through that backend while each partition’s E-step threads independent model work. The full model remains resident in every process. The M-step (estimate) and the accumulator’s value/combine contract are the wrapped estimator’s, unchanged.

Parameters:

• inner (ParameterEstimator)

• num_workers (int | None)

accumulator_factory()

Return a factory that wraps the inner estimator’s accumulator factory.

Return type:

ModelParallelAccumulatorFactory

estimate(nobs, suff_stat)

Delegate the M-step to the wrapped estimator.

Parameters:

• nobs (float | None)

• suff_stat (Any)

Return type:

Any

class ModelParallelAccumulator(inner, num_workers=None)

Bases: SequenceEncodableStatisticAccumulator

Wrap an accumulator so its E-step threads independent component work.

All sufficient-statistic methods delegate to the wrapped (inner) accumulator unchanged, so the value/combine/from_value/key-merge contract – and thus every data backend’s reduce – is preserved; only seq_update is replaced with the distributed fold. Holding inner in vars() keeps the accumulator key-validator’s recursion transparent.

Parameters:

• inner (SequenceEncodableStatisticAccumulator)

• num_workers (int | None)

update(x, weight, estimate)

Delegate scalar accumulation to the wrapped accumulator.

Parameters:

• x (Any)

• weight (float)

• estimate (Any)

Return type:

None

initialize(x, weight, rng)

Delegate scalar initialization to the wrapped accumulator.

Parameters:

• x (Any)

• weight (float)

• rng (Any)

Return type:

None

seq_update(x, weights, estimate)

Run the wrapped accumulator’s sequence update through the model-parallel fold.

Parameters:

• x (Any)

• weights (ndarray)

• estimate (Any)

Return type:

None

seq_initialize(x, weights, rng)

Delegate encoded initialization to the wrapped accumulator.

Parameters:

• x (Any)

• weights (ndarray)

• rng (Any)

Return type:

None

combine(suff_stat)

Merge sufficient statistics into the wrapped accumulator.

Parameters:

suff_stat (Any)

Return type:

ModelParallelAccumulator

value()

Return the wrapped accumulator’s sufficient-statistic value.

Return type:

Any

from_value(x)

Replace the wrapped accumulator from a sufficient-statistic value.

Parameters:

x (Any)

Return type:

ModelParallelAccumulator

key_merge(stats_dict)

Delegate keyed statistic merging to the wrapped accumulator.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

key_replace(stats_dict)

Delegate keyed statistic replacement to the wrapped accumulator.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

acc_to_encoder()

Return the wrapped accumulator’s compatible data encoder.

Return type:

Any

class ModelParallelAccumulatorFactory(inner_factory, num_workers=None)

Bases: StatisticAccumulatorFactory

Factory that wraps another accumulator factory with model-parallel updates.

Parameters:

• inner_factory (Any)

• num_workers (int | None)

make()

Create a fresh model-parallel accumulator wrapper.

Return type:

ModelParallelAccumulator

ComponentParallelEncodedData

alias of ModelParallelEncodedData

ComponentParallelEstimator

alias of ModelParallelEstimator

ComponentParallelAccumulator

alias of ModelParallelAccumulator

ComponentParallelAccumulatorFactory

alias of ModelParallelAccumulatorFactory

model_parallel_fold(acc, model, enc, weights, num_workers=None, *, parallel_ids=None)

Run a component-parallel E-step over a fully resident model.

Parameters:

• acc (Any)

• model (Any)

• enc (Any)

• weights (ndarray)

• num_workers (int | None)

• parallel_ids (frozenset[int] | None)

Return type:

None

component_parallel_fold(acc, model, enc, weights, num_workers=None, *, parallel_ids=None)

Run a component-parallel E-step over a fully resident model.

Parameters:

• acc (Any)

• model (Any)

• enc (Any)

• weights (ndarray)

• num_workers (int | None)

• parallel_ids (frozenset[int] | None)

Return type:

None

auto_parallel_estimator(estimator, model, resources=None, *, n_data=None, min_components_per_shard=1)

Return a plain estimator when no distributed model placement is required.

The structural planner is advisory. If it selects cuts that require model storage to be partitioned across devices, this helper raises UnrealizedModelPlacementError: the local component-thread runtime cannot truthfully realize those cuts. Callers may explicitly choose ComponentParallelEstimator when the full model fits in each process.

Parameters:

• estimator (Any)

• model (Any)

• resources (Any)

• n_data (int | None)

• min_components_per_shard (int)

Return type:

tuple[Any, Any]

exception UnrealizedModelPlacementError(plan)

Bases: RuntimeError

A planner selected model shards for which this runtime has no placement executor.

Parameters:

plan (Any)

Return type:

None