mixle.utils.parallel.em_observability module

Runtime observability for distributed EM backends (MP/MPI/Spark).

Companion to mixle.utils.parallel.balance: balance produces a static worker grid from a FLOPs-under-memory model before the fit starts; this module captures what actually happened per rank per round once the fit is running, so stragglers and data-volume skew that the static plan could not see (heterogeneous hardware, cold caches, a lopsided shard) are caught within a single round instead of silently eating wall-clock for the whole fit.

The unlock is the same one that makes shard recomputation a corruption audit and accumulators checkpoints: EM’s sufficient statistics are additive, so per-rank records are additive too – a round’s RankRecord list is a complete, replayable receipt of what each worker did, and folding/analyzing them needs no coordination beyond what already flows back to the driver.

Pieces:

  • RankRecord – one worker-rank’s timing/bytes/accumulator-size receipt for one round.

  • record_rank_round() – append a RankRecord as a mixle.telemetry event.

  • detect_stragglers() – robust (median/MAD) outlier test over one round’s rank times.

  • imbalance_receipt() – quantitative data-volume skew across ranks for one round.

  • plan_rebalance_weights() – turn observed timings into a feedback signal for mixle.utils.parallel.balance.balance_plan()’s next static plan.

  • fit_report() – one human/machine-readable summary of a whole fit’s collected records.

class FitReport(n_rounds, n_ranks, total_seconds, total_bytes, total_obs, rounds, stragglers_by_round, imbalance_by_round, worst_straggler_ratio, worst_imbalance_ratio)[source]

Bases: object

A whole-fit summary rolled up from per-rank, per-round records – the EM-side run.report().

Mirrors the shape F4’s training-health run.report() is expected to have (per-round health, an overall verdict, a printable render) without depending on F4, which had not landed when this was written; aligning field names/shape once F4 exists is a natural, tracked follow-up.

Parameters:

  • n_rounds (int)

  • n_ranks (int)

  • total_seconds (float)

  • total_bytes (int)

  • total_obs (int)

  • rounds (tuple[int, ...])

  • stragglers_by_round (dict[int, StragglerReport])

  • imbalance_by_round (dict[int, ImbalanceReceipt])

  • worst_straggler_ratio (float)

  • worst_imbalance_ratio (float)

render()[source]

A short human-readable summary, e.g. for a CLI/log line.

Return type:

str

as_dict()[source]

Machine-readable form (JSON-friendly).

Return type:

dict[str, Any]

class ImbalanceReceipt(round, bytes_per_rank, n_obs_per_rank, mean_bytes, max_bytes_ratio, skew_by_rank)[source]

Bases: object

Quantitative data-volume skew across ranks for one round (bytes and observation counts).

Parameters:
class RankRecord(rank, round, e_step_seconds, m_step_seconds=0.0, bytes_processed=0, accumulator_bytes=0, n_obs=0, extra=<factory>)[source]

Bases: object

One worker-rank’s structured receipt for one EM round.

e_step_seconds/m_step_seconds are wall-clock durations measured on the rank itself (E-step = the accumulate/score pass over the shard, M-step = folding+``estimate`` when a rank does its own partial M-step, 0.0 when the M-step is driver-side only). bytes_processed is the raw/encoded data volume the rank touched this round; accumulator_bytes is the size of the sufficient-statistics payload it shipped back – both additive across ranks, so a round’s totals are a plain sum.

Parameters:
as_features()[source]

Flat dict form used as a Event features payload.

Return type:

dict[str, Any]

classmethod from_features(features)[source]

Inverse of as_features() (round-trips through a Event).

Parameters:

features (dict[str, Any])

Return type:

RankRecord

class StragglerReport(round, rank_seconds, median_seconds, mad_seconds, ratios, z_scores, threshold_ratio, z_threshold, slow_ranks)[source]

Bases: object

Straggler/imbalance verdict for one EM round, from a robust median/MAD outlier test.

Parameters:
detect_stragglers(records, *, round=None, threshold_ratio=1.5, z_threshold=3.0)[source]

Flag ranks that are meaningfully slower than the rest for one round.

A robust (median / median-absolute-deviation) outlier test rather than a mean/stddev test, since a single straggler should not be allowed to inflate the scale it is measured against. A rank is flagged when it is BOTH threshold_ratio``x slower than the median rank time AND (when the round has enough spread to have a nonzero MAD) ``z_threshold robust-z above the median – the ratio catches the practically-significant case, the z-score guards against flagging normal jitter when every rank is close together.

Parameters:
Return type:

StragglerReport

fit_report(records, *, threshold_ratio=1.5, z_threshold=3.0)[source]

Build one summary report from a completed (or in-progress) distributed EM fit’s records.

Groups records by round, runs detect_stragglers() and imbalance_receipt() per round, and rolls the results into one FitReport. Safe to call mid-fit with only the rounds collected so far – an empty records sequence yields a zeroed report rather than raising.

Parameters:
Return type:

FitReport

imbalance_receipt(records, *, round=None)[source]

Measure how unevenly data volume was actually spread across ranks for one round.

skew_by_rank[rank] is that rank’s bytes_processed divided by the mean across ranks – a planted “rank 0 gets 10x the data” skew shows up directly as skew_by_rank[0] ~= 10 * (P / (P - 1 + 10)) scaling (exactly proportional to the planted ratio for the standard case of one heavy rank among otherwise-equal peers; see the test for the exact algebra), so this is a real correctness check against the planted skew, not just a boolean “imbalanced” flag.

Parameters:
Return type:

ImbalanceReceipt

plan_rebalance_weights(records, *, round=None)[source]

Turn one round’s observed rank timings into a feedback signal for the next static plan.

balance_plan (mixle.utils.parallel.balance) picks a static (D, M) worker grid and, when M > 1, splits the model’s units across shards with _balance_units() weighted by a predicted per-unit FLOP cost (best.unit_works). This function produces the observed analogue: a per-rank weight inversely proportional to how long that rank actually took, normalized to sum to the rank count – a rank that ran 2x slower than average gets ~0.5x the weight, meaning “give this rank half as much data/model next round”. The weights are in exactly the shape _balance_units consumes (a sequence of relative per-unit works, here one “unit” per rank), so a re-planning loop can plug them in directly; wiring that end-to-end (re-running balance_plan with per-worker throughput overrides derived from these weights) is future work – this function only produces the signal, documented here as the integration point.

Parameters:
Return type:

dict[int, float]

record_rank_round(telemetry, *, rank, round, e_step_seconds, m_step_seconds=0.0, bytes_processed=0, accumulator_bytes=0, n_obs=0, run_id=None, extra=None)[source]

Append one rank’s per-round receipt to telemetry as an "em_round" event.

Thin wrapper around Telemetry.record() – reuses the existing typed-event/JSONL recorder rather than inventing a parallel logging object. run_id (if given) tags the event so records from concurrent fits sharing one recorder can be told apart.

Parameters:

  • telemetry (Telemetry)

  • rank (int)

  • round (int)

  • e_step_seconds (float)

  • m_step_seconds (float)

  • bytes_processed (int)

  • accumulator_bytes (int)

  • n_obs (int)

  • run_id (str | None)

  • extra (dict[str, Any] | None)

Return type:

Event

records_from_events(events)[source]

Recover RankRecord objects from "em_round" Event rows.

Parameters:

events (Sequence[Event])

Return type:

list[RankRecord]