mixle.utils.parallel.planner module

Resource and placement planning helpers.

The planner is advisory: it estimates memory pressure from an encoder/model pair and produces a printable, editable placement. Orchestrators still own actual data movement and sufficient-statistic folding.

class CalibrationCatalog(records=())[source]

Bases: object

Append-only catalog of resource calibration measurements.

Parameters:

records (tuple[CalibrationRecord, ...])

records: tuple[CalibrationRecord, ...] = ()
add(record)[source]

Append record and return it.

Parameters:

record (CalibrationRecord)

Return type:

CalibrationRecord

latest(model_type=None, workload=None, precision=None)[source]

Return the newest record matching the supplied filters.

Parameters:

  • model_type (str | None)

  • workload (str | None)

  • precision (str | None)

Return type:

CalibrationRecord | None

resources_for(model_type=None, workload=None, precision=None)[source]

Return calibrated resources from the newest matching record.

Parameters:

  • model_type (str | None)

  • workload (str | None)

  • precision (str | None)

Return type:

Resources | None

to_dict()[source]

Return a JSON-serializable calibration catalog.

Return type:

dict[str, Any]

classmethod from_dict(payload)[source]

Build a catalog from to_dict output.

Parameters:

payload (dict[str, Any])

Return type:

CalibrationCatalog

to_json(**kwargs)[source]

Serialize the calibration catalog to JSON.

Parameters:

kwargs (Any)

Return type:

str

classmethod from_json(text)[source]

Deserialize a calibration catalog from JSON.

Parameters:

text (str)

Return type:

CalibrationCatalog

save(path, **kwargs)[source]

Persist the catalog to a JSON file.

Parameters:
Return type:

None

classmethod load(path)[source]

Load a calibration catalog from disk.

Parameters:

path (Any)

Return type:

CalibrationCatalog

class CalibrationRecord(model_type, workload, resources, sample_size, repeats, precision=None, estimator_type=None, row_count=None, model_bytes=None, statistic_bytes=None, timestamp=0.0)[source]

Bases: object

One persisted calibration measurement for a model/workload/resource set.

Parameters:

  • model_type (str)

  • workload (str)

  • resources (Resources)

  • sample_size (int)

  • repeats (int)

  • precision (str | None)

  • estimator_type (str | None)

  • row_count (int | None)

  • model_bytes (int | None)

  • statistic_bytes (int | None)

  • timestamp (float)

model_type: str
workload: str
resources: Resources
sample_size: int
repeats: int
precision: str | None = None
estimator_type: str | None = None
row_count: int | None = None
model_bytes: int | None = None
statistic_bytes: int | None = None
timestamp: float = 0.0
to_dict()[source]

Return a JSON-serializable calibration record.

Return type:

dict[str, Any]

classmethod from_dict(payload)[source]

Build a calibration record from to_dict output.

Parameters:

payload (dict[str, Any])

Return type:

CalibrationRecord

class DeviceSpec(name, kind='cpu', memory_bytes=None, engine='numpy', throughput=1.0, precision=None)[source]

Bases: object

Description of one compute placement target.

Parameters:

  • name (str)

  • kind (str)

  • memory_bytes (int | None)

  • engine (str)

  • throughput (float)

  • precision (str | None)

name: str
kind: str = 'cpu'
memory_bytes: int | None = None
engine: str = 'numpy'
throughput: float = 1.0
precision: str | None = None
property is_gpu: bool

Return true for CUDA/MPS/GPU-like devices.

to_dict()[source]

Return a JSON-serializable device description.

Return type:

dict[str, Any]

classmethod from_dict(payload)[source]

Build a device description from to_dict output.

Parameters:

payload (dict[str, Any])

Return type:

DeviceSpec

class DaskEncodedData(data, estimator=None, model=None, encoder=None, client=None, num_partitions=None, sub_chunks=1, materialize=True)[source]

Bases: EncodedDataHandle

dask.distributed encoded-data handle implementing the orchestrator protocol.

Parameters:

  • data (Any)

  • estimator (Any | None)

  • model (Any | None)

  • encoder (DataSequenceEncoder | None)

  • client (Any | None)

  • num_partitions (int | None)

  • sub_chunks (int)

  • materialize (bool)

pysp_seq_log_density_sum(estimate)[source]

Return dask-folded (count, summed_log_density) for estimate.

Parameters:

