mixle.utils.parallel package

Distributed-estimation backends (multiprocessing, MPI, torchrun).

class ModelParallelEstimator(inner, num_workers=None)[source]

Bases: ParameterEstimator

Wrap an estimator so EM distributes the model axis – composing with any backend= for the data.

optimize(ModelParallelEstimator(est), backend="spark"|"mpi"|"mp"|"local") shards the data through that backend while each partition’s E-step distributes the model axis. 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()[source]
Return type:

ModelParallelAccumulatorFactory

estimate(nobs, suff_stat)[source]
Parameters:
Return type:

Any

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

Bases: EncodedDataHandle

Replicate the data, distribute the model’s shardable axis across threads (single machine).

Parameters:

  • data (Any)

  • estimator (Any | None)

  • model (Any | None)

  • encoder (Any | None)

  • num_workers (int | None)

  • _ (Any)

pysp_seq_log_density_sum(estimate)[source]

Return (num_observations, summed_log_density) for estimate.

Parameters:

estimate (Any)

Return type:

tuple[float, float]

pysp_seq_estimate(estimator, prev_estimate)[source]

Run one distributed/local sufficient-statistic fold and M-step.

Parameters:

  • estimator (Any)

  • prev_estimate (Any)

Return type:

Any

pysp_seq_initialize(estimator, rng, p)[source]

Initialize a model through the handle’s resident encoded data.

Parameters:
Return type:

Any

pysp_stream_accumulate(estimator, model)[source]

Return folded sufficient statistics for streaming/incremental EM.

Parameters:
Return type:

tuple[float, Any]

model_parallel_fold(acc, model, enc, weights, num_workers=None)[source]

Run a model-parallel E-step of model over encoded enc into accumulator acc (in place).

Parameters:
Return type:

None

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

Consult the C2 planner (decompose_model()) and return (estimator, decomposition).

When the planner picks model-parallelism for model on resources, the estimator is wrapped in ModelParallelEstimator sized to the planner’s cuts; otherwise the plain estimator is returned (replicate the model, shard the data – already optimal when N dominates). Either way, run the returned estimator through optimize(data, est, backend=<data backend>): the data axis is handled by backend and the model axis, if any, by the wrapper – composing into the data x model split. The decomposition’s rationale explains the choice. resources defaults to the local CPU slots (use Resources.from_spark(sc) / Resources.from_mpi() to size the split to a real cluster).

Parameters:

  • estimator (Any)

  • model (Any)

  • resources (Any)

  • n_data (int | None)

  • min_components_per_shard (int)

Return type:

tuple[Any, Any]

class Resources(devices)[source]

Bases: object

A collection of placement targets.

Parameters:

devices (tuple[DeviceSpec, ...])

devices: tuple[DeviceSpec, ...]
classmethod single_cpu(memory_bytes=None, throughput=1.0, precision=None)[source]

Return a one-device CPU resource description.

Parameters:

  • memory_bytes (int | None)

  • throughput (float)

  • precision (Any | None)

Return type:

Resources

classmethod local(num_cpus=None, memory_bytes=None, precision=None)[source]

Return local CPU resources split into logical worker slots.

Parameters:

  • num_cpus (int | None)

  • memory_bytes (int | None)

  • precision (Any | None)

Return type:

Resources

classmethod discover(include_torch=True, cpu_workers=None, precision=None)[source]

Best-effort local resource discovery with no required extras.

Parameters:

  • include_torch (bool)

  • cpu_workers (int | None)

  • precision (Any | None)

Return type:

Resources

classmethod from_specs(specs)[source]

Build resources from an iterable of device specifications.

Parameters:

specs (Iterable[DeviceSpec])

Return type:

Resources

classmethod from_mpi(comm=None, memory_bytes=None, throughput=1.0, precision=None)[source]

Return CPU resource slots for an MPI world without importing mpi4py.

comm may be an mpi4py-style communicator exposing Get_size. When it is omitted, common MPI launcher environment variables are used as a best-effort size hint, falling back to one slot.

Parameters:

  • comm (Any | None)

  • memory_bytes (int | None)

  • throughput (float)

  • precision (Any | None)

