mixle.utils.parallel.sdc_audit module

Silent-data-corruption (SDC) audit for the resilient mp EM backend (K5).

Status audit this module leans on (see K4’s docstring in resilient_em.py, which this module builds directly on top of): fits are deterministic given seed (#115) and sufficient statistics are ADDITIVE. K4 already turns those two facts into exact, cheap failure recovery (a crashed worker’s shard is just redone). K5’s contribution is different: it turns the SAME two facts into a cheap corruption detector.

The idea. Silent data corruption (SDC) is a bit-flip in memory or compute that produces a WRONG result with no exception and no crash – the fit just silently converges to a corrupted answer. For gradient-based training this is expensive to catch (you’d have to redo an entire gradient step, or use redundant hardware). mixle’s additive-stat EM makes it cheap: recompute the SAME shard’s accumulator TWICE, from the SAME raw shard bytes / same estimator / same model / same sub_chunks split, once on the shard’s normal (“primary”) rank and once on a DIFFERENT (“audit”) rank. Determinism-given-seed + additive stats + fixed chunking means the two (count, value()) payloads MUST be bitwise identical if nothing is corrupted – so any byte-level difference is proof of a real hardware/software fault, not numerical noise (there is no tolerance band to tune, unlike a gradient re-check).

Why zero false positives holds “by construction”, not just empirically:

1. Both the primary and audit recompute start from self._shard_raw[shard_id] – the IDENTICAL pickled raw observations the driver has held since construction. Neither rank is given a different view of the data.

2. Both recomputes run the SAME code path: _worker_main’s "update_shard" handler, which pickles the shard with the SAME sub_chunks value (see _encode_shard), so floating-point summation order – which is NOT associative – is identical on both ranks. This is exactly the “fixed chunking” half of the guarantee: without pinning sub_chunks, two honest ranks could legitimately disagree in the last bit.

3. seq_update (the E-step) and seq_initialize are pure, deterministic functions of (encoded data, weights, model[, seed]) – no per-rank RNG state, no nondeterministic parallel reduction inside a single call. IEEE-754 arithmetic is deterministic given a fixed sequence of operations, so replaying the identical operation sequence on different physical hardware still produces the identical bit pattern.

4. Therefore an uncorrupted primary and an uncorrupted audit recompute of the same shard are the same pure computation evaluated twice, and must agree bit-for-bit. A mismatch can only arise if at least one of the two computations was NOT the pure computation – i.e. something (memory fault, cosmic ray, buggy kernel) perturbed it. This is verified empirically too, at scale, in mixle/tests/sdc_audit_test.py.

NaN/Inf watchdog. mixle.models._neural_serial.check_finite already guards individual density evaluations. K5 extends that same “fail loud, immediately, with the offending location named” philosophy to the accumulator-fold boundary: finite_guarded_fold() is a drop-in replacement for checkpointed_fold() that checks finiteness of the running accumulator’s value() immediately after EVERY combine() call, not just once at the end – so a NaN/Inf introduced while folding payload i is caught at payload i’s combine() boundary, before it silently propagates into payload i+1..n. Scope note: this wraps combine() calls made by THIS module’s own fold loop only (mirroring checkpointed_fold’s loop) – it deliberately does NOT touch SequenceEncodableStatisticAccumulator.combine()’s contract itself, which every other caller in the codebase still uses unguarded, exactly as before.

class AuditedMPEncodedData(data, estimator=None, encoder=None, num_workers=None, sub_chunks=1, max_retries=2, audit_rate=0.1, rng=None)

Bases: ResilientMPEncodedData

ResilientMPEncodedData (K4) plus a continuous SDC audit (K5).

Every round (pysp_seq_estimate / pysp_stream_accumulate), a random audit_rate fraction of shards are ALSO recomputed on a second, different rank via the same ad hoc "update_shard" wire command K4 already uses for shard recovery (see ResilientMPEncodedData._recover_shard) – no new worker-side machinery. The primary and audit (count, value()) payloads are compared BYTE-FOR-BYTE. A mismatch:

1. is recorded as a structured SDCAuditReceipt (self.audit_receipts, self.last_round_audit_mismatches);

2. quarantines BOTH ranks that produced the disagreeing payloads via K4’s existing _blacklist / _retire_worker / _migrate_shard_permanently machinery – this is a deliberately conservative policy: a single 2-way mismatch cannot, by itself, prove which of the two ranks is the corrupted one (that needs a third witness / majority vote), so both are treated as suspect and the shard is migrated to a clean survivor rather than risk silently trusting either.

The main EM round itself (retry, blacklisting on repeated failure, checkpointed fold) is entirely K4’s, untouched, reused via inheritance – K5 only adds the audit phase (run BEFORE the main round each call, so a quarantine decided by the audit is already reflected in that same round’s live-worker set) and swaps K4’s plain checkpointed_fold for finite_guarded_fold() via the _fold_fn hook K4 exposes for exactly this purpose.

Parameters:

• audit_rate (float) – fraction of shards (0..1) redundantly recomputed each round.

• rng (np.random.RandomState | None) – drives which shards are audited each round and which live rank is picked as the second (“audit”) rank. Not used for anything that needs to be reproducible bit-for-bit across ranks – only for which shards get the (always bit-exact) double-check.

• data (Sequence[Any])

• estimator (Any | None)

• encoder (Any | None)

• num_workers (int | None)

• sub_chunks (int)

• max_retries (int)

Testing hook:

arm_corruption() registers a one-shot-per-round hook that can mutate a payload right after it comes off the wire from an ad hoc "update_shard" recompute, letting a test inject a deterministic, reproducible corruption (see inject_bit_flip()) into a chosen (rank, shard, role) combination without touching worker internals – the same “observe the wire, mutate deterministically” pattern K4’s arm_kill uses.

arm_corruption(hook)

Register hook(worker_id, shard_id, role, payload_bytes) -> payload_bytes (role is "primary" or "audit"), applied to every ad hoc "update_shard" payload this instance receives until cleared. Unlike arm_kill this is NOT one-shot by default (an SDC fault is typically persistent, e.g. a stuck bit in one DIMM) – clear it explicitly with arm_corruption(None) to simulate a transient fault.

Parameters:

hook (Any)

Return type:

None

pysp_seq_estimate(estimator, prev_estimate)

One distributed EM step, tolerant of a worker dying mid-accumulation.

Parameters:

• estimator (Any)

• prev_estimate (Any)

Return type:

Any

pysp_stream_accumulate(estimator, model)

Return globally folded batch sufficient statistics for streaming EM, chaos-tolerant.

Parameters:

• estimator (Any)

• model (Any)

Return type:

tuple[float, Any]

class SDCAuditReceipt(round, shard_id, primary_worker, audit_worker, primary_nbytes, audit_nbytes, first_diff_byte_offset, primary_sha256, audit_sha256, primary_value_repr='', audit_value_repr='')

Bases: object

A structured record of one detected primary-vs-audit mismatch: which shard, which two ranks, and a summary of what actually differed (never just a bare “mismatch” boolean).

Parameters:

• round (int)

• shard_id (int)

• primary_worker (int)

• audit_worker (int)

• primary_nbytes (int)

• audit_nbytes (int)

• first_diff_byte_offset (int | None)

• primary_sha256 (str)

• audit_sha256 (str)

• primary_value_repr (str)

• audit_value_repr (str)

finite_guarded_fold(estimator, payloads, where='AuditedMPEncodedData.combine')

checkpointed_fold, but check finiteness of the running accumulator immediately after EVERY combine() call (see module docstring for why this is a wrapper around this module’s own fold loop rather than a change to combine()’s contract).

Parameters:

• estimator (Any)

• payloads (Sequence[bytes])

• where (str)

Return type:

tuple[float, Any]

inject_bit_flip(payload, bit_offset=None)

Flip exactly one bit of payload, deterministically, and return the corrupted bytes.

A real, reproducible stand-in for a hardware/software bit-flip: this is the corruption primitive the acceptance tests inject via AuditedMPEncodedData.arm_corruption(). bit_offset defaults to the middle bit of the payload (an arbitrary but fixed choice – determinism of the test does not depend on which bit, only that the two payloads it is applied to differ afterward).

Parameters:

• payload (bytes)

• bit_offset (int | None)

Return type:

bytes