estimate (Any)

Return type:

tuple[float, float]

pysp_seq_estimate(estimator, prev_estimate)[source]

Run one dask distributed E-step fold and estimator M-step.

Parameters:

  • estimator (Any)

  • prev_estimate (Any)

Return type:

Any

pysp_seq_initialize(estimator, rng, p)[source]

Initialize a model over persisted dask encoded partitions.

Parameters:
Return type:

Any

pysp_stream_accumulate(estimator, model)[source]

Return dask-folded sufficient statistics for streaming EM.

Parameters:
Return type:

tuple[float, Any]

property num_chunks: int

Return the number of persisted dask encoded partitions.

close()[source]

Cancel persisted partitions and close an owned dask client.

Return type:

None

class EncodedDataHandle[source]

Bases: object

Duck-typed orchestrator contract consumed by mixle.stats.

Local, multiprocessing, MPI, Spark, dask, or future worker handles can implement these methods without sharing inheritance. The base class exists to document the contract and to give local code a common type to return.

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]

close()[source]

Release worker resources owned by this handle, if any.

Return type:

None

class EncodedFold(*args, **kwargs)[source]

Bases: Protocol

Structural contract for the parallel/distributed fold consumed by mixle.stats.

This formalizes the duck-typed orchestrator contract that the concrete EncodedDataHandle base class documents. Local, multiprocessing, MPI, Spark, dask, Ray, Lightning, or future worker handles satisfy this Protocol structurally without sharing inheritance. Membership is decided by isinstance() against the four pysp_seq_*/pysp_stream_* methods.

pysp_seq_log_density_sum(estimate)[source]
Parameters:

estimate (Any)

Return type:

tuple[float, float]

pysp_seq_estimate(estimator, prev_estimate)[source]
Parameters:

  • estimator (Any)

  • prev_estimate (Any)

Return type:

Any

pysp_seq_initialize(estimator, rng, p)[source]
Parameters:
Return type:

Any

pysp_stream_accumulate(estimator, model)[source]
Parameters:
Return type:

tuple[float, Any]

class LocalEncodedData(data, estimator=None, model=None, encoder=None, placement=None, resources=None, engine=None, precision=None, num_chunks=None, sub_chunks=1, parallel_chunks=False, chunk_workers=None)[source]

Bases: EncodedDataHandle

In-process encoded-data handle implementing the orchestrator protocol.

The handle encodes raw data once according to an advisory Placement and then exposes the pysp_seq_* methods recognized by mixle.stats. It is intentionally small: distributions/estimators own math, kernels own backend scoring/accumulation, and this object owns only local movement and global sufficient-statistic folding.