Return type:

Resources

classmethod from_dask(client, precision=None)[source]

Return resource slots from a dask.distributed-like client.

The method relies only on client.scheduler_info() and therefore does not introduce a dask dependency.

Parameters:

  • client (Any)

  • precision (Any | None)

Return type:

Resources

classmethod from_spark(spark_context, memory_bytes=None, precision=None)[source]

Return CPU resource slots from a SparkContext-like object.

Parameters:

  • spark_context (Any)

  • memory_bytes (int | None)

  • precision (Any | None)

Return type:

Resources

classmethod from_torchrun(memory_bytes=None, precision=None)[source]

Return torchrun rank/device slots from launcher environment hints.

Parameters:

  • memory_bytes (int | None)

  • precision (Any | None)

Return type:

Resources

fastest()[source]

Return the device with the largest advisory throughput weight.

Return type:

DeviceSpec

to_dict()[source]

Return a JSON-serializable resource description.

Return type:

dict[str, Any]

classmethod from_dict(payload)[source]

Build resources from to_dict output.

Parameters:

payload (dict[str, Any])

Return type:

Resources

to_json(**kwargs)[source]

Serialize resources, including calibrated throughput, to JSON.

Parameters:

kwargs (Any)

Return type:

str

classmethod from_json(text)[source]

Deserialize resources from JSON produced by to_json.

Parameters:

text (str)

Return type:

Resources

save(path, **kwargs)[source]

Persist resources to a JSON file for reuse by later plans.

Parameters:
Return type:

None

classmethod load(path)[source]

Load resources from a JSON file created by save.

Parameters:

path (Any)

Return type:

Resources

encoded_data(data, estimator=None, model=None, encoder=None, placement=None, resources=None, engine=None, precision=None, num_chunks=None, sub_chunks=1, backend='local', num_workers=None, client=None, comm=None, root=0, root_only=False, parallel_chunks=False, chunk_workers=None)[source]

Return an encoded-data handle, preserving existing compatible handles.

Backend dispatch goes through a registry (see register_encoded_data_backend()) rather than a hard-coded branch, so a new distributed framework (Lightning, Ray, JAX, …) plugs in by registering a factory – the same “register, don’t branch” pattern the compute engines use – without editing this function.

Parameters:

  • data (Any)

  • estimator (Any | None)

  • model (Any | None)

  • encoder (DataSequenceEncoder | None)

  • placement (Placement | None)

  • resources (Resources | None)

  • engine (Any | None)

  • precision (Any | None)

  • num_chunks (int | None)

  • sub_chunks (int)

  • backend (str)

  • num_workers (int | None)

  • client (Any | None)

  • comm (Any | None)

  • root (int)

  • root_only (bool)

  • parallel_chunks (bool)

  • chunk_workers (int | None)

Return type:

EncodedDataHandle

is_encoded_data_handle(obj)[source]

Return true when obj exposes the sequence-orchestrator contract.

Parameters:

obj (Any)

Return type:

bool

model_sharding_plan(model, resources, estimator=None, axis='components', min_components_per_shard=1)[source]

Return an advisory component-axis sharding plan.

This is the planning half of model parallelism. It does not create DTensor placements or move tensors; it answers which component range each device should own when a stacked/generated component kernel is available.

Parameters:

  • model (Any)

  • resources (Resources)

  • estimator (Any | None)

  • axis (str)

  • min_components_per_shard (int)

Return type:

tuple[ModelShard, …]

plan(data=None, model=None, estimator=None, encoder=None, resources=None, engine=None, precision=None, num_chunks=None, sub_chunks=1, sample_size=256, safety_factor=1.25)[source]

Return an advisory placement for local or distributed orchestration.

Parameters:

  • data (Sequence[Any] | None)

  • model (Any | None)

  • estimator (Any | None)

  • encoder (DataSequenceEncoder | None)

  • resources (Resources | None)

  • engine (Any | None)

  • precision (Any | None)

  • num_chunks (int | None)

  • sub_chunks (int)

  • sample_size (int)

  • safety_factor (float)

Return type:

Placement

Submodules