mixle.utils.parallel.balance module

Automatic compute / memory / load-balancing planner for EM estimation of any mixle model.

The currency is FLOPs per iteration, with memory as a hard constraint – memory decides whether the model fits, compute decides how long the iteration takes, and the planner balances compute across the cluster subject to memory.

An EM iteration is a fixed amount of work W = N * C FLOPs (C = per-observation model cost from compute_cost()). A worker grid of D data-parallel replicas x M model-shards gives every worker W / (D*M) FLOPs, so the iteration time is max_worker_FLOPs / throughput + coupling. To balance the load we therefore:

  • prefer data parallelism (D) – it has no cross-worker coupling (a data point’s whole model lives on one worker) and balances trivially by equal row counts;

  • use model parallelism (M) only as forced – by memory (the model does not fit: M >= ceil(bytes/mem)) or by compute concurrency (too few data points to fill the cluster: with N points only N data-replicas exist, so the rest of the cluster can only be used by splitting the model);

  • balance the model split by FLOPs, not bytes – a memory-light but compute-heavy leaf (a GP, a big quadratic form) must not become the straggler everyone waits on.

This covers the whole spectrum the same way: a tiny model on lots of data -> M=1, D=P (data-parallel); a model too big for one worker -> M from memory, D fills the rest; a huge model on a single observation (N=1) -> D=1, M= as many model shards as the model exposes; and an unbalanced heterogeneous nest -> the FLOP cost model finds where the work is and the split equalizes it.

class BalancePlan(data_parallel, model_parallel, workers_used, workers_total, axis, model_cuts, model_flops, model_bytes, per_worker_flops, fits, rationale, extra=<factory>)[source]

Bases: object

A worker-grid assignment for one EM estimation, balanced by compute under a memory constraint.

Parameters:

  • data_parallel (int)

  • model_parallel (int)

  • workers_used (int)

  • workers_total (int)

  • axis (DecompAxis)

  • model_cuts (tuple[ModelCut, ...])

  • model_flops (float)

  • model_bytes (int)

  • per_worker_flops (float)

  • fits (bool)

  • rationale (str)

  • extra (dict[str, Any])

data_parallel: int
model_parallel: int
workers_used: int
workers_total: int
axis: DecompAxis
model_cuts: tuple[ModelCut, ...]
model_flops: float
model_bytes: int
per_worker_flops: float
fits: bool
rationale: str
extra: dict[str, Any]
property is_model_parallel: bool
property workers_idle: int
balance_plan(model, resources, *, n_data)[source]

Choose the (D data-parallel) x (M model-parallel) worker grid that balances compute under memory.

Searches M from the memory-required minimum up to the model’s splittable units, picking the grid that keeps the most workers busy (ties broken toward smaller M – less coupling). Works for any model: a model with no splittable axis simply gets M=1 (data-parallel / single worker).

Parameters:

  • model (Any)

  • resources (Resources)

  • n_data (int)

Return type:

BalancePlan

auto_balanced_estimator(estimator, model, resources=None, *, n_data)[source]

Realize balance_plan() – return (estimator, plan) ready to drive optimize.

When the plan is model-parallel the estimator is wrapped in ModelParallelEstimator sized to plan.model_parallel (the model axis distributes inside each worker / across the model shards); otherwise the plain estimator is returned (pure data-parallel). The data degree plan.data_parallel is realized by the data backend you pass to optimize ("local"|"mp"|"spark"|"mpi"), so optimize(data, returned_estimator, backend=...) runs the full D x M grid. resources defaults to the local CPU slots; pass Resources.from_spark(sc) / from_mpi() to plan for a real cluster.

Parameters:
Return type:

tuple[Any, BalancePlan]