Parameters:

  • data (Sequence[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)

  • parallel_chunks (bool)

  • chunk_workers (int | None)

shards: tuple[_LocalShard, ...]
pysp_seq_log_density_sum(estimate)[source]

Return total count and summed log density over resident chunks.

Parameters:

estimate (Any)

Return type:

tuple[float, float]

pysp_seq_estimate(estimator, prev_estimate)[source]

Run one local EM E-step/M-step through the unified handle contract.

Parameters:

  • estimator (Any)

  • prev_estimate (Any)

Return type:

Any

pysp_seq_initialize(estimator, rng, p)[source]

Randomized initialization over resident encoded chunks.

Parameters:
Return type:

Any

pysp_stream_accumulate(estimator, model)[source]

Return globally tied batch sufficient statistics for streaming EM.

Parameters:
Return type:

tuple[float, Any]

property num_chunks: int

Return the number of local encoded chunks across all shards.

close()[source]

No-op for parity with process-backed handles.

Return type:

None

class ModelShard(device, component_start, component_stop, parameter_bytes=0, statistic_bytes=0)[source]

Bases: object

Advisory component-axis shard for a large model.

Parameters:

  • device (DeviceSpec)

  • component_start (int)

  • component_stop (int)

  • parameter_bytes (int)

  • statistic_bytes (int)

device: DeviceSpec
component_start: int
component_stop: int
parameter_bytes: int = 0
statistic_bytes: int = 0
property num_components: int

Return the number of mixture components in this model shard.

property total_bytes: int

Return estimated parameter plus statistic bytes for this shard.

class Placement(shards, total_rows, encoded_row_bytes, transient_row_bytes, model_bytes, statistic_bytes, dtype_bytes)[source]

Bases: object

Printable placement returned by plan.

Parameters:

  • shards (tuple[PlacementShard, ...])

  • total_rows (int)

  • encoded_row_bytes (float)

  • transient_row_bytes (float)

  • model_bytes (int)

  • statistic_bytes (int)

  • dtype_bytes (int)

shards: tuple[PlacementShard, ...]
total_rows: int
encoded_row_bytes: float
transient_row_bytes: float
model_bytes: int
statistic_bytes: int
dtype_bytes: int
for_device(name)[source]

Return all placement shards assigned to a named device.

Parameters:

name (str)

Return type:

tuple[PlacementShard, …]

to_dict()[source]

Return a JSON-friendly placement summary.

Return type:

dict[str, Any]

class PlacementShard(device, start, stop, sub_chunks=1, encoded_bytes=0, transient_bytes=0)[source]

Bases: object

A contiguous row range assigned to one device.

Parameters:

  • device (DeviceSpec)

  • start (int)

  • stop (int)

  • sub_chunks (int)

  • encoded_bytes (int)

  • transient_bytes (int)

device: DeviceSpec
start: int
stop: int
sub_chunks: int = 1
encoded_bytes: int = 0
transient_bytes: int = 0
property size: int

Return the number of rows assigned to this shard.

property total_bytes: int

Return estimated encoded plus transient bytes for this shard.

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

class SparkEncodedData(data, estimator=None, model=None, encoder=None, materialize=True)[source]

Bases: EncodedDataHandle

Spark RDD encoded-data handle implementing the orchestrator protocol.

Parameters:

  • data (Any)

  • estimator (Any | None)

  • model (Any | None)

  • encoder (DataSequenceEncoder | None)

  • materialize (bool)

pysp_seq_log_density_sum(estimate)[source]

Return Spark-folded (count, summed_log_density) for estimate.

Parameters:

estimate (Any)

Return type:

tuple[float, float]

pysp_seq_estimate(estimator, prev_estimate)[source]

Run one Spark distributed E-step fold and estimator M-step.

Parameters:

  • estimator (Any)

  • prev_estimate (Any)

Return type:

Any

pysp_seq_initialize(estimator, rng, p)[source]

Initialize a model over the resident Spark encoded RDD.

Parameters:
Return type:

Any

pysp_stream_accumulate(estimator, model)[source]

Return Spark-folded sufficient statistics for streaming EM.

Parameters:
Return type:

tuple[float, Any]

property num_chunks: int

Return the number of Spark partitions in the encoded RDD.

close()[source]

Unpersist the encoded RDD from Spark storage.

Return type:

None

available_encoded_data_backends()[source]

Return the sorted names of all registered encoded-data backends.

Return type:

list[str]

calibrate_resources(data, model, resources=None, estimator=None, encoder=None, sample_size=512, repeats=3, workload='score', precision=None, catalog=None, catalog_path=None)[source]

Time model scoring on each resource and return updated throughputs.

Calibration is deliberately advisory and local. It runs a small scoring or E-step pass through model.kernel(engine=...) where possible and leaves a device’s previous throughput unchanged if that device cannot run the requested workload. workload may be score, estep / accumulate, or em. The em workload includes the estimator’s M-step on the sampled sufficient statistics. Pass catalog and/or catalog_path to append a persisted model/workload calibration record.

Parameters:

  • data (Sequence[Any])

  • model (Any)

  • resources (Resources | None)

  • estimator (Any | None)

  • encoder (DataSequenceEncoder | None)

  • sample_size (int)

  • repeats (int)

  • workload (str)

  • precision (Any | None)

  • catalog (CalibrationCatalog | None)

  • catalog_path (Any | None)

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

estimate_estimator_stat_nbytes(estimator)[source]

Approximate bytes in a zero accumulator value for an estimator.

Parameters:

estimator (Any)

Return type:

int

estimate_model_nbytes(model)[source]

Approximate bytes held by a model’s public parameter payload.

Parameters:

model (Any)

Return type:

int

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

register_encoded_data_backend(name, factory, aliases=())[source]

Register an encoded-data backend factory under name (and any aliases).

factory is called as factory(data, **params) with the keyword arguments of encoded_data(); it should accept **_ for the parameters it ignores and return an EncodedDataHandle. This is the extension point for new parallel/distributed frameworks – registering is all that is needed, no core edits.

Parameters:
Return type:

None