diff --git a/docs/PROTOCOL.md b/docs/PROTOCOL.md index 579db093..db8e47e9 100644 --- a/docs/PROTOCOL.md +++ b/docs/PROTOCOL.md @@ -32,7 +32,7 @@ or allocation overflows are errors rather than panics. | type | name | direction | payload | |---:|---|---|---| -| 1 | HELLO | learner → syncer | protocol_version:u16 (=4), learner_id:u32, connection_generation:u64, dtype:u8 (1=f32, 2=bf16, 3=q4), num_fragments:u32, per-fragment layout, layout_fingerprint:[u8;32], session_contract_hash:[u8;32], optional syncer_profile_hash:[u8;32], num_streams:u16, optional compat_group trailer (`b"YCG1"`, u32 length, 1–64 bytes utf-8) | +| 1 | HELLO | learner → syncer | protocol_version:u16 (=4), learner_id:u32, connection_generation:u64, dtype:u8 (1=f32, 2=bf16, 3=q4, 4=f32 with bf16 pushes, 5=f32 with int8 pushes; flag 0x80 = bf16 downlink deltas), num_fragments:u32, per-fragment layout, layout_fingerprint:[u8;32], session_contract_hash:[u8;32], optional syncer_profile_hash:[u8;32], num_streams:u16, optional compat_group trailer (`b"YCG1"`, u32 length, 1–64 bytes utf-8) | | 2 | INIT_PARAMS | learner → syncer | fragment_id:u32, full tensor bytes; only learner 0 may initialize | | 3 | PULL_REQ | syncer → learner | fragment_id:u32, global_step:u64, round_attempt:u32 | | 4 | PUSH_FRAGMENT | learner → syncer | learner_id:u32, fragment_id:u32, global_step:u64, round_attempt:u32, base_version:u64, local_step:u64, c_steps:u32, c_tokens:u64, base-relative learner-delta bytes | @@ -223,6 +223,37 @@ The coordinator is one sequential scheduler actor. Pipelined rounds target distinct fragments, and merge + version update + broadcast for a completion is serialized. No two completions mutate one fragment concurrently. +## bf16 pushes over an f32 session (dtype 4) + +Dtype 4 keeps INIT, BCAST and FINAL exact f32 and sends only PUSH deltas as +bf16 (round-to-nearest-even). The syncer decodes and negates them like dtype 2 +and keeps f32 masters and momentum. The full-parameter policy sweep accepts +dtype 1 or 4 (S26 E3, `--rl-sync-push-dtype bf16`). + +## int8 pushes over an f32 session (dtype 5) + +Dtype 5 sends PUSH deltas as block-scaled int8: u32 block size B +(32..1048576), then per block of B values (the last block may be shorter) an +f32 scale (block absmax / 127) and one i8 per value (`-128` is invalid). +A value decodes to `q * scale`, so its error is at most absmax / 254. The +learner picks B (launcher `--rl-sync-push-int8-block`, default 128); it is +bound into the role-stream session contract. INIT, BCAST and FINAL stay f32. + +## bf16 downlink deltas (dtype flag 0x80) + +On an f32-bulk session (dtype 1, 4 or 5) the flag 0x80 changes BCAST_FRAGMENT +to `fragment_id:u32, version:u64, kind:u8, base_version:u64, +anchor_sha256:[u8;32]`, then the tensor: kind 0 (FULL) carries f32 values, +kind 1 (DELTA) carries `bf16(theta - anchor)`. Syncer and learner hold the +same f32 anchor per fragment: FULL sets it, DELTA adds the decoded delta in +f32; the digest is checked after every update. Every N-th version of a +fragment is FULL (`--bcast-delta-full-every N`, N >= 2; the syncer refuses a +flagged session without it and the reverse). Initial and reconnect re-sends +are always FULL with the shared anchor. A learner with a missing or wrong +anchor drops its connection, and the reconnect re-send restores it. The +syncer's masters, momentum, checkpoints and FINAL are unchanged and exact; +anchors are not checkpointed, so a resumed syncer starts with FULL. + ## Q4 pushes Q4 uses the same learner-delta semantics. INIT and BCAST stay bf16; PUSH is diff --git a/syncer/src/downlink.rs b/syncer/src/downlink.rs new file mode 100644 index 00000000..32e6f82d --- /dev/null +++ b/syncer/src/downlink.rs @@ -0,0 +1,302 @@ +//! bf16 downlink deltas for BCAST_FRAGMENT (S26 E3, session flag +//! `DTYPE_FLAG_BCAST_DELTA`). +//! +//! Every learner of a session holds, per fragment, an f32 *anchor* that is +//! bit-identical to the syncer's copy here. A broadcast of a new version is +//! either +//! +//! * `FULL`: the f32 values; the anchor becomes exactly those values, or +//! * `DELTA`: `d = bf16(θ − anchor)`; both sides set `anchor += d` in f32 +//! (IEEE addition is exact-rounded, so both sides get the same bits). +//! +//! Each delta is taken against the current anchor, so the rounding error does +//! not accumulate: `|θ − anchor| <= |θ − anchor_prev| * 2^-9` per element +//! after every delta. The learner trains from the anchor; its next push is +//! relative to the anchor, and the syncer adds that push to its exact θ, so +//! θ stays the exact sum of all pushes. +//! +//! Drift guard: every `full_every`-th broadcast of a fragment is FULL; a +//! re-sent fragment (HELLO/reconnect recovery) is always FULL; and every +//! header carries the SHA-256 of the f32 anchor after the update. A learner +//! whose anchor version or digest does not match drops its connection; the +//! reconnect recovery then sends FULL. +//! +//! Anchors live only in memory: a resumed syncer starts with none, so its +//! first broadcast of every fragment is FULL. +//! +//! Header after the usual `fragment_id u32, version u64` prefix: +//! `kind u8 (0 FULL, 1 DELTA), base_version u64, anchor_sha256 [32]`. +//! The tensor that follows is f32 (FULL) or bf16 (DELTA). + +use std::sync::Arc; + +use anyhow::{bail, Result}; +use sha2::{Digest, Sha256}; + +use crate::protocol::{DTYPE_BF16, DTYPE_F32}; + +pub const KIND_FULL: u8 = 0; +pub const KIND_DELTA: u8 = 1; +/// fragment_id + version + kind + base_version + digest. +pub const PREFIX_LEN: usize = 4 + 8 + 1 + 8 + 32; + +struct Anchor { + version: u64, + values: Arc>, + digest: [u8; 32], + deltas_since_full: u64, +} + +/// One encoded BCAST: header prefix plus tensor values in `dtype` (DELTA +/// values are already bf16-representable, so encoding them is exact). +pub struct EncodedBcast { + pub prefix: Vec, + pub dtype: u8, + pub values: Arc>, +} + +impl EncodedBcast { + #[cfg(test)] + pub fn kind(&self) -> u8 { + self.prefix[12] + } +} + +pub struct DownlinkAnchors { + full_every: u64, + anchors: Vec>, +} + +pub fn digest_f32(values: &[f32]) -> [u8; 32] { + let mut hasher = Sha256::new(); + // Hash in bounded slices to avoid a second full-size byte buffer. + let mut buf = Vec::with_capacity(4 * 65536); + for chunk in values.chunks(65536) { + buf.clear(); + for v in chunk { + buf.extend_from_slice(&v.to_le_bytes()); + } + hasher.update(&buf); + } + hasher.finalize().into() +} + +fn prefix(fragment: usize, version: u64, kind: u8, base_version: u64, digest: &[u8; 32]) -> Vec { + let mut out = Vec::with_capacity(PREFIX_LEN); + out.extend_from_slice(&(fragment as u32).to_le_bytes()); + out.extend_from_slice(&version.to_le_bytes()); + out.push(kind); + out.extend_from_slice(&base_version.to_le_bytes()); + out.extend_from_slice(digest); + out +} + +impl DownlinkAnchors { + pub fn new(fragments: usize, full_every: u64) -> Result { + if full_every < 2 { + bail!("--bcast-delta-full-every must be at least 2"); + } + Ok(Self { + full_every, + anchors: (0..fragments).map(|_| None).collect(), + }) + } + + /// Encode the broadcast of a newly committed `version` of fragment `p` + /// and advance the shared anchor. Call exactly once per new version. + pub fn publish(&mut self, p: usize, version: u64, params: &[f32]) -> Result { + let slot = self + .anchors + .get_mut(p) + .ok_or_else(|| anyhow::anyhow!("downlink anchor for unknown fragment {p}"))?; + let delta_ok = matches!( + slot, + Some(a) if a.deltas_since_full + 1 < self.full_every + && a.values.len() == params.len() + && a.version < version + ); + if !delta_ok { + return Ok(Self::reset(slot, p, version, params)); + } + let anchor = slot.as_mut().unwrap(); + let base_version = anchor.version; + let values = Arc::make_mut(&mut anchor.values); + let mut delta = Vec::new(); + delta + .try_reserve_exact(params.len()) + .map_err(|error| anyhow::anyhow!("cannot allocate downlink delta: {error}"))?; + for (a, t) in values.iter_mut().zip(params) { + let d = half::bf16::from_f32(t - *a).to_f32(); + if !d.is_finite() { + // bf16 overflow of a huge step: fall back to FULL. + return Ok(Self::reset(slot, p, version, params)); + } + delta.push(d); + *a += d; + } + anchor.version = version; + anchor.digest = digest_f32(values); + anchor.deltas_since_full += 1; + Ok(EncodedBcast { + prefix: prefix(p, version, KIND_DELTA, base_version, &anchor.digest), + dtype: DTYPE_BF16, + values: Arc::new(delta), + }) + } + + /// Encode a newly committed `version` of fragment `p` as FULL (exact f32) + /// and reset the shared anchor to it. + pub fn publish_full(&mut self, p: usize, version: u64, params: &[f32]) -> EncodedBcast { + match self.anchors.get_mut(p) { + Some(slot) => Self::reset(slot, p, version, params), + None => EncodedBcast { + prefix: prefix(p, version, KIND_FULL, version, &digest_f32(params)), + dtype: DTYPE_F32, + values: Arc::new(params.to_vec()), + }, + } + } + + /// FULL re-send for recovery (HELLO / reconnect): the anchor every other + /// learner holds, when it is at `version`; otherwise the exact params. + pub fn resend(&self, p: usize, version: u64, params: &[f32]) -> EncodedBcast { + match self.anchors.get(p).and_then(Option::as_ref) { + Some(anchor) if anchor.version == version => EncodedBcast { + prefix: prefix(p, version, KIND_FULL, version, &anchor.digest), + dtype: DTYPE_F32, + values: Arc::clone(&anchor.values), + }, + _ => EncodedBcast { + prefix: prefix(p, version, KIND_FULL, version, &digest_f32(params)), + dtype: DTYPE_F32, + values: Arc::new(params.to_vec()), + }, + } + } + + fn reset(slot: &mut Option, p: usize, version: u64, params: &[f32]) -> EncodedBcast { + let values = Arc::new(params.to_vec()); + let digest = digest_f32(&values); + *slot = Some(Anchor { + version, + values: Arc::clone(&values), + digest, + deltas_since_full: 0, + }); + EncodedBcast { + prefix: prefix(p, version, KIND_FULL, version, &digest), + dtype: DTYPE_F32, + values, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A learner-side mirror: applies broadcasts exactly as yeto/protocol.py. + fn apply(anchor: &mut Option<(u64, Vec)>, enc: &EncodedBcast) -> Result<()> { + let version = u64::from_le_bytes(enc.prefix[4..12].try_into().unwrap()); + let base = u64::from_le_bytes(enc.prefix[13..21].try_into().unwrap()); + let digest: [u8; 32] = enc.prefix[21..53].try_into().unwrap(); + let next = match enc.kind() { + KIND_FULL => enc.values.to_vec(), + KIND_DELTA => { + let Some((held, values)) = anchor.as_ref() else { bail!("no anchor") }; + if *held != base { + bail!("anchor version {held} != base {base}"); + } + values.iter().zip(enc.values.iter()).map(|(a, d)| a + d).collect() + } + _ => bail!("bad kind"), + }; + if digest_f32(&next) != digest { + bail!("digest mismatch"); + } + *anchor = Some((version, next)); + Ok(()) + } + + fn params(seed: u64, n: usize) -> Vec { + (0..n) + .map(|i| ((i as f64 * 0.37 + seed as f64 * 1.3).sin() * 3.0) as f32) + .collect() + } + + #[test] + fn deltas_track_params_within_bf16_error_and_refresh_every_m() { + let mut syncer = DownlinkAnchors::new(1, 4).unwrap(); + let mut learner = None; + let mut theta = params(0, 1000); + let mut kinds = vec![]; + for version in 1..=9u64 { + let step = params(version, 1000); + let prev_anchor = learner.as_ref().map(|(_, v): &(u64, Vec)| v.clone()); + for (t, s) in theta.iter_mut().zip(&step) { + *t += 0.01 * s; + } + let enc = syncer.publish(0, version, &theta).unwrap(); + kinds.push(enc.kind()); + apply(&mut learner, &enc).unwrap(); + let (_, held) = learner.as_ref().unwrap(); + if enc.kind() == KIND_FULL { + assert_eq!(held, &theta); + } else { + let prev = prev_anchor.unwrap(); + for ((h, t), a) in held.iter().zip(&theta).zip(&prev) { + // one bf16 rounding of (θ − anchor_prev): rel. 2^-8 bound + assert!((h - t).abs() <= (t - a).abs() / 256.0 + 2.0 * f32::EPSILON * t.abs()); + } + } + } + assert_eq!( + kinds, + vec![KIND_FULL, KIND_DELTA, KIND_DELTA, KIND_DELTA, KIND_FULL, KIND_DELTA, KIND_DELTA, KIND_DELTA, KIND_FULL] + ); + } + + #[test] + fn resend_is_the_shared_anchor_and_lets_a_reconnected_learner_continue() { + let mut syncer = DownlinkAnchors::new(1, 100).unwrap(); + let mut a = None; + let theta1 = params(1, 64); + apply(&mut a, &syncer.publish(0, 1, &theta1).unwrap()).unwrap(); + let theta2: Vec = theta1.iter().map(|v| v + 0.123).collect(); + apply(&mut a, &syncer.publish(0, 2, &theta2).unwrap()).unwrap(); + // A fresh learner (reconnect) gets the anchor, not θ, as FULL. + let mut b = None; + let full = syncer.resend(0, 2, &theta2); + assert_eq!(full.kind(), KIND_FULL); + apply(&mut b, &full).unwrap(); + assert_eq!(a, b); + let theta3: Vec = theta2.iter().map(|v| v - 0.05).collect(); + let enc = syncer.publish(0, 3, &theta3).unwrap(); + assert_eq!(enc.kind(), KIND_DELTA); + apply(&mut a, &enc).unwrap(); + apply(&mut b, &enc).unwrap(); + assert_eq!(a, b); + } + + #[test] + fn a_learner_with_a_wrong_anchor_detects_it() { + let mut syncer = DownlinkAnchors::new(1, 100).unwrap(); + let mut a = None; + apply(&mut a, &syncer.publish(0, 1, ¶ms(1, 32)).unwrap()).unwrap(); + // Corrupt the learner's anchor bits: the digest check must fail. + if let Some((_, v)) = a.as_mut() { + v[3] += 1e-3; + } + let enc = syncer.publish(0, 2, ¶ms(2, 32)).unwrap(); + assert!(apply(&mut a, &enc).unwrap_err().to_string().contains("digest")); + // A learner that skipped a version fails the base-version check. + let mut stale = Some((0u64, params(1, 32))); + let enc = syncer.publish(0, 3, ¶ms(3, 32)).unwrap(); + assert!(apply(&mut stale, &enc).unwrap_err().to_string().contains("base")); + } + + #[test] + fn full_every_below_two_is_refused() { + assert!(DownlinkAnchors::new(1, 1).is_err()); + } +} diff --git a/syncer/src/main.rs b/syncer/src/main.rs index 513caabe..58ddc964 100644 --- a/syncer/src/main.rs +++ b/syncer/src/main.rs @@ -2,6 +2,7 @@ mod elastic; mod elastic_server; mod iso_worker; mod merge; +mod downlink; mod protocol; mod server; mod state; @@ -98,6 +99,15 @@ struct Args { /// Write the snapshot every N outer steps (0 disables). #[arg(long, default_value_t = 8)] checkpoint_every: u64, + /// Policy sweep: write per-round checkpoints off the scheduler; each + /// round is broadcast only after its checkpoint is durable. + #[arg(long, default_value_t = false)] + async_checkpoint: bool, + /// Learners whose session asks for bf16 downlink deltas get a full f32 + /// BCAST every N-th version of a fragment (N >= 2). Unset: such + /// sessions are refused. + #[arg(long)] + bcast_delta_full_every: Option, /// Resume from --checkpoint-path if it exists. #[arg(long, default_value_t = false)] resume: bool, @@ -273,6 +283,8 @@ fn main() -> anyhow::Result<()> { final_state: args.final_state, checkpoint_path: args.checkpoint_path, checkpoint_every: args.checkpoint_every, + async_checkpoint: args.async_checkpoint, + bcast_delta_full_every: args.bcast_delta_full_every, resume: args.resume, mark_final_checkpoint: args.mark_final_checkpoint, learner_budget_steps: args.learner_budget_steps, diff --git a/syncer/src/protocol.rs b/syncer/src/protocol.rs index bf4b530e..e214faf2 100644 --- a/syncer/src/protocol.rs +++ b/syncer/src/protocol.rs @@ -45,6 +45,10 @@ pub const MSG_ELASTIC_BASE: u8 = 23; pub const MSG_SAMPLE_VERDICT: u8 = 24; /// Elastic: run finished (final outer_version + base policy hash). pub const MSG_FINISHED: u8 = 25; +/// S26 E3b: link-bandwidth probe. First and only frame of its own connection +/// (never part of a learner session): the island uploads N bytes, the syncer +/// acknowledges with the u64 byte count, then streams M bytes back. +pub const MSG_BW_PROBE: u8 = 26; /// Version of the final-artifact handshake carried inside FINAL_MANIFEST and /// FINAL_ACK. Keeping this in the payload makes incompatible peers fail with @@ -57,17 +61,85 @@ pub const DTYPE_BF16: u8 = 2; /// base-relative learner deltas; INIT_PARAMS and BCAST_FRAGMENT travel as /// bf16 (see `bulk_dtype`, docs/PROTOCOL.md v4). pub const DTYPE_Q4: u8 = 3; +/// Session dtype 4 (S26 E3): PUSH_FRAGMENT deltas travel as bf16 while +/// INIT_PARAMS, BCAST_FRAGMENT and FINAL stay exact f32. The syncer keeps f32 +/// masters and momentum; only the learner's outer gradient is rounded. +pub const DTYPE_F32_PUSH_BF16: u8 = 4; + +/// Session dtype 5 (S26 E3): PUSH_FRAGMENT deltas travel as block-scaled +/// int8 (see `decode_i8`); INIT_PARAMS, BCAST_FRAGMENT and FINAL stay f32. +pub const DTYPE_F32_PUSH_I8: u8 = 5; +/// Session-dtype flag (S26 E3): BCAST_FRAGMENT carries a 41-byte downlink +/// header and, between periodic f32 refreshes, a bf16 delta against the +/// anchor every learner holds (see `crate::downlink`). Valid only on f32-bulk +/// sessions (dtypes 1, 4, 5). +pub const DTYPE_FLAG_BCAST_DELTA: u8 = 0x80; /// Values per q4 scale block (f32 absmax scale + 128 nibble bytes each). pub const Q4_BLOCK: usize = 256; +/// Bounds of the int8 push block size the learner declares per payload. +pub const I8_BLOCK_MIN: usize = 32; +pub const I8_BLOCK_MAX: usize = 1 << 20; + +/// The session dtype without its flags. +pub fn session_base(dtype: u8) -> u8 { + dtype & !DTYPE_FLAG_BCAST_DELTA +} + +/// Whether the session receives bf16 downlink deltas. +pub fn bcast_delta(dtype: u8) -> bool { + dtype & DTYPE_FLAG_BCAST_DELTA != 0 +} + +/// Whether `dtype` is a session dtype this syncer accepts. +pub fn valid_session_dtype(dtype: u8) -> bool { + match session_base(dtype) { + DTYPE_F32 | DTYPE_F32_PUSH_BF16 | DTYPE_F32_PUSH_I8 => true, + DTYPE_BF16 | DTYPE_Q4 => !bcast_delta(dtype), + _ => false, + } +} /// Dtype of INIT_PARAMS/BCAST_FRAGMENT tensors for a session dtype: q4 /// applies only to push deltas; full parameter payloads stay bf16. pub fn bulk_dtype(dtype: u8) -> u8 { - if dtype == DTYPE_Q4 { - DTYPE_BF16 - } else { - dtype + match session_base(dtype) { + DTYPE_Q4 => DTYPE_BF16, + DTYPE_F32_PUSH_BF16 | DTYPE_F32_PUSH_I8 => DTYPE_F32, + other => other, + } +} + +/// Dtype of PUSH_FRAGMENT deltas for an unquantized session dtype (q4 and +/// int8 have their own block codecs). +pub fn push_dtype(dtype: u8) -> u8 { + match session_base(dtype) { + DTYPE_F32_PUSH_BF16 => DTYPE_BF16, + other => other, + } +} + +/// Largest PUSH_FRAGMENT delta payload a session dtype may carry. +pub fn max_push_nbytes(dtype: u8, numel: usize) -> Result { + match session_base(dtype) { + DTYPE_Q4 => q4_nbytes(numel), + DTYPE_F32_PUSH_I8 => i8_nbytes(numel, I8_BLOCK_MIN), + other => tensor_nbytes(push_dtype(other), numel), + } +} + +/// Decode a PUSH_FRAGMENT delta payload of any session dtype into f32s. +pub fn decode_push(dtype: u8, bytes: &[u8], numel: usize, out: &mut Vec) -> Result<()> { + match session_base(dtype) { + DTYPE_Q4 => decode_q4(bytes, numel, out), + DTYPE_F32_PUSH_I8 => decode_i8(bytes, numel, out), + other => { + let expected = tensor_nbytes(push_dtype(other), numel)?; + if bytes.len() != expected { + bail!("push delta has {} bytes, expected {expected}", bytes.len()); + } + decode_tensor(push_dtype(other), bytes, out) + } } } @@ -257,6 +329,80 @@ pub fn encode_tensor(dtype: u8, vals: &[f32], out: &mut Vec) -> Result<()> { Ok(()) } +/// Exact encoded byte length of an int8 push payload with block size `block`. +pub fn i8_nbytes(numel: usize, block: usize) -> Result { + if !(I8_BLOCK_MIN..=I8_BLOCK_MAX).contains(&block) { + bail!("int8 block size {block} outside {I8_BLOCK_MIN}..={I8_BLOCK_MAX}"); + } + numel + .div_ceil(block) + .checked_mul(4) + .and_then(|scales| scales.checked_add(numel)) + .and_then(|body| body.checked_add(4)) + .ok_or_else(|| anyhow::anyhow!("int8 tensor byte length overflow")) +} + +/// Decode an int8 delta payload (yeto/protocol.py `f32_bytes_to_i8`). +/// +/// Layout: u32 block size B, then per block of B values (the last block may +/// be shorter): f32 scale (absmax / 127), then one i8 per value. A value +/// decodes to `q as f32 * scale`, so |error| <= scale / 2 = absmax / 254 per +/// element (plus one f32 rounding). The scale is non-negative and finite. +pub fn decode_i8(bytes: &[u8], numel: usize, out: &mut Vec) -> Result<()> { + if bytes.len() < 4 { + bail!("int8 payload is shorter than its block header"); + } + let block = u32::from_le_bytes(bytes[..4].try_into().unwrap()) as usize; + let expected = i8_nbytes(numel, block)?; + if bytes.len() != expected { + bail!( + "int8 payload has {} bytes, expected {expected} for {numel} values in blocks of {block}", + bytes.len() + ); + } + out.clear(); + out.try_reserve(numel) + .map_err(|error| anyhow::anyhow!("cannot allocate int8 tensor: {error}"))?; + let mut at = 4usize; + let mut remaining = numel; + let mut b = 0usize; + while remaining > 0 { + let len = remaining.min(block); + let scale = f32::from_le_bytes(bytes[at..at + 4].try_into().unwrap()); + if !scale.is_finite() || scale < 0.0 { + bail!("int8 block {b} has invalid scale {scale}"); + } + at += 4; + for &q in &bytes[at..at + len] { + let q = q as i8; + if q == i8::MIN { + bail!("int8 block {b} uses the reserved value -128"); + } + out.push(q as f32 * scale); + } + at += len; + remaining -= len; + b += 1; + } + Ok(()) +} + +/// Encode f32s as an int8 push payload (test counterpart of the Python codec). +#[cfg(test)] +pub fn encode_i8(vals: &[f32], block: usize) -> Vec { + let mut out = (block as u32).to_le_bytes().to_vec(); + for chunk in vals.chunks(block) { + let absmax = chunk.iter().fold(0f32, |m, v| m.max(v.abs())); + let scale = absmax / 127.0; + out.extend_from_slice(&scale.to_le_bytes()); + for v in chunk { + let q = if scale > 0.0 { (v / scale).round_ties_even().clamp(-127.0, 127.0) } else { 0.0 }; + out.push(q as i8 as u8); + } + } + out +} + /// Decode a q4 delta payload (yeto/tensor_io.py `quantize_q4`) into f32s. /// /// Per block of `Q4_BLOCK` values: f32 absmax scale, then 128 bytes of @@ -312,6 +458,59 @@ mod tests { use super::*; use tokio::io::AsyncWriteExt; + #[test] + fn int8_push_round_trips_within_half_a_scale_step() { + let vals: Vec = (0..1000) + .map(|i| ((i as f32) * 0.731).sin() * if i % 300 < 150 { 1e-3 } else { 2.5 }) + .collect(); + for block in [I8_BLOCK_MIN, 128, 333, 4096] { + let wire = encode_i8(&vals, block); + assert_eq!(wire.len(), i8_nbytes(vals.len(), block).unwrap()); + assert!(wire.len() <= max_push_nbytes(DTYPE_F32_PUSH_I8, vals.len()).unwrap()); + let mut out = Vec::new(); + decode_push(DTYPE_F32_PUSH_I8, &wire, vals.len(), &mut out).unwrap(); + for (chunk, decoded) in vals.chunks(block).zip(out.chunks(block)) { + let absmax = chunk.iter().fold(0f32, |m, v| m.max(v.abs())); + for (v, d) in chunk.iter().zip(decoded) { + // |error| <= scale/2 = absmax/254, plus f32 rounding. + assert!((v - d).abs() <= absmax / 254.0 * (1.0 + 1e-5), "{v} {d} {absmax}"); + } + } + } + // All-zero blocks decode to exact zeros. + let mut out = Vec::new(); + decode_i8(&encode_i8(&[0.0; 40], 32), 40, &mut out).unwrap(); + assert!(out.iter().all(|v| *v == 0.0)); + } + + #[test] + fn int8_push_rejects_malformed_payloads() { + let vals = vec![1.0f32; 64]; + let mut out = Vec::new(); + let wire = encode_i8(&vals, 32); + assert!(decode_i8(&wire[..wire.len() - 1], 64, &mut out).is_err()); + assert!(decode_i8(&encode_i8(&vals, 16), 64, &mut out).is_err()); + let mut bad_scale = wire.clone(); + bad_scale[4..8].copy_from_slice(&(-1.0f32).to_le_bytes()); + assert!(decode_i8(&bad_scale, 64, &mut out).is_err()); + let mut reserved = wire; + reserved[8] = 0x80; + assert!(decode_i8(&reserved, 64, &mut out).is_err()); + } + + #[test] + fn session_dtype_flags_compose_only_with_f32_bulk() { + assert!(valid_session_dtype(DTYPE_F32_PUSH_I8)); + assert!(valid_session_dtype(DTYPE_F32 | DTYPE_FLAG_BCAST_DELTA)); + assert!(valid_session_dtype(DTYPE_F32_PUSH_BF16 | DTYPE_FLAG_BCAST_DELTA)); + assert!(valid_session_dtype(DTYPE_F32_PUSH_I8 | DTYPE_FLAG_BCAST_DELTA)); + assert!(!valid_session_dtype(DTYPE_BF16 | DTYPE_FLAG_BCAST_DELTA)); + assert!(!valid_session_dtype(DTYPE_Q4 | DTYPE_FLAG_BCAST_DELTA)); + assert!(!valid_session_dtype(6)); + assert_eq!(bulk_dtype(DTYPE_F32_PUSH_I8 | DTYPE_FLAG_BCAST_DELTA), DTYPE_F32); + assert_eq!(push_dtype(DTYPE_F32_PUSH_BF16 | DTYPE_FLAG_BCAST_DELTA), DTYPE_BF16); + } + #[test] fn reserved_message_ids_are_stable() { assert_eq!(MSG_ERROR, 10); @@ -417,6 +616,9 @@ mod tests { assert_eq!(bulk_dtype(DTYPE_Q4), DTYPE_BF16); assert_eq!(bulk_dtype(DTYPE_BF16), DTYPE_BF16); assert_eq!(bulk_dtype(DTYPE_F32), DTYPE_F32); + assert_eq!(bulk_dtype(DTYPE_F32_PUSH_BF16), DTYPE_F32); + assert_eq!(push_dtype(DTYPE_F32_PUSH_BF16), DTYPE_BF16); + assert_eq!(push_dtype(DTYPE_F32), DTYPE_F32); } #[test] diff --git a/syncer/src/server.rs b/syncer/src/server.rs index a5ed5b67..c03011a3 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -91,6 +91,16 @@ pub struct Config { /// the quiescent cut between rounds, resumed from when `resume` is set. pub checkpoint_path: Option, pub checkpoint_every: u64, + /// Policy-sweep only: persist each per-round checkpoint on a blocking + /// writer instead of the scheduler, and broadcast that round only after + /// the write is durable. Merge results, versions, and checkpoint bytes + /// are unchanged, so this is deliberately not part of the semantic + /// profile hash. + pub async_checkpoint: bool, + /// S26 E3: full-f32 refresh period of bf16 downlink deltas (sessions + /// with `DTYPE_FLAG_BCAST_DELTA`). Wire encoding only: merges, versions + /// and checkpoints are unchanged, so it is not in the profile hash. + pub bcast_delta_full_every: Option, pub resume: bool, /// Create the adjacent final marker only for a completed terminal cut. pub mark_final_checkpoint: bool, @@ -702,8 +712,11 @@ fn validate_config(cfg: &Config) -> Result<()> { if cfg.total_steps == 0 { bail!("--policy-sweep-fragments requires positive --total-steps"); } - if cfg.pipeline != 1 { - bail!("--policy-sweep-fragments requires --pipeline 1"); + // Fragment rounds of one sweep are independent (H=1, plain average, + // no momentum), so up to P of them may gather concurrently; commits + // stay in strict step order. + if cfg.pipeline == 0 || cfg.pipeline > fragments { + bail!("--policy-sweep-fragments requires --pipeline between 1 and the fragment count"); } if !cfg.total_steps.is_multiple_of(u64::from(fragments)) { bail!("--total-steps must be divisible by --policy-sweep-fragments"); @@ -726,13 +739,28 @@ fn validate_config(cfg: &Config) -> Result<()> { // --resume is not required here: a fresh sweep has no checkpoint yet and // --resume refuses a missing file (see check_resume_checkpoint). The // launcher adds --resume only when the checkpoint exists. - if cfg.checkpoint_path.is_none() || cfg.checkpoint_every != 1 { + if cfg.checkpoint_path.is_none() + || cfg.checkpoint_every == 0 + || (cfg.checkpoint_every != 1 + && !cfg.checkpoint_every.is_multiple_of(u64::from(fragments))) + { bail!( "--policy-sweep-fragments requires --checkpoint-path and \ - --checkpoint-every 1 for crash-safe sweeps" + --checkpoint-every 1 or a multiple of the fragment count" ); } } + if matches!(cfg.bcast_delta_full_every, Some(n) if n < 2) { + bail!("--bcast-delta-full-every must be at least 2"); + } + if cfg.async_checkpoint { + if cfg.policy_sweep_fragments.is_none() { + bail!("--async-checkpoint requires --policy-sweep-fragments"); + } + if cfg.learner_budget_steps.is_some() { + bail!("--async-checkpoint does not support --learner-budget-steps"); + } + } Ok(()) } @@ -743,6 +771,20 @@ fn checkpoint_before_broadcast(cfg: &Config) -> bool { cfg.checkpoint_path.is_some() && cfg.checkpoint_every == 1 } +/// Whether step `t` must be durable before its BCAST. Per-round checkpoints +/// always are; a policy sweep with `--checkpoint-every N` (N a multiple of +/// the fragment count) persists only every N-th step, at sweep boundaries, +/// and broadcasts the other steps without a durable record (a crash resumes +/// at the newest checkpoint; islands ahead of it fail closed on their next +/// push's local-step check). +fn persist_before_broadcast(cfg: &Config, t: u64) -> bool { + checkpoint_before_broadcast(cfg) + || (cfg.policy_sweep_fragments.is_some() + && cfg.checkpoint_path.is_some() + && cfg.checkpoint_every > 0 + && t.is_multiple_of(cfg.checkpoint_every)) +} + fn validate_policy_sweep_push(push: &Push, fragments: u32) -> Result<()> { if fragments == 0 { bail!("policy sweep fragment count must be positive"); @@ -980,11 +1022,7 @@ fn negotiated_payload_limits(layout: &Layout, dtype: u8) -> Result<(u64, u64)> { for fragment in &layout.fragments { let numel = fragment.numel()?; let init_bytes = tensor_nbytes(bulk_dtype(dtype), numel)? as u64; - let delta_bytes = if dtype == DTYPE_Q4 { - q4_nbytes(numel)? - } else { - tensor_nbytes(dtype, numel)? - } as u64; + let delta_bytes = max_push_nbytes(dtype, numel)? as u64; max_init = max_init.max( 4u64.checked_add(init_bytes) .context("INIT_PARAMS payload limit overflow")?, @@ -1129,7 +1167,7 @@ fn parse_hello(payload: &[u8], expected_learners: u32) -> Result { bail!("connection generation must be nonzero"); } let dtype = r.u8()?; - if !matches!(dtype, DTYPE_F32 | DTYPE_BF16 | DTYPE_Q4) { + if !valid_session_dtype(dtype) { bail!("unsupported session dtype {dtype}"); } let num_fragments = r.u32()?; @@ -1222,8 +1260,8 @@ async fn handle_connection( stream.set_nodelay(true)?; let (mut rd, mut wr) = stream.into_split(); let first = match read_frame_limited(&mut rd, |msg_type| match msg_type { - MSG_HELLO | MSG_DATA_HELLO => Ok(MAX_HELLO_FRAME), - other => bail!("first frame must be HELLO/DATA_HELLO, got {other}"), + MSG_HELLO | MSG_DATA_HELLO | MSG_BW_PROBE => Ok(MAX_HELLO_FRAME), + other => bail!("first frame must be HELLO/DATA_HELLO/BW_PROBE, got {other}"), }) .await { @@ -1428,8 +1466,98 @@ async fn handle_connection( tokio::spawn(writer_task(wr, rx)); read_loop(&mut rd, &group, &event_tx, &budget_cutoff).await } - t => bail!("first frame must be HELLO/DATA_HELLO, got {t}"), + MSG_BW_PROBE => { + let request = open_hello(hello_auth.key.as_deref(), &first.payload) + .and_then(parse_bw_probe); + let (upload, download) = match request { + Ok(sizes) => sizes, + Err(error) => { + let message = format!("invalid BW_PROBE: {error:#}"); + let _ = send_direct(&mut wr, MSG_ERROR, message.as_bytes()).await; + return Err(error); + } + }; + let Some(_slot) = BwProbeSlot::acquire() else { + let message = "too many concurrent BW_PROBE connections"; + send_direct(&mut wr, MSG_ERROR, message.as_bytes()).await?; + bail!(message); + }; + tokio::time::timeout(BW_PROBE_TIMEOUT, serve_bw_probe(&mut rd, &mut wr, upload, download)) + .await + .context("BW_PROBE timed out")? + } + t => bail!("first frame must be HELLO/DATA_HELLO/BW_PROBE, got {t}"), + } +} + +/// Largest upload or download one BW_PROBE connection may ask for. +const BW_PROBE_MAX_BYTES: u64 = 1 << 30; +const BW_PROBE_MAX_CONNECTIONS: usize = 64; +const BW_PROBE_TIMEOUT: Duration = Duration::from_secs(300); +static BW_PROBE_ACTIVE: AtomicUsize = AtomicUsize::new(0); + +/// One of at most BW_PROBE_MAX_CONNECTIONS concurrent probe connections. +struct BwProbeSlot; + +impl BwProbeSlot { + fn acquire() -> Option { + let previous = BW_PROBE_ACTIVE.fetch_add(1, Ordering::AcqRel); + if previous >= BW_PROBE_MAX_CONNECTIONS { + BW_PROBE_ACTIVE.fetch_sub(1, Ordering::AcqRel); + return None; + } + Some(BwProbeSlot) + } +} + +impl Drop for BwProbeSlot { + fn drop(&mut self) { + BW_PROBE_ACTIVE.fetch_sub(1, Ordering::AcqRel); + } +} + +/// BW_PROBE body (after the HELLO-style HMAC trailer is removed): +/// u64 upload bytes, u64 download bytes, both at most BW_PROBE_MAX_BYTES. +fn parse_bw_probe(body: &[u8]) -> Result<(u64, u64)> { + if body.len() != 16 { + bail!("BW_PROBE body has {} bytes, expected 16", body.len()); + } + let upload = u64::from_le_bytes(body[0..8].try_into().unwrap()); + let download = u64::from_le_bytes(body[8..16].try_into().unwrap()); + if upload > BW_PROBE_MAX_BYTES || download > BW_PROBE_MAX_BYTES { + bail!("BW_PROBE asks for more than {BW_PROBE_MAX_BYTES} bytes per direction"); + } + Ok((upload, download)) +} + +/// Read and discard `upload` raw bytes, acknowledge with their count, then +/// write `download` zero bytes and close. Nothing touches session state. +async fn serve_bw_probe( + rd: &mut (impl tokio::io::AsyncReadExt + Unpin), + wr: &mut OwnedWriteHalf, + upload: u64, + download: u64, +) -> Result<()> { + let mut buf = vec![0u8; 1 << 20]; + let mut remaining = upload; + while remaining > 0 { + let want = remaining.min(buf.len() as u64) as usize; + let got = rd.read(&mut buf[..want]).await?; + if got == 0 { + bail!("BW_PROBE upload ended after {} of {upload} bytes", upload - remaining); + } + remaining -= got as u64; } + wr.write_all(&upload.to_le_bytes()).await?; + buf.fill(0); + let mut remaining = download; + while remaining > 0 { + let take = remaining.min(buf.len() as u64) as usize; + wr.write_all(&buf[..take]).await?; + remaining -= take as u64; + } + wr.shutdown().await?; + Ok(()) } async fn send_direct(wr: &mut OwnedWriteHalf, msg_type: u8, payload: &[u8]) -> Result<()> { @@ -1700,23 +1828,9 @@ async fn dispatch_inner( .with_context(|| format!("PUSH_FRAGMENT for unknown fragment {fragment_id}"))?; let numel = fragment.numel()?; let bytes = r.rest(); - let expected = if group.dtype == DTYPE_Q4 { - q4_nbytes(numel)? - } else { - tensor_nbytes(group.dtype, numel)? - }; - if bytes.len() != expected { - bail!( - "PUSH_FRAGMENT fragment {fragment_id} has {} delta bytes, expected {expected}", - bytes.len() - ); - } let mut outer_gradient = Vec::new(); - if group.dtype == DTYPE_Q4 { - decode_q4(bytes, numel, &mut outer_gradient)?; - } else { - decode_tensor(group.dtype, bytes, &mut outer_gradient)?; - } + decode_push(group.dtype, bytes, numel, &mut outer_gradient) + .with_context(|| format!("PUSH_FRAGMENT fragment {fragment_id}"))?; // Every wire dtype carries learner_delta = local - raw_anchor. // Merge math consumes the corresponding outer gradient with the // opposite sign and never reconstructs a full learner parameter. @@ -1862,6 +1976,17 @@ async fn send_fragment_until_cutoff( dtype: u8, budget_cutoff: &BudgetCutoff, ) -> Result { + if let Some(encoded) = recovery_bcast(st, p, msg_type) { + return send_tensor_until_cutoff( + group, + msg_type, + &encoded.prefix, + encoded.dtype, + &encoded.values, + budget_cutoff, + ) + .await; + } let mut prefix = [0u8; 12]; prefix[..4].copy_from_slice(&(p as u32).to_le_bytes()); prefix[4..].copy_from_slice(&st.versions[p].to_le_bytes()); @@ -1876,6 +2001,17 @@ async fn send_fragment_until_cutoff( .await } +/// A BCAST outside publication (initial cut, HELLO/reconnect recovery) on a +/// bf16-downlink session: always FULL, carrying the shared anchor. +fn recovery_bcast(st: &GlobalState, p: usize, msg_type: u8) -> Option { + if msg_type != MSG_BCAST_FRAGMENT { + return None; + } + st.downlink + .as_ref() + .map(|anchors| anchors.lock().unwrap().resend(p, st.versions[p], &st.params[p])) +} + async fn send_pull_until_cutoff( groups: &[Arc], pull: &bytes::Bytes, @@ -2085,7 +2221,7 @@ async fn scheduler( // Persist a fresh version zero before its first BCAST. A resumed cut is // already committed and only needs to be rebroadcast. - if checkpoint_each_round { + if checkpoint_each_round || cfg.policy_sweep_fragments.is_some() { if let Some(path) = cfg.checkpoint_path.as_ref() { if !cfg.resume || !path.exists() { st.save_checkpoint(path)?; @@ -2116,12 +2252,15 @@ async fn scheduler( let mut computing = tokio::task::JoinSet::>::new(); let mut ready: BTreeMap = BTreeMap::new(); let mut busy_fragments: HashSet = HashSet::new(); + // --async-checkpoint: the committed round whose checkpoint is in flight. + let mut publishing: Option = None; let mut cutoff_established = false; let mut first_budget_report: Option<(Member, u64)> = None; 'outer: while next_t <= cfg.total_steps || !inflight.is_empty() || !computing.is_empty() || !ready.is_empty() + || publishing.is_some() { if budget_cutoff.deadline().is_some() { cutoff_established = true; @@ -2138,17 +2277,53 @@ async fn scheduler( let completed = take_next_commit(&mut ready, next_commit_t) .context("ready commit disappeared after linearization")?; let p = completed.round.p; - commit_round( - &cfg, - &mut st, - ®istry, - &mut last_sync_secs, - completed, - &budget_cutoff, - ) - .await?; - if !busy_fragments.remove(&p) { - bail!("fragment {p} lost its busy ownership at commit"); + if cfg.async_checkpoint { + // At most one checkpoint write is outstanding: the snapshot + // of t must be durable before t+1 is snapshotted, which also + // bounds host memory to one extra copy of the state. + if let Some(pending) = publishing.take() { + settle_pending_publish( + &cfg, + &st, + ®istry, + &mut last_sync_secs, + pending, + &budget_cutoff, + &mut busy_fragments, + ) + .await?; + } + let committed = apply_round_commit(&cfg, &mut st, completed)?; + if persist_before_broadcast(&cfg, committed.t) { + publishing = Some(start_async_checkpoint(&cfg, &st, committed)?); + } else { + publish_round( + &cfg, + &st, + ®istry, + &mut last_sync_secs, + committed, + None, + &budget_cutoff, + ) + .await?; + if !busy_fragments.remove(&p) { + bail!("fragment {p} lost its busy ownership at commit"); + } + } + } else { + commit_round( + &cfg, + &mut st, + ®istry, + &mut last_sync_secs, + completed, + &budget_cutoff, + ) + .await?; + if !busy_fragments.remove(&p) { + bail!("fragment {p} lost its busy ownership at commit"); + } } next_commit_t = next_commit_t .checked_add(1) @@ -2349,7 +2524,19 @@ async fn scheduler( if next_fragment_available { earliest = Some(earliest.map_or(next_launch, |d| d.min(next_launch))); } - let wake = match earliest { + let wake = if let Some(pending) = publishing.as_mut() { + // Same event-first bias; a durable checkpoint outranks compute + // so the oldest committed round is exposed first. + let sleep_until = earliest.map(tokio::time::Instant::from_std); + tokio::select! { + biased; + event = events.recv() => SchedulerWake::Event(event), + written = &mut pending.write => SchedulerWake::Durable(written), + result = computing.join_next(), if !computing.is_empty() => SchedulerWake::Computed(result), + _ = sleep_until_opt(sleep_until) => SchedulerWake::Deadline, + } + } else { + match earliest { Some(deadline) if !computing.is_empty() => tokio::select! { biased; wake = prefer_event_over_compute(events.recv(), computing.join_next()) => wake, @@ -2364,9 +2551,27 @@ async fn scheduler( prefer_event_over_compute(events.recv(), computing.join_next()).await } None => continue, + } }; match wake { SchedulerWake::Deadline => continue, + SchedulerWake::Durable(written) => { + let pending = publishing + .take() + .context("durable checkpoint without a pending publish")?; + finish_pending_publish( + &cfg, + &st, + ®istry, + &mut last_sync_secs, + pending.committed, + pending.write_started, + written, + &budget_cutoff, + &mut busy_fragments, + ) + .await?; + } SchedulerWake::Computed(Some(result)) => { let completed = result.context("iso compute task panicked")??; let t = completed.round.t; @@ -2469,6 +2674,20 @@ async fn scheduler( // slots in existing rounds. The helper revalidates the // generation because catch-up awaits socket writers and // a newer HELLO may supersede this one meanwhile. + // Catch-up must never expose a round whose checkpoint + // is not yet durable. + if let Some(pending) = publishing.take() { + settle_pending_publish( + &cfg, + &st, + ®istry, + &mut last_sync_secs, + pending, + &budget_cutoff, + &mut busy_fragments, + ) + .await?; + } if st.global_step < cfg.total_steps { send_all_fragments(&st, &group, &budget_cutoff).await; } @@ -2553,6 +2772,18 @@ async fn scheduler( // Catch-up must precede a replay even when the // disconnect event wins the race with the queued // replacement HELLO. + if let Some(pending) = publishing.take() { + settle_pending_publish( + &cfg, + &st, + ®istry, + &mut last_sync_secs, + pending, + &budget_cutoff, + &mut busy_fragments, + ) + .await?; + } if st.global_step < cfg.total_steps { send_all_fragments(&st, &group, &budget_cutoff).await; } @@ -2640,6 +2871,7 @@ async fn scheduler( if next_commit_t != cfg.total_steps.saturating_add(1) || st.global_step != cfg.total_steps || !busy_fragments.is_empty() + || publishing.is_some() { bail!( "non-quiescent terminal cut: next_commit_t={next_commit_t} global_step={} busy={:?}", @@ -2935,6 +3167,14 @@ enum SchedulerWake { Deadline, Event(Option), Computed(Option, tokio::task::JoinError>>), + Durable(std::result::Result, tokio::task::JoinError>), +} + +async fn sleep_until_opt(deadline: Option) { + match deadline { + Some(deadline) => tokio::time::sleep_until(deadline).await, + None => std::future::pending().await, + } } type ComputeJoin = Option, tokio::task::JoinError>>; @@ -3189,6 +3429,24 @@ fn prepare_round_compute(cfg: &Config, st: &GlobalState, round: &Round) -> Resul st.prepare_merge(round.p, round.base_version, &outer_gradients, &weights) } +/// A round whose merge is applied to `GlobalState` but not yet externally +/// visible: no BCAST, no tape record, no log line. +struct CommittedRound { + t: u64, + p: usize, + base_version: u64, + attempt: u32, + started: Instant, + expected_members: Vec, + quorum_size: usize, + quorum_ms: Option, + grace_ms: Option, + pushes: HashMap, + gnorm: f64, + compute_secs: f64, + sync_start: Instant, +} + /// Commit an already-computed round. Called only for the contiguous t prefix, /// so Nesterov, versions, ledger, event tape, checkpoint, and broadcast all /// share one deterministic order even when SVD completion is out of order. @@ -3200,6 +3458,26 @@ async fn commit_round( completed: ComputedRound, budget_cutoff: &BudgetCutoff, ) -> Result<()> { + let committed = apply_round_commit(cfg, st, completed)?; + let mut checkpoint_ms = None; + if persist_before_broadcast(cfg, committed.t) { + if let Some(path) = cfg.checkpoint_path.as_ref() { + let checkpoint_start = Instant::now(); + st.save_checkpoint(path)?; + checkpoint_ms = Some(checkpoint_start.elapsed().as_millis() as u64); + info!(step = committed.t, path = %path.display(), "checkpoint committed"); + } + } + publish_round(cfg, st, registry, last_sync_secs, committed, checkpoint_ms, budget_cutoff).await +} + +/// Apply the merge, version, step, and ledger updates of the next contiguous +/// round. Nothing leaves the process here. +fn apply_round_commit( + cfg: &Config, + st: &mut GlobalState, + completed: ComputedRound, +) -> Result { let ComputedRound { round, merge, @@ -3237,7 +3515,6 @@ async fn commit_round( st.global_step = t; let mut ordered_pushes: Vec<_> = pushes.iter().collect(); ordered_pushes.sort_unstable_by_key(|(member, _)| **member); - let responders: Vec = ordered_pushes.iter().map(|(member, _)| **member).collect(); for (_, push) in ordered_pushes { if let Some(fragments) = cfg.policy_sweep_fragments { st.record_fragment_merge( @@ -3252,15 +3529,52 @@ async fn commit_round( st.record_merge(push.learner_id, push.c_steps, push.c_tokens); } } + Ok(CommittedRound { + t, + p, + base_version, + attempt, + started, + expected_members, + quorum_size, + quorum_ms, + grace_ms, + pushes, + gnorm, + compute_secs, + sync_start, + }) +} - let checkpoint_each_round = checkpoint_before_broadcast(cfg); - if checkpoint_each_round { - if let Some(path) = cfg.checkpoint_path.as_ref() { - st.save_checkpoint(path)?; - info!(step = t, path = %path.display(), "checkpoint committed"); - } - } - +/// Expose a committed round: BCAST its fragment, log it, append its tape +/// record, and write the periodic (non-per-round) checkpoint. A per-round +/// checkpoint must already be durable when this is called. +async fn publish_round( + cfg: &Config, + st: &GlobalState, + registry: &Registry, + last_sync_secs: &mut f64, + committed: CommittedRound, + checkpoint_ms: Option, + budget_cutoff: &BudgetCutoff, +) -> Result<()> { + let CommittedRound { + t, + p, + base_version, + attempt, + started, + expected_members, + quorum_size, + quorum_ms, + grace_ms, + pushes, + gnorm, + compute_secs, + sync_start, + } = committed; + let mut responders: Vec = pushes.keys().copied().collect(); + responders.sort_unstable(); // A disconnected learner can catch up after reconnecting. Deterministic // local encoding/allocation failures must terminate the coordinator // instead of silently committing a cut that was never broadcastable. @@ -3268,16 +3582,43 @@ async fn commit_round( // its exact f32 cut immediately after the scheduler reaches quiescence. if t < cfg.total_steps { let dtype = bulk_dtype(st.wire_dtype); + // bf16 downlink: encode this version once (advancing the shared + // anchor) and send the same bytes to every learner. + let published = match &st.downlink { + Some(anchors) => { + // The last version of each fragment is also its FINAL version. + // Islands key stored fragments by (fragment, version) and + // refuse a second copy with other bytes, so that broadcast + // must carry the exact f32 masters (FULL), not a bf16 delta. + let last_version = t + st.params.len() as u64 > cfg.total_steps; + let mut anchors = anchors.lock().unwrap(); + Some(if last_version { + anchors.publish_full(p, st.versions[p], &st.params[p]) + } else { + anchors.publish(p, st.versions[p], &st.params[p])? + }) + } + None => None, + }; for group in current_groups(registry) { - match send_fragment_until_cutoff( - st, - &group, - p, - MSG_BCAST_FRAGMENT, - dtype, - budget_cutoff, - ) - .await + let sent = match &published { + Some(encoded) => { + send_tensor_until_cutoff( + &group, + MSG_BCAST_FRAGMENT, + &encoded.prefix, + encoded.dtype, + &encoded.values, + budget_cutoff, + ) + .await + } + None => { + send_fragment_until_cutoff(st, &group, p, MSG_BCAST_FRAGMENT, dtype, budget_cutoff) + .await + } + }; + match sent { Ok(true) => {} Ok(false) => break, @@ -3318,6 +3659,7 @@ async fn commit_round( quorum_ms, grace_ms, sync_ms, + checkpoint_ms, &pushes, gnorm, ms, @@ -3330,7 +3672,8 @@ async fn commit_round( // have also been broadcast; the terminal cut is delivered by FINAL. // A crash-resume loses those gathers; their fragments simply merge on // a later cycle, which the quorum design already tolerates. - if !checkpoint_each_round { + // Policy sweeps persist only before broadcast (see persist_before_broadcast). + if cfg.policy_sweep_fragments.is_none() && !checkpoint_before_broadcast(cfg) { if let Some(path) = &cfg.checkpoint_path { if cfg.checkpoint_every > 0 && t % cfg.checkpoint_every == 0 { st.save_checkpoint(path)?; @@ -3341,6 +3684,97 @@ async fn commit_round( Ok(()) } +/// A committed round whose per-round checkpoint is being written by a +/// blocking task (`--async-checkpoint`). Its fragment stays busy and its +/// BCAST is withheld until the write is durable, so no learner can observe +/// a version a crash could lose. +struct PendingPublish { + committed: CommittedRound, + write_started: Instant, + write: tokio::task::JoinHandle>, +} + +/// Snapshot the just-committed state and start persisting it off the +/// scheduler thread. +fn start_async_checkpoint( + cfg: &Config, + st: &GlobalState, + committed: CommittedRound, +) -> Result { + let path = cfg + .checkpoint_path + .clone() + .context("--async-checkpoint requires --checkpoint-path")?; + let write_started = Instant::now(); + let bytes = st.encode_checkpoint()?; + let write = tokio::task::spawn_blocking(move || { + crate::state::write_checkpoint_bytes(&path, &bytes) + }); + Ok(PendingPublish { + committed, + write_started, + write, + }) +} + +/// Wait for a pending per-round checkpoint, then publish its round. A failed +/// write fails the coordinator before the round becomes visible; the +/// previous checkpoint stays the resume point. +async fn settle_pending_publish( + cfg: &Config, + st: &GlobalState, + registry: &Registry, + last_sync_secs: &mut f64, + pending: PendingPublish, + budget_cutoff: &BudgetCutoff, + busy_fragments: &mut HashSet, +) -> Result<()> { + let PendingPublish { + committed, + write_started, + write, + } = pending; + finish_pending_publish( + cfg, + st, + registry, + last_sync_secs, + committed, + write_started, + write.await, + budget_cutoff, + busy_fragments, + ) + .await +} + +#[allow(clippy::too_many_arguments)] +async fn finish_pending_publish( + cfg: &Config, + st: &GlobalState, + registry: &Registry, + last_sync_secs: &mut f64, + committed: CommittedRound, + write_started: Instant, + written: std::result::Result, tokio::task::JoinError>, + budget_cutoff: &BudgetCutoff, + busy_fragments: &mut HashSet, +) -> Result<()> { + let (t, p) = (committed.t, committed.p); + written + .context("async checkpoint task panicked")? + .with_context(|| format!("async checkpoint for step {t} failed; step {t} not broadcast"))?; + let checkpoint_ms = Some(write_started.elapsed().as_millis() as u64); + if let Some(path) = cfg.checkpoint_path.as_ref() { + info!(step = t, path = %path.display(), "checkpoint committed"); + } + publish_round(cfg, st, registry, last_sync_secs, committed, checkpoint_ms, budget_cutoff).await?; + if !busy_fragments.remove(&p) { + bail!("fragment {p} lost its busy ownership at publish"); + } + Ok(()) +} + #[cfg(test)] async fn complete_round( cfg: &Config, @@ -3361,10 +3795,38 @@ async fn complete_round( commit_round(cfg, st, registry, last_sync_secs, completed, &cutoff).await } +/// Test twin of `complete_round` through the `--async-checkpoint` path. +#[cfg(test)] +async fn complete_round_async( + cfg: &Config, + st: &mut GlobalState, + registry: &Registry, + last_sync_secs: &mut f64, + round: Round, +) -> Result<()> { + let prepared = prepare_round_compute(cfg, st, &round)?; + let compute_start = Instant::now(); + let merge = prepared.compute().await?; + let completed = ComputedRound { + round, + merge, + compute_secs: compute_start.elapsed().as_secs_f64(), + }; + let cutoff = BudgetCutoff::new(false, Duration::ZERO); + let mut busy = HashSet::from([completed.round.p]); + let committed = apply_round_commit(cfg, st, completed)?; + if !persist_before_broadcast(cfg, committed.t) { + return publish_round(cfg, st, registry, last_sync_secs, committed, None, &cutoff).await; + } + let pending = start_async_checkpoint(cfg, st, committed)?; + settle_pending_publish(cfg, st, registry, last_sync_secs, pending, &cutoff, &mut busy).await +} + fn new_state_for(group: &Arc, cfg: &Config) -> Result { if let Some(fragments) = cfg.policy_sweep_fragments { - if group.dtype != DTYPE_F32 { - bail!("--policy-sweep-fragments requires an FP32 HELLO session"); + // Masters and every BCAST/FINAL stay f32; only push deltas may be bf16. + if bulk_dtype(group.dtype) != DTYPE_F32 { + bail!("--policy-sweep-fragments requires an FP32 HELLO session (optionally with bf16/int8 pushes)"); } if group.layout.fragments.len() != fragments as usize { bail!( @@ -3383,6 +3845,17 @@ fn new_state_for(group: &Arc, cfg: &Config) -> Result { )?; st.policy_sweep_fragments = cfg.policy_sweep_fragments; st.session_contract_hash = Some(group.session_contract_hash); + match (bcast_delta(group.dtype), cfg.bcast_delta_full_every) { + (true, Some(full_every)) => { + st.downlink = Some(Mutex::new(crate::downlink::DownlinkAnchors::new( + group.layout.fragments.len(), + full_every, + )?)); + } + (true, None) => bail!("session asks for bf16 downlink deltas, but --bcast-delta-full-every is unset"), + (false, Some(_)) => bail!("--bcast-delta-full-every is set, but the session has no bf16 downlink deltas"), + (false, None) => {} + } if cfg.delta_correction { st.delta_correction = Some(crate::merge::Heloco::default()); } @@ -3721,6 +4194,11 @@ async fn send_fragment( msg_type: u8, dtype: u8, ) -> Result<()> { + if let Some(encoded) = recovery_bcast(st, p, msg_type) { + return group + .send_tensor_large(msg_type, &encoded.prefix, encoded.dtype, &encoded.values) + .await; + } let mut prefix = [0u8; 12]; prefix[..4].copy_from_slice(&(p as u32).to_le_bytes()); prefix[4..].copy_from_slice(&st.versions[p].to_le_bytes()); @@ -3856,6 +4334,7 @@ fn append_tape( quorum_ms: Option, grace_ms: Option, sync_ms: u64, + checkpoint_ms: Option, pushes: &HashMap, gnorm: f64, ms: u64, @@ -3932,9 +4411,10 @@ fn append_tape( .map(|byte| format!("{byte:02x}")) .collect(); let merge_seconds = sync_ms as f64 / 1000.0; + let checkpoint_field = checkpoint_tape_field(checkpoint_ms); let sweep_fields = policy_sweep_event_fields(step, policy_sweep_fragments); let line = format!( - "{{\"protocol_version\":{PROTOCOL_VERSION},\"delta_semantics\":\"local_minus_raw_anchor\",\"sync/layout_hash\":\"{layout_hash}\",\"sync/base_version\":{launch_base_version},\"sync/responders\":{},\"sync/quorum\":{quorum},\"sync/rejected_stale_updates\":0,\"sync/merge_seconds\":{merge_seconds},\"sync/global_delta_norm\":{gnorm},\"step\":{step},\"fragment\":{fragment}{sweep_fields},\"launch_base_version\":{launch_base_version},\"attempt\":{attempt},\"gnorm\":{gnorm},\"ms\":{ms},\"quorum\":{quorum},\"expected\":{},\"expected_members\":{},\"responded\":{},\"responded_members\":{},\"missed_grace\":{},\"missed_members\":{},\"quorum_ms\":{quorum_ms},\"grace_ms\":{grace_ms},\"sync_ms\":{sync_ms},\"time_unix\":{:?},\"responders\":[{}]}}\n", + "{{\"protocol_version\":{PROTOCOL_VERSION},\"delta_semantics\":\"local_minus_raw_anchor\",\"sync/layout_hash\":\"{layout_hash}\",\"sync/base_version\":{launch_base_version},\"sync/responders\":{},\"sync/quorum\":{quorum},\"sync/rejected_stale_updates\":0,\"sync/merge_seconds\":{merge_seconds},\"sync/global_delta_norm\":{gnorm},\"step\":{step},\"fragment\":{fragment}{sweep_fields},\"launch_base_version\":{launch_base_version},\"attempt\":{attempt},\"gnorm\":{gnorm},\"ms\":{ms},\"quorum\":{quorum},\"expected\":{},\"expected_members\":{},\"responded\":{},\"responded_members\":{},\"missed_grace\":{},\"missed_members\":{},\"quorum_ms\":{quorum_ms},\"grace_ms\":{grace_ms},\"sync_ms\":{sync_ms}{checkpoint_field},\"time_unix\":{:?},\"responders\":[{}]}}\n", responded.len(), json_ids(&expected_learners), json_members(expected_members), @@ -3954,6 +4434,12 @@ fn append_tape( Ok(()) } +/// Optional top-level tape field: wall time of the per-round checkpoint +/// (synchronous save, or snapshot + async write until durable). +fn checkpoint_tape_field(checkpoint_ms: Option) -> String { + checkpoint_ms.map_or_else(String::new, |ms| format!(",\"checkpoint_ms\":{ms}")) +} + fn json_ids(ids: &[u32]) -> String { let mut ids = ids.to_vec(); ids.sort_unstable(); @@ -4014,6 +4500,28 @@ mod hello_auth_tests { out } + #[test] + fn bw_probe_body_is_two_bounded_sizes() { + let mut body = Vec::new(); + body.extend_from_slice(&5u64.to_le_bytes()); + body.extend_from_slice(&7u64.to_le_bytes()); + assert_eq!(parse_bw_probe(&body).unwrap(), (5, 7)); + assert!(parse_bw_probe(&body[..15]).is_err()); + let mut too_big = Vec::new(); + too_big.extend_from_slice(&(BW_PROBE_MAX_BYTES + 1).to_le_bytes()); + too_big.extend_from_slice(&0u64.to_le_bytes()); + assert!(parse_bw_probe(&too_big).is_err()); + } + + #[test] + fn bw_probe_slots_are_bounded() { + let slots: Vec<_> = (0..BW_PROBE_MAX_CONNECTIONS).map(|_| BwProbeSlot::acquire()).collect(); + assert!(slots.iter().all(Option::is_some)); + assert!(BwProbeSlot::acquire().is_none()); + drop(slots); + assert!(BwProbeSlot::acquire().is_some()); + } + #[test] fn open_hello_accepts_the_right_key_only() { let body = b"hello body bytes"; @@ -4354,6 +4862,8 @@ mod tests { final_state: None, checkpoint_path: None, checkpoint_every: 0, + async_checkpoint: false, + bcast_delta_full_every: None, resume: false, mark_final_checkpoint: false, learner_budget_steps: None, @@ -4553,7 +5063,7 @@ mod tests { } let mut stream = stream.expect("legacy syncer did not start"); // A correctly sealed elastic payload (epoch 0) still must not be - // accepted: legacy only knows HELLO/DATA_HELLO as a first frame. + // accepted: legacy only knows HELLO/DATA_HELLO/BW_PROBE as a first frame. let payload = crate::elastic::seal(&key, msg_type, 0u64.to_le_bytes().to_vec()); let mut frame = MAGIC.to_le_bytes().to_vec(); frame.push(msg_type); @@ -4570,7 +5080,7 @@ mod tests { assert_eq!(reply.get(4), Some(&MSG_ERROR), "type {msg_type}"); let text = String::from_utf8_lossy(&reply[13..]); assert!( - text.contains(&format!("first frame must be HELLO/DATA_HELLO, got {msg_type}")), + text.contains(&format!("first frame must be HELLO/DATA_HELLO/BW_PROBE, got {msg_type}")), "type {msg_type}: {text}" ); assert!(!server.is_finished(), "type {msg_type}: server exited"); @@ -4618,10 +5128,30 @@ mod tests { assert!(format!("{:#}", validate_config(&invalid).unwrap_err()) .contains("requires positive --total-steps")); + let mut streamed = sweep_test_config(2, 4); + streamed.pipeline = 2; + streamed.async_checkpoint = true; + assert!(validate_config(&streamed).is_ok()); + for pipeline in [0, 3] { + invalid = sweep_test_config(2, 4); + invalid.pipeline = pipeline; + assert!(format!("{:#}", validate_config(&invalid).unwrap_err()) + .contains("requires --pipeline between 1 and the fragment count")); + } + let mut legacy_async = round_test_config(3); + legacy_async.async_checkpoint = true; + assert!(format!("{:#}", validate_config(&legacy_async).unwrap_err()) + .contains("--async-checkpoint requires --policy-sweep-fragments")); invalid = sweep_test_config(2, 4); - invalid.pipeline = 2; + invalid.async_checkpoint = true; + invalid.learner_budget_steps = Some(1); assert!(format!("{:#}", validate_config(&invalid).unwrap_err()) - .contains("requires --pipeline 1")); + .contains("does not support --learner-budget-steps")); + // Async persistence is not semantic: the profile hash ignores it. + let mut hashed = sweep_test_config(2, 4); + let before = hashed.semantic_profile_hash(); + hashed.async_checkpoint = true; + assert_eq!(before, hashed.semantic_profile_hash()); invalid = sweep_test_config(3, 4); assert!( @@ -4654,13 +5184,21 @@ mod tests { for mutate in [ |config: &mut Config| config.checkpoint_path = None, - |config: &mut Config| config.checkpoint_every = 2, + |config: &mut Config| config.checkpoint_every = 0, + |config: &mut Config| config.checkpoint_every = 3, ] { invalid = sweep_test_config(2, 4); mutate(&mut invalid); assert!(format!("{:#}", validate_config(&invalid).unwrap_err()) .contains("requires --checkpoint-path")); } + // Sweep-aligned checkpoint intervals are accepted. + let mut interval = sweep_test_config(2, 4); + interval.checkpoint_every = 4; + assert!(validate_config(&interval).is_ok()); + assert!(persist_before_broadcast(&interval, 4)); + assert!(!persist_before_broadcast(&interval, 2)); + assert!(persist_before_broadcast(&sweep_test_config(2, 4), 1)); } #[test] @@ -5244,6 +5782,197 @@ mod tests { std::fs::remove_dir_all(&directory).ok(); } + fn unique_test_directory(tag: &str) -> std::path::PathBuf { + let directory = std::env::temp_dir().join(format!( + "yeto-{tag}-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + std::fs::create_dir_all(&directory).unwrap(); + directory + } + + #[test] + fn bf16_push_session_rounds_only_the_delta_within_bf16_error() { + // dtype 4 decodes pushes as bf16 and keeps every bulk payload f32. + assert_eq!(bulk_dtype(DTYPE_F32_PUSH_BF16), DTYPE_F32); + let delta: Vec = (0..4096) + .map(|i| ((i as f32) * 0.37).sin() * 10f32.powi((i % 9) - 4)) + .collect(); + let mut wire = Vec::new(); + encode_tensor(push_dtype(DTYPE_F32_PUSH_BF16), &delta, &mut wire).unwrap(); + assert_eq!(wire.len(), delta.len() * 2); + let mut decoded = Vec::new(); + decode_tensor(push_dtype(DTYPE_F32_PUSH_BF16), &wire, &mut decoded).unwrap(); + for (exact, rounded) in delta.iter().zip(&decoded) { + // bf16 keeps 8 significant bits: round-to-nearest error <= 2^-8 relative. + assert!((exact - rounded).abs() <= exact.abs() * (1.0 / 256.0)); + } + // A merge of bf16 deltas into f32 masters stays within the same bound + // per learner, so a two-learner average is off by at most that much. + let master = 1.0f32; + for (exact, rounded) in delta.iter().zip(&decoded) { + let merged_exact = master + exact; + let merged = master + rounded; + assert!((merged_exact - merged).abs() <= exact.abs() / 256.0 + f32::EPSILON * 2.0); + } + } + + #[tokio::test] + async fn sweep_checkpoint_interval_persists_only_sweep_aligned_steps() { + let directory = unique_test_directory("sweep-checkpoint-interval"); + let learner = member(0, 30); + for asynchronous in [false, true] { + let checkpoint = directory.join(format!("state-{asynchronous}.ckpt")); + let (group, mut frames) = streaming_test_group(DTYPE_F32, false); + let registry = registry_with_group(group); + let mut config = sweep_test_config(2, 8); + config.checkpoint_path = Some(checkpoint.clone()); + config.checkpoint_every = 4; + config.async_checkpoint = asynchronous; + let mut state = two_fragment_sweep_state(); + state.save_checkpoint(&checkpoint).unwrap(); // version zero + let mut sync_seconds = 0.0; + for step in 1..=6u64 { + let fragment = ((step - 1) % 2) as usize; + let mut round = sweep_test_round(step, fragment, step.div_ceil(2), 1, learner); + let base = state.versions[fragment]; + round.base_version = base; + for push in round.pushes.values_mut() { + push.base_version = base; + } + if asynchronous { + complete_round_async(&config, &mut state, ®istry, &mut sync_seconds, round) + .await + .unwrap(); + } else { + complete_round(&config, &mut state, ®istry, &mut sync_seconds, round) + .await + .unwrap(); + } + assert_eq!( + chunked_inner_type(&frames.try_recv().unwrap()), + MSG_BCAST_FRAGMENT + ); + let mut on_disk = two_fragment_sweep_state(); + on_disk.load_checkpoint(&checkpoint).unwrap(); + // Only step 4 (a sweep boundary and a multiple of 4) is durable. + let expected = if step >= 4 { 4 } else { 0 }; + assert_eq!(on_disk.global_step, expected, "step {step}"); + } + // The newest checkpoint resumes exactly at its sweep boundary. + let mut resumed = two_fragment_sweep_state(); + resumed.load_checkpoint(&checkpoint).unwrap(); + validate_resumed_policy_sweep(&config, &resumed).unwrap(); + assert_eq!(resumed.versions, vec![3, 4]); + } + std::fs::remove_dir_all(&directory).ok(); + } + + #[tokio::test] + async fn async_checkpoint_matches_the_synchronous_path_bit_for_bit() { + let directory = unique_test_directory("async-checkpoint-equal"); + let learner = member(0, 30); + let mut outputs = Vec::new(); + for asynchronous in [false, true] { + let checkpoint = directory.join(format!("state-{asynchronous}.ckpt")); + let tape = directory.join(format!("tape-{asynchronous}.jsonl")); + let (group, mut frames) = streaming_test_group(DTYPE_F32, false); + let registry = registry_with_group(group); + let mut config = round_test_config(4); + config.pipeline = 2; + config.checkpoint_path = Some(checkpoint.clone()); + config.checkpoint_every = 1; + config.async_checkpoint = asynchronous; + config.event_tape = Some(tape.clone()); + let mut state = two_fragment_streaming_state(); + let mut sync_seconds = 0.0; + for step in 1..=2u64 { + let round = sweep_test_round(step, (step - 1) as usize, 1, step, learner); + if asynchronous { + complete_round_async(&config, &mut state, ®istry, &mut sync_seconds, round) + .await + .unwrap(); + } else { + complete_round(&config, &mut state, ®istry, &mut sync_seconds, round) + .await + .unwrap(); + } + // Every non-terminal round is broadcast exactly once. + assert_eq!( + chunked_inner_type(&frames.try_recv().unwrap()), + MSG_BCAST_FRAGMENT + ); + } + let tape_text = std::fs::read_to_string(&tape).unwrap(); + assert_eq!(tape_text.lines().count(), 2); + assert!(tape_text.lines().all(|line| line.contains("\"checkpoint_ms\":"))); + outputs.push((state.params.clone(), std::fs::read(&checkpoint).unwrap())); + // The snapshot encoder writes exactly the streaming file format. + let mut on_disk = two_fragment_streaming_state(); + on_disk.load_checkpoint(&checkpoint).unwrap(); + assert_eq!(on_disk.encode_checkpoint().unwrap(), std::fs::read(&checkpoint).unwrap()); + } + assert_eq!(outputs[0].0, outputs[1].0); + assert_eq!(outputs[0].1, outputs[1].1); + std::fs::remove_dir_all(&directory).ok(); + } + + #[tokio::test] + async fn async_checkpoint_failure_withholds_broadcast_and_keeps_previous_checkpoint() { + let directory = unique_test_directory("async-checkpoint-failure"); + let checkpoint = directory.join("state.ckpt"); + let learner = member(0, 30); + let (group, mut frames) = streaming_test_group(DTYPE_F32, false); + let registry = registry_with_group(group); + let mut config = round_test_config(4); + config.pipeline = 2; + config.checkpoint_path = Some(checkpoint.clone()); + config.checkpoint_every = 1; + config.async_checkpoint = true; + let mut state = two_fragment_streaming_state(); + let mut sync_seconds = 0.0; + complete_round_async( + &config, + &mut state, + ®istry, + &mut sync_seconds, + sweep_test_round(1, 0, 1, 1, learner), + ) + .await + .unwrap(); + assert_eq!( + chunked_inner_type(&frames.try_recv().unwrap()), + MSG_BCAST_FRAGMENT + ); + // The atomic writer creates this exact temporary path; a directory + // there makes the next write fail deterministically. + std::fs::create_dir(checkpoint.with_extension("tmp")).unwrap(); + let error = complete_round_async( + &config, + &mut state, + ®istry, + &mut sync_seconds, + sweep_test_round(2, 1, 1, 1, learner), + ) + .await + .unwrap_err(); + assert!(format!("{error:#}").contains("step 2 not broadcast")); + assert!(matches!( + frames.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Empty) + )); + // The resume point is still the last durable, broadcast cut. + let mut recovered = two_fragment_streaming_state(); + recovered.load_checkpoint(&checkpoint).unwrap(); + assert_eq!(recovered.global_step, 1); + assert_eq!(recovered.versions, vec![1, 0]); + std::fs::remove_dir_all(&directory).ok(); + } + #[tokio::test] async fn policy_sweep_checkpoint_resume_mid_sweep_does_not_double_account() { let directory = std::env::temp_dir().join(format!( @@ -5410,13 +6139,13 @@ mod tests { #[tokio::test] async fn updated_broadcast_propagates_local_errors_but_tolerates_closed_streams() { - let invalid_state = broadcast_test_state(255); + let invalid_state = broadcast_test_state(9); let (connected, _receiver) = streaming_test_group(DTYPE_F32, false); let error = broadcast_updated_fragment(&invalid_state, ®istry_with_group(connected), 0) .await .unwrap_err(); assert!( - format!("{error:#}").contains("unknown tensor dtype 255"), + format!("{error:#}").contains("unknown tensor dtype 9"), "unexpected error: {error:#}" ); @@ -6227,6 +6956,7 @@ mod tests { Some(11), Some(22), 33, + None, &pushes, 0.5, 44, @@ -6275,6 +7005,7 @@ mod tests { Some(1), Some(0), 2, + None, &pushes, 0.5, 3, @@ -6303,6 +7034,7 @@ mod tests { Some(1), Some(0), 2, + None, &pushes, 0.5, 3, @@ -6410,6 +7142,7 @@ mod tests { Some(1), Some(0), 2, + None, &pushes, 0.5, 3, diff --git a/syncer/src/state.rs b/syncer/src/state.rs index f19ce4c9..3b54cb9e 100644 --- a/syncer/src/state.rs +++ b/syncer/src/state.rs @@ -146,6 +146,8 @@ pub struct GlobalState { pub delta_correction: Option, /// Spectrum flattener. Torch mode owns one persistent child process. iso_backend: IsoBackend, + /// bf16 downlink anchors (S26 E3); in memory only, never checkpointed. + pub downlink: Option>, } /// A merge whose deterministic learner reduction has completed but whose @@ -344,6 +346,7 @@ impl GlobalState { wire_dtype, delta_correction: None, iso_backend: IsoBackend::start(iso_backend_config)?, + downlink: None, }) } @@ -628,6 +631,36 @@ pub fn remove_final_marker(path: &std::path::Path) -> Result<()> { } } +/// Little-endian f32 values in 64 KiB blocks (same bytes as per-value +/// `to_le_bytes`, without one writer call per element). +fn write_f32s(f: &mut W, values: &[f32]) -> Result<()> { + let mut block = Vec::with_capacity(16 * 1024 * 4); + for chunk in values.chunks(16 * 1024) { + block.clear(); + for v in chunk { + block.extend_from_slice(&v.to_le_bytes()); + } + f.write_all(&block)?; + } + Ok(()) +} + +/// Atomically persist checkpoint bytes from `encode_checkpoint`: write a +/// sibling temporary file, fsync it, rename it over `path`, fsync the +/// directory. A failure at any point leaves the previous checkpoint intact. +pub fn write_checkpoint_bytes(path: &std::path::Path, bytes: &[u8]) -> Result<()> { + use std::io::Write; + let tmp = path.with_extension("tmp"); + { + let mut f = std::fs::File::create(&tmp) + .with_context(|| format!("create checkpoint temporary {}", tmp.display()))?; + f.write_all(bytes)?; + f.sync_all()?; + } + std::fs::rename(&tmp, path)?; + sync_parent_directory(path) +} + fn sync_parent_directory(path: &std::path::Path) -> Result<()> { let parent = path.parent().unwrap_or_else(|| std::path::Path::new(".")); std::fs::File::open(parent) @@ -660,47 +693,11 @@ impl GlobalState { /// Written to `.tmp` then renamed, so a crash mid-write never /// corrupts the previous checkpoint. pub fn save_checkpoint(&self, path: &std::path::Path) -> Result<()> { - use std::io::Write; - let layout_fingerprint = self.layout_fingerprint; let tmp = path.with_extension("tmp"); { let mut f = std::io::BufWriter::new(std::fs::File::create(&tmp)?); - f.write_all(&CKPT_MAGIC_V3.to_le_bytes())?; - f.write_all(&[self.iso_backend_kind() as u8])?; - f.write_all(&layout_fingerprint)?; - f.write_all(&self.global_step.to_le_bytes())?; - f.write_all(&(self.params.len() as u32).to_le_bytes())?; - for p in 0..self.params.len() { - f.write_all(&self.versions[p].to_le_bytes())?; - f.write_all(&(self.params[p].len() as u64).to_le_bytes())?; - for v in &self.params[p] { - f.write_all(&v.to_le_bytes())?; - } - for v in &self.momentum[p] { - f.write_all(&v.to_le_bytes())?; - } - } - f.write_all(&(self.ledger.len() as u32).to_le_bytes())?; - for (id, l) in &self.ledger { - f.write_all(&id.to_le_bytes())?; - f.write_all(&l.merges.to_le_bytes())?; - f.write_all(&l.steps.to_le_bytes())?; - f.write_all(&l.tokens.to_le_bytes())?; - } - f.write_all(&self.layout_fingerprint)?; - if let Some(fragments) = self.policy_sweep_fragments { - f.write_all(&POLICY_SWEEP_CKPT_MAGIC.to_le_bytes())?; - f.write_all(&fragments.to_le_bytes())?; - f.write_all( - &self - .session_contract_hash - .context("policy-sweep checkpoint requires a session contract hash")?, - )?; - } else if let Some(contract_hash) = self.session_contract_hash { - f.write_all(&SESSION_CONTRACT_CKPT_MAGIC.to_le_bytes())?; - f.write_all(&contract_hash)?; - } - f.flush()?; + self.write_checkpoint_to(&mut f)?; + std::io::Write::flush(&mut f)?; f.get_ref().sync_all()?; } std::fs::rename(&tmp, path)?; @@ -708,6 +705,54 @@ impl GlobalState { Ok(()) } + /// The exact bytes `save_checkpoint` writes, as an owned snapshot. + /// + /// The asynchronous checkpoint path encodes on the scheduler (one memory + /// copy of the consistent cut) and hands the bytes to a blocking writer, + /// so later commits may mutate `self` while the snapshot is persisted. + pub fn encode_checkpoint(&self) -> Result> { + let state_floats: usize = self.params.iter().map(Vec::len).sum::() * 2; + let mut bytes = Vec::with_capacity(state_floats * 4 + 4096); + self.write_checkpoint_to(&mut bytes)?; + Ok(bytes) + } + + fn write_checkpoint_to(&self, f: &mut W) -> Result<()> { + let layout_fingerprint = self.layout_fingerprint; + f.write_all(&CKPT_MAGIC_V3.to_le_bytes())?; + f.write_all(&[self.iso_backend_kind() as u8])?; + f.write_all(&layout_fingerprint)?; + f.write_all(&self.global_step.to_le_bytes())?; + f.write_all(&(self.params.len() as u32).to_le_bytes())?; + for p in 0..self.params.len() { + f.write_all(&self.versions[p].to_le_bytes())?; + f.write_all(&(self.params[p].len() as u64).to_le_bytes())?; + write_f32s(f, &self.params[p])?; + write_f32s(f, &self.momentum[p])?; + } + f.write_all(&(self.ledger.len() as u32).to_le_bytes())?; + for (id, l) in &self.ledger { + f.write_all(&id.to_le_bytes())?; + f.write_all(&l.merges.to_le_bytes())?; + f.write_all(&l.steps.to_le_bytes())?; + f.write_all(&l.tokens.to_le_bytes())?; + } + f.write_all(&self.layout_fingerprint)?; + if let Some(fragments) = self.policy_sweep_fragments { + f.write_all(&POLICY_SWEEP_CKPT_MAGIC.to_le_bytes())?; + f.write_all(&fragments.to_le_bytes())?; + f.write_all( + &self + .session_contract_hash + .context("policy-sweep checkpoint requires a session contract hash")?, + )?; + } else if let Some(contract_hash) = self.session_contract_hash { + f.write_all(&SESSION_CONTRACT_CKPT_MAGIC.to_le_bytes())?; + f.write_all(&contract_hash)?; + } + Ok(()) + } + /// Restore params/momentum/versions/step/ledger from a snapshot. /// /// V3 binds the snapshot to both the ISO backend and exact semantic HELLO diff --git a/tests/test_e3_overlap_miles_role_stream.py b/tests/test_e3_overlap_miles_role_stream.py new file mode 100644 index 00000000..6a69f85f --- /dev/null +++ b/tests/test_e3_overlap_miles_role_stream.py @@ -0,0 +1,504 @@ +"""S26 E3: sync/train overlap on the real Miles role stream (CPU tensors, no GPU). + +``tests/test_e3_sync_overlap.py`` drives ``FullParameterSync`` with a numeric +fake role. Here the coordinator drives the real ``MilesFullParameterRoleStream`` +and ``MilesChunkedFullParameterAdapter`` (``rebase_fragment``, held +submissions, ``finish_rebase``, ``mark_round_base``, retained replays). Only +the Megatron trainer group, Ray and the syncer connection are fakes: + +* the trainer keeps one f32 tensor per fragment and adds a fixed progress per + round, so every pushed delta must equal that round's progress exactly; +* the syncer merges the island's delta with a second, virtual island's delta + (plain average) so the global weights differ from the local ones and a + wrong rebase shows up in the numbers. + +The syncer-restart tests cut the run while the island trains the overlapped +round: the syncer resumes from its last checkpoint (uncommitted pushes lost, +PULLs re-sent with a new attempt, committed fragments re-broadcast) and the +run must still count every round exactly once. + +Needs Miles (``importorskip``): run it in the Miles image. +""" + +from __future__ import annotations + +# ruff: noqa: I001 -- the Miles test helpers install a bridge stub before import. + +import asyncio +import threading +from types import SimpleNamespace + +import pytest +import torch + +pytest.importorskip("miles") + +from test_rl_miles_chunked_full_parameter import FakeRay, _evidence, component, manifest # noqa: E402 + +from miles.backends.megatron_utils.full_parameter_state import ( # noqa: E402 + FullParameterChunkDescriptor, + FullParameterChunkedShardState, + FullParameterFragmentDescriptor, + _chunk_hash, + _fragment_hasher, + _tensor_bytes, +) +from miles.ray.full_parameter_transport import FullParameterChunkedCut # noqa: E402 +from yeto.protocol import BcastFragment, FinalFragment, FinalManifest, PullRequest # noqa: E402 +from yeto.rl.adapters.miles.full_parameter import role_stream as role_stream_module # noqa: E402 +from yeto.rl.adapters.miles.full_parameter.role_stream import ( # noqa: E402 + MilesFullParameterRoleStream, + RoleStreamConfig, +) +from yeto.rl.adapters.miles.models.chunked_full_parameter import ( # noqa: E402 + MilesChunkedFullParameterAdapter, +) +from yeto.rl.engine.full_parameter_sync import ( # noqa: E402 + FullParameterSync, + FullParameterSyncProfile, + RolePairing, +) + +P = 2 +SIZES = (3, 2) # one parameter per owner shard; one fragment per shard +OTHER_WEIGHT = 0.5 # plain average of this island and one virtual peer + + +def _adapter() -> MilesChunkedFullParameterAdapter: + return MilesChunkedFullParameterAdapter.create( + (manifest(0, (SIZES[0],)), manifest(1, (SIZES[1],))), + algorithm="grpo", + components=(component(),), + minimum_fragments=P, + max_fragment_bytes=64, + max_chunk_bytes=8, # two floats per chunk: fragment 0 spans two chunks + ) + + +def _build_cut(adapter, values, *, policy_version, generation, ray) -> FullParameterChunkedCut: + """A reference-backed Miles cut holding ``values[fragment_id]``.""" + shards = [] + for plan in adapter.plans: + descriptors, refs = [], [] + for fragment in plan.fragments: + value = values[fragment.fragment_id].to(torch.float32).contiguous() + assert value.numel() == fragment.numel + digest = _fragment_hasher(fragment.fragment_id, adapter.layout.layout_hash, plan.plan_hash) + chunks, chunk_refs = [], [] + step = adapter.max_chunk_bytes // 4 + for index, start in enumerate(range(0, value.numel(), step)): + chunk = value[start:start + step].clone().contiguous() + digest.update(_tensor_bytes(chunk)) + chunks.append(FullParameterChunkDescriptor( + index, start, chunk.numel(), + _chunk_hash(fragment.fragment_id, index, start, chunk, + parameter_layout_hash=adapter.layout.layout_hash, + plan_hash=plan.plan_hash))) + chunk_refs.append(ray.put(chunk)) + descriptors.append(FullParameterFragmentDescriptor( + fragment.fragment_id, value.numel(), digest.hexdigest(), tuple(chunks))) + refs.append(chunk_refs) + shards.append(FullParameterChunkedShardState( + policy_version, generation, plan.topology, adapter.layout.layout_hash, + plan.plan_hash, tuple(descriptors), refs)) + return FullParameterChunkedCut(policy_version, generation, adapter.layout.layout_hash, + tuple(sorted(shards, key=lambda state: state.topology))) + + +class _Trainer: + """The Megatron trainer group: one f32 tensor per fragment, fixed progress per round.""" + + def __init__(self, adapter, ray, initial, progress) -> None: + self.adapter = adapter + self.ray = ray + self.values = {f: value.clone() for f, value in initial.items()} + self.progress = progress + self.installed = 0 + self.generation = 0 + self.scheduler_steps = 0 + self.trained_from: list[dict[int, torch.Tensor]] = [] + self.on_train = None # hook: runs while a round trains (cut point) + + async def record_full_parameter_local_step(self, *, base_policy_version, rollout_id): + assert base_policy_version == self.installed + self.trained_from.append({f: v.clone() for f, v in self.values.items()}) + for f, delta in self.progress(rollout_id).items(): + self.values[f] = self.values[f] + delta + self.generation += 1 + start, self.scheduler_steps = self.scheduler_steps, self.scheduler_steps + 1 + if self.on_train is not None: + self.on_train(rollout_id) + return tuple( + SimpleNamespace(topology=shard.topology, role="actor", base_policy_version=base_policy_version, + local_step_generation=self.generation, rollout_id=rollout_id, + optimizer_steps=1, scheduler_start_steps=start, + scheduler_end_steps=self.scheduler_steps) + for shard in self.adapter.manifests) + + async def export_full_parameter_chunked_cut(self, policy_version, local_step_generation=0, *, + max_chunk_bytes): + assert max_chunk_bytes == self.adapter.max_chunk_bytes + return _build_cut(self.adapter, self.values, policy_version=policy_version, + generation=local_step_generation, ray=self.ray) + + async def apply_full_parameter_chunked_cut(self, cut, *, commit_token=None, max_commit_attempts=3): + assert commit_token + for shard in cut.shards: + for descriptor, refs in zip(shard.fragments, shard.chunk_refs, strict=True): + assert all(ref is not None for ref in refs), "applied a released chunk" + self.values[descriptor.fragment_id] = torch.cat([self.ray.get(ref) for ref in refs]).clone() + self.installed = cut.policy_version + self.generation = 0 + return self.adapter.expected_parameter_tensor_count + + async def full_parameter_optimizer_states(self): + return tuple(SimpleNamespace(role="actor", installed_policy_version=self.installed, + local_step_generation=self.generation, + scheduler_num_steps=self.scheduler_steps) + for _ in self.adapter.manifests) + + +class _Syncer: + """In-process strict sweep (``--pipeline P``, checkpoint every round). + + ``G_f <- G_f + OTHER_WEIGHT * (island delta + peer delta)``; committed state + is checkpointed at once, so a restart loses only uncommitted pushes.""" + + def __init__(self, initial, sweeps, peer) -> None: + self.global_ = {f: v.clone() for f, v in initial.items()} + self.versions = [0] * P + self.sweeps = sweeps + self.peer = peer # rollout_id -> {fragment: peer delta} + self.received: dict[int, torch.Tensor] = {} # uncommitted pushes by step + self.pushed: list[tuple[int, int, torch.Tensor]] = [] # (step, attempt, delta), every PUSH seen + self.committed: list[int] = [] + self.attempt = 1 + self.client = None + + def commit_ready(self) -> None: + """Commit fragments in strict step order once their PUSH arrived.""" + while True: + step = len(self.committed) + 1 + if step > self.sweeps * P or step not in self.received: + return + f, k = (step - 1) % P, (step - 1) // P + delta = self.received.pop(step) + self.global_[f] = self.global_[f] + OTHER_WEIGHT * (delta + self.peer(k)[f]) + self.versions[f] = step + self.committed.append(step) + if step == self.sweeps * P: + self.client.publish_final() + elif step % P == 0: + # The sweep is committed: broadcast it, then PULL the next one. + for fragment in range(P): + self.client.broadcast(fragment, self.versions[fragment], self.global_[fragment]) + self.client.pull_sweep(step // P, self.attempt) + + def restart_from_checkpoint(self) -> None: + """Crash + ``--resume``: uncommitted pushes are lost; the reconnect + re-broadcasts every committed fragment and re-sends the open PULLs.""" + self.received.clear() + self.attempt += 1 + self.client.reconnect() + for fragment in range(P): + if self.versions[fragment]: + self.client.broadcast(fragment, self.versions[fragment], self.global_[fragment]) + if len(self.committed) % P == 0 and len(self.committed) < self.sweeps * P: + self.client.pull_sweep(len(self.committed) // P, self.attempt) + else: + # A partly committed sweep: only its uncommitted fragments are open. + sweep = len(self.committed) // P + for step in range(len(self.committed) + 1, (sweep + 1) * P + 1): + self.client.pulls.append(PullRequest((step - 1) % P, step, self.attempt)) + + +class _Client: + """The ``SyncerClient`` surface the role stream uses, backed by ``_Syncer``.""" + + def __init__(self, syncer: _Syncer) -> None: + self.syncer = syncer + syncer.client = self + self.connection_generation = 1 + self.finalizing = threading.Event() + self.pulls: list[PullRequest] = [] + self.updates: list[BcastFragment] = [] + self.inbound = None + self.final: tuple[FinalManifest, list[FinalFragment]] | None = None + self.acknowledged: FinalManifest | None = None + self.closed = False + + # -- syncer side + def _stored(self, fragment_id, version, value): + return self.inbound.sink(fragment_id, version, _tensor_bytes(value.to(torch.float32).contiguous())) + + def broadcast(self, fragment_id, version, value) -> None: + self.updates.append(BcastFragment(fragment_id, version, self._stored(fragment_id, version, value), + stored=True)) + + def pull_sweep(self, sweep, attempt) -> None: + self.pulls.extend(PullRequest(f, sweep * P + f + 1, attempt) for f in range(P)) + + def reconnect(self) -> None: + self.connection_generation += 1 + self.pulls.clear() + + def publish_final(self) -> None: + last = self.syncer.sweeps * P + versions = tuple(range(last - P + 1, last + 1)) + fragments = [FinalFragment(f, versions[f], self._stored(f, versions[f], self.syncer.global_[f]), + stored=True) for f in range(P)] + self.final = (FinalManifest(last, versions), fragments) + self.finalizing.set() + + # -- island side (SyncerClient API) + def install_inbound_chunk_sink(self, sink) -> None: + self.inbound = sink + + def start(self) -> None: + pass + + def send_init_parts(self, fragment_id, parts) -> bool: + payload = b"".join(bytes(part) for part in parts) + value = torch.frombuffer(bytearray(payload), dtype=torch.float32) + assert torch.equal(value, self.syncer.global_[fragment_id]) + self.broadcast(fragment_id, 0, value) + if fragment_id == P - 1: + self.pull_sweep(0, 1) + return True + + def send_initial_fragments(self, parts_for) -> bool: + return all(self.send_init_parts(f, parts_for(f)) for f in range(P)) + + def init_replay_due(self) -> bool: + return False + + def mark_initial_cut_received(self) -> None: + pass + + def check_health(self) -> None: + assert not self.closed + + def drain_updates(self): + out, self.updates = self.updates, [] + return out + + def drain_pulls(self): + out, self.pulls = self.pulls, [] + return out + + def push_fragment_parts(self, fragment_id, global_step, round_attempt, base_version, local_step, + c_steps, c_tokens, parts) -> bool: + assert (global_step - 1) % P == fragment_id and c_steps == 1 and c_tokens >= 1 + assert base_version == self.syncer.versions[fragment_id], "PUSH on a stale base" + delta = torch.frombuffer(bytearray(b"".join(bytes(part) for part in parts)), dtype=torch.float32) + self.syncer.pushed.append((global_step, round_attempt, delta.clone())) + self.syncer.received[global_step] = delta.clone() + self.syncer.commit_ready() + return True + + def wait_for_final_fragments(self, timeout): + assert self.final is not None, "no FINAL before the terminal barrier" + return self.final + + def acknowledge_finalization(self, manifest) -> None: + self.acknowledged = manifest + + def close(self) -> None: + self.closed = True + + +class _Driver: + def __init__(self) -> None: + self.events = [] + + def emit(self, event, **fields) -> None: + self.events.append((event, fields)) + + +def _config(sweeps: int) -> RoleStreamConfig: + adapter = _adapter() + return RoleStreamConfig( + algorithm="grpo", role="actor", component=component(), syncer_addr=("127.0.0.1", 30100), + learner_id=0, learner_generation=0, learner_generations=(0, 0), + total_fragment_steps=sweeps * P, expected_fragments=P, + expected_layout_hash=adapter.layout.layout_hash, local_horizon=1, optimizer_steps_per_round=1, + training_contract_hash="2" * 64, syncer_profile_hash="3" * 64, pipeline_depth=1, + algorithm_spec_sha256="a1" * 32, wait_timeout=5.0, poll_seconds=0.001, + max_fragment_bytes=64, max_chunk_bytes=8, overlap=True) + + +def _run(monkeypatch, sweeps: int, *, cut=None): + """Run ``sweeps`` overlapped rounds; ``cut(rollout_id, syncer)`` runs while + that round trains (the mid-overlap cut point).""" + generator = torch.Generator().manual_seed(11) + + def rand(scale=1.0): + return {f: torch.randn(n, generator=generator) * scale for f, n in enumerate(SIZES)} + + initial = rand() + steps = {k: rand(0.1) for k in range(sweeps)} + peer = {k: rand(0.1) for k in range(sweeps)} + ray = FakeRay() + adapter = _adapter() + trainer = _Trainer(adapter, ray, initial, lambda k: steps[k]) + syncer = _Syncer(initial, sweeps, lambda k: peer[k]) + client = _Client(syncer) + + async def initialize(cls, group, *, policy_version, **_kwargs): + cut = await group.export_full_parameter_chunked_cut(policy_version, 0, + max_chunk_bytes=adapter.max_chunk_bytes) + return adapter, adapter.validate_cut(cut, policy_version, 0) + + monkeypatch.setattr(MilesChunkedFullParameterAdapter, "initialize", classmethod(initialize)) + monkeypatch.setattr(role_stream_module, "SyncerClient", lambda *_args, **_kwargs: client) + + stream = MilesFullParameterRoleStream(_config(sweeps), ray_module=ray) + profile = FullParameterSyncProfile.from_configs({"actor": stream.config}, sweep=True) + assert profile.overlap + session = FullParameterSync( + RolePairing({"actor": stream}, profile), groups={"actor": trainer}, run=asyncio.run, + trajectories=lambda rollout_id: _evidence(rollout_id=rollout_id, + behavior_policy_hash=stream.current.policy_hash)) + if cut is not None: + trainer.on_train = lambda rollout_id: cut(rollout_id, syncer) + driver = _Driver() + session.start(driver) + stops = [] + for rollout_id in range(sweeps): + boundary = session.boundary(driver, rollout_id=rollout_id, stats=None) + stops.append(boundary.stop) + if boundary.stop: + session.published(driver, rollout_id=rollout_id, policy_hash=boundary.state.policy_tensor_hash()) + break + session.finish(driver) + return SimpleNamespace(initial=initial, steps=steps, peer=peer, trainer=trainer, syncer=syncer, + client=client, stream=stream, stops=stops, ray=ray) + + +def _expected_global(run, sweeps): + return {f: run.initial[f] + sum(OTHER_WEIGHT * (run.steps[k][f] + run.peer[k][f]) + for k in range(sweeps)) for f in range(P)} + + +def _committed_deltas(run): + """The delta of each committed step (the last PUSH seen for that step).""" + last = {} + for step, attempt, delta in run.syncer.pushed: + if attempt >= last.get(step, (0, None))[0]: + last[step] = (attempt, delta) + return {step: delta for step, (_attempt, delta) in last.items()} + + +def _check_exactly_once(run, sweeps): + assert run.stops == [False] * (sweeps - 1) + [True] + assert run.syncer.committed == list(range(1, sweeps * P + 1)) + deltas = _committed_deltas(run) + for k in range(sweeps): + for f in range(P): + # Each push is exactly that round's progress: no round is pushed + # twice, none is lost, the rebase kept the local progress. + torch.testing.assert_close(deltas[k * P + f + 1], run.steps[k][f], rtol=0, atol=1e-6) + want = _expected_global(run, sweeps) + for f in range(P): + torch.testing.assert_close(run.syncer.global_[f], want[f], rtol=0, atol=1e-5) + torch.testing.assert_close(run.trainer.values[f], want[f], rtol=0, atol=1e-5) + assert run.stream.final_applied and run.stream.final_acknowledged + assert run.client.acknowledged == run.stream.terminal_manifest + + +def test_overlap_on_the_miles_role_stream_rebases_and_counts_every_round_once(monkeypatch): + sweeps = 4 + run = _run(monkeypatch, sweeps) + _check_exactly_once(run, sweeps) + trained = run.trainer.trained_from + for f in range(P): + # Round 0 starts at the initial weights, round 1 at its own round-0 + # result (sweep 0 in flight), round k >= 2 at G_{k-2} + progress of + # round k-1: the rebase ran in the trainer through rebase_fragment. + torch.testing.assert_close(trained[0][f], run.initial[f]) + torch.testing.assert_close(trained[1][f], run.initial[f] + run.steps[0][f]) + for k in range(2, sweeps): + g = run.initial[f] + sum(OTHER_WEIGHT * (run.steps[j][f] + run.peer[j][f]) for j in range(k - 1)) + torch.testing.assert_close(trained[k][f], g + run.steps[k - 1][f], rtol=0, atol=1e-5) + # Exactly one PUSH per step without a restart. + assert sorted(step for step, _attempt, _delta in run.syncer.pushed) == list(range(1, sweeps * P + 1)) + + +def test_overlap_role_stream_releases_every_superseded_cut(monkeypatch): + run = _run(monkeypatch, 3) + stream = run.stream + assert not stream._pinned_cuts and not stream._pinned_anchors + assert not stream._staged + + +@pytest.mark.parametrize("cut_round", [1, 2]) +def test_syncer_restart_before_the_commit_mid_overlap_replays_the_retained_push(monkeypatch, cut_round): + """Cut while round ``cut_round`` trains: sweep ``cut_round - 1`` was pushed + but not committed. The resumed syncer re-PULLs it with attempt 2; the + island replays the retained bytes at its next boundary.""" + sweeps = 4 + + def cut(rollout_id, syncer): + if rollout_id == cut_round and syncer.attempt == 1: + # Undo the commit of the in-flight sweep: the crash came before it. + sweep = cut_round - 1 + for step in range(sweep * P + 1, (sweep + 1) * P + 1): + syncer.committed.remove(step) + syncer.global_ = {f: g.clone() for f, g in checkpoints[sweep].items()} + syncer.versions = [sweep * P - P + f + 1 if sweep else 0 for f in range(P)] + syncer.client.updates.clear() # its BCAST never left the syncer + syncer.restart_from_checkpoint() + + checkpoints = {} + original = _Syncer.commit_ready + + def commit_ready(self): + if len(self.committed) % P == 0: + checkpoints.setdefault(len(self.committed) // P, {f: g.clone() for f, g in self.global_.items()}) + original(self) + + monkeypatch.setattr(_Syncer, "commit_ready", commit_ready) + run = _run(monkeypatch, sweeps, cut=cut) + _check_exactly_once(run, sweeps) + replayed = [(step, attempt) for step, attempt, _ in run.syncer.pushed if attempt == 2] + first = (cut_round - 1) * P + 1 + assert [step for step, _ in replayed][:P] == list(range(first, first + P)) + + +@pytest.mark.parametrize("cut_round", [1, 2]) +def test_syncer_restart_after_the_commit_mid_overlap_rebases_on_the_resent_broadcast(monkeypatch, cut_round): + """Cut while round ``cut_round`` trains: sweep ``cut_round - 1`` committed + and checkpointed, but its BCAST was lost. The resumed syncer re-sends it + (FULL) and PULLs the next sweep; the island rebases on the re-sent bytes.""" + sweeps = 4 + + def cut(rollout_id, syncer): + if rollout_id == cut_round and syncer.attempt == 1: + syncer.client.updates.clear() # the BCAST was lost with the connection + syncer.restart_from_checkpoint() + + run = _run(monkeypatch, sweeps, cut=cut) + _check_exactly_once(run, sweeps) + assert all(attempt == 1 for step, attempt, _ in run.syncer.pushed if step <= cut_round * P) + + +def test_syncer_restart_with_a_partly_committed_sweep_mid_overlap(monkeypatch): + """Cut while round 2 trains: fragment 0 of sweep 1 committed (and is + re-broadcast), fragment 1 was lost. Only fragment 1 is re-PULLed and + replayed; the island applies the sweep only once both fragments arrived.""" + sweeps = 3 + lost = 2 * P # sweep 1, fragment 1 + + def cut(rollout_id, syncer): + if rollout_id == 2 and syncer.attempt == 1: + f = (lost - 1) % P + pushed = [delta for step, _attempt, delta in syncer.pushed if step == lost][-1] + syncer.committed.remove(lost) + syncer.global_[f] = syncer.global_[f] - OTHER_WEIGHT * (pushed + syncer.peer(1)[f]) + syncer.versions[f] = lost - P + syncer.client.updates.clear() + syncer.restart_from_checkpoint() + + run = _run(monkeypatch, sweeps, cut=cut) + _check_exactly_once(run, sweeps) + assert [step for step, attempt, _ in run.syncer.pushed if attempt == 2][0] == lost diff --git a/tests/test_e3_sync_options.py b/tests/test_e3_sync_options.py new file mode 100644 index 00000000..a3256227 --- /dev/null +++ b/tests/test_e3_sync_options.py @@ -0,0 +1,318 @@ +"""S26 E3 sync options: bf16 push transcoding, session binding, preset scoping. CPU only.""" + +from __future__ import annotations + +import hashlib +import itertools +from dataclasses import replace + +import numpy as np +import pytest +import torch +from test_rl_algorithm_provenance import ( # noqa: F401 + _fake_sky, + _island_run, + _launcher_args, + _no_modal_listing, +) +from test_rl_dense_sweep_real_syncer import syncer_binary # noqa: F401 +from test_rl_full_parameter_sync import _role_stream_config + +from yeto.downlink import ( + HEADER, + KIND_DELTA, + KIND_FULL, + DownlinkAnchors, + DownlinkChunkSink, + DownlinkResync, +) +from yeto.protocol import ( + DTYPE_BF16, + DTYPE_F32, + DTYPE_F32_PUSH_BF16, + DTYPE_F32_PUSH_I8, + DTYPE_FLAG_BCAST_DELTA, + _f32_parts_to_bf16, + _f32_parts_to_i8, + bulk_dtype, + f32_bytes_to_bf16, + f32_bytes_to_i8, + i8_nbytes, + push_wire_dtype, + valid_session_dtype, +) + + +def test_push_bf16_session_keeps_bulk_payloads_f32(): + assert bulk_dtype(DTYPE_F32_PUSH_BF16) == DTYPE_F32 + assert push_wire_dtype(DTYPE_F32_PUSH_BF16) == DTYPE_BF16 + assert push_wire_dtype(DTYPE_F32) == DTYPE_F32 + + +def test_f32_to_bf16_matches_torch_round_to_nearest_even_across_split_parts(): + values = torch.randn(50_001) * torch.logspace(-20, 20, 50_001) + values[:3] = torch.tensor([0.0, -0.0, float("inf")]) + raw = values.numpy().tobytes() + want = values.to(torch.bfloat16).view(torch.int16).numpy().astype(" np.ndarray: + block = int(np.frombuffer(wire, "= 2"), + (("--rl-sync-preset", "decoupled", "--local-rl-rounds-per-sync", "2", + "--rl-sync-bcast-delta-full-every", "4"), "only applies to --rl-sync-preset strict-avg"), + (("--rl-sync-push-dtype", "int8"), "--rl-sync-push-dtype only applies"), +)) +def test_e3_downlink_and_int8_scope(extra, match): + from yeto.launcher import _prepare_rl_args + + args = _launcher_args("ports", extra, gpu="modal:2xh100") + args.controller = "local" + with pytest.raises(ValueError, match=match): + _prepare_rl_args(args) + + +def _run_bridge_pair(binary, tmp_path, rounds, *, full_every): + """Two fake islands through ``StrictRlBridge`` (the verl/Miles LoRA path).""" + import subprocess + + from test_rl_integration import _BridgeThread, _FakeMiles, _port, _start + + from yeto.rl.bridge import BridgeConfig, StrictRlBridge + + port = _port() + process = _start(binary, port, tmp_path / "state.ckpt", rounds=rounds) + if full_every: + process.kill() + process.wait() + argv = [*process.args, "--bcast-delta-full-every", str(full_every)] + process = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) + runtimes = (_FakeMiles([0.1, 0.3], 0), _FakeMiles([0.3, 0.5], 1)) + threads = [] + try: + for learner_id, runtime in enumerate(runtimes): + config = BridgeConfig( + syncer_addr=("127.0.0.1", port), learner_id=learner_id, global_rounds=rounds, + groups_per_round=1, samples_per_group=2, local_optimizer_steps=1, wan_streams=0, + expected_specs=runtime.current.specs, + base_model_revision=runtime.current.base_model_revision, + lora_config_hash=runtime.current.lora_config_hash, + layout_hash=runtime.current.layout_hash, + event_tape=str(tmp_path / f"island-{learner_id}.jsonl"), + bcast_delta=bool(full_every)) + threads.append(_BridgeThread(StrictRlBridge(runtime, config))) + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=30) + assert not thread.is_alive() and thread.error is None, thread.error + assert process.wait(timeout=10) == 0 + return [next(iter(t.result.tensors.values())) for t in threads], runtimes + finally: + for thread in threads: + thread.bridge.client.close() + if process.poll() is None: + process.kill() + process.wait() + + +def test_strict_bridge_with_bf16_downlink_finishes_on_the_exact_f32_cut(syncer_binary, tmp_path): # noqa: F811 + """E3, verl/LoRA path: islands train from the bf16-delta anchor, the syncer + keeps exact masters, and the final cut matches the exact f32 run (up to the + f32 rounding of ``local - anchor`` on a different anchor).""" + binary = syncer_binary + exact_dir, delta_dir = tmp_path / "f32", tmp_path / "delta" + exact_dir.mkdir() + delta_dir.mkdir() + exact, _ = _run_bridge_pair(binary, exact_dir, 4, full_every=0) + delta, runtimes = _run_bridge_pair(binary, delta_dir, 4, full_every=3) + for want, got in zip(exact, delta, strict=True): + assert torch.allclose(want, got, rtol=0, atol=1e-6), (want, got) + assert all(runtime.applied[-1] == 4 for runtime in runtimes) + + +def test_verl_island_accepts_the_downlink_flag(): + from yeto.rl.adapters.verl.island_entry import parse_args as verl_parse_args + + args = verl_parse_args([ + "--model", "m", "--model-revision", "r", "--data", "d", "--learner-id", "0", + "--reward-function", "f", "--global-rounds", "1", "--groups-per-round", "1", + "--samples-per-group", "1", "--rollout-max-response-len", "8", "--event-tape", "t", + "--inner-lr", "1e-6", "--seq-len", "8", "--seed", "0", "--bcast-delta"]) + assert args[0].bcast_delta is True diff --git a/tests/test_e3_sync_overlap.py b/tests/test_e3_sync_overlap.py new file mode 100644 index 00000000..6421d6ce --- /dev/null +++ b/tests/test_e3_sync_overlap.py @@ -0,0 +1,370 @@ +"""S26 E3 item 3: sync/train overlap of the full-parameter sweep (CPU only). + +A numeric fake role stream and an in-process strict-sweep syncer drive the +real ``FullParameterSync``/``RolePairing``: the island trains round k+1 while +sweep k is in flight and rebases ``theta <- G_k + (local - base)``. The tests +check that every round's progress reaches the global weights exactly once +(no double counting), that the trainer keeps its progress across a rebase, +and the launcher rules (rollout policy age 0, fragment pipeline P). +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import numpy as np +import pytest +from test_rl_algorithm_provenance import ( # noqa: F401 + _fake_sky, + _launcher_args, + _no_modal_listing, +) + +from yeto.protocol import PullRequest +from yeto.rl.adapters.miles.full_parameter.role_stream import RoleSubmissionCandidate +from yeto.rl.engine.full_parameter_sync import ( + FullParameterSync, + FullParameterSyncProfile, + RolePairing, +) +from yeto.rl.engine.sync_overlap import OUTER_STEP_LAG, overlap_option_problems, rebase_values + +P = 2 # fragments +NUMEL = 4 + + +class _SweepSyncer: + """One-island strict sweep: PULLs a sweep once the previous one committed; + ``G_f += pushed delta``; the terminal step publishes FINAL.""" + + def __init__(self, initial: list[np.ndarray], sweeps: int) -> None: + self.global_ = [value.copy() for value in initial] + self.versions = [0] * P + self.sweeps = sweeps + self.step = 0 # last committed global step + self.pushed: list[tuple[int, np.ndarray]] = [] + self.bcasts: list[tuple[int, int, np.ndarray]] = [] + self.finalizing = False + self._since_commit = 0 + + def permits(self) -> list[PullRequest]: + sweep = self.step // P # completed sweeps + if self.finalizing or sweep >= self.sweeps or self.step % P: + return [] + pushed = {step for step, _ in self.pushed} + return [PullRequest(f, sweep * P + f + 1, 1) for f in range(P) if sweep * P + f + 1 not in pushed] + + def push(self, step: int, fragment_id: int, delta: np.ndarray) -> None: + assert self.step < step <= self.step + P and (step - 1) % P == fragment_id + self.pushed.append((step, delta.copy())) + self.global_[fragment_id] = self.global_[fragment_id] + delta + self.versions[fragment_id] = step + self._since_commit += 1 + if self._since_commit < P: + return + self._since_commit = 0 + self.step += P + if self.step == self.sweeps * P: + self.finalizing = True + else: + self.bcasts.extend((f, self.versions[f], self.global_[f].copy()) for f in range(P)) + + +class _NumericRole: + """A RoleStream whose cuts are lists of per-fragment f32 arrays.""" + + role = "actor" + + def __init__(self, syncer: _SweepSyncer, progress) -> None: + self.syncer = syncer + self.progress = progress # rollout_id -> list of per-fragment progress + self.config = SimpleNamespace(overlap=True) + self.current = [g.copy() for g in syncer.global_] + self.anchors = [g.copy() for g in syncer.global_] + self._versions = [0] * P + self.submission = None + self.submission_versions = None + self.local_rounds = self.submission_rounds = 0 + self.rounds_at_anchor = [0] * P + self.round_base = None + self.rounds_at_base = 0 + self._pending: dict[int, PullRequest] = {} + self._staged: dict[tuple[int, int], np.ndarray] = {} + self.terminal_submitted = False + self.final_acknowledged = False + self.closed = False + self.trained_from: list[list[np.ndarray]] = [] + self.push_rounds: list[tuple[int, int]] = [] # (global step, local rounds trained then) + + @property + def fragment_versions(self): + return tuple(self._versions) + + @property + def finalizing(self): + return self.syncer.finalizing + + async def initialize(self, group): + del group + + def resident_state(self, policy_version): + return SimpleNamespace(policy_version=policy_version, policy_tensor_hash=lambda: "cut") + + async def complete_local_round(self, *, rollout_id, trajectories, hold_submission=False): + del trajectories + self.trained_from.append([c.copy() for c in self.current]) + self.current = [c + d for c, d in zip(self.current, self.progress(rollout_id), strict=True)] + self.local_rounds += 1 + if not hold_submission: + self.submission = [c.copy() for c in self.current] + self.submission_versions = self.fragment_versions + self.submission_rounds = self.local_rounds + + def ready_submission_candidates(self): + for permit in self.syncer.permits(): + self._pending.setdefault(permit.fragment_id, permit) + out = [] + for f, permit in sorted(self._pending.items()): + base = max(0, permit.global_step - P) + if (self.submission is None or self._versions[f] != base + or self.submission_versions[f] != base + or self.submission_rounds - self.rounds_at_anchor[f] < 1): + continue + out.append(RoleSubmissionCandidate( + "actor", permit, base, self.submission_rounds, 1, 1, NUMEL * 4, + self.submission, self.anchors[f])) + return tuple(out) + + def retain_candidate(self, candidate): + del candidate + + def release_candidate(self, candidate): + del candidate + + def push_candidate(self, candidate): + f = candidate.permit.fragment_id + self.push_rounds.append((candidate.permit.global_step, self.local_rounds)) + self.syncer.push(candidate.permit.global_step, f, + candidate.submission_cut[f] - candidate.anchor) + self._pending.pop(f) + self.terminal_submitted |= candidate.permit.global_step == self.syncer.sweeps * P + return SimpleNamespace(role="actor", fragment_id=f, global_step=candidate.permit.global_step) + + async def stage_available_broadcasts(self): + for f, version, value in self.syncer.bcasts: + if version > self._versions[f]: + self._staged[(f, version)] = value + self.syncer.bcasts.clear() + return tuple(self._staged) + + @property + def staged_fragment_versions(self): + return tuple(sorted(self._staged)) + + def has_staged_broadcast(self, fragment_id, version): + return (fragment_id, version) in self._staged + + async def apply_staged_broadcast(self, fragment_id, version, *, rebase=False): + assert rebase, "overlap applies every sweep by rebasing" + g = self._staged.pop((fragment_id, version)) + self.current[fragment_id] = rebase_values( + g, self.current[fragment_id], self.round_base[fragment_id]) + self.anchors[fragment_id] = g + self._versions[fragment_id] = version + self.rounds_at_anchor[fragment_id] = self.rounds_at_base + return (fragment_id,) + + def finish_rebase(self): + self.submission = [c.copy() for c in self.current] + self.submission_versions = self.fragment_versions + self.submission_rounds = self.local_rounds + + def mark_round_base(self): + self.round_base = [c.copy() for c in self.current] + self.rounds_at_base = self.local_rounds + + async def wait_and_apply_final(self): + self.current = [g.copy() for g in self.syncer.global_] + self._versions = [self.syncer.sweeps * P - P + f + 1 for f in range(P)] + + def acknowledge_final(self): + self.final_acknowledged = True + + def discard_detached_submission_cut(self): + pass + + def close(self): + self.closed = True + + +class _Driver: + def __init__(self): + self.events = [] + + def emit(self, event, **fields): + self.events.append((event, fields)) + + +class _LatePullSyncer(_SweepSyncer): + """Shows the last fragment's PULL of each sweep only after a few polls.""" + + LATE_POLLS = 3 + + def __init__(self, *args) -> None: + super().__init__(*args) + self._polls: dict[int, int] = {} + + def permits(self): + permits = super().permits() + late = [p for p in permits if p.fragment_id == P - 1] + if late: + seen = self._polls[late[0].global_step] = self._polls.get(late[0].global_step, 0) + 1 + if seen <= self.LATE_POLLS: + return [p for p in permits if p.fragment_id != P - 1] + return permits + + +def _run(sweeps: int, overlap: bool = True, syncer_type=_SweepSyncer): + rng = np.random.default_rng(7) + initial = [rng.standard_normal(NUMEL).astype(np.float32) for _ in range(P)] + steps = {k: [rng.standard_normal(NUMEL).astype(np.float32) * 0.1 for _ in range(P)] + for k in range(sweeps)} + syncer = syncer_type(initial, sweeps) + role = _NumericRole(syncer, lambda rollout_id: steps[rollout_id]) + profile = FullParameterSyncProfile(roles=("actor",), total_fragment_steps=sweeps * P, fragments=P, + local_horizon=1, pipeline_depth=1, wait_timeout=5.0, + poll_seconds=0.001, sweep=True, overlap=overlap) + session = FullParameterSync(RolePairing({"actor": role}, profile), groups={"actor": "g"}, + run=asyncio.run, trajectories=lambda rollout_id: rollout_id) + driver = _Driver() + session.start(driver) + stops = [] + for rollout_id in range(sweeps): + boundary = session.boundary(driver, rollout_id=rollout_id, stats=None) + stops.append(boundary.stop) + if boundary.stop: + session.published(driver, rollout_id=rollout_id, policy_hash="x") + break + return initial, steps, syncer, role, stops + + +def test_overlap_counts_every_round_exactly_once(): + sweeps = 5 + initial, steps, syncer, role, stops = _run(sweeps) + assert stops == [False] * (sweeps - 1) + [True] + # Pushed deltas are exactly the per-round progress, in round order. + pushed = sorted(syncer.pushed, key=lambda row: row[0]) + for k in range(sweeps): + for f in range(P): + np.testing.assert_allclose(pushed[k * P + f][1], steps[k][f], rtol=0, atol=1e-6) + # Global = initial + sum of every round's progress, once. + for f in range(P): + want = initial[f] + sum(steps[k][f] for k in range(sweeps)) + np.testing.assert_allclose(syncer.global_[f], want, rtol=0, atol=1e-5) + np.testing.assert_allclose(role.current[f], want, rtol=0, atol=1e-5) + assert role.final_acknowledged + + +def test_overlap_boundary_pushes_the_whole_sweep_even_when_a_pull_arrives_late(): + sweeps = 3 + _initial, _steps, syncer, role, stops = _run(sweeps, syncer_type=_LatePullSyncer) + assert stops == [False] * (sweeps - 1) + [True] + # Every sweep was complete before the next round trained: sweep k's last + # fragment is pushed before round k+1's progress exists (no sweep spans + # two boundaries). + assert role.push_rounds == [(k * P + f + 1, k + 1) for k in range(sweeps) for f in range(P)] + assert role.final_acknowledged + assert all(polls > _LatePullSyncer.LATE_POLLS for polls in syncer._polls.values()) + + +def test_overlap_trains_one_outer_step_stale_and_keeps_local_progress(): + sweeps = 4 + initial, steps, _syncer, role, _ = _run(sweeps) + # Round 0 starts from the initial cut, round 1 from its own round-0 result + # (sweep 0 still in flight), round k>=2 from G_{k-2} + progress of round k-1. + for f in range(P): + np.testing.assert_allclose(role.trained_from[0][f], initial[f]) + np.testing.assert_allclose(role.trained_from[1][f], initial[f] + steps[0][f], atol=1e-6) + for k in range(2, sweeps): + g_prev = initial[f] + sum(steps[j][f] for j in range(k - 1)) + np.testing.assert_allclose(role.trained_from[k][f], g_prev + steps[k - 1][f], atol=1e-5) + + +def test_rebase_values_is_anchor_plus_local_progress(): + anchor = np.array([1.0, -2.0], dtype=np.float32) + current = np.array([0.5, 0.25], dtype=np.float32) + base = np.array([0.25, 0.5], dtype=np.float32) + np.testing.assert_array_equal(rebase_values(anchor, current, base), [1.25, -2.25]) + with pytest.raises(ValueError, match="float32"): + rebase_values(anchor.astype(np.float64), current, base) + with pytest.raises(ValueError, match="non-finite"): + rebase_values(np.array([np.inf, 0], np.float32), current, base) + + +def test_overlap_profile_needs_the_strict_sweep(): + with pytest.raises(ValueError, match="overlap needs the dense strict sweep"): + FullParameterSyncProfile(roles=("actor",), total_fragment_steps=4, fragments=2, local_horizon=2, + pipeline_depth=1, wait_timeout=1.0, poll_seconds=0.01, overlap=True) + + +def _dense(*extra): + from yeto.launcher import _prepare_rl_args + + args = _launcher_args("ports", ( + "--rl-sync-preset", "full-parameter", "--rl-placement", "fixed-partition", + "--rl-rollout-gpus", "1", "--rl-parameter-layout-sha256", "e" * 64, + "--fragments", "2", "--total-steps", "2", *extra), gpu="modal:2xh100") + args.controller = "local" + _prepare_rl_args(args) + return args + + +def test_overlap_keeps_the_rollout_policy_age_zero_and_reaches_the_island(): + from yeto.launch.rl_island_task import _rl_trainable_island_flags + from yeto.launch.rl_ports import _ports_infra_flags + + args = _dense("--rl-sync-overlap", "--rl-sync-fragment-pipeline", "2") + # Overlap is not rollout staleness: no policy age, so Miles gets neither + # partial rollout (carry-over) nor fully_async generation. + assert int(args.rl_max_policy_age or 0) == 0 + assert OUTER_STEP_LAG == 1 + assert "--sweep-overlap" in _rl_trainable_island_flags(args, 0, 1) + assert "--rl-max-policy-age" not in " ".join(_ports_infra_flags(args)) + plain = _dense() + assert "--sweep-overlap" not in _rl_trainable_island_flags(plain, 0, 1) + + +def test_overlap_launch_refuses_a_policy_age_or_a_partial_fragment_pipeline(): + with pytest.raises(ValueError, match="--rl-max-policy-age 0 with --rl-sync-overlap"): + _dense("--rl-sync-overlap", "--rl-sync-fragment-pipeline", "2", "--rl-max-policy-age", "1") + with pytest.raises(ValueError, match="--rl-sync-fragment-pipeline equal to --fragments"): + _dense("--rl-sync-overlap") + assert overlap_option_problems(max_policy_age=0, fragment_pipeline=4, fragments=4) == [] + + +def test_island_builds_overlap_without_a_policy_age(): + from yeto.rl.adapters.miles.island_entry import _dense_full_sync_config + + args = SimpleNamespace( + sweep_overlap=True, rl_max_policy_age=0, model_revision="A" * 40, syncer="127.0.0.1:30100", + learner_id=0, num_learners=2, total_fragment_steps=4, fragments=2, + parameter_layout_sha256="1" * 64, optimizer_steps=1, syncer_profile_sha256="3" * 64, + wan_streams=4) + config = _dense_full_sync_config(args, model_config_sha256="0" * 64, + training_contract_sha256="0" * 64, algorithm_spec_sha256="0" * 64) + assert config.roles[0].overlap + + +def test_overlap_is_bound_into_the_role_stream_contract(): + from dataclasses import replace + + from test_rl_full_parameter_sync import _role_stream_config + + from yeto.rl.adapters.miles.full_parameter.role_stream import ( + role_stream_session_contract_hash, + ) + + exact = _role_stream_config("actor", 30100, sweep=True) + assert (role_stream_session_contract_hash("1" * 64, exact) + != role_stream_session_contract_hash("1" * 64, replace(exact, overlap=True))) diff --git a/tests/test_rl_dense_sweep_real_syncer.py b/tests/test_rl_dense_sweep_real_syncer.py index fd14e997..76351037 100644 --- a/tests/test_rl_dense_sweep_real_syncer.py +++ b/tests/test_rl_dense_sweep_real_syncer.py @@ -13,6 +13,7 @@ import shutil import socket import subprocess +import threading import time from pathlib import Path @@ -21,7 +22,13 @@ from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: F401 from yeto.fragments import MERGE_RDA, Fragment, FragmentLayout -from yeto.protocol import DTYPE_F32, SyncerClient +from yeto.protocol import ( + DTYPE_F32, + DTYPE_F32_PUSH_BF16, + DTYPE_F32_PUSH_I8, + DTYPE_FLAG_BCAST_DELTA, + SyncerClient, +) ROOT = Path(__file__).resolve().parent.parent CARGO = shutil.which("cargo") or str(Path.home() / ".cargo/bin/cargo") @@ -41,23 +48,25 @@ def _hmac_key(monkeypatch): monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", "dense-sweep-test-key") -def _dense_args(): +def _dense_args(*extra: str, total_steps: int = 1): from yeto.launcher import _prepare_rl_args args = _launcher_args("ports", ( "--rl-sync-preset", "full-parameter", "--rl-placement", "fixed-partition", "--rl-rollout-gpus", "1", "--rl-parameter-layout-sha256", "e" * 64, - "--fragments", str(FRAGMENTS), "--total-steps", "1"), gpu="modal:2xh100") + "--fragments", str(FRAGMENTS), "--total-steps", str(total_steps), *extra), + gpu="modal:2xh100") args.controller = "local" _prepare_rl_args(args) return args -def _start_syncer(binary: Path, tmp_path: Path) -> tuple[subprocess.Popen, int, str]: +def _start_syncer(binary: Path, tmp_path: Path, *extra: str, + total_steps: int = 1) -> tuple[subprocess.Popen, int, str]: """The launcher's dense actor syncer, relocated to a free port and tmp_path.""" from yeto import launcher - args = _dense_args() + args = _dense_args(*extra, total_steps=total_steps) command = launcher.role_syncer_commands(args, 1, binary=str(binary))["actor"] with socket.socket() as probe: probe.bind(("127.0.0.1", 0)) @@ -76,10 +85,11 @@ def _layout() -> FragmentLayout: return FragmentLayout([Fragment(MERGE_RDA, [(f"actor.w{i}", NUMEL)]) for i in range(FRAGMENTS)]) -def _client(port: int, profile_hash: str) -> SyncerClient: - return SyncerClient(("127.0.0.1", port), 0, _layout(), dtype=DTYPE_F32, num_streams=1, - connect_timeout=10, max_reconnects=0, finalization_timeout=20, - syncer_profile_hash=bytes.fromhex(profile_hash)) +def _client(port: int, profile_hash: str, dtype: int = DTYPE_F32, *, + max_reconnects: int = 0, push_i8_block: int = 128) -> SyncerClient: + return SyncerClient(("127.0.0.1", port), 0, _layout(), dtype=dtype, num_streams=1, + connect_timeout=10, max_reconnects=max_reconnects, finalization_timeout=20, + syncer_profile_hash=bytes.fromhex(profile_hash), push_i8_block=push_i8_block) def _f32(values: list[float]) -> bytes: @@ -153,3 +163,583 @@ def test_launcher_dense_syncer_refuses_another_profile_hash(syncer_binary, tmp_p client.close() log = _stop(process) assert "syncer semantic profile does not match" in log + + +STREAMED = ("--rl-sync-fragment-pipeline", str(FRAGMENTS), "--rl-sync-async-checkpoint") + + +def _run_sweeps(binary: Path, tmp_path: Path, extra: tuple[str, ...], sweeps: int, + dtype: int = DTYPE_F32, *, after_sweep=None, max_reconnects: int = 0, + push_i8_block: int = 128, updates: dict | None = None, sink=None): + """Drive ``sweeps`` strict sweeps; return (max PULLs seen at once, finals, checkpoint). + + ``updates`` collects every BCAST the client accepted, as + ``{(fragment_id, version): f32 bytes}``; ``after_sweep(client, sweep)`` + runs once all pushes of a sweep are queued.""" + process, port, profile_hash = _start_syncer(binary, tmp_path, *extra, total_steps=sweeps) + client = _client(port, profile_hash, dtype, max_reconnects=max_reconnects, + push_i8_block=push_i8_block) + if sink is not None: + client.install_inbound_chunk_sink(sink) + concurrent = 0 + try: + client.start() + for fid in range(FRAGMENTS): + client.send_init(fid, _f32([float(fid) + 0.125 * i for i in range(NUMEL)])) + for sweep in range(1, sweeps + 1): + permits: dict[int, object] = {} + + def gather(): + client.check_health() + for permit in client.drain_pulls(): + permits[permit.fragment_id] = permit + return any(not getattr(p, "pushed", False) for p in permits.values()) + + while len(permits) < FRAGMENTS: + _wait(gather, f"no PULL in sweep {sweep}") + time.sleep(0.2) # let every PULL the profile permits arrive + gather() + concurrent = max(concurrent, sum(not getattr(p, "pushed", False) + for p in permits.values())) + # Push only what is permitted now; the serial profile then + # releases the next fragment's PULL after this commit. + for fid, permit in sorted(permits.items()): + if getattr(permit, "pushed", False): + continue + base = 0 if sweep == 1 else (sweep - 2) * FRAGMENTS + fid + 1 + delta = [0.5 + fid + 0.03125 * sweep + 0.001 * i for i in range(NUMEL)] + client.push_fragment(fid, permit.global_step, permit.round_attempt, + base, sweep, 1, 7, _f32(delta)) + permit.pushed = True + if after_sweep is not None: + after_sweep(client, sweep) + manifest, finals = client.wait_for_final_fragments(timeout=30) + if updates is not None: + updates.update({(u.fragment_id, u.version): bytes(u.data) for u in client.drain_updates()}) + result = {f.fragment_id: bytes(f.data) for f in finals} + client.acknowledge_finalization(manifest) + assert process.wait(timeout=15) == 0 + finally: + client.close() + log = _stop(process) + assert "syncer semantic profile" not in log, log + return concurrent, list(manifest.versions), result, (tmp_path / "state.ckpt").read_bytes(), log + + +def test_streamed_sweep_matches_the_serial_sweep_bit_for_bit(syncer_binary, tmp_path): + """E3: fragment pipelining + async checkpoint changes timing, not results.""" + serial_dir, streamed_dir = tmp_path / "serial", tmp_path / "streamed" + serial_dir.mkdir() + streamed_dir.mkdir() + serial = _run_sweeps(syncer_binary, serial_dir, (), sweeps=3) + streamed = _run_sweeps(syncer_binary, streamed_dir, STREAMED, sweeps=3) + assert serial[0] == 1, "the serial profile must permit one fragment at a time" + assert streamed[0] == FRAGMENTS, "the streamed profile must permit a whole sweep at once" + assert serial[1] == streamed[1] == [5, 6] + assert serial[2] == streamed[2] + assert serial[3] == streamed[3] + assert "--async-checkpoint" not in serial[4] + + +def test_streamed_flags_bind_the_profile_and_are_full_parameter_only(): + serial = _dense_args() + streamed = _dense_args(*STREAMED) + assert serial.pipeline == 1 and streamed.pipeline == FRAGMENTS + from yeto import launcher + + command = launcher.role_syncer_commands(streamed, 1, binary="yeto-syncer")["actor"] + assert f"--pipeline {FRAGMENTS}" in command and "--async-checkpoint" in command + assert (launcher.rl_syncer_profile(serial, 1).sha256 + != launcher.rl_syncer_profile(streamed, 1).sha256) + with pytest.raises(ValueError, match="--rl-sync-fragment-pipeline between 1 and --fragments"): + _dense_args("--rl-sync-fragment-pipeline", str(FRAGMENTS + 1)) + + +def test_bf16_pushes_stay_within_bf16_error_of_the_f32_sweep(syncer_binary, tmp_path): + """E3: bf16 outer deltas change results only by bf16 rounding of each delta.""" + exact_dir, bf16_dir = tmp_path / "f32", tmp_path / "bf16" + exact_dir.mkdir() + bf16_dir.mkdir() + exact = _run_sweeps(syncer_binary, exact_dir, STREAMED, sweeps=3) + rounded = _run_sweeps(syncer_binary, bf16_dir, (*STREAMED, "--rl-sync-push-dtype", "bf16"), + sweeps=3, dtype=DTYPE_F32_PUSH_BF16) + assert exact[1] == rounded[1] == [5, 6] # versions are unchanged + for fid, data in exact[2].items(): + want = torch.frombuffer(bytearray(data), dtype=torch.float32)[-NUMEL:] + got = torch.frombuffer(bytearray(rounded[2][fid]), dtype=torch.float32)[-NUMEL:] + # Three sweeps, each delta < 2 in magnitude, bf16 error <= 2^-8 relative. + assert torch.allclose(got, want, rtol=0, atol=3 * 2.0 / 256), fid + assert not torch.equal(got, want), "bf16 pushes must actually round" + + +def test_sweep_checkpoint_interval_reaches_the_syncer_and_its_profile(): + every = 2 * FRAGMENTS + args = _dense_args("--rl-sync-checkpoint-every", str(every), total_steps=4) + from yeto import launcher + + command = launcher.role_syncer_commands(args, 1, binary="yeto-syncer")["actor"] + assert f"--checkpoint-every {every}" in command + assert launcher.rl_syncer_profile(args, 1).checkpoint_every == every + with pytest.raises(ValueError, match="--rl-sync-checkpoint-every 1 or a multiple"): + _dense_args("--rl-sync-checkpoint-every", str(FRAGMENTS + 1)) + + +def test_bf16_push_flag_reaches_the_island_command(): + from yeto.launch.rl_island_task import _rl_trainable_island_flags + + args = _dense_args("--rl-sync-push-dtype", "bf16") + assert "--sweep-push-dtype bf16" in _rl_trainable_island_flags(args, 0, 1) + assert "--sweep-push-dtype" not in _rl_trainable_island_flags(_dense_args(), 0, 1) + + +def _tail(data: bytes) -> torch.Tensor: + return torch.frombuffer(bytearray(data), dtype=torch.float32) + + +def test_int8_pushes_stay_within_their_block_error_of_the_f32_sweep(syncer_binary, tmp_path): + """E3: int8 outer deltas change results only by half a scale step per delta.""" + exact_dir, int8_dir = tmp_path / "f32", tmp_path / "int8" + exact_dir.mkdir() + int8_dir.mkdir() + exact = _run_sweeps(syncer_binary, exact_dir, STREAMED, sweeps=3) + quantized = _run_sweeps(syncer_binary, int8_dir, + (*STREAMED, "--rl-sync-push-dtype", "int8", "--rl-sync-push-int8-block", "32"), + sweeps=3, dtype=DTYPE_F32_PUSH_I8, push_i8_block=32) + assert exact[1] == quantized[1] == [5, 6] + for fid, data in exact[2].items(): + want, got = _tail(data)[-NUMEL:], _tail(quantized[2][fid])[-NUMEL:] + # Three sweeps; each delta's absmax < 2, so error <= 3 * 2 / 254. + assert torch.allclose(got, want, rtol=0, atol=3 * 2.0 / 254), fid + assert not torch.equal(got, want), "int8 pushes must actually quantize" + + +def _downlink_run(binary, tmp_path, full_every: int, sweeps: int, **kwargs): + updates: dict = {} + result = _run_sweeps(binary, tmp_path, + (*STREAMED, "--rl-sync-bcast-delta-full-every", str(full_every)), + sweeps=sweeps, dtype=DTYPE_F32 | DTYPE_FLAG_BCAST_DELTA, + updates=updates, **kwargs) + return result, updates + + +def test_bf16_downlink_tracks_the_exact_broadcasts_and_leaves_the_merge_exact(syncer_binary, tmp_path): + """E3: BCAST deltas change only what islands hold between f32 refreshes; + the syncer's masters, versions, FINAL and checkpoint stay bit-identical.""" + exact_dir, delta_dir = tmp_path / "f32", tmp_path / "delta" + exact_dir.mkdir() + delta_dir.mkdir() + exact_updates: dict = {} + exact = _run_sweeps(syncer_binary, exact_dir, STREAMED, sweeps=4, updates=exact_updates) + delta, delta_updates = _downlink_run(syncer_binary, delta_dir, 3, sweeps=4) + assert exact[1] == delta[1] == [7, 8] + assert exact[2] == delta[2], "FINAL must be the exact f32 masters" + assert exact[3] == delta[3], "checkpoints must not depend on the downlink" + assert set(exact_updates) == set(delta_updates) + # Fragment 0 publishes versions 1, 3, 5, 7: FULL, DELTA, DELTA, FULL (M=3). + kinds = {} + for key, data in sorted(delta_updates.items()): + want, got = _tail(exact_updates[key]), _tail(data) + kinds[key] = "full" if torch.equal(got, want) else "delta" + # one bf16 rounding of (theta - anchor), |theta - anchor| < 4 here + assert torch.allclose(got, want, rtol=0, atol=4.0 / 256), key + assert [kinds[(0, v)] for v in (0, 1, 3, 5, 7)] == ["full", "full", "delta", "delta", "full"] + assert [kinds[(1, v)] for v in (0, 2, 4, 6)] == ["full", "full", "delta", "delta"] + + +def test_bf16_downlink_anchor_mismatch_reconnects_and_resyncs_with_f32(syncer_binary, tmp_path): + """E3 fallback: a learner whose anchor drifted fails the digest check, + reconnects, receives the shared anchor as FULL, and the run finishes + with the same exact result.""" + exact_dir, delta_dir = tmp_path / "f32", tmp_path / "delta" + exact_dir.mkdir() + delta_dir.mkdir() + exact = _run_sweeps(syncer_binary, exact_dir, STREAMED, sweeps=4) + corrupted, reconnects = [], [] + + def corrupt(client, sweep): + if sweep == 4: + reconnects.append(client._reconnects_used) + if sweep != 1: + return + _wait(lambda: client._downlink.version(0) == 1, "fragment 0 version 1 never arrived") + held = client._downlink._held[0] + held.values = held.values.copy() + held.values[0] += 1e-3 + corrupted.append(client.connection_generation) + + delta, _ = _downlink_run(syncer_binary, delta_dir, 100, sweeps=4, after_sweep=corrupt, + max_reconnects=5) + assert corrupted and reconnects and reconnects[0] >= 1, "the digest mismatch must reconnect" + assert exact[1] == delta[1] and exact[2] == delta[2] and exact[3] == delta[3] + + +def test_bf16_downlink_and_int8_flags_reach_syncer_and_island_commands(): + from yeto import launcher + from yeto.launch.rl_island_task import _rl_trainable_island_flags + + args = _dense_args("--rl-sync-push-dtype", "int8", "--rl-sync-push-int8-block", "64", + "--rl-sync-bcast-delta-full-every", "4") + island = _rl_trainable_island_flags(args, 0, 1) + assert "--sweep-push-dtype int8" in island and "--sweep-push-int8-block 64" in island + command = launcher.role_syncer_commands(args, 1, binary="yeto-syncer")["actor"] + assert "--bcast-delta-full-every 4" in command + plain = _dense_args() + assert "--bcast-delta" not in launcher.role_syncer_commands(plain, 1, binary="yeto-syncer")["actor"] + with pytest.raises(ValueError, match="--rl-sync-bcast-delta-full-every 0"): + _dense_args("--rl-sync-bcast-delta-full-every", "1") + with pytest.raises(ValueError, match="--rl-sync-push-int8-block only with"): + _dense_args("--rl-sync-push-int8-block", "64") + + +class _BytesSink: + """Minimal transactional inbound sink: assembles the f32 fragment bytes.""" + + def begin_message(self, message_id, total_bytes): + return {"buf": bytearray(), "bound": None} + + def bind_fragment(self, tx, msg_type, fragment_id, version, payload_bytes): + tx["bound"] = payload_bytes + + def stage_chunk(self, tx, payload): + return bytes(payload) + + def consume_staged_chunk(self, tx, offset, token, nbytes): + self.consume_chunk(tx, offset, token) + + def consume_chunk(self, tx, offset, payload): + assert offset == len(tx["buf"]) + tx["buf"] += bytes(payload) + + def finish_message(self, tx): + import hashlib + + from yeto.protocol import StreamedInboundPayload + + assert len(tx["buf"]) == tx["bound"] + return StreamedInboundPayload(bytes(tx["buf"]), hashlib.sha256(tx["buf"]).digest()) + + def abort_message(self, tx, staged_tokens): + tx["buf"].clear() + + +def test_bf16_downlink_through_a_streaming_sink_matches_the_buffered_path(syncer_binary, tmp_path): + """E3: the Miles path streams BCAST chunks into a transactional sink; the + downlink wrapper must hand it the same f32 anchor bytes.""" + buffered_dir, streamed_dir = tmp_path / "buffered", tmp_path / "streamed" + buffered_dir.mkdir() + streamed_dir.mkdir() + buffered, buffered_updates = _downlink_run(syncer_binary, buffered_dir, 3, sweeps=3) + streamed, streamed_updates = _downlink_run(syncer_binary, streamed_dir, 3, sweeps=3, + sink=_BytesSink()) + assert buffered[2] == streamed[2] + assert buffered_updates == streamed_updates + + +def _overlap_island_run(binary: Path, tmp_path: Path, extra: tuple[str, ...], dtype: int, + sweeps: int, progress: dict[tuple[int, int], list[float]]): + """One overlap island (E3 item 3) against the real syncer. + + Round s trains before sweep s-1 comes back; at boundary s the island waits + for G_{s-1}, rebases ``theta <- G_{s-1} + (theta - base)`` and pushes + ``theta - G_{s-1}``. Returns (FINAL bytes, the island's anchors per sweep, + redials used).""" + process, port, profile_hash = _start_syncer(binary, tmp_path, *extra, total_steps=sweeps) + client = _client(port, profile_hash, dtype) + init = {fid: torch.tensor([float(fid) + 0.125 * i for i in range(NUMEL)]) for fid in range(FRAGMENTS)} + anchors: dict[int, dict[int, torch.Tensor]] = {} + held: dict[tuple[int, int], torch.Tensor] = {} + try: + client.start() + for fid in range(FRAGMENTS): + client.send_init(fid, _f32(init[fid].tolist())) + anchor = dict(init) # what the last push was measured against + theta = {fid: init[fid] + torch.tensor(progress[(1, fid)]) for fid in range(FRAGMENTS)} + for sweep in range(1, sweeps + 1): + if sweep > 1: + # Round `sweep` trained from `base` while sweep-1 was in flight. + base = dict(theta) + theta = {fid: base[fid] + torch.tensor(progress[(sweep, fid)]) for fid in range(FRAGMENTS)} + + def got_sweep(): + client.check_health() + held.update({(u.fragment_id, u.version): _tail(bytes(u.data)) + for u in client.drain_updates()}) + return all((fid, (sweep - 2) * FRAGMENTS + fid + 1) in held for fid in range(FRAGMENTS)) + + _wait(got_sweep, f"sweep {sweep - 1} never came back") + anchor = {fid: held[(fid, (sweep - 2) * FRAGMENTS + fid + 1)] for fid in range(FRAGMENTS)} + anchors[sweep - 1] = anchor + theta = {fid: anchor[fid] + (theta[fid] - base[fid]) for fid in range(FRAGMENTS)} + permits: dict[int, object] = {} + + def gather(): + client.check_health() + permits.update({p.fragment_id: p for p in client.drain_pulls()}) + return len(permits) == FRAGMENTS + + _wait(gather, f"no PULLs for sweep {sweep}") + for fid, permit in sorted(permits.items()): + base_version = 0 if sweep == 1 else (sweep - 2) * FRAGMENTS + fid + 1 + client.push_fragment(fid, permit.global_step, permit.round_attempt, base_version, + sweep, 1, 7, _f32((theta[fid] - anchor[fid]).tolist())) + manifest, finals = client.wait_for_final_fragments(timeout=30) + result = {f.fragment_id: _tail(bytes(f.data)) for f in finals} + client.acknowledge_finalization(manifest) + redials = client._reconnects_used + assert process.wait(timeout=15) == 0 + finally: + client.close() + log = _stop(process) + assert "syncer semantic profile" not in log, log + return init, result, anchors, redials + + +def test_overlap_with_bf16_downlink_counts_each_round_once_without_a_resync(syncer_binary, tmp_path): + """E3: overlap + bf16 BCAST deltas. The island rebases on a bf16-tracked + anchor, but it pushes ``theta - anchor`` against that same anchor, so the + syncer's f32 masters still get every round's progress exactly once; the + downlink error stays out of the masters and never triggers a resync.""" + sweeps = 4 + progress = {(s, fid): [0.5 + fid + 0.03125 * s + 0.001 * i for i in range(NUMEL)] + for s in range(1, sweeps + 1) for fid in range(FRAGMENTS)} + exact_dir, delta_dir = tmp_path / "f32", tmp_path / "delta" + exact_dir.mkdir() + delta_dir.mkdir() + init, exact, exact_anchors, exact_redials = _overlap_island_run( + syncer_binary, exact_dir, STREAMED, DTYPE_F32, sweeps, progress) + _, delta, delta_anchors, delta_redials = _overlap_island_run( + syncer_binary, delta_dir, (*STREAMED, "--rl-sync-bcast-delta-full-every", "3"), + DTYPE_F32 | DTYPE_FLAG_BCAST_DELTA, sweeps, progress) + assert exact_redials == delta_redials == 0, "no DownlinkResync expected" + for fid in range(FRAGMENTS): + want = init[fid] + sum(torch.tensor(progress[(s, fid)]) for s in range(1, sweeps + 1)) + # (anchor + p) - anchor in f32 is p up to an ulp of the anchor: no bf16 error. + assert torch.allclose(exact[fid], want, rtol=0, atol=1e-5), fid + assert torch.allclose(delta[fid], want, rtol=0, atol=1e-5), fid + for sweep, anchors in exact_anchors.items(): + got = delta_anchors[sweep][fid] + # one bf16 rounding of (G - previous anchor); |G - previous anchor| < 4 + assert torch.allclose(got, anchors[fid], rtol=0, atol=4.0 / 256), (sweep, fid) + assert any(not torch.equal(delta_anchors[s][f], exact_anchors[s][f]) + for s in delta_anchors for f in range(FRAGMENTS)), "the downlink must actually round" + + +def test_bf16_downlink_sends_each_fragments_final_version_exact(syncer_binary, tmp_path): + """E3 (s26-e3-p06-D8): the last BCAST of a fragment has the same version as + its FINAL. Islands keep one copy per (fragment, version) and refuse a + second copy with other bytes, so that BCAST must be the exact f32 masters + even when the FULL cadence (here M=100) would send a bf16 delta.""" + delta, updates = _downlink_run(syncer_binary, tmp_path, 100, sweeps=4) + final_versions, finals = delta[1], delta[2] + assert final_versions == [7, 8] + # Fragment 0's FINAL version 7 was also broadcast; fragment 1's version 8 + # is the terminal step, which only FINAL carries. + assert updates[(0, 7)] == finals[0] + assert (1, 8) not in updates + # Earlier versions still travel as bf16 deltas. + assert any(updates[(0, v)] != finals[0] for v in (3, 5)) + + +class _StallingSocket: + """Proxy whose ``send`` stops making progress once ``armed`` is set, + until the socket is closed.""" + + def __init__(self, sock: socket.socket, armed: threading.Event) -> None: + self._sock = sock + self._armed = armed + + def send(self, data) -> int: + if not self._armed.is_set(): + return self._sock.send(data) + while self._sock.fileno() != -1: + time.sleep(0.02) + raise OSError("stalled stream closed") + + def __getattr__(self, name): + return getattr(self._sock, name) + + +class _StallOnceClient(SyncerClient): + """Data stream 1 of the first connection group stalls once ``armed``.""" + + stalled_generation: int | None = None + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.armed = threading.Event() + + def _send_loop(self, gen, sock, q, meter): + if meter is self._meters[1] and self.stalled_generation in (None, gen): + self.stalled_generation = gen + sock = _StallingSocket(sock, self.armed) + return super()._send_loop(gen, sock, q, meter) + + +def _sweep_until_final(client: SyncerClient, sweeps: int): + """Answer every PULL (also PULLs the syncer replays after a redial) with + the deterministic delta of its (fragment, global_step) until FINAL.""" + answered: set[tuple[int, int, int]] = set() + deadline = time.monotonic() + 60 + while not client.finalizing.is_set(): + assert time.monotonic() < deadline, "sweeps did not finish" + client.check_health() + for permit in client.drain_pulls(): + key = (client.connection_generation, permit.global_step, permit.round_attempt) + if key in answered: + continue + fid = permit.fragment_id + sweep = (permit.global_step - 1) // FRAGMENTS + 1 + assert permit.global_step == (sweep - 1) * FRAGMENTS + fid + 1 and sweep <= sweeps + base = 0 if sweep == 1 else permit.global_step - FRAGMENTS + delta = [0.5 + fid + 0.03125 * sweep + 0.001 * i for i in range(NUMEL)] + client.push_fragment(fid, permit.global_step, permit.round_attempt, + base, sweep, 1, 7, _f32(delta)) + answered.add(key) + time.sleep(0.01) + return client.wait_for_final_fragments(timeout=30) + + +@pytest.mark.parametrize("schedule", ["rr", "adaptive"]) +def test_stalled_stream_redials_and_the_replayed_push_keeps_the_exact_result( + syncer_binary, tmp_path, schedule): + """E3b: a data stream that stops moving bytes (C32 first run) is dropped + after the stall timeout; the syncer replays its PULL on the new group and + the island pushes that fragment again, so versions and finals match a run + without the stall bit for bit.""" + exact = _run_sweeps(syncer_binary, tmp_path, STREAMED, sweeps=3) + stall_dir = tmp_path / "stall" + stall_dir.mkdir() + process, port, profile_hash = _start_syncer(syncer_binary, stall_dir, *STREAMED, total_steps=3) + client = _StallOnceClient(("127.0.0.1", port), 0, _layout(), dtype=DTYPE_F32, num_streams=2, + connect_timeout=10, max_reconnects=5, finalization_timeout=20, + syncer_profile_hash=bytes.fromhex(profile_hash), + stream_schedule=schedule, stream_stall_timeout=1.0) + try: + client.start() + for fid in range(FRAGMENTS): + client.send_init(fid, _f32([float(fid) + 0.125 * i for i in range(NUMEL)])) + # Stall only the pushes: INIT is sent once and not replayed after a redial. + _wait(lambda: client.send_progress().queued_chunks == 0 and client._bcasts.qsize() >= FRAGMENTS, + "initial broadcast") + client.armed.set() + manifest, finals = _sweep_until_final(client, sweeps=3) + result = {f.fragment_id: bytes(f.data) for f in finals} + client.acknowledge_finalization(manifest) + assert process.wait(timeout=15) == 0 + progress = client.send_progress() + finally: + client.close() + _stop(process) + assert progress.stalls >= 1 and progress.reconnects >= 1 + assert list(manifest.versions) == exact[1] + assert result == exact[2] + + +def _initial_values(fid: int) -> list[float]: + return [float(fid) + 0.125 * i for i in range(NUMEL)] + + +def test_init_interrupted_by_a_redial_is_replayed_and_the_result_stays_exact(syncer_binary, tmp_path): + """E3b (s26-e3b-p06-C32a): learner 0's group died after INIT of fragment 0 + and before fragment 1, so the syncer never finished initializing and no + island got the initial cut. The client now replays INIT on the next group + until the initial cut arrives; the result matches an undisturbed run. A + link probe on the same port (yeto.wan_probe) leaves the session alone.""" + from yeto.protocol import island_hmac_key_from_env + from yeto.wan_probe import probe_link + + exact = _run_sweeps(syncer_binary, tmp_path, STREAMED, sweeps=3) + run_dir = tmp_path / "replay" + run_dir.mkdir() + process, port, profile_hash = _start_syncer(syncer_binary, run_dir, *STREAMED, total_steps=3) + client = SyncerClient(("127.0.0.1", port), 0, _layout(), dtype=DTYPE_F32, num_streams=2, + connect_timeout=10, max_reconnects=5, finalization_timeout=20, + syncer_profile_hash=bytes.fromhex(profile_hash), stream_stall_timeout=30.0) + killed = [] + + def parts_for(fid: int): + if fid == 1 and not killed: + killed.append(client.connection_generation) + client._poison_outbound_generation(client._gen, ConnectionError("test: group died mid-INIT")) + return [_f32(_initial_values(fid))] + + try: + def listening() -> bool: + try: + socket.create_connection(("127.0.0.1", port), timeout=1).close() + except OSError: + return False + return True + + _wait(listening, "syncer listening") + good = probe_link(("127.0.0.1", port), streams=2, bytes_per_direction=8 << 20, + key=island_hmac_key_from_env(), timeout=30) + assert good.error is None and good.upload_mbps > 0 and good.download_mbps > 0 + bad = probe_link(("127.0.0.1", port), streams=1, bytes_per_direction=1 << 20, + key=b"wrong key", timeout=10) + assert bad.error is not None and not bad.accepted(0.0, 0.0) + client.start() + assert client.send_initial_fragments(parts_for) is False + replays = 0 + while client._bcasts.qsize() < FRAGMENTS: + client.check_health() + if client.init_replay_due(): + replays += client.send_initial_fragments(parts_for) + time.sleep(0.02) + client.mark_initial_cut_received() + assert killed and replays == 1 and client.connection_generation != killed[0] + manifest, finals = _sweep_until_final(client, sweeps=3) + result = {f.fragment_id: bytes(f.data) for f in finals} + client.acknowledge_finalization(manifest) + assert process.wait(timeout=15) == 0 + finally: + client.close() + _stop(process) + assert list(manifest.versions) == exact[1] + assert result == exact[2] + + +def test_one_socket_lost_mid_init_is_redialed_and_the_result_stays_exact(syncer_binary, tmp_path): + """E3c: a single data socket of learner 0's group fails between INIT of + fragment 0 and fragment 1 (as in s26-e3b-p06-C32a, where one connection + reported BrokenPipe). The client replaces the whole group, replays INIT, + and the final versions and bytes equal an undisturbed run.""" + exact = _run_sweeps(syncer_binary, tmp_path, STREAMED, sweeps=3) + run_dir = tmp_path / "one-socket" + run_dir.mkdir() + process, port, profile_hash = _start_syncer(syncer_binary, run_dir, *STREAMED, total_steps=3) + client = SyncerClient(("127.0.0.1", port), 0, _layout(), dtype=DTYPE_F32, num_streams=2, + connect_timeout=10, max_reconnects=5, finalization_timeout=20, + syncer_profile_hash=bytes.fromhex(profile_hash)) + broken = [] + + def parts_for(fid: int): + if fid == 1 and not broken: + with client._lock: + sock = client._socks[-1] + broken.append(client.connection_generation) + sock.shutdown(socket.SHUT_RDWR) # one real socket dies; the others stay open + return [_f32(_initial_values(fid))] + + try: + client.start() + client.send_initial_fragments(parts_for) + replays = 0 + deadline = time.monotonic() + 60 + while client._bcasts.qsize() < FRAGMENTS: + assert time.monotonic() < deadline, "initial cut never arrived" + client.check_health() + if client.init_replay_due(): + replays += client.send_initial_fragments(parts_for) + time.sleep(0.02) + client.mark_initial_cut_received() + assert broken and replays >= 1 and client.connection_generation != broken[0] + manifest, finals = _sweep_until_final(client, sweeps=3) + result = {f.fragment_id: bytes(f.data) for f in finals} + client.acknowledge_finalization(manifest) + assert process.wait(timeout=15) == 0 + finally: + client.close() + _stop(process) + assert list(manifest.versions) == exact[1] + assert result == exact[2] diff --git a/tests/test_rl_sao_streaming_runtime.py b/tests/test_rl_sao_streaming_runtime.py index ced8e08e..88ec57b2 100644 --- a/tests/test_rl_sao_streaming_runtime.py +++ b/tests/test_rl_sao_streaming_runtime.py @@ -279,7 +279,7 @@ def test_load_or_role_pairing_rejects_role_or_lockstep_drift(tmp_path, mutation, ("checkpoint_every", 2, "checkpointing every round"), ("resume", False, "checkpointing every round"), ("max_base_lag", 1, "max_base_lag=0"), - ("policy_sweep_fragments", 4, "policy sweep"), + ("policy_sweep_fragments", 4, "policy.sweep"), ("learner_budget_steps", 2, "learner-budget cutoff mode"), ), ) diff --git a/tests/test_run_flags.py b/tests/test_run_flags.py index 9ef5c8ea..9cdba9d4 100644 --- a/tests/test_run_flags.py +++ b/tests/test_run_flags.py @@ -250,7 +250,7 @@ def test_forward_tables_name_flags_of_both_layers(): # flag renamed in the schema but still spelled the old way in the launcher fails. _NON_RL_LITERALS = frozenset({ "--abbrev-ref", "--adapter-dir", "--add", "--address", "--allow-unattested-legacy-adapter", - "--allow-unauthenticated-islands", "--bind", "--cap-add", "--carry-gamma", + "--allow-unauthenticated-islands", "--async-checkpoint", "--bcast-delta-full-every", "--bind", "--cap-add", "--carry-gamma", "--checkpoint-every", "--checkpoint-path", "--codex-harness-binary", "--config-json", "--create", "--dashboard-host", "--dearmor", "--depth", "--detach", "--dtype", "--entrypoint", "--expected-island-contract", "--final-ack-timeout-s", "--final-grace-s", "--follow", "--force", diff --git a/tests/test_wan_probe.py b/tests/test_wan_probe.py new file mode 100644 index 00000000..ad8a19df --- /dev/null +++ b/tests/test_wan_probe.py @@ -0,0 +1,150 @@ +"""S26 E3b: island link probe, its island hook, and the launcher's relaunch of +an island whose probe rejected the container's link.""" + +from __future__ import annotations + +import json +from argparse import Namespace + +import pytest + +from yeto.wan_probe import EXIT_SLOW_LINK, ProbeResult, probe_link + + +def _fake_stream(up_s: float, down_s: float): + def run(addr, upload, download, key, timeout, out): + out["up"], out["down"] = up_s, down_s + return run + + +def test_probe_rates_use_the_slowest_stream_per_direction(): + result = probe_link(("h", 1), streams=4, bytes_per_direction=40_000_000, key=None, + stream_fn=_fake_stream(2.0, 0.5)) + assert (result.upload_mbps, result.download_mbps) == (20.0, 80.0) + assert result.accepted(20.0, 80.0) and not result.accepted(20.1, 0.0) + assert result.fields()["wan_probe/upload_mbps"] == 20.0 + + +def test_a_failed_stream_rejects_the_link(): + def broken(addr, upload, download, key, timeout, out): + raise ConnectionError("reset") + result = probe_link(("h", 1), streams=2, bytes_per_direction=2 << 20, key=None, stream_fn=broken) + assert result.error == "ConnectionError: reset" and not result.accepted(0.0, 0.0) + with pytest.raises(ValueError): + probe_link(("h", 1), streams=0, bytes_per_direction=1, key=None) + + +def _island_args(tmp_path, **kw): + return Namespace(syncer="10.0.0.1:29400", wan_streams=8, wan_probe_mib=16, + event_tape=str(tmp_path / "tape.jsonl"), learner_id=1, **kw) + + +def test_island_probe_is_off_by_default_and_exits_75_on_a_slow_link(tmp_path, capsys): + from yeto.rl.adapters.miles.island_entry import probe_wan_link + + calls = [] + + def probe(addr, **kw): + calls.append((addr, kw["streams"], kw["bytes_per_direction"])) + return ProbeResult(8, 16 << 20, 16 << 20, 4.0, 1.0) + + probe_wan_link(_island_args(tmp_path), probe=probe) + assert calls == [] + probe_wan_link(_island_args(tmp_path, wan_probe_min_mbps=2.0, wan_probe_min_down_mbps=10.0), probe=probe) + with pytest.raises(SystemExit) as stop: + probe_wan_link(_island_args(tmp_path, wan_probe_min_mbps=8.0), probe=probe) + assert stop.value.code == EXIT_SLOW_LINK + assert calls == [(("10.0.0.1", 29400), 8, 16 << 20)] * 2 + records = [json.loads(line) for line in (tmp_path / "tape.jsonl").read_text().splitlines()] + assert [r["accepted"] for r in records if r["event"] == "rl_wan_probe"] == [True, False] + assert "[yeto-wan-probe] rejected: up 4.2 MB/s" in capsys.readouterr().out + + +def test_launcher_flags_reach_the_island_and_bound_the_relaunches(): + from yeto.launch.modal_islands import wan_probe_relaunches + from yeto.launch.rl_island_task import _wan_transport_flags + + on = Namespace(wan_probe_min_mbps=20.0, wan_probe_min_down_mbps=0.0, wan_probe_mib=128, + wan_probe_relaunches=3) + assert _wan_transport_flags(on) == (" --wan-probe-min-mbps 20.0 --wan-probe-min-down-mbps 0.0" + " --wan-probe-mib 128") + assert wan_probe_relaunches(on) == 3 + assert wan_probe_relaunches(Namespace(wan_probe_relaunches=3)) == 0 # probe off + with pytest.raises(ValueError): + _wan_transport_flags(Namespace(wan_probe_min_mbps=1.0, wan_probe_mib=0)) + + +def test_echo_collector_counts_rejected_probes(tmp_path): + from yeto.rl.event_echo import PREFIX, TapeCollector + + collector = TapeCollector(tmp_path / "tape.jsonl", fresh=True) + for accepted, t in ((False, 1), (True, 2), (False, 3)): + collector.feed(PREFIX + json.dumps({"event": "rl_wan_probe", "accepted": accepted, "island_id": 1, "time_unix": t}) + "\n") + assert collector.wan_rejects == 2 + + +def _fleet(ops, rejects, relaunches, **kw): + from test_controller import RUNNING, SYNCER, ImmediateThread + + from yeto.launcher import FleetController + + ops.status_seq.setdefault(SYNCER, [RUNNING]) + return FleetController(learners={"l0": ("task-l0", 1), "l1": ("task-l1", 2)}, + syncer=(SYNCER, "task-syncer", 1), sky_ops=ops, poll_interval=30, + recover_timeout=0, thread_cls=ImmediateThread, fixed_roster=True, + no_island_relaunch=True, wan_reject_probe=lambda name: rejects[name], + wan_probe_relaunches=relaunches, **kw) + + +def test_rejected_link_is_relaunched_even_without_island_relaunch(): + from test_controller import FAILED, RUNNING, SUCCEEDED, FakeOps + + ops = FakeOps() + ops.status_seq["l0"] = [RUNNING, RUNNING, RUNNING, SUCCEEDED] + ops.status_seq["l1"] = [FAILED] + ops.relaunch_results["l1"] = [7] + ops.after_relaunch["l1"] = [RUNNING, SUCCEEDED] + rejects = {"l0": 0, "l1": 1} + ctl = _fleet(ops, rejects, relaunches=2) + ctl.run() + assert ops.relaunch_calls == ["l1"] + assert ctl.learners["l1"]["wan_relaunches"] == 1 and ctl.learners["l1"]["state"] == "done" + + +def test_rejected_link_budget_is_bounded(): + from test_controller import FAILED, RUNNING, FakeOps + + from yeto.launcher import FixedRosterIslandAbandoned + + ops = FakeOps() + ops.status_seq["l0"] = [RUNNING] + ops.status_seq["l1"] = [FAILED] + rejects = {"l0": 0, "l1": 1} + ctl = _fleet(ops, rejects, relaunches=0) + with pytest.raises(FixedRosterIslandAbandoned): + ctl.run() + assert ops.relaunch_calls == [] + + +@pytest.mark.parametrize( + ("up_s", "down_s", "error", "min_up", "min_down", "accepted"), + [ + (1.0, 1.0, None, 100.0, 100.0, True), # 100 MB in 1 s each way: exactly at the minimum + (1.0, 1.0, None, 100.1, 0.0, False), # upload just below the minimum + (1.0, 4.0, None, 50.0, 50.0, False), # fast upload, slow download + (1.0, 4.0, None, 50.0, 0.0, True), # download not checked when its minimum is 0 + (0.1, 0.1, "OSError: x", 1.0, 1.0, False), # any stream error rejects a fast link + ], +) +def test_slow_container_decision(up_s, down_s, error, min_up, min_down, accepted): + result = ProbeResult(streams=8, upload_bytes=100_000_000, download_bytes=100_000_000, + upload_seconds=up_s, download_seconds=down_s, error=error) + assert result.accepted(min_up, min_down) is accepted + + +def test_probe_is_refused_with_the_elastic_syncer(): + from yeto.launch.rl_island_task import _wan_transport_flags + + with pytest.raises(ValueError, match="elastic"): + _wan_transport_flags(Namespace(wan_probe_min_mbps=20.0, rl_elastic=True)) + assert _wan_transport_flags(Namespace(rl_elastic=True)) == "" # probe off: no error diff --git a/tests/test_wan_streams.py b/tests/test_wan_streams.py new file mode 100644 index 00000000..99f3d4dc --- /dev/null +++ b/tests/test_wan_streams.py @@ -0,0 +1,261 @@ +"""S26 E3b: per-stream accounting, stall watchdog and rate-aware striping.""" + +from __future__ import annotations + +import queue +import socket +import threading +import time + +import pytest + +from yeto.protocol import SyncerClient +from yeto.rl.engine.sync_timing import SyncTiming, upload_fields +from yeto.wan_streams import ( + StreamMeter, + StreamSnapshot, + UploadWindow, + pick_stream, + upload_window, +) + + +class _Clock: + def __init__(self) -> None: + self.now = 100.0 + + def __call__(self) -> float: + return self.now + + +def _meter_with_rate(clock: _Clock, rate_bps: float, backlog: int = 0) -> StreamMeter: + meter = StreamMeter(clock) + meter.enqueued(int(rate_bps)) + meter.begin() + clock.now += 1.0 + meter.progress(int(rate_bps)) + meter.end() + if backlog: + meter.enqueued(backlog) + return meter + + +def test_round_robin_ignores_rates(): + clock = _Clock() + meters = [_meter_with_rate(clock, 1e6), _meter_with_rate(clock, 1e3)] + assert [pick_stream("rr", meters, ticket) for ticket in range(4)] == [0, 1, 0, 1] + + +def test_adaptive_spreads_an_idle_start_over_every_stream(): + meters = [StreamMeter() for _ in range(4)] + picked = [] + for ticket in range(4): + index = pick_stream("adaptive", meters, ticket) + meters[index].enqueued(4 << 20) + picked.append(index) + assert sorted(picked) == [0, 1, 2, 3] + + +def test_adaptive_gives_a_slow_stream_fewer_chunks_in_proportion_to_its_rate(): + clock = _Clock() + fast, slow = 40e6, 5e6 + meters = [_meter_with_rate(clock, fast), _meter_with_rate(clock, slow)] + counts = [0, 0] + for ticket in range(90): + index = pick_stream("adaptive", meters, ticket) + meters[index].enqueued(4 << 20) + counts[index] += 1 + # Equal estimated finish time: chunk counts follow the 8:1 rate ratio. + assert counts[1] == pytest.approx(90 / 9, abs=1) + assert counts[0] == pytest.approx(90 * 8 / 9, abs=1) + + +def test_adaptive_starves_a_stream_whose_backlog_is_stuck(): + clock = _Clock() + meters = [_meter_with_rate(clock, 10e6), _meter_with_rate(clock, 10e6, backlog=64 << 20)] + assert {pick_stream("adaptive", meters, ticket) for ticket in range(10)} == {0} + + +def test_unknown_schedule_is_refused(): + with pytest.raises(ValueError, match="schedule"): + pick_stream("random", [StreamMeter()], 0) + with pytest.raises(ValueError, match="stream_schedule"): + SyncerClient(("127.0.0.1", 1), 0, _layout(), stream_schedule="random") + with pytest.raises(ValueError, match="stall_timeout"): + SyncerClient(("127.0.0.1", 1), 0, _layout(), stream_stall_timeout=0) + + +def test_meter_reports_a_stall_only_while_bytes_wait(): + clock = _Clock() + meter = StreamMeter(clock) + clock.now += 50 + assert meter.stalled_for() == 0.0 # idle is not stalled + meter.enqueued(10) + meter.begin() + meter.progress(4) + clock.now += 7 + assert meter.stalled_for() == pytest.approx(7) + meter.progress(6) + meter.end() + clock.now += 30 + assert meter.stalled_for() == 0.0 and meter.idle + meter.enqueued(5) + meter.drop_pending() + assert meter.idle + + +def test_upload_window_and_fields_report_rates_and_the_wait_after_upload(): + before = (StreamSnapshot(100, 1.0, 0, None), StreamSnapshot(0, 0.0, 0, None)) + after = (StreamSnapshot(100 + 40_000_000, 3.0, 0, 2e7), StreamSnapshot(5_000_000, 1.0, 0, 5e6)) + window = upload_window(before, after) + assert window == UploadWindow((40_000_000, 5_000_000), (2.0, 1.0)) + assert window.stream_mbps() == [20.0, 5.0] + fields = upload_fields(window, drained_ts=1010.0, timing=SyncTiming(1000.0, 1025.0), stalls=1) + assert fields["sync/upload_mbps"] == pytest.approx(22.5) + assert fields["sync/upload_min_stream_mbps"] == 5.0 + assert fields["sync/upload_active_streams"] == 2 + assert fields["sync/wait_after_upload_seconds"] == pytest.approx(15.0) + assert fields["sync/stream_stalls"] == 1 + # The upload drained before this boundary's submit: only the submit counts. + early = upload_fields(window, drained_ts=990.0, timing=SyncTiming(1000.0, 1025.0), stalls=0) + assert early["sync/wait_after_upload_seconds"] == pytest.approx(25.0) + unknown = upload_fields(window, drained_ts=None, timing=SyncTiming(1000.0, None), stalls=0) + assert unknown["sync/wait_after_upload_seconds"] is None + + +def test_sync_round_diagnostics_reach_wandb_on_the_policy_axis(): + from yeto.rl.wandb_rl import event_metrics + + metrics = event_metrics({ + "event": "rl_full_parameter_round", "policy_version": 3, + "sync/sync_wait_seconds": 40.0, "sync/upload_mbps": 22.5, + "sync/upload_min_stream_mbps": 5.0, "sync/upload_stream_mbps": [20.0, 5.0], + "sync/wait_after_upload_seconds": 15.0, "sync/global_policy_hash": "ab", + }) + assert metrics == {"rl/step": 3, "sync/sync_wait_seconds": 40.0, "sync/upload_mbps": 22.5, + "sync/upload_min_stream_mbps": 5.0, "sync/wait_after_upload_seconds": 15.0} + assert event_metrics({"event": "rl_full_parameter_round", "policy_version": 3}) is None + + +# -- socket-level behaviour --------------------------------------------------- + + +def _layout(): + from yeto.fragments import MERGE_RDA, Fragment, FragmentLayout + + return FragmentLayout([Fragment(MERGE_RDA, [("w", 4)])]) + + +def test_send_loop_meters_bytes_and_marks_the_upload_drained(): + left, right = socket.socketpair() + client = SyncerClient(("127.0.0.1", 1), 0, _layout(), num_streams=1) + q: queue.Queue = queue.Queue() + meter = client._meters[1] + worker = threading.Thread(target=client._send_loop, args=(client._gen, left, q, meter)) + worker.start() + try: + meter.enqueued(3 << 20) + q.put(type("Item", (), {"data": b"x" * (3 << 20), "sent": None})()) + received = 0 + while received < 3 << 20: + received += len(right.recv(1 << 20)) + q.put(None) + worker.join(timeout=5) + snap = client.stream_snapshots()[0] + assert snap.bytes_sent == 3 << 20 and snap.backlog_bytes == 0 + assert client.upload_drained_at is not None + finally: + left.close() + right.close() + + +def test_watchdog_fails_the_group_when_a_stream_stops_moving(): + client = SyncerClient(("127.0.0.1", 1), 0, _layout(), num_streams=2, stream_stall_timeout=0.2) + a, b = socket.socketpair() + try: + client._socks = [a] + client._connected.set() + meter = client._meters[2] + meter.enqueued(10) + meter.begin() + assert client._stalled_stream() is None + time.sleep(0.3) + assert "outbound stream 2" in client._stalled_stream() + client._watchdog = threading.Thread(target=client._watch_stalls, daemon=True) + client._watchdog.start() + assert client._failure.wait(3), "the watchdog must fail the live group" + assert client.send_progress().stalls == 1 + assert "no progress" in str(client._last_err) + finally: + client._closed.set() + a.close() + b.close() + + +def test_watchdog_notices_an_inbound_message_that_stops_growing(): + client = SyncerClient(("127.0.0.1", 1), 0, _layout(), num_streams=1, stream_stall_timeout=0.2) + client._reasm_touched[5] = time.monotonic() - 1.0 + assert "inbound message 5" in client._stalled_stream() + client._clear_partial_reassemblies() + assert client._stalled_stream() is None + + +def test_stall_timeout_sets_tcp_user_timeout(): + if not hasattr(socket, "TCP_USER_TIMEOUT"): + pytest.skip("Linux-only socket option") + from yeto.protocol import set_tcp_user_timeout + + with socket.socket() as sock: + set_tcp_user_timeout(sock, 2.5) + assert sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_USER_TIMEOUT) == 2500 + + +def test_transport_flags_reach_the_island_command_only_when_set(): + from argparse import Namespace + + from yeto.launch.rl_island_task import _wan_transport_flags + + assert _wan_transport_flags(Namespace()) == "" + assert _wan_transport_flags(Namespace(wan_stream_schedule="adaptive", wan_stall_timeout=60.0)) == ( + " --wan-stream-schedule adaptive --wan-stall-timeout 60.0") + with pytest.raises(ValueError, match="wan-stall-timeout"): + _wan_transport_flags(Namespace(wan_stall_timeout=-1.0)) + + +def test_island_parser_accepts_the_transport_flags(): + from yeto.rl.engine.run_flags import ISLAND, rl_option + + assert rl_option("wan_stream_schedule", ISLAND) == "--wan-stream-schedule" + assert rl_option("wan_stall_timeout", ISLAND) == "--wan-stall-timeout" + + +def test_tcp_user_timeout_waits_for_the_initial_cut(): + if not hasattr(socket, "TCP_USER_TIMEOUT"): + pytest.skip("Linux-only socket option") + with socket.create_server(("127.0.0.1", 0)) as server: + addr = server.getsockname() + client = SyncerClient(addr, 0, _layout(), num_streams=1, stream_stall_timeout=2.0) + early = client._dial(5) + try: + assert early.getsockopt(socket.IPPROTO_TCP, socket.TCP_USER_TIMEOUT) == 0 + client._socks = [early] + client.mark_initial_cut_received() + assert early.getsockopt(socket.IPPROTO_TCP, socket.TCP_USER_TIMEOUT) == 2000 + late = client._dial(5) + assert late.getsockopt(socket.IPPROTO_TCP, socket.TCP_USER_TIMEOUT) == 2000 + late.close() + finally: + early.close() + + +def test_init_replay_is_due_only_on_a_new_group_before_the_initial_cut(): + client = SyncerClient(("127.0.0.1", 1), 0, _layout(), num_streams=1) + client._connected.set() + client.connection_generation = 11 + assert client.init_replay_due() + client._init_generation = 11 + assert not client.init_replay_due() + client.connection_generation = 12 + assert client.init_replay_due() + client.mark_initial_cut_received() + assert not client.init_replay_due() diff --git a/tests/test_wandb_tape.py b/tests/test_wandb_tape.py index 0a9abf08..2305d7fc 100644 --- a/tests/test_wandb_tape.py +++ b/tests/test_wandb_tape.py @@ -176,13 +176,15 @@ def test_every_rust_tape_field_is_consumed_or_deliberately_skipped(): # The responders literal is built first, the record literal last. responder_fmt, record_fmt = body.split("let line = format!(", 1) top_level = set(re.findall(r'\\"(\w+)\\":', record_fmt)) + optional = source.split("fn checkpoint_tape_field(", 1)[1].split("\nfn ", 1)[0] + top_level |= set(re.findall(r'\\"(\w+)\\":', optional)) responder_keys = set(re.findall(r'\\"(\w+)\\":', responder_fmt)) assert responder_keys, "could not find the responder literal in append_tape" consumed_top = { "step", "fragment", "attempt", "gnorm", "ms", "sync_ms", "quorum", "expected", "responded", "missed_grace", "launch_base_version", - "quorum_ms", "grace_ms", "responders", + "quorum_ms", "grace_ms", "responders", "checkpoint_ms", } # Constants and the member-identity mirrors of the id lists: nothing a # time series can be drawn from. diff --git a/tools/bench_sync_boundary.py b/tools/bench_sync_boundary.py new file mode 100644 index 00000000..a55fb906 --- /dev/null +++ b/tools/bench_sync_boundary.py @@ -0,0 +1,175 @@ +"""CPU micro-benchmark of the full-parameter sync boundary (S26 E3). + +Runs the real Rust syncer on localhost with two in-process learners and a +synthetic f32 model, and times each strict sweep from the first push to the +last updated fragment received by every learner. It compares the serial +profile (pipeline 1, checkpoint on the scheduler), the streamed profile +(pipeline P, ``--async-checkpoint``), and streamed plus bf16 pushes and a +checkpoint every second sweep. + +These are local CPU numbers: loopback transfer and a local NVMe checkpoint. +They show where the boundary time goes on this machine; they do not predict +GPU-cluster or WAN timings. + + python tools/bench_sync_boundary.py --floats-per-fragment 33554432 --sweeps 4 +""" + +from __future__ import annotations + +import argparse +import json +import socket +import statistics +import subprocess +import tempfile +import threading +import time +from pathlib import Path + +import numpy as np + +from yeto.fragments import MERGE_RDA, Fragment, FragmentLayout +from yeto.protocol import DTYPE_F32, DTYPE_F32_PUSH_BF16, SyncerClient + +ROOT = Path(__file__).resolve().parent.parent + + +def _free_port() -> int: + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + return probe.getsockname()[1] + + +def _syncer(binary: Path, workdir: Path, *, port: int, learners: int, fragments: int, + sweeps: int, streamed: bool, checkpoint_every: int) -> subprocess.Popen: + command = [ + str(binary), "--port", str(port), "--learners", str(learners), + "--quorum", str(learners), "--grace-ms", "0", + "--pipeline", str(fragments if streamed else 1), + "--sync-interval-steps", "0", "--delta-correction", "none", + "--total-steps", str(sweeps * fragments), "--policy-sweep-fragments", str(fragments), + "--outer-lr", "1", "--outer-momentum", "0", "--max-base-lag", "0", + "--learner-weight", "equal", "--checkpoint-path", str(workdir / "state.ckpt"), + "--checkpoint-every", str(checkpoint_every), "--event-tape", str(workdir / "tape.jsonl"), + "--allow-unauthenticated-islands", # loopback-only benchmark + ] + if streamed: + command.append("--async-checkpoint") + return subprocess.Popen(command, stdout=open(workdir / "syncer.log", "w"), + stderr=subprocess.STDOUT) + + +def _learner(learner_id: int, port: int, layout: FragmentLayout, numel: int, sweeps: int, + timings: list[list[float]], errors: list[BaseException], dtype: int) -> None: + fragments = len(layout.fragments) + client = SyncerClient(("127.0.0.1", port), learner_id, layout, dtype=dtype, + num_streams=1, connect_timeout=30, max_reconnects=0, + finalization_timeout=600) + rng = np.random.default_rng(learner_id) + try: + client.start() + if learner_id == 0: + for fid in range(fragments): + client.send_init(fid, rng.standard_normal(numel, dtype=np.float32).tobytes()) + deltas = [rng.standard_normal(numel, dtype=np.float32).tobytes() for _ in range(fragments)] + carried: list = [] # a PULL of the next sweep can arrive before this sweep's BCAST + for sweep in range(1, sweeps + 1): + pushed: set[int] = set() + updated: set[int] = set() + first_push = None + while len(updated) < fragments or len(pushed) < fragments: + client.check_health() + permits, carried = carried + client.drain_pulls(), [] + for permit in permits: + if permit.fragment_id in pushed: + carried.append(permit) + continue + base = 0 if sweep == 1 else (sweep - 2) * fragments + permit.fragment_id + 1 + first_push = first_push or time.monotonic() + client.push_fragment(permit.fragment_id, permit.global_step, + permit.round_attempt, base, sweep, 1, 1, + deltas[permit.fragment_id]) + pushed.add(permit.fragment_id) + for update in client.drain_updates(): + if update.version > (sweep - 1) * fragments: + updated.add(update.fragment_id) + if sweep == sweeps and len(pushed) == fragments: + break # the terminal sweep is delivered by FINAL, not BCAST + time.sleep(0.001) + if sweep == sweeps: + manifest, _ = client.wait_for_final_fragments(timeout=600) + timings[learner_id].append(time.monotonic() - first_push) + client.acknowledge_finalization(manifest) + else: + timings[learner_id].append(time.monotonic() - first_push) + except BaseException as error: # noqa: BLE001 - reported by the caller + errors.append(error) + finally: + client.close() + + +PROFILES = { + "serial": {"streamed": False, "dtype": DTYPE_F32, "sweeps_per_checkpoint": 0}, + "streamed": {"streamed": True, "dtype": DTYPE_F32, "sweeps_per_checkpoint": 0}, + "streamed+bf16+ckpt/2sweeps": {"streamed": True, "dtype": DTYPE_F32_PUSH_BF16, + "sweeps_per_checkpoint": 2}, +} + + +def run(binary: Path, *, profile: str, learners: int, fragments: int, numel: int, + sweeps: int) -> dict: + options = PROFILES[profile] + streamed, dtype = options["streamed"], options["dtype"] + every = options["sweeps_per_checkpoint"] * fragments or 1 + layout = FragmentLayout([Fragment(MERGE_RDA, [(f"w{i}", numel)]) for i in range(fragments)]) + with tempfile.TemporaryDirectory(prefix="e3-bench-", dir="/tmp") as tmp: + workdir = Path(tmp) + port = _free_port() + process = _syncer(binary, workdir, port=port, learners=learners, fragments=fragments, + sweeps=sweeps, streamed=streamed, checkpoint_every=every) + timings: list[list[float]] = [[] for _ in range(learners)] + errors: list[BaseException] = [] + threads = [threading.Thread(target=_learner, args=(i, port, layout, numel, sweeps, + timings, errors, dtype)) + for i in range(learners)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + code = process.wait(timeout=120) + if errors or code != 0: + raise RuntimeError(f"benchmark failed (syncer exit {code}): {errors!r}; " + f"log: {(workdir / 'syncer.log').read_text()[-2000:]}") + tape = [json.loads(line) for line in (workdir / "tape.jsonl").read_text().splitlines() + if '"sync_ms"' in line] + # Sweep 1 includes connection set-up; the last sweep ends with FINAL. + steady = [max(per[i] for per in timings) for i in range(1, sweeps - 1)] or \ + [max(per[i] for per in timings) for i in range(sweeps)] + return { + "profile": profile, + "boundary_s_per_sweep": [round(x, 3) for x in steady], + "boundary_s_median": round(statistics.median(steady), 3), + "checkpoints_written": sum("checkpoint_ms" in r for r in tape), + "sync_ms_median": statistics.median(r["sync_ms"] for r in tape), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--binary", type=Path, default=ROOT / "syncer/target/release/yeto-syncer") + parser.add_argument("--learners", type=int, default=2) + parser.add_argument("--fragments", type=int, default=2) + parser.add_argument("--floats-per-fragment", type=int, default=32 * 1024 * 1024) + parser.add_argument("--sweeps", type=int, default=4) + args = parser.parse_args() + results = [run(args.binary, profile=profile, learners=args.learners, + fragments=args.fragments, numel=args.floats_per_fragment, sweeps=args.sweeps) + for profile in PROFILES] + print(json.dumps({"note": "CPU localhost numbers; not GPU or WAN timings", + "learners": args.learners, "fragments": args.fragments, + "floats_per_fragment": args.floats_per_fragment, "results": results}, + indent=2)) + + +if __name__ == "__main__": + main() diff --git a/yeto/cli.py b/yeto/cli.py index 7c68497f..ab16fb4b 100644 --- a/yeto/cli.py +++ b/yeto/cli.py @@ -438,6 +438,45 @@ def int_or_auto(value: str): "deltas (~4x less learner egress; broadcasts stay bf16)", ) sync.add_argument("--wan-streams", type=int, default=4, help="parallel TCP streams per learner") + sync.add_argument( + "--wan-stream-schedule", + choices=("rr", "adaptive"), + default="rr", + help="full-parameter RL islands: rr = round-robin chunks over the streams; " + "adaptive = send each chunk to the stream with the earliest estimated " + "finish (backlog / measured rate), so slow streams carry less", + ) + sync.add_argument( + "--wan-stall-timeout", + type=float, + default=0.0, + help="full-parameter RL islands: drop and redial the connection group when " + "one stream holds bytes with no progress for this many seconds (also sets " + "TCP_USER_TIMEOUT); the fragment is pushed again on the replayed PULL. 0 = off", + ) + sync.add_argument( + "--wan-probe-min-mbps", + type=float, + default=0.0, + help="Miles RL islands: before loading the model, measure the upload rate to the " + "actor syncer over --wan-streams connections; below this MB/s the island exits " + "(code 75) and the launcher requests another container (--wan-probe-relaunches). 0 = off", + ) + sync.add_argument( + "--wan-probe-min-down-mbps", + type=float, + default=0.0, + help="like --wan-probe-min-mbps, for the download rate from the syncer. 0 = off", + ) + sync.add_argument("--wan-probe-mib", type=int, default=256, + help="bytes per direction of the link probe, in MiB") + sync.add_argument( + "--wan-probe-relaunches", + type=int, + default=2, + help="containers the launcher may request again per island after a rejected link " + "probe (also under --no-island-relaunch: the island never joined the session)", + ) infra = p.add_argument_group("infrastructure") infra.add_argument( diff --git a/yeto/downlink.py b/yeto/downlink.py new file mode 100644 index 00000000..d3a613c0 --- /dev/null +++ b/yeto/downlink.py @@ -0,0 +1,235 @@ +"""Learner side of bf16 downlink deltas (S26 E3; syncer ``src/downlink.rs``). + +A session whose dtype carries ``DTYPE_FLAG_BCAST_DELTA`` receives every +BCAST_FRAGMENT with a 41-byte downlink header after ``fragment_id, version``:: + + kind u8 (0 FULL, 1 DELTA) | base_version u64 | anchor_sha256 [32] + +FULL carries f32 values; DELTA carries ``bf16(theta - anchor)``. The client +keeps one f32 *anchor* per fragment, bit-identical to the syncer's: FULL sets +it, DELTA adds the decoded delta in f32 (numpy float32 addition rounds like +Rust's). The digest is checked after every update. Callers above the client +(``drain_updates``, inbound chunk sinks) only ever see the reconstructed f32 +anchor, so adapters, N7.6 cuts and manifests are unchanged. + +A missing or mismatched anchor raises :class:`DownlinkResync` (a +``ConnectionError``): the client drops its connection group, and the +syncer's reconnect recovery re-sends every fragment as FULL. +""" + +from __future__ import annotations + +import hashlib +import struct +import threading +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +if TYPE_CHECKING: # numpy is imported lazily: yeto.protocol stays import-light + import numpy as np + +KIND_FULL = 0 +KIND_DELTA = 1 +HEADER = struct.Struct(" int: + if kind == KIND_FULL: + return 4 * numel + if kind == KIND_DELTA: + return 2 * numel + raise ValueError(f"unknown downlink kind {kind}") + + +@dataclass +class _Held: + version: int + values: np.ndarray # float32, never mutated after commit + + +class DownlinkDecode: + """Streams one BCAST's tensor bytes into the next f32 anchor.""" + + def __init__(self, fragment_id: int, version: int, kind: int, base: np.ndarray | None, + digest: bytes, numel: int) -> None: + import numpy as np + + self.fragment_id = fragment_id + self.version = version + self.kind = kind + self._base = base + self._digest = digest + self.values = np.empty(numel, dtype=np.float32) + self._width = 4 if kind == KIND_FULL else 2 + self._carry = b"" + self._filled = 0 + self._hasher = hashlib.sha256() + + def feed(self, data: bytes | bytearray | memoryview) -> bytes: + """Decode the next wire bytes; return the f32 bytes they complete.""" + import numpy as np + + raw = self._carry + bytes(data) + whole = len(raw) - len(raw) % self._width + self._carry = raw[whole:] + if not whole: + return b"" + count = whole // self._width + start, end = self._filled, self._filled + count + if end > self.values.size: + raise ValueError("downlink BCAST has more values than its fragment") + out = self.values[start:end] + if self.kind == KIND_FULL: + out[:] = np.frombuffer(raw, dtype=" None: + if self._carry or self._filled != self.values.size: + raise ValueError("downlink BCAST ended before its fragment was complete") + if self._hasher.digest() != self._digest: + raise DownlinkResync( + f"downlink anchor digest mismatch for fragment {self.fragment_id} " + f"version {self.version}") + + +class DownlinkAnchors: + """Per-fragment f32 anchors shared (bit for bit) with the syncer.""" + + def __init__(self, num_fragments: int) -> None: + self._held: list[_Held | None] = [None] * num_fragments + self._lock = threading.Lock() + + def version(self, fragment_id: int) -> int | None: + with self._lock: + held = self._held[fragment_id] + return None if held is None else held.version + + def begin(self, fragment_id: int, version: int, header: bytes, numel: int) -> DownlinkDecode: + kind, base_version, digest = HEADER.unpack(header) + if kind == KIND_FULL: + return DownlinkDecode(fragment_id, version, kind, None, digest, numel) + if kind != KIND_DELTA: + raise ValueError(f"unknown downlink kind {kind}") + with self._lock: + held = self._held[fragment_id] + if held is None or held.version != base_version or held.values.size != numel: + raise DownlinkResync( + f"downlink delta for fragment {fragment_id} version {version} needs anchor " + f"version {base_version}; learner holds " + f"{None if held is None else held.version}") + return DownlinkDecode(fragment_id, version, kind, held.values, digest, numel) + + def commit(self, decode: DownlinkDecode) -> None: + """Adopt a finished decode as the anchor unless a newer one is held.""" + with self._lock: + held = self._held[decode.fragment_id] + if held is None or decode.version >= held.version: + self._held[decode.fragment_id] = _Held(decode.version, decode.values) + + def decode_whole(self, fragment_id: int, version: int, payload: bytes | memoryview, + numel: int) -> bytes: + """Non-streamed BCAST: header + tensor -> committed f32 anchor bytes.""" + if len(payload) < HEADER_LEN: + raise ValueError("downlink BCAST header truncated") + decode = self.begin(fragment_id, version, bytes(payload[:HEADER_LEN]), numel) + if len(payload) - HEADER_LEN != tensor_nbytes(decode.kind, numel): + raise ValueError(f"downlink BCAST {fragment_id} has {len(payload) - HEADER_LEN} " + f"tensor bytes, expected {tensor_nbytes(decode.kind, numel)}") + data = decode.feed(payload[HEADER_LEN:]) + decode.finish() + self.commit(decode) + return data + + +@dataclass +class _Transaction: + inner: object + passthrough: bool = True + fragment_id: int | None = None + version: int | None = None + numel: int = 0 + header: bytearray = field(default_factory=bytearray) + decode: DownlinkDecode | None = None + out_offset: int = 0 + + +class DownlinkChunkSink: + """Wraps a transactional inbound chunk sink for a downlink session. + + BCAST payload bytes (downlink header + f32/bf16 tensor) are decoded in + order and handed to the wrapped sink as the plain f32 fragment; FINAL + passes through. Out-of-order chunks are staged here as bytes (the + wrapped sink sees only in-order f32 bytes).""" + + def __init__(self, inner: object, anchors: DownlinkAnchors, numels: list[int], + bcast_type: int) -> None: + self.inner = inner + self.anchors = anchors + self.numels = numels + self.bcast_type = bcast_type + + def begin_message(self, message_id: int, total_bytes: int) -> _Transaction: + return _Transaction(self.inner.begin_message(message_id, total_bytes)) + + def bind_fragment(self, tx: _Transaction, msg_type: int, fragment_id: int, version: int, + payload_bytes: int) -> None: + if msg_type != self.bcast_type: + self.inner.bind_fragment(tx.inner, msg_type, fragment_id, version, payload_bytes) + return + tx.passthrough = False + tx.fragment_id, tx.version = fragment_id, version + tx.numel = self.numels[fragment_id] + self.inner.bind_fragment(tx.inner, msg_type, fragment_id, version, 4 * tx.numel) + + def stage_chunk(self, tx: _Transaction, payload) -> bytes: + del tx + return bytes(payload) + + def consume_staged_chunk(self, tx: _Transaction, offset: int, token: bytes, nbytes: int) -> None: + if len(token) != nbytes: + raise ValueError("staged downlink chunk is malformed") + self.consume_chunk(tx, offset, token) + + def consume_chunk(self, tx: _Transaction, offset: int, payload) -> None: + if tx.passthrough: + self.inner.consume_chunk(tx.inner, offset, payload) + return + data = memoryview(payload).cast("B") + if tx.decode is None: + take = min(HEADER_LEN - len(tx.header), data.nbytes) + tx.header += data[:take] + data = data[take:] + if len(tx.header) < HEADER_LEN: + return + tx.decode = self.anchors.begin(tx.fragment_id, tx.version, bytes(tx.header), tx.numel) + out = tx.decode.feed(data) + if out: + self.inner.consume_chunk(tx.inner, tx.out_offset, out) + tx.out_offset += len(out) + + def finish_message(self, tx: _Transaction): + if tx.passthrough: + return self.inner.finish_message(tx.inner) + if tx.decode is None: + raise ValueError("downlink BCAST ended inside its header") + tx.decode.finish() + committed = self.inner.finish_message(tx.inner) + self.anchors.commit(tx.decode) + return committed + + def abort_message(self, tx: _Transaction, staged_tokens) -> None: + del staged_tokens # plain bytes held here + self.inner.abort_message(tx.inner, ()) diff --git a/yeto/launch/fleet.py b/yeto/launch/fleet.py index b7051a46..d1d9ce63 100644 --- a/yeto/launch/fleet.py +++ b/yeto/launch/fleet.py @@ -48,6 +48,8 @@ def __init__(self): # time spent recovering in the window. FIXED_ROSTER_MAX_RELAUNCHES = 2 RECOVERY_STABLE_S = 300.0 +# S26 E3b: how long one rejected-link relaunch may take to get a new container. +WAN_RECOVERY_S = 900.0 # How long a failed-looking island is given for its rl_learner_finalized record # to arrive over the log stream before recovery starts. FINALIZE_GRACE_S = 60.0 @@ -154,6 +156,8 @@ def __init__( elastic: bool = False, no_island_relaunch: bool = False, replacement_planner=None, + wan_reject_probe=None, + wan_probe_relaunches: int = 0, ): """`learners` maps cluster name -> (task, job_id); `syncer` is (name, task, job_id) for a cluster syncer, or None with @@ -169,6 +173,12 @@ def __init__( self.left_pool_probe = left_pool_probe self.elastic = elastic self.no_island_relaunch = no_island_relaunch + # S26 E3b: name -> number of rejected link probes the island logged + # (rl_wan_probe accepted=false, exit 75). Such an island never joined + # the syncer session, so it is relaunched in a fresh container up to + # wan_probe_relaunches times, also under --no-island-relaunch. + self.wan_reject_probe = wan_reject_probe + self.wan_probe_relaunches = wan_probe_relaunches # rl-spot-cost-saving 2.3: yeto.cloud.replace.ReplacementPlanner or None. # Asked once per recovery window before the same-place relaunch. Auto # launch off (default): it only writes spot_replace_advice and the @@ -372,6 +382,8 @@ def _poll(self, rec, is_syncer: bool) -> None: self._fleet("island_done", rec["name"], reason="syncer finished") print(f"[launcher] {rec['name']}: {verdict} after the syncer finished; " "counted as a normal end, not relaunched", file=sys.stderr) + elif not is_syncer and self._rejected_link(rec): + self._relaunch_rejected_link(rec) elif not is_syncer and self.elastic and self._left_pool(rec): # 0.27: re-JOIN exhausted (exit 7): the island left the elastic pool. # Not the fixed-roster abandon path (exit 4): relaunch unless @@ -418,6 +430,36 @@ def _syncer_finished(self) -> bool: except Exception: return False + def _rejected_link(self, rec) -> bool: + """A new rejected link probe since the island's last (re)launch.""" + if self.wan_reject_probe is None: + return False + try: + seen = int(self.wan_reject_probe(rec["name"])) + except Exception: + return False + return seen > rec.get("wan_rejects_handled", 0) + + def _relaunch_rejected_link(self, rec) -> None: + rec["wan_rejects_handled"] = rec.get("wan_rejects_handled", 0) + 1 + used = rec.get("wan_relaunches", 0) + if used >= self.wan_probe_relaunches: + self._fleet("island_lost", rec["name"], reason="link probe rejected; relaunch budget used") + print(f"[launcher] {rec['name']}: link probe rejected again after {used} new " + "container(s); not relaunched", file=sys.stderr) + self._abandon(rec, 0.0, reason="link probe rejected") + return + rec["wan_relaunches"] = used + 1 + self._fleet("island_lost", rec["name"], + reason=f"link probe rejected; new container {used + 1}/{self.wan_probe_relaunches}") + print(f"[launcher] {rec['name']}: link probe rejected; requesting a new container " + f"({used + 1}/{self.wan_probe_relaunches})", file=sys.stderr) + rec["state"] = RECOVERING + rec["failed_at"] = self.ops.now() + rec["wan_recovery"] = True + rec.pop("recovered_at", None) + self._drive_recovery(rec, False) + def _left_pool(self, rec) -> bool: if self.left_pool_probe is None: return False @@ -577,6 +619,7 @@ def _drive_recovery(self, rec, is_syncer: bool) -> None: self.ops.now() - rec["failed_at"]) rec["failed_at"] = None rec["replace_planned"] = False + rec.pop("wan_recovery", None) rec["recovered_at"] = self.ops.now() print( f"[launcher] {rec['name']} recovered: relaunched as job " @@ -597,7 +640,9 @@ def _drive_recovery(self, rec, is_syncer: bool) -> None: if rec.get("failures", 0) > FIXED_ROSTER_MAX_RELAUNCHES: self._abandon(rec, elapsed) return - if self.recover_timeout <= 0 or elapsed > self.recover_timeout: + # A rejected-link relaunch has its own budget (the island never joined). + timeout = WAN_RECOVERY_S if rec.get("wan_recovery") else self.recover_timeout + if timeout <= 0 or elapsed > timeout: if is_syncer: # The syncer is never abandoned: without it no learner can # make outer progress, so keep trying and complain loudly. diff --git a/yeto/launch/modal_islands.py b/yeto/launch/modal_islands.py index a7f12af6..976211e3 100644 --- a/yeto/launch/modal_islands.py +++ b/yeto/launch/modal_islands.py @@ -315,3 +315,15 @@ def effective_recover_timeout(args) -> float: if getattr(args, "modal_retries", None) == 0 and not getattr(args, "modal_launcher_relaunch", False): return 0 return args.recover_timeout + + +def wan_probe_relaunches(args) -> int: + """New containers per island after a rejected link probe (S26 E3b); + 0 when no island runs the probe.""" + if not (float(getattr(args, "wan_probe_min_mbps", 0.0) or 0.0) > 0 + or float(getattr(args, "wan_probe_min_down_mbps", 0.0) or 0.0) > 0): + return 0 + count = int(getattr(args, "wan_probe_relaunches", 2)) + if count < 0: + raise ValueError("--wan-probe-relaunches must be >= 0") + return count diff --git a/yeto/launch/rl_args.py b/yeto/launch/rl_args.py index 0c21669b..e623e6dc 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -148,6 +148,40 @@ def sao_contract_mounts(args, learner_id: int) -> dict[str, str]: for shipped in (island.runtime, island.layout_attestation)} +_FULL_PARAMETER_SYNC_DEFAULTS = ( + ("rl_sync_fragment_pipeline", 1, "--rl-sync-fragment-pipeline"), + ("rl_sync_async_checkpoint", False, "--rl-sync-async-checkpoint"), + ("rl_sync_push_dtype", "f32", "--rl-sync-push-dtype"), + ("rl_sync_checkpoint_every", 1, "--rl-sync-checkpoint-every"), + ("rl_sync_push_int8_block", 128, "--rl-sync-push-int8-block"), + ("rl_sync_overlap", False, "--rl-sync-overlap"), +) + + +def _refuse_full_parameter_sync_options(args) -> None: + """The S26 E3 sync options exist only for the full-parameter preset + (the bf16 downlink also for strict-avg, see ``_check_bcast_delta``).""" + for dest, default, option in _FULL_PARAMETER_SYNC_DEFAULTS: + if getattr(args, dest, default) != default: + raise ValueError(f"{option} only applies to --rl-sync-preset {FULL_PARAMETER_SYNC_PRESET}") + + +def _check_bcast_delta(args) -> None: + """S26 E3 bf16 downlink: a legacy-scheduled strict-avg or full-parameter + syncer only (the elastic server and SAO have no downlink anchors).""" + full_every = int(getattr(args, "rl_sync_bcast_delta_full_every", 0) or 0) + if full_every == 0: + return + if full_every < 2: + raise ValueError("--rl-sync-bcast-delta-full-every 0 (off) or >= 2") + if args.rl_sync_preset not in ("strict-avg", FULL_PARAMETER_SYNC_PRESET): + raise ValueError("--rl-sync-bcast-delta-full-every only applies to --rl-sync-preset " + f"strict-avg or {FULL_PARAMETER_SYNC_PRESET}") + if island_scheduling_mode(args) != "legacy" or args.experimental_rl_sync: + raise ValueError("--rl-sync-bcast-delta-full-every needs --rl-island-scheduling legacy " + "and no --experimental-rl-sync") + + def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: """full-parameter preset (P0-D8): one actor role swept as whole-model fragments. @@ -169,11 +203,30 @@ def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: problems.append("--rl-island-scheduling legacy (no fragment frames in elastic mode)") if rl_needs_critic(args): problems.append("a critic-free algorithm (the strict sweep trains the actor only)") + fragment_pipeline = int(getattr(args, "rl_sync_fragment_pipeline", 1)) + if not 1 <= fragment_pipeline <= max(args.fragments, 1): + problems.append("--rl-sync-fragment-pipeline between 1 and --fragments") + checkpoint_every = int(getattr(args, "rl_sync_checkpoint_every", 1)) + if checkpoint_every != 1 and (checkpoint_every < 1 or checkpoint_every % max(args.fragments, 1)): + problems.append("--rl-sync-checkpoint-every 1 or a multiple of --fragments") + int8_block = int(getattr(args, "rl_sync_push_int8_block", 128)) + if not 32 <= int8_block <= 1 << 20: + problems.append("--rl-sync-push-int8-block between 32 and 1048576") + elif int8_block != 128 and getattr(args, "rl_sync_push_dtype", "f32") != "int8": + problems.append("--rl-sync-push-int8-block only with --rl-sync-push-dtype int8") + if getattr(args, "rl_sync_overlap", False): + from ..rl.engine.sync_overlap import overlap_option_problems + + problems += overlap_option_problems( + max_policy_age=getattr(args, "rl_max_policy_age", 0) or 0, + fragment_pipeline=fragment_pipeline, fragments=args.fragments) if problems: raise ValueError(f"--rl-sync-preset {FULL_PARAMETER_SYNC_PRESET} requires " + "; ".join(problems)) args.quorum = num_islands args.grace_ms = 0 - args.pipeline = 1 + # Fragment rounds of one sweep are independent (H=1, plain average), so the + # syncer may gather several at once; it still commits them in step order. + args.pipeline = fragment_pipeline args.sync_interval_steps = 0.0 args.delta_correction = "none" args.outer_lr = 1.0 @@ -723,6 +776,9 @@ def _prepare_rl_args( f"--rl-sync-preset {FULL_PARAMETER_SYNC_PRESET}") if args.rl_sync_preset != SAO_STREAMING_SYNC_PRESET and getattr(args, "rl_sao_contracts", None): raise ValueError("--rl-sao-contracts only applies to --rl-sync-preset sao-streaming-full") + if args.rl_sync_preset != FULL_PARAMETER_SYNC_PRESET: + _refuse_full_parameter_sync_options(args) + _check_bcast_delta(args) if args.rl_sync_preset == "decoupled": if args.experimental_rl_sync: raise ValueError("decoupled RL does not accept --experimental-rl-sync") diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index 1185f4c9..ce0ecea7 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -51,6 +51,34 @@ def submitted_unix() -> float: return time.time() +def _wan_transport_flags(args) -> str: + """S26 E3b island transport options; emitted only when not the default.""" + flags = "" + schedule = getattr(args, "wan_stream_schedule", "rr") + if schedule != "rr": + flags += f" --wan-stream-schedule {schedule}" + stall = float(getattr(args, "wan_stall_timeout", 0.0) or 0.0) + if stall < 0: + raise ValueError("--wan-stall-timeout must be >= 0 (0 = off)") + if stall > 0: + flags += f" --wan-stall-timeout {stall!r}" + min_up = float(getattr(args, "wan_probe_min_mbps", 0.0) or 0.0) + min_down = float(getattr(args, "wan_probe_min_down_mbps", 0.0) or 0.0) + if min_up < 0 or min_down < 0: + raise ValueError("--wan-probe-min-mbps / --wan-probe-min-down-mbps must be >= 0 (0 = off)") + if min_up > 0 or min_down > 0: + if getattr(args, "rl_elastic", False): + # Only the legacy syncer serves MSG_BW_PROBE; on the elastic server + # every probe would fail and every island would be rejected. + raise ValueError("--wan-probe-min-mbps / --wan-probe-min-down-mbps need the " + "legacy syncer; the elastic syncer has no link probe") + mib = int(getattr(args, "wan_probe_mib", 256)) + if not 1 <= mib <= 1024: + raise ValueError("--wan-probe-mib must be between 1 and 1024") + flags += f" --wan-probe-min-mbps {min_up!r} --wan-probe-min-down-mbps {min_down!r} --wan-probe-mib {mib}" + return flags + + def _rl_trainable_island_flags(args, learner_id: int, num_learners: int) -> str: """LoRA rank/targets; for the full-parameter preset the sweep identity (the probed layout hash, the frozen roster size and the syncer profile hash); @@ -72,6 +100,12 @@ def _rl_trainable_island_flags(args, learner_id: int, num_learners: int) -> str: sweep = (f" --parameter-layout-sha256 {args.rl_parameter_layout_sha256}" f" --num-learners {num_learners}" f" --syncer-profile-sha256 {rl_syncer_profile(args, num_learners).sha256}") + if getattr(args, "rl_sync_push_dtype", "f32") != "f32": + sweep += f" --sweep-push-dtype {args.rl_sync_push_dtype}" + if getattr(args, "rl_sync_overlap", False): + sweep += " --sweep-overlap" + if getattr(args, "rl_sync_push_dtype", "f32") == "int8": + sweep += f" --sweep-push-int8-block {int(getattr(args, 'rl_sync_push_int8_block', 128))}" if getattr(args, "expert_full_count", 0): return lora + sweep return " --parameter-mode full" + sweep @@ -150,6 +184,7 @@ def make_miles_island_task( f" --rollout-num-gpus-per-engine {args.rollout_num_gpus_per_engine}" f" --sglang-mem-fraction-static {args.sglang_mem_fraction_static}" + _rl_trainable_island_flags(args, learner_id, num_learners) + + (" --bcast-delta" if int(getattr(args, "rl_sync_bcast_delta_full_every", 0) or 0) else "") + f" --inner-lr {args.inner_lr}" f" --seq-len {args.seq_len}" + (" --rl-activation-recompute" if getattr(args, "rl_activation_recompute", False) else "") @@ -159,7 +194,8 @@ def make_miles_island_task( if float(getattr(args, "rl_optimizer_offload_fraction", 0.0) or 0.0) > 0 else "") + f" --seed {args.seed}" f" --wan-streams {args.wan_streams}" - " --miles-root ~/miles" + + _wan_transport_flags(args) + + " --miles-root ~/miles" ) if getattr(args, "wandb", False): flags += " --wandb" diff --git a/yeto/launch/syncer.py b/yeto/launch/syncer.py index 93b1aa6a..2836e1b0 100644 --- a/yeto/launch/syncer.py +++ b/yeto/launch/syncer.py @@ -145,6 +145,17 @@ def _syncer_quorum_timeout(args) -> str: return "" if value is None else f" --quorum-timeout-s {value}" +def _syncer_checkpoint_every(args) -> int: + """Per-round checkpoints (1) unless the full-parameter preset sets an interval.""" + return int(getattr(args, "rl_sync_checkpoint_every", 1)) if rl_dense_sweep(args) else 1 + + +def _syncer_bcast_delta(args) -> str: + """S26 E3 bf16 downlink deltas (checked by rl_args._check_bcast_delta).""" + full_every = int(getattr(args, "rl_sync_bcast_delta_full_every", 0) or 0) + return f" --bcast-delta-full-every {full_every}" if full_every else "" + + def rl_dense_sweep(args) -> bool: return (getattr(args, "training_mode", "sft") == "rl" and getattr(args, "rl_sync_preset", None) == FULL_PARAMETER_SYNC_PRESET) @@ -161,7 +172,8 @@ def _syncer_profile_binding(args) -> str: full-parameter adds strict P-fragment policy sweeps; sao-streaming-full adds the frozen contract's timeouts and the publishable final checkpoint.""" if rl_dense_sweep(args): - return f" --policy-sweep-fragments {args.fragments} --require-profile-binding" + async_checkpoint = " --async-checkpoint" if getattr(args, "rl_sync_async_checkpoint", False) else "" + return f" --policy-sweep-fragments {args.fragments} --require-profile-binding{async_checkpoint}" if rl_sao_streaming(args): profile = args.rl_sao_manifest.profile return (f" --quorum-timeout-s {profile.quorum_timeout_s}" @@ -205,7 +217,7 @@ def rl_syncer_profile(args, num_learners: int): "outer_lr": args.outer_lr, "outer_momentum": args.outer_momentum, "checkpoint_enabled": True, - "checkpoint_every": 1, + "checkpoint_every": _syncer_checkpoint_every(args), "resume": False, # not hashed (v2): fresh start and restart agree "mark_final_checkpoint": mark_final, "learner_budget_steps": None, @@ -442,11 +454,12 @@ def _rl_role_syncer_command(args, num_learners: int, binary: str, f"{_syncer_island_scheduling(args)}" f"{_syncer_expected_contract(args)}" f"{_syncer_profile_binding(args)}" + f"{_syncer_bcast_delta(args)}" f" --outer-lr {args.outer_lr}" f" --outer-momentum {args.outer_momentum}" " --max-base-lag 0 --learner-weight equal" f" --checkpoint-path {syncer.checkpoint}" - " --checkpoint-every 1" + f" --checkpoint-every {_syncer_checkpoint_every(args)}" f" {_resume_if_exists(syncer.checkpoint)}" f" --event-tape {syncer.event_tape}" ) diff --git a/yeto/launcher.py b/yeto/launcher.py index 0abd2f4b..7819fbc5 100644 --- a/yeto/launcher.py +++ b/yeto/launcher.py @@ -61,7 +61,7 @@ from .launch.modal_islands import ( _modal_tape_run_dir, _tail_modal, _verify_modal_app_stopped, build_modal_island_config, container_set_poller_for, effective_recover_timeout, pull_modal_tapes, - require_modal_for_gpu_exact, + require_modal_for_gpu_exact, wan_probe_relaunches, ) from .launch.fleet import ( CONTAINER_CHANGED_EXIT, DEFAULT_RL_STALL_TIMEOUT_S, FAILED_RUN_DRAIN_S, ISLAND_FAILED_EXIT, @@ -660,6 +660,11 @@ def container_changed() -> bool: ), elastic=island_scheduling_mode(args) == "elastic", no_island_relaunch=bool(getattr(args, "no_island_relaunch", False)), + wan_reject_probe=( + (lambda name: event_collectors[name].wan_rejects if name in event_collectors else 0) + if echo_names else None + ), + wan_probe_relaunches=wan_probe_relaunches(args), ) def drain_tapes(limit: float = NO_SYNC_EVENT_DRAIN_S) -> None: if echo_names: diff --git a/yeto/protocol.py b/yeto/protocol.py index 7beb2254..e55444b5 100644 --- a/yeto/protocol.py +++ b/yeto/protocol.py @@ -34,10 +34,14 @@ import struct import threading import time -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Iterator from dataclasses import dataclass, field +from .downlink import HEADER_LEN as DOWNLINK_HEADER_LEN +from .downlink import DownlinkAnchors, DownlinkChunkSink +from .downlink import tensor_nbytes as downlink_tensor_nbytes from .fragments import MERGE_ISO, FragmentLayout +from .wan_streams import STREAM_SCHEDULES, StreamMeter, StreamSnapshot, StreamStalledError, pick_stream MAGIC = 0xD170C0DE PROTOCOL_VERSION = 4 @@ -56,6 +60,8 @@ MSG_FINAL_ACK = 12 MSG_FINAL_FRAGMENT = 13 MSG_BUDGET_DONE = 14 +# S26 E3b: link-bandwidth probe connection (yeto/wan_probe.py; syncer serve_bw_probe). +MSG_BW_PROBE = 26 FINALIZATION_REVISION = 1 FINALIZATION_TIMEOUT = 900.0 @@ -66,6 +72,20 @@ # base-relative learner deltas; INIT_PARAMS and BCAST_FRAGMENT stay bf16 # (see bulk_dtype and docs/PROTOCOL.md v4). DTYPE_Q4 = 3 +# Session dtype 4 (S26 E3): PUSH_FRAGMENT deltas travel as bf16; INIT_PARAMS, +# BCAST_FRAGMENT and FINAL stay exact f32. Callers still hand the client f32 +# delta bytes; the client rounds them (round-to-nearest-even, as the syncer's +# ``half::bf16::from_f32``) on the way out. +DTYPE_F32_PUSH_BF16 = 4 +# Session dtype 5 (S26 E3): PUSH_FRAGMENT deltas travel as block-scaled int8 +# (``f32_bytes_to_i8``); INIT/BCAST/FINAL stay f32. Callers hand f32 bytes. +DTYPE_F32_PUSH_I8 = 5 +# Session-dtype flag (S26 E3): BCAST_FRAGMENT carries bf16 downlink deltas +# between f32 refreshes (yeto/downlink.py). Only on f32-bulk sessions. +DTYPE_FLAG_BCAST_DELTA = 0x80 +I8_BLOCK_MIN = 32 +I8_BLOCK_MAX = 1 << 20 +DEFAULT_I8_BLOCK = 128 CHUNK_SIZE = 4 * 1024 * 1024 MAX_HANDSHAKE_FRAME = 16 * 1024 * 1024 @@ -87,13 +107,127 @@ _BUDGET_DONE = struct.Struct(" int: + """The session dtype without its flags.""" + return dtype & ~DTYPE_FLAG_BCAST_DELTA + + +def bcast_delta(dtype: int) -> bool: + """Whether the session receives bf16 downlink deltas.""" + return bool(dtype & DTYPE_FLAG_BCAST_DELTA) + + +def valid_session_dtype(dtype: int) -> bool: + base = session_base(dtype) + if base in (DTYPE_F32, DTYPE_F32_PUSH_BF16, DTYPE_F32_PUSH_I8): + return 0 <= dtype <= 0xFF + return base in (DTYPE_BF16, DTYPE_Q4) and not bcast_delta(dtype) + + def bulk_dtype(dtype: int) -> int: """Dtype of INIT_PARAMS/BCAST_FRAGMENT tensors for a session dtype. Q4 applies only to push deltas (small dynamic range); full parameter - payloads would not survive 4 bits, so they travel as bf16. + payloads would not survive 4 bits, so they travel as bf16. A downlink + session's BCAST is reconstructed to f32 before callers see it. + """ + base = session_base(dtype) + if base == DTYPE_Q4: + return DTYPE_BF16 + if base in (DTYPE_F32_PUSH_BF16, DTYPE_F32_PUSH_I8): + return DTYPE_F32 + return base + + +def push_wire_dtype(dtype: int) -> int: + """Dtype of unquantized PUSH_FRAGMENT deltas on the wire for a session dtype.""" + base = session_base(dtype) + return DTYPE_BF16 if base == DTYPE_F32_PUSH_BF16 else base + + +def i8_nbytes(numel: int, block: int) -> int: + """Encoded size of an int8 push: u32 block, then per block f32 scale + i8s.""" + if not I8_BLOCK_MIN <= block <= I8_BLOCK_MAX: + raise ValueError(f"int8 block size {block} outside {I8_BLOCK_MIN}..{I8_BLOCK_MAX}") + return 4 + 4 * -(-numel // block) + numel + + +def f32_bytes_to_i8(data: bytes | bytearray | memoryview, block: int, *, header: bool = True) -> bytes: + """Quantize little-endian f32 values to block-scaled int8 (syncer ``decode_i8``). + + Per block of ``block`` values (the last may be shorter): f32 scale = + absmax / 127, then ``round_half_even(x / scale)`` clipped to [-127, 127]. + Decoding ``q * scale`` errs by at most absmax / 254 per element (plus f32 + rounding). Blocks must start at multiples of ``block`` values; with + ``header`` the u32 block size is prepended. """ - return DTYPE_BF16 if dtype == DTYPE_Q4 else dtype + import numpy as np + + i8_nbytes(0, block) # validates the block size + values = np.frombuffer(data, dtype=" None: # chunk: (blocks, width) + absmax = np.abs(chunk).max(axis=1) + scale = (absmax / np.float32(127)).astype(np.float32) + safe = np.where(scale > 0, scale, np.float32(1)) + q = np.clip(np.rint(chunk / safe[:, None]), -127, 127).astype(np.int8) + q[scale == 0] = 0 + rows = np.empty((chunk.shape[0], 4 + chunk.shape[1]), dtype=np.uint8) + rows[:, :4] = scale.astype(" full: + encode(values[full:].reshape(1, -1)) + return b"".join(parts) + + +def _f32_parts_to_i8(parts: Iterable[bytes | bytearray | memoryview], block: int) -> Iterator[bytes]: + """Stream ``f32_bytes_to_i8`` over parts; blocks may span part boundaries.""" + yield struct.pack(" bytes: + """Round little-endian f32 values to bf16 (nearest, ties to even).""" + import numpy as np # noqa: PLC0415 - keep the protocol import light + + bits = np.frombuffer(data, dtype="> 16) & 1)) >> 16 + nan = (bits & 0x7FFF_FFFF) > 0x7F80_0000 + rounded[nan] = (bits[nan] >> 16) | 0x40 # keep NaN quiet, never round to inf + return rounded.astype(" Iterator[bytes]: + """Stream ``f32_bytes_to_bf16`` over parts whose boundaries may split a value.""" + carry = b"" + for part in parts: + data = carry + bytes(part) + whole = len(data) - len(data) % 4 + carry = data[whole:] + if whole: + yield f32_bytes_to_bf16(data[:whole]) + if carry: + raise ValueError("f32 push delta ends mid-value") def layout_fingerprint(layout: FragmentLayout) -> bytes: @@ -168,8 +302,8 @@ def seal_hello(body: bytes, key: bytes | None) -> bytes: The legacy syncer with a key refuses a HELLO without this trailer.""" if not key: return body - import hmac import hashlib + import hmac return body + hmac.new(key, HELLO_MAC_DOMAIN + body, hashlib.sha256).digest() @@ -440,6 +574,7 @@ class SendProgress: queued_chunks: int reconnects: int connected: bool + stalls: int = 0 @dataclass @@ -463,6 +598,10 @@ class SyncerClient: forever); once exhausted, check_health raises. """ + # Transport defaults (S26 E3b); __init__ overrides them per client. + stream_schedule: str = "rr" + stream_stall_timeout: float | None = None + def __init__( self, addr: tuple[str, int], @@ -477,11 +616,19 @@ def __init__( syncer_profile_hash: bytes | None = None, hmac_key: bytes | None = None, compat_group: str | None = None, + push_i8_block: int = DEFAULT_I8_BLOCK, + stream_schedule: str = "rr", + stream_stall_timeout: float | None = None, ): + if stream_schedule not in STREAM_SCHEDULES: + raise ValueError(f"stream_schedule must be one of {STREAM_SCHEDULES}, got {stream_schedule!r}") + if stream_stall_timeout is not None and not stream_stall_timeout > 0: + raise ValueError("stream_stall_timeout must be positive or None") if not 0 <= learner_id <= 0xFFFF_FFFF: raise ValueError(f"learner_id must fit u32, got {learner_id}") - if dtype not in (DTYPE_F32, DTYPE_BF16, DTYPE_Q4): + if not valid_session_dtype(dtype): raise ValueError(f"unsupported session dtype {dtype}") + i8_nbytes(0, push_i8_block) # validates the int8 block size if layout.num_fragments < 1: raise ValueError("layout must contain at least one fragment") if not 0 <= num_streams <= 256: @@ -499,7 +646,27 @@ def __init__( self.learner_id = learner_id self.layout = layout self.dtype = dtype + self._push_base = session_base(dtype) + self.push_i8_block = push_i8_block + # S26 E3: f32 anchors of a bf16-downlink session (yeto/downlink.py). + self._downlink = DownlinkAnchors(layout.num_fragments) if bcast_delta(dtype) else None + self._bcast_prefix = 12 + (DOWNLINK_HEADER_LEN if self._downlink is not None else 0) self.num_streams = num_streams + # S26 E3b: chunk placement over the data streams and the per-stream + # no-progress timeout (None = off, the historical behaviour). + self.stream_schedule = stream_schedule + self.stream_stall_timeout = stream_stall_timeout + # Index 0 is the control socket; the meters outlive connection groups + # so their counters are cumulative for the whole session. + self._meters = [StreamMeter() for _ in range(num_streams + 1)] + self._upload_drained_at: float | None = None + self._stalls = 0 + # S26 E3b: set once the initial cut arrived. Before that a redial must + # replay INIT_PARAMS (learner 0), and TCP_USER_TIMEOUT stays unset: + # the multi-GB INIT upload and initial broadcast must not be reset by + # the kernel on a slow link (s26-e3b-p06-C32a). + self._initial_cut_received = threading.Event() + self._init_generation: int | None = None self.session_contract_hash = session_contract_hash self.syncer_profile_hash = syncer_profile_hash # secret-handling-hardening: HELLO carries an HMAC when the island has @@ -510,7 +677,7 @@ def __init__( self.max_reconnects = max_reconnects self.finalization_timeout = finalization_timeout self._max_bcast_payload = max( - 12 + _tensor_nbytes(bulk_dtype(dtype), fragment.numel) + self._bcast_prefix + _tensor_nbytes(bulk_dtype(dtype), fragment.numel) for fragment in layout.fragments ) self._max_final_payload = max( @@ -544,6 +711,7 @@ def __init__( | _StreamedReassembly, ] = {} self._reasm_lock = threading.Lock() + self._reasm_touched: dict[int, float] = {} # msg_id -> last chunk (monotonic) self._inbound_chunk_sink: object | None = None self._msg_id = itertools.count() self._rr = itertools.count() @@ -564,6 +732,7 @@ def __init__( self._bytes_sent_lock = threading.Lock() self._reset_bcasts_on_reconnect = False self._supervisor: threading.Thread | None = None + self._watchdog: threading.Thread | None = None # -- lifecycle ------------------------------------------------------------- @@ -572,6 +741,8 @@ def _dial(self, timeout: float) -> socket.socket: sock.settimeout(None) sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) enable_tcp_keepalive(sock) + if self.stream_stall_timeout is not None and self._initial_cut_received.is_set(): + set_tcp_user_timeout(sock, self.stream_stall_timeout) return sock def _connect_one(self) -> socket.socket: @@ -602,6 +773,9 @@ def start(self) -> None: self._connect_group(patient=True) self._supervisor = threading.Thread(target=self._supervise, name="yeto-supervisor", daemon=True) self._supervisor.start() + if self.stream_stall_timeout is not None: + self._watchdog = threading.Thread(target=self._watch_stalls, name="yeto-stall-watchdog", daemon=True) + self._watchdog.start() def _connect_group(self, patient: bool) -> None: """Dial control + data sockets, send HELLO/DATA_HELLO, start workers. @@ -659,7 +833,7 @@ def _connect_group(self, patient: bool) -> None: q: queue.Queue[_Outbound | None] = queue.Queue(maxsize=256) self._queues.append(q) ts = threading.Thread( - target=self._send_loop, args=(gen, s, q), name=f"yeto-send-{i}", daemon=True + target=self._send_loop, args=(gen, s, q, self._meters[i]), name=f"yeto-send-{i}", daemon=True ) tr = threading.Thread(target=self._recv_loop, args=(gen, s), name=f"yeto-recv-{i}", daemon=True) self._threads += [ts, tr] @@ -684,6 +858,8 @@ def _teardown_group(self) -> None: pass for s in socks: _close_socket(s) + for meter in self._meters: + meter.drop_pending() self._clear_partial_reassemblies() self._drain(self._pulls) # stale pull requests; answering them is pointless with self._lock: @@ -776,10 +952,72 @@ def send_progress(self) -> SendProgress: queued_chunks=queued, reconnects=self._reconnects_used, connected=connected, + stalls=self._stalls, ) + def stream_snapshots(self) -> tuple[StreamSnapshot, ...]: + """Cumulative counters of each data stream (control stream excluded; + with ``num_streams == 0`` the control stream carries the data).""" + meters = self._meters[1:] if self.num_streams else self._meters[:1] + return tuple(meter.snapshot() for meter in meters) + + @property + def upload_drained_at(self) -> float | None: + """Unix time when the data streams last became idle with nothing + queued, i.e. when this island's latest upload left the host.""" + return self._upload_drained_at + # -- learner-facing API ------------------------------------------------------ + def send_initial_fragments( + self, + parts_for: Callable[[int], Iterable[bytes | bytearray | memoryview]], + ) -> bool: + """Learner 0: stream INIT_PARAMS of every fragment on the live group. + + ``parts_for(fragment_id)`` returns a fresh part iterable each call, so + the same INIT can be replayed after a redial. Returns True when every + fragment was queued on one connection generation; False on an outage + (nothing to undo: the syncer ignores INIT it already applied, and a + partial INIT of a dead generation is dropped by the syncer).""" + generation = self.connection_generation + try: + for fragment_id in range(self.layout.num_fragments): + if self.send_init_parts(fragment_id, parts_for(fragment_id)) is not True: + return False + except PartialMessageGenerationLost: + return False + if generation != self.connection_generation: + return False + self._init_generation = generation + return True + + def init_replay_due(self) -> bool: + """True when the live group never carried a complete INIT and the + initial cut has not arrived yet (learner 0 must send INIT again).""" + with self._lock: + live = self._connected.is_set() + generation = self.connection_generation + return ( + live + and not self._initial_cut_received.is_set() + and generation != self._init_generation + ) + + def mark_initial_cut_received(self) -> None: + """The initial cut arrived: INIT is no longer replayed, and the stall + timeout now also bounds unacknowledged data in the kernel.""" + self._initial_cut_received.set() + if self.stream_stall_timeout is None: + return + with self._lock: + socks = list(self._socks) + for sock in socks: + try: + set_tcp_user_timeout(sock, self.stream_stall_timeout) + except OSError: + pass # a dying socket; the group is being replaced anyway + def send_init(self, fragment_id: int, tensor_bytes: bytes) -> None: if not 0 <= fragment_id < self.layout.num_fragments: raise ValueError(f"INIT_PARAMS for unknown fragment {fragment_id}") @@ -832,8 +1070,9 @@ def push_fragment( numel = self.layout.fragments[fragment_id].numel expected = ( _q4_nbytes(numel) - if self.dtype == DTYPE_Q4 - else _tensor_nbytes(self.dtype, numel) + if self._push_base == DTYPE_Q4 + # dtypes 4/5 callers pass f32 bytes; the client encodes them below + else _tensor_nbytes(bulk_dtype(self.dtype), numel) ) if len(tensor_bytes) != expected: raise ValueError( @@ -851,6 +1090,10 @@ def push_fragment( c_steps, c_tokens, ) + if self._push_base == DTYPE_F32_PUSH_BF16: + tensor_bytes = f32_bytes_to_bf16(tensor_bytes) + elif self._push_base == DTYPE_F32_PUSH_I8: + tensor_bytes = f32_bytes_to_i8(tensor_bytes, self.push_i8_block) self._send_large(MSG_PUSH_FRAGMENT, head + tensor_bytes) def push_fragment_parts( @@ -882,11 +1125,15 @@ def push_fragment_parts( if c_steps < 1: raise ValueError("c_steps must be positive") numel = self.layout.fragments[fragment_id].numel - expected = ( - _q4_nbytes(numel) - if self.dtype == DTYPE_Q4 - else _tensor_nbytes(self.dtype, numel) - ) + if self._push_base == DTYPE_Q4: + expected = _q4_nbytes(numel) + elif self._push_base == DTYPE_F32_PUSH_I8: + expected = i8_nbytes(numel, self.push_i8_block) + tensor_parts = _f32_parts_to_i8(tensor_parts, self.push_i8_block) + else: + expected = _tensor_nbytes(push_wire_dtype(self.dtype), numel) + if self._push_base == DTYPE_F32_PUSH_BF16: + tensor_parts = _f32_parts_to_bf16(tensor_parts) head = struct.pack( " None: raise TypeError( "inbound chunk sink does not implement the transaction API" ) + if sink is not None and self._downlink is not None: + sink = DownlinkChunkSink( + sink, self._downlink, [f.numel for f in self.layout.fragments], MSG_BCAST_FRAGMENT) with self._reasm_lock: if self._reasm: raise RuntimeError("cannot replace inbound chunk sink mid-message") @@ -1129,6 +1379,8 @@ def _enqueue( if gen is not None and gen != self._gen: return False q = self._queues[stream] + meter = self._meters[stream] + meter.enqueued(len(data)) # before put: the sender may drain it at once try: q.put( _Outbound(data, sent), @@ -1136,8 +1388,16 @@ def _enqueue( ) # bounded wait, then re-check the group return True except queue.Full: + meter.enqueued(-len(data)) continue + def _next_data_stream(self) -> int: + """Socket index (1-based; 0 only without data streams) for the next + CHUNK envelope, by ``stream_schedule``.""" + if self.num_streams == 0: + return 0 + return 1 + pick_stream(self.stream_schedule, self._meters[1:], next(self._rr)) + def _send_large(self, msg_type: int, payload: bytes) -> None: inner = self._frame(msg_type, payload) if self.num_streams == 0: @@ -1152,7 +1412,7 @@ def _send_large(self, msg_type: int, payload: bytes) -> None: for offset in range(0, total, CHUNK_SIZE): chunk = inner[offset : offset + CHUNK_SIZE] envelope = self._frame(MSG_CHUNK, _CHUNK_HEAD.pack(msg_id, total, offset) + chunk) - stream = 1 + next(self._rr) % self.num_streams + stream = self._next_data_stream() if not self._enqueue(stream, envelope, gen=gen): return # group died mid-message; drop the remainder @@ -1273,11 +1533,7 @@ def _send_large_parts( before_last_enqueue = None callback() - stream = ( - 0 - if self.num_streams == 0 - else 1 + next(self._rr) % self.num_streams - ) + stream = self._next_data_stream() if not self._enqueue(stream, envelope, gen=gen): if enqueued_chunks: raise PartialMessageGenerationLost( @@ -1337,20 +1593,85 @@ def _protocol_failed(self, gen: int, exc: BaseException) -> None: with self._final_cond: self._final_cond.notify_all() - def _send_loop(self, gen: int, sock: socket.socket, q: queue.Queue) -> None: + def _send_loop( + self, + gen: int, + sock: socket.socket, + q: queue.Queue, + meter: StreamMeter | None = None, + ) -> None: + if meter is None: # a sender outside a connection group (tests) + meter = StreamMeter() try: while True: item = q.get() if item is None: return - sock.sendall(item.data) + # send() instead of sendall() so the meter sees byte progress + # inside one 4 MiB envelope (stall watchdog, per-stream rate). + view = memoryview(item.data) + meter.begin() + try: + while view: + sent = sock.send(view) + meter.progress(sent) + view = view[sent:] + finally: + meter.end() + view.release() with self._bytes_sent_lock: self._bytes_sent += len(item.data) + if all(m.idle for m in self._meters): + self._upload_drained_at = time.time() if item.sent is not None: item.sent.set() except BaseException as e: self._socket_failed(gen, e) + def _stalled_stream(self) -> str | None: + """Describe the first stream without progress for the stall timeout.""" + limit = self.stream_stall_timeout + now = time.monotonic() + for index, meter in enumerate(self._meters): + idle_for = meter.stalled_for(now) + if idle_for > limit: + return f"outbound stream {index} made no progress for {idle_for:.0f}s" + with self._reasm_lock: + touched = dict(self._reasm_touched) + for msg_id, last in touched.items(): + if now - last > limit: + return f"inbound message {msg_id} got no chunk for {now - last:.0f}s" + return None + + def _watch_stalls(self) -> None: + """Fail the live group when one stream stops moving bytes. + + A stalled TCP stream can sit for many minutes before the kernel gives + up (C32: 900 s with every socket still open). The failure goes + through the normal reconnect path, so the syncer replays its PULLs + and the island pushes the interrupted fragment again.""" + period = min(1.0, self.stream_stall_timeout / 4) + while not self._closed.wait(period): + if self.shutdown.is_set(): + return + with self._lock: + gen = self._gen + live = self._connected.is_set() + if not live: + continue + reason = self._stalled_stream() + if reason is None: + continue + with self._lock: + if gen != self._gen: + continue + self._stalls += 1 + socks = list(self._socks) + _connect_log(f"{reason} (learner {self.learner_id}); dropping the connection group") + self._socket_failed(gen, StreamStalledError(reason)) + for sock in socks: # unblock senders stuck in send() + _close_socket(sock) + def _recv_loop(self, gen: int, sock: socket.socket) -> None: try: while True: @@ -1419,6 +1740,7 @@ def _reassemble( sink=sink, transaction=sink.begin_message(msg_id, total), ) + self._reasm_touched[msg_id] = time.monotonic() partial = self._reasm[msg_id] if isinstance(partial, _StreamedReassembly): if partial.total != total: @@ -1438,6 +1760,7 @@ def _reassemble( raise if completed is not None: self._reasm.pop(msg_id, None) + self._reasm_touched.pop(msg_id, None) return completed buf, filled, ranges = partial if len(buf) != total: @@ -1452,6 +1775,7 @@ def _reassemble( if filled[0] != total: raise ValueError("chunk byte count exceeds frame length") del self._reasm[msg_id] + self._reasm_touched.pop(msg_id, None) magic, msg_type, length = _HEADER.unpack_from(buf) if magic != MAGIC or length != total - _HEADER.size: raise ValueError("corrupt reassembled frame") @@ -1499,6 +1823,13 @@ def _stream_chunk_locked( dtype, self.layout.fragments[fragment_id].numel, ) + if msg_type == MSG_BCAST_FRAGMENT and self._downlink is not None: + # The downlink header (kind, base version, digest) is part of + # the payload the wrapping sink decodes; its kind fixes the size. + if len(data) < tensor_start + DOWNLINK_HEADER_LEN: + raise ValueError("first downlink fragment chunk is header-truncated") + expected_payload = DOWNLINK_HEADER_LEN + downlink_tensor_nbytes( + data[tensor_start], self.layout.fragments[fragment_id].numel) if length != 12 + expected_payload: raise ValueError( f"streamed fragment {fragment_id} has {length - 12} tensor " @@ -1612,6 +1943,7 @@ def _clear_partial_reassemblies(self) -> None: with self._reasm_lock: partials = list(self._reasm.values()) self._reasm.clear() + self._reasm_touched.clear() for partial in partials: if isinstance(partial, _StreamedReassembly): try: @@ -1767,15 +2099,19 @@ def _dispatch_live( fid, version = struct.unpack_from("= self.layout.num_fragments: raise ValueError(f"BCAST_FRAGMENT for unknown fragment {fid}") - expected = _tensor_nbytes( - bulk_dtype(self.dtype), self.layout.fragments[fid].numel - ) - if len(payload) - 12 != expected: - raise ValueError( - f"BCAST_FRAGMENT {fid} has {len(payload) - 12} tensor bytes, " - f"expected {expected}" + if self._downlink is not None: + data = self._downlink.decode_whole( + fid, version, payload[12:], self.layout.fragments[fid].numel) + else: + expected = _tensor_nbytes( + bulk_dtype(self.dtype), self.layout.fragments[fid].numel ) - data = payload[12:] + if len(payload) - 12 != expected: + raise ValueError( + f"BCAST_FRAGMENT {fid} has {len(payload) - 12} tensor bytes, " + f"expected {expected}" + ) + data = payload[12:] digest = hashlib.sha256(data).digest() with self._bcast_lock: seen = self._bcast_seen[fid] @@ -1882,6 +2218,14 @@ def enable_tcp_keepalive(sock: socket.socket) -> None: sock.setsockopt(socket.IPPROTO_TCP, option, value) +def set_tcp_user_timeout(sock: socket.socket, seconds: float) -> None: + """Bound how long unacknowledged data may sit before the kernel resets + the connection (Linux TCP_USER_TIMEOUT); other platforms keep defaults.""" + option = getattr(socket, "TCP_USER_TIMEOUT", None) + if option is not None: + sock.setsockopt(socket.IPPROTO_TCP, option, max(1, int(seconds * 1000))) + + def _close_socket(sock: socket.socket) -> None: try: sock.shutdown(socket.SHUT_RDWR) diff --git a/yeto/rl/adapters/miles/full_parameter/role_stream.py b/yeto/rl/adapters/miles/full_parameter/role_stream.py index e1c0f6b2..549743f7 100644 --- a/yeto/rl/adapters/miles/full_parameter/role_stream.py +++ b/yeto/rl/adapters/miles/full_parameter/role_stream.py @@ -20,18 +20,28 @@ from typing import Any from yeto.protocol import ( + DEFAULT_I8_BLOCK, DTYPE_F32, + DTYPE_F32_PUSH_BF16, + DTYPE_F32_PUSH_I8, + DTYPE_FLAG_BCAST_DELTA, + I8_BLOCK_MAX, + I8_BLOCK_MIN, FinalManifest, PullRequest, SyncerClient, ) +from yeto.rl.adapters.miles.full_parameter.inbound_sink import MilesRayInboundChunkSink from yeto.rl.adapters.miles.models.chunked_full_parameter import ( MilesChunkedFullParameterAdapter, ReferencedPolicyCut, StoredAuthoritativeFragment, ) -from yeto.rl.adapters.miles.full_parameter.inbound_sink import MilesRayInboundChunkSink -from yeto.rl.engine.policy_digest import PolicyDigest, TrainerResidentState, WeightsChanged +from yeto.rl.engine.policy_digest import ( + PolicyDigest, + TrainerResidentState, + WeightsChanged, +) from yeto.rl.engine.trainable_state import LAYOUT_FRAGMENTS, UnsupportedLayoutError from yeto.rl.local_learner import ComponentIdentity, layout_stream_role from yeto.rl.trajectory_evidence import TrajectoryBatchEvidence @@ -78,10 +88,36 @@ class RoleStreamConfig: # P0-4.8: the island's AlgorithmSpec hash; placement capabilities stay out. algorithm_spec_sha256: str wan_streams: int = 4 + # S26 E3b transport options (yeto/wan_streams.py); not part of the session + # contract: they change how bytes travel, never which bytes. + wan_stream_schedule: str = "rr" + wan_stall_timeout: float = 0.0 # seconds; 0 = off wait_timeout: float = 900.0 poll_seconds: float = 0.01 max_fragment_bytes: int = 2 << 30 max_chunk_bytes: int = 256 << 20 + # S26 E3 wire options (default exact f32): push outer deltas as bf16 + # (session dtype 4) or block-scaled int8 (dtype 5); receive BCAST as bf16 + # deltas against a shared f32 anchor. FINAL always stays f32. + push_dtype: str = "f32" + push_int8_block: int = DEFAULT_I8_BLOCK + bcast_delta: bool = False + # S26 E3: train round k+1 while sweep k syncs, then rebase the trainer + # (see FullParameterSync._overlap_boundary); the local weights lag the + # global policy by one outer step (yeto.rl.engine.sync_overlap). + overlap: bool = False + + @property + def session_dtype(self) -> int: + base = {"f32": DTYPE_F32, "bf16": DTYPE_F32_PUSH_BF16, "int8": DTYPE_F32_PUSH_I8}[self.push_dtype] + return base | (DTYPE_FLAG_BCAST_DELTA if self.bcast_delta else 0) + + @property + def wire_dtype_label(self) -> str: + """Session-contract label; plain f32 keeps the historical ``fp32``.""" + label = {"f32": "fp32", "bf16": "fp32-push-bf16", + "int8": f"fp32-push-int8-b{self.push_int8_block}"}[self.push_dtype] + return label + ("+bcast-bf16-delta" if self.bcast_delta else "") def __post_init__(self) -> None: if not self.algorithm.isidentifier(): @@ -126,6 +162,14 @@ def __post_init__(self) -> None: value = getattr(self, name) if type(value) is not int or value < 1: raise ValueError(f"role-stream {name} must be positive") + if self.wan_stream_schedule not in ("rr", "adaptive"): + raise ValueError("role-stream wan stream schedule must be rr or adaptive") + if not self.wan_stall_timeout >= 0: + raise ValueError("role-stream wan stall timeout must be >= 0") + if self.push_dtype not in ("f32", "bf16", "int8"): + raise ValueError("role-stream push dtype must be f32, bf16 or int8") + if type(self.push_int8_block) is not int or not I8_BLOCK_MIN <= self.push_int8_block <= I8_BLOCK_MAX: + raise ValueError("role-stream int8 block size is invalid") if self.wait_timeout <= 0 or not 0 < self.poll_seconds <= 1: raise ValueError("role-stream wait bounds are invalid") if ( @@ -219,7 +263,8 @@ def role_stream_session_contract_hash( "pipeline_depth": config.pipeline_depth, "local_horizon": config.local_horizon, "optimizer_steps_per_round": config.optimizer_steps_per_round, - "wire_dtype": "fp32", + "wire_dtype": config.wire_dtype_label, + **({"sync_train_overlap": True} if config.overlap else {}), }, } encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() @@ -261,6 +306,11 @@ def __init__(self, config: RoleStreamConfig, *, ray_module=None) -> None: self._staged: dict[tuple[int, int], StoredAuthoritativeFragment] = {} self._pinned_cuts: dict[int, tuple[ReferencedPolicyCut, int]] = {} self._pinned_anchors: dict[int, tuple[StoredAuthoritativeFragment, int]] = {} + # Overlap: the cut the current local round started from, and the + # round/token counters at that point (the next push's c_steps base). + self._round_base: ReferencedPolicyCut | None = None + self._rounds_at_base = 0 + self._tokens_at_base = 0 self.terminal_manifest: FinalManifest | None = None self.terminal_submitted = False self.final_applied = False @@ -334,8 +384,11 @@ async def initialize(self, group) -> None: self.config.syncer_addr, self.config.learner_id, adapter.layout.fragments, - dtype=DTYPE_F32, + dtype=self.config.session_dtype, + push_i8_block=self.config.push_int8_block, num_streams=self.config.wan_streams, + stream_schedule=self.config.wan_stream_schedule, + stream_stall_timeout=self.config.wan_stall_timeout or None, max_reconnects=None, session_contract_hash=role_stream_session_contract_hash( adapter.layout.layout_hash, @@ -347,27 +400,23 @@ async def initialize(self, group) -> None: client.install_inbound_chunk_sink(MilesRayInboundChunkSink(sink)) try: client.start() + init_parts = None if self.config.learner_id == 0: - for fragment_id in range(adapter.layout.fragments.num_fragments): - queued = client.send_init_parts( - fragment_id, - adapter.fragment_parts( - initial, - fragment_id, - ray_module=self.ray_module, - )(), - ) - if queued is not True: - raise RuntimeError( - f"{self.config.algorithm} {self.config.role} initial fragment was not queued" - ) + def init_parts(fragment_id: int): + return adapter.fragment_parts( + initial, fragment_id, ray_module=self.ray_module + )() + + # An outage here is not fatal: the wait below replays INIT on + # the next connection group until the initial cut arrives. + client.send_initial_fragments(init_parts) # Wait off the event loop: an island initializes its roles # concurrently, and each role's initial cut needs every island's # INIT for that role. A blocking wait here would stop this # island's other role from connecting, and two islands blocked # on different roles deadlock (s23-g3-sao-main-20261014a). authoritative = await asyncio.to_thread( - self._wait_for_initial_fragments, client, adapter + self._wait_for_initial_fragments, client, adapter, init_parts ) adapter.verify_initial_fragments(initial, authoritative) except BaseException: @@ -390,6 +439,7 @@ def _wait_for_initial_fragments( self, client, adapter: MilesChunkedFullParameterAdapter, + init_parts=None, ) -> tuple[StoredAuthoritativeFragment, ...]: started = time.monotonic() deadline = started + self.config.wait_timeout @@ -399,6 +449,11 @@ def _wait_for_initial_fragments( self._report_initial_cut(client, observed, count, started, "waiting") while len(observed) < count: client.check_health() + if init_parts is not None and client.init_replay_due(): + # S26 E3b (s26-e3b-p06-C32a): the group carrying INIT died + # before the syncer applied it; send it again on this group. + self._report_initial_cut(client, observed, count, started, "replaying INIT") + client.send_initial_fragments(init_parts) for update in client.drain_updates(): stored = self._stored_update(update) if update.version != 0 or stored.fragment_id in observed: @@ -418,6 +473,7 @@ def _wait_for_initial_fragments( ) time.sleep(self.config.poll_seconds) self._report_initial_cut(client, observed, count, started, "received") + client.mark_initial_cut_received() return tuple(observed[index] for index in range(count)) def _report_initial_cut( @@ -471,7 +527,14 @@ async def complete_local_round( *, rollout_id: int, trajectories: TrajectoryBatchEvidence, + hold_submission: bool = False, ) -> RoleRoundResult: + """Record one trained round and capture its cut. + + ``hold_submission`` (overlap): the previous sweep is still in flight, + so the new cut is not yet submittable; it becomes the submission only + after :meth:`apply_staged_broadcast` rebased it and + :meth:`finish_rebase`. The round base stays pinned until then.""" if ( self.adapter is None or self.current is None @@ -500,14 +563,19 @@ async def complete_local_round( previous = self.current previous_submission = self.submission_cut self.current = local - self.submission_cut = local - self.submission_fragment_versions = self.fragment_versions self.local_generation = expected_generation self.local_rounds += 1 self.action_tokens += trajectories.trained_tokens - self.submission_rounds = self.local_rounds - self.submission_tokens = self.action_tokens - self._release_cuts(previous, previous_submission, keep=(local,)) + if hold_submission: + if not self.config.overlap or self._round_base is None: + raise RuntimeError(f"{self.config.role} holds a submission without an overlap round base") + self._release_cuts(previous, keep=(local, self._round_base, previous_submission)) + else: + self.submission_cut = local + self.submission_fragment_versions = self.fragment_versions + self.submission_rounds = self.local_rounds + self.submission_tokens = self.action_tokens + self._release_cuts(previous, previous_submission, keep=(local, self._round_base)) # The policy-level coordinator owns permit matching and release. A # role must never answer a PULL here: the peer role may not yet have @@ -559,8 +627,15 @@ async def apply_staged_broadcast( self, fragment_id: int, version: int, + *, + rebase: bool = False, ) -> tuple[int, ...]: - """Apply one coordinator-approved staged fragment at a safe boundary.""" + """Apply one coordinator-approved staged fragment at a safe boundary. + + ``rebase`` (overlap): the trainer keeps its progress since the round + base, ``theta_f <- G_f + (current_f - base_f)``; ``G_f`` becomes the + anchor, and the counters at the round base become the anchor's, so + the next push carries exactly that progress.""" key = (fragment_id, version) stored = self._staged.get(key) @@ -573,14 +648,60 @@ async def apply_staged_broadcast( raise RuntimeError( f"{self.config.algorithm} {self.config.role} staged fragment skipped its base version" ) - await self._apply_authoritative({fragment_id: stored}) + if not rebase: + await self._apply_authoritative({fragment_id: stored}) + else: + if not self.config.overlap or self._round_base is None or self.current is None: + raise RuntimeError(f"{self.config.role} rebase without an overlap round base") + rebased = self.adapter.rebase_fragment( + stored, self.current, self._round_base, ray_module=self.ray_module) + try: + await self._apply_authoritative( + {fragment_id: rebased}, + anchors={fragment_id: stored}, + counters=(self._rounds_at_base, self._tokens_at_base), + ) + finally: + rebased.release() # the applied target holds its own references self._staged.pop(key, None) return (fragment_id,) + def finish_rebase(self) -> None: + """Overlap: the rebased cut is the next submission (``G + progress``). + + The applied cut has generation 0; the submission is the same weights + labeled with the rounds trained since the round base.""" + + if not self.config.overlap or self.current is None or self.adapter is None: + raise RuntimeError(f"{self.config.role} finish_rebase outside overlap") + previous_submission = self.submission_cut + self.submission_cut = self.adapter.relabel_local_progress( + self.current, local_rounds=self.local_rounds - self._rounds_at_base) + self.submission_fragment_versions = self.fragment_versions + self.submission_rounds = self.local_rounds + self.submission_tokens = self.action_tokens + self._release_cuts(previous_submission, keep=(self.current, self.submission_cut, self._round_base)) + + def mark_round_base(self) -> None: + """Overlap: pin the cut the next local round starts from.""" + + if not self.config.overlap or self.current is None: + raise RuntimeError(f"{self.config.role} mark_round_base outside overlap") + previous = self._round_base + self._round_base = self.current + self._rounds_at_base = self.local_rounds + self._tokens_at_base = self.action_tokens + self._release_cuts(previous, keep=(self.current, self.submission_cut)) + async def _apply_authoritative( self, selected: Mapping[int, StoredAuthoritativeFragment], + *, + anchors: Mapping[int, StoredAuthoritativeFragment] | None = None, + counters: tuple[int, int] | None = None, ) -> None: + """Apply ``selected`` fragments; ``anchors``/``counters`` override the + recorded anchor and its round/token counters (overlap rebase).""" if self.adapter is None or self.current is None or self.group is None: raise RuntimeError(f"{self.config.algorithm} {self.config.role} stream is unavailable") target_version = self.transport_version + 1 @@ -610,16 +731,18 @@ async def _apply_authoritative( self.local_generation = 0 self._release_cuts( previous_cut, - keep=(self.current, self.submission_cut), + keep=(self.current, self.submission_cut, self._round_base), ) - for fragment_id, stored in selected.items(): + rounds, tokens = counters if counters is not None else (self.local_rounds, self.action_tokens) + for fragment_id, applied in selected.items(): + stored = applied if anchors is None else anchors[fragment_id] previous_anchor = self._anchors[fragment_id] if previous_anchor is not None and previous_anchor is not stored: self._release_anchor(previous_anchor) self._anchors[fragment_id] = stored self._fragment_versions[fragment_id] = stored.version - self._rounds_at_anchor[fragment_id] = self.local_rounds - self._tokens_at_anchor[fragment_id] = self.action_tokens + self._rounds_at_anchor[fragment_id] = rounds + self._tokens_at_anchor[fragment_id] = tokens def _apply_token( self, @@ -911,8 +1034,9 @@ def close(self) -> None: retained_cuts = tuple(value for value, _count in self._pinned_cuts.values()) self._pinned_cuts.clear() self._pinned_anchors.clear() - self._release_cuts(self.current, self.submission_cut, *retained_cuts) + self._release_cuts(self.current, self.submission_cut, self._round_base, *retained_cuts) self._staged.clear() + self._round_base = None self.current = None self.submission_cut = None self.submission_fragment_versions = None @@ -934,6 +1058,7 @@ def _unpin_cut(self, cut: ReferencedPolicyCut) -> None: remaining == 0 and cut is not self.current and cut is not self.submission_cut + and cut is not self._round_base and self.adapter is not None ): self.adapter.release(cut) diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 48585224..cf98a201 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -838,6 +838,14 @@ def _dense_full_sync_config( pipeline_depth=1, algorithm_spec_sha256=algorithm_spec_sha256, wan_streams=args.wan_streams, + wan_stream_schedule=getattr(args, "wan_stream_schedule", "rr"), + wan_stall_timeout=float(getattr(args, "wan_stall_timeout", 0.0) or 0.0), + push_dtype=getattr(args, "sweep_push_dtype", "f32"), + push_int8_block=int(getattr(args, "sweep_push_int8_block", 128)), + bcast_delta=bool(getattr(args, "bcast_delta", False)), + # The rollout policy age stays 0 (``_check_dense_full_sweep``); the + # overlap's outer-step lag is bound into the session contract. + overlap=bool(getattr(args, "sweep_overlap", False)), ) return FullParameterSyncConfig((actor,), sweep=True) @@ -1625,6 +1633,40 @@ def probe_syncers(args, *, connect=None, attempts: int = 10, delay: float = 3.0, return addrs +def probe_wan_link(args, *, probe=None) -> None: + """S26 E3b: measure this container's link to the actor syncer before any + model work; exit EXIT_SLOW_LINK below the configured rates so the launcher + can request another container. Off unless a minimum rate is set.""" + min_up = float(getattr(args, "wan_probe_min_mbps", 0.0) or 0.0) + min_down = float(getattr(args, "wan_probe_min_down_mbps", 0.0) or 0.0) + if min_up <= 0 and min_down <= 0: + return + from yeto.protocol import island_hmac_key_from_env + from yeto.wan_probe import EXIT_SLOW_LINK, PROBE_EVENT, PROBE_MARKER, probe_link + + result = (probe or probe_link)( + _syncer_address(args.syncer), + streams=max(1, int(args.wan_streams)), + bytes_per_direction=int(getattr(args, "wan_probe_mib", 256)) << 20, + key=island_hmac_key_from_env(), + ) + accepted = result.accepted(min_up, min_down) + record = {"event": PROBE_EVENT, "accepted": accepted, + "wan_probe/min_upload_mbps": min_up, "wan_probe/min_download_mbps": min_down, + **result.fields()} + print(f"{PROBE_MARKER} {'accepted' if accepted else 'rejected'}: up {result.upload_mbps:.1f} MB/s " + f"(min {min_up}), down {result.download_mbps:.1f} MB/s (min {min_down}), " + f"{result.streams} stream(s){'' if result.error is None else '; ' + result.error}", flush=True) + if getattr(args, "event_tape", None): + from yeto.rl.engine.driver import EventTape + + if getattr(args, "rl_echo_events", False): + install_event_echo() # the launcher reads the verdict from the log stream + EventTape(args.event_tape, int(getattr(args, "learner_id", 0) or 0)).append(record) + if not accepted: + raise SystemExit(EXIT_SLOW_LINK) + + def require_ports_router_mode(miles_args) -> None: """Upstream Miles owns the SGLang router. @@ -2098,6 +2140,7 @@ def run_miles( algorithm_semantic_sha256=algorithm_semantic_sha256, code_pack_digest=code_pack_digest, compat_group=_island_compat_group(args), + bcast_delta=bool(getattr(args, "bcast_delta", False)), ) if getattr(args, "rl_island_scheduling", "legacy") == "elastic": # rl-inter-island-scheduling 0.15 (legacy sets nothing) @@ -2385,6 +2428,7 @@ def main(argv=None) -> None: verify_miles_revision(miles_root, expected=pins) _preflight_codex_harness(args) probe_syncers(args) + probe_wan_link(args) from miles.utils.misc import load_function diff --git a/yeto/rl/adapters/miles/models/chunked_full_parameter.py b/yeto/rl/adapters/miles/models/chunked_full_parameter.py index 9b7e6c99..14a675d6 100644 --- a/yeto/rl/adapters/miles/models/chunked_full_parameter.py +++ b/yeto/rl/adapters/miles/models/chunked_full_parameter.py @@ -786,6 +786,96 @@ def assemble_mixed_target( ) return self.validate_cut(target, target_policy_version, 0) + def rebase_fragment( + self, + stored: StoredAuthoritativeFragment, + current: ReferencedPolicyCut, + base: ReferencedPolicyCut, + *, + ray_module=None, + ) -> StoredAuthoritativeFragment: + """Sync/train overlap: ``stored + (current - base)`` for one fragment. + + The result is stored like an inbound authoritative fragment (same + id, version and provenance, fresh Ray refs) so ``assemble_mixed_target`` + validates and applies it unchanged. The caller keeps ``stored`` as + the anchor and releases the result after the apply.""" + + import numpy as np + from miles.backends.megatron_utils.full_parameter_state import ( + _resolve_ray_module, + ) + + from yeto.rl.engine.sync_overlap import rebase_values + + fragment_id = stored.fragment_id + for cut in (current, base): + if cut.layout_hash != self.layout.layout_hash: + raise ValueError("Miles rebase cut layout changed") + owner = _fragment_owner_shard(current.transport_cut, fragment_id) + if ( + stored.parameter_layout_hash != self.layout.layout_hash + or stored.topology != owner.topology + or stored.plan_hash != owner.plan_hash + ): + raise ValueError("authoritative fragment provenance changed") + ray = _resolve_ray_module(ray_module) + chunks = sorted(zip(stored.descriptor.chunks, stored.refs, strict=True), + key=lambda row: row[0].flat_offset) + anchor_parts, offset = [], 0 + for descriptor, reference in chunks: + if reference is None or descriptor.flat_offset != offset: + raise ValueError("authoritative fragment chunks are released or not contiguous") + value = ray.get(reference) + anchor_parts.append(value.detach().to("cpu").reshape(-1).numpy()) + offset += descriptor.numel + + def flat(cut: ReferencedPolicyCut) -> np.ndarray: + raw = b"".join(bytes(part) for part in self.fragment_parts( + cut, fragment_id, ray_module=ray_module)()) + return np.frombuffer(raw, dtype=" ReferencedPolicyCut: + """Sync/train overlap: label a rebased cut with its local progress. + + ``apply`` resets the local step generation to 0, but a rebased cut + (``G + (current - base)``) still holds ``local_rounds`` rounds of + progress over its anchor ``G``, and Miles pushes only cuts with + progress (generation >= 1). Values, chunk hashes and refs are shared + with ``cut`` (fresh ref lists, so releasing one leaves the other); + only the generation label changes.""" + + from dataclasses import replace + + from miles.ray.full_parameter_transport import FullParameterChunkedCut + + if isinstance(local_rounds, bool) or not isinstance(local_rounds, int) or local_rounds < 1: + raise ValueError("a rebased cut needs at least one local round") + if cut.layout_hash != self.layout.layout_hash or cut.local_step_generation != 0: + raise ValueError("only an applied (generation 0) cut is relabeled") + transport = cut.transport_cut + relabeled = FullParameterChunkedCut( + policy_version=transport.policy_version, + local_step_generation=local_rounds, + parameter_layout_hash=transport.parameter_layout_hash, + shards=tuple( + replace(shard, local_step_generation=local_rounds, + chunk_refs=[list(references) for references in shard.chunk_refs]) + for shard in transport.shards + ), + ) + return self.validate_cut(relabeled, cut.policy_version, local_rounds) + async def apply( self, group: MilesChunkedFullParameterGroup, diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index 8a9f5202..803b19e1 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -313,6 +313,30 @@ ), island_flag("--fragments", type=int, default=1), island_flag("--parameter-layout-sha256", default=None), + island_flag( + "--sweep-push-dtype", + choices=("f32", "bf16", "int8"), + default="f32", + help="full-parameter sweep: wire dtype of pushed outer deltas (syncer masters stay f32)", + ), + island_flag( + "--sweep-push-int8-block", + type=int, + default=128, + help="full-parameter sweep, int8 pushes: values per f32 absmax scale", + ), + island_flag( + "--sweep-overlap", + action="store_true", + help="full-parameter sweep: train the next round while the sweep syncs, then " + "rebase onto the new global weights (rollout policy age stays 0)", + ), + island_flag( + "--bcast-delta", + action="store_true", + help="strict-avg / full-parameter sweep: accept bf16 downlink deltas against the " + "shared f32 anchor (the syncer runs --bcast-delta-full-every)", + ), island_flag( "--syncer-profile-sha256", default=None, @@ -392,6 +416,11 @@ island_flag("--seq-len", type=int, required=True), island_flag("--seed", type=int, required=True), island_flag("--wan-streams", type=int, default=4), + island_flag("--wan-stream-schedule", choices=("rr", "adaptive"), default="rr"), + island_flag("--wan-stall-timeout", type=float, default=0.0), + island_flag("--wan-probe-min-mbps", type=float, default=0.0), + island_flag("--wan-probe-min-down-mbps", type=float, default=0.0), + island_flag("--wan-probe-mib", type=int, default=256), island_flag("--miles-root", required=True), island_flag( "--miles-commit", diff --git a/yeto/rl/adapters/verl/island_entry.py b/yeto/rl/adapters/verl/island_entry.py index ab6f640f..5fec3a02 100644 --- a/yeto/rl/adapters/verl/island_entry.py +++ b/yeto/rl/adapters/verl/island_entry.py @@ -60,6 +60,8 @@ def parse_args(argv): p.add_argument("--seq-len", type=int, required=True) p.add_argument("--seed", type=int, required=True) p.add_argument("--wan-streams", type=int, default=4) + # S26 E3: bf16 downlink deltas (yeto/downlink.py; the syncer runs --bcast-delta-full-every) + p.add_argument("--bcast-delta", action="store_true") p.add_argument("--apply-chat-template-kwargs", default=None) p.add_argument("--rl-algorithm-spec", default=None) p.add_argument("--rl-island-scheduling", default="legacy") @@ -357,6 +359,7 @@ def main(argv=None) -> int: "model_path": model_path, "model_revision": args.model_revision, "lora_rank": args.lora_r, "lora_targets": args.lora_targets, "event_tape": os.path.expanduser(args.event_tape), "out_dir": str(OUTPUT), "tis_upper": algo["tis_upper"], "wan_streams": args.wan_streams, + "bcast_delta": bool(args.bcast_delta), "thresholds_key": ["verl", "fsdp2", "vllm-" + str(manifest["versions"].get("vllm")), _threshold_card(manifest)], "algorithm_spec": spec_dict, "syncer_epoch": args.rl_syncer_epoch, "island_lease_s": args.rl_island_lease_s, diff --git a/yeto/rl/adapters/verl/trainer.py b/yeto/rl/adapters/verl/trainer.py index 6be18c9e..08c72e83 100644 --- a/yeto/rl/adapters/verl/trainer.py +++ b/yeto/rl/adapters/verl/trainer.py @@ -186,7 +186,8 @@ def build_sync(plan: dict, specs, *, lora_hash: str, layout_hash: str, identity_ expected_specs=tuple(specs), base_model_revision=plan["model_revision"], lora_config_hash=lora_hash, layout_hash=layout_hash, event_tape=plan["event_tape"], wan_streams=int(plan.get("wan_streams", 4)), - backend_identity_sha256=identity_sha256, compat_group=compat_group) + backend_identity_sha256=identity_sha256, compat_group=compat_group, + bcast_delta=bool(plan.get("bcast_delta", False))) if mode == "strict": sync = StrictAvgSync(bridge, progress=None) elif mode == "elastic": diff --git a/yeto/rl/bridge.py b/yeto/rl/bridge.py index 9cf0e388..7cae9256 100644 --- a/yeto/rl/bridge.py +++ b/yeto/rl/bridge.py @@ -13,8 +13,13 @@ import torch -from .engine.backend_identity import config_island_contract_sha256 -from ..protocol import DTYPE_F32, FinalManifest, PullRequest, SyncerClient +from ..protocol import ( + DTYPE_F32, + DTYPE_FLAG_BCAST_DELTA, + FinalManifest, + PullRequest, + SyncerClient, +) from ..tensor_io import pack_tensor, unpack_fragment from .core import ( CanonicalLoraState, @@ -29,6 +34,7 @@ policy_delta, tensors_from_flat_owned, ) +from .engine.backend_identity import config_island_contract_sha256 class IslandRuntime(Protocol): @@ -77,6 +83,12 @@ class BridgeConfig: # 7.7c: "-" (yeto.hw.catalog); sent in HELLO / JOIN so the # syncer refuses another card type by name. None = not declared. compat_group: str | None = None + # S26 E3: receive BCAST as bf16 deltas against the shared f32 anchor + # (session flag; the syncer runs --bcast-delta-full-every). The client + # hands this bridge the reconstructed f32 anchor, so nothing else changes. + # An anchor mismatch drops the connection; with max_reconnects=0 the + # island exits and its restart re-reads the committed cut as f32. + bcast_delta: bool = False def _write_round_audit( @@ -278,7 +290,7 @@ def __init__(self, runtime: IslandRuntime, config: BridgeConfig) -> None: config.syncer_addr, config.learner_id, self.layout, - dtype=DTYPE_F32, + dtype=DTYPE_F32 | (DTYPE_FLAG_BCAST_DELTA if config.bcast_delta else 0), num_streams=config.wan_streams, session_contract_hash=_session_contract( self.layout, config_island_contract_sha256(config)), diff --git a/yeto/rl/engine/full_parameter_sync.py b/yeto/rl/engine/full_parameter_sync.py index f4b2141d..af2a1a1d 100644 --- a/yeto/rl/engine/full_parameter_sync.py +++ b/yeto/rl/engine/full_parameter_sync.py @@ -33,9 +33,10 @@ from yeto.protocol import PartialMessageGenerationLost from yeto.rl.elastic_client import LeaveReason +from yeto.wan_streams import upload_window from .driver import SyncBoundary, SyncStart -from .sync_timing import SyncTiming, unix_now +from .sync_timing import SyncTiming, unix_now, upload_fields if TYPE_CHECKING: from .driver import IslandDriver @@ -81,6 +82,11 @@ async def stage_available_broadcasts(self) -> tuple[tuple[int, int], ...]: ... def staged_fragment_versions(self) -> tuple[tuple[int, int], ...]: ... def has_staged_broadcast(self, fragment_id: int, version: int) -> bool: ... async def apply_staged_broadcast(self, fragment_id: int, version: int) -> tuple[int, ...]: ... + # S26 E3 sync/train overlap (profile.overlap) only: + # complete_local_round(..., hold_submission=True) keeps the round base and + # the previous submission; apply_staged_broadcast(..., rebase=True) applies + # G + (current - base) with G as the anchor; finish_rebase() makes the + # rebased cut the submission; mark_round_base() pins the next round base. async def wait_and_apply_final(self) -> None: ... def acknowledge_final(self) -> None: ... def discard_detached_submission_cut(self) -> None: ... @@ -99,6 +105,8 @@ class FullParameterSyncProfile: wait_timeout: float poll_seconds: float sweep: bool = False + # S26 E3: train round k+1 while sweep k syncs; rebase at the next boundary. + overlap: bool = False def __post_init__(self) -> None: if not self.roles or len(set(self.roles)) != len(self.roles): @@ -107,6 +115,13 @@ def __post_init__(self) -> None: raise ValueError("the actor is the published role and must come first") if self.sweep and (self.roles != ("actor",) or self.local_horizon != 1): raise ValueError("a dense strict sweep has one actor role and H=1") + if self.overlap and not self.sweep: + raise ValueError("sync/train overlap needs the dense strict sweep") + + @property + def total_rounds(self) -> int: + """Local rounds of a sweep profile (one sweep per round).""" + return self.total_fragment_steps // self.fragments @property def effective_pipeline_depth(self) -> int: @@ -131,6 +146,7 @@ def from_configs( wait_timeout=min(config.wait_timeout for config in configs.values()), poll_seconds=min(config.poll_seconds for config in configs.values()), sweep=sweep, + overlap=bool(getattr(first, "overlap", False)), ) @@ -389,9 +405,12 @@ async def _apply_pair( step, fragment_id, _base = key if not pair.complete: raise RuntimeError("broadcast arrived before every paired PUSH was retained") - results = await asyncio.gather( - *(stream.apply_staged_broadcast(fragment_id, step) for stream in self.streams.values()) - ) + if self.profile.overlap: + calls = (stream.apply_staged_broadcast(fragment_id, step, rebase=True) + for stream in self.streams.values()) + else: + calls = (stream.apply_staged_broadcast(fragment_id, step) for stream in self.streams.values()) + results = await asyncio.gather(*calls) for role, fragments in zip(self.streams, results, strict=True): applied[role].extend(fragments) if self.streams[role].fragment_versions[fragment_id] != step: @@ -509,6 +528,10 @@ def __init__( self.closed = False self._driver: IslandDriver | None = None self._timing = SyncTiming(None, None) # L5.6, set by each boundary + # Overlap: fragment versions before the sweep still in flight. + self._pending_sweep: tuple[int, ...] | None = None + # S26 E3b: actor data-stream counters at the previous boundary. + self._upload_before: tuple | None = None @property def actor(self) -> RoleStream: @@ -517,6 +540,7 @@ def actor(self) -> RoleStream: def start(self, driver: IslandDriver) -> SyncStart: self._driver = driver self.run(self.pairing.gather_roles(lambda role, stream: stream.initialize(self.groups[role]))) + self._upload_before = self._stream_snapshots() return SyncStart(state=self.actor.resident_state(0), rollout_id=0) def boundary(self, driver: IslandDriver, *, rollout_id: int, stats: Any) -> SyncBoundary: @@ -534,11 +558,32 @@ def boundary(self, driver: IslandDriver, *, rollout_id: int, stats: Any) -> Sync "sync/global_policy_hash": state.policy_tensor_hash(), "sync/inflight_pairs": len(self.pairing.inflight), **self._timing.fields(), + **self._upload_fields(), }, ) return SyncBoundary(state=state, stop=self.terminal) + def _stream_snapshots(self) -> tuple | None: + snapshots = getattr(getattr(self.actor, "client", None), "stream_snapshots", None) + return snapshots() if callable(snapshots) else None + + def _upload_fields(self) -> dict[str, object]: + """Actor upload rates of this boundary and the wait after the upload.""" + after = self._stream_snapshots() + before, self._upload_before = self._upload_before, after + if after is None or before is None: + return {} + client = self.actor.client + return upload_fields( + upload_window(before, after), + drained_ts=client.upload_drained_at, + timing=self._timing, + stalls=client.send_progress().stalls, + ) + async def _boundary(self, rollout_id: int) -> bool: + if self.pairing.profile.overlap: + return await self._overlap_boundary(rollout_id) pairing = self.pairing trajectories = self.trajectories(rollout_id) await pairing.gather_roles( @@ -564,6 +609,79 @@ async def _boundary(self, rollout_id: int) -> bool: stream.discard_detached_submission_cut() return False + async def _overlap_boundary(self, rollout_id: int) -> bool: + """Streaming-DiLoCo overlap (S26 E3; one more outer step of staleness). + + Round k's delta (``current - base``) is pushed at boundary k and the + island trains round k+1 at once. At boundary k+1 the island first + waits for sweep k (``G_k``), then rebases its trainer in place: + ``theta <- G_k + (theta_local_{k+1} - base_k)`` with ``G_k`` as the + anchor, so the next push is exactly round k+1's progress and round + k's delta is never counted twice. The terminal round waits for the + final cut like the strict sweep.""" + + pairing = self.pairing + trajectories = self.trajectories(rollout_id) + hold = self._pending_sweep is not None + await pairing.gather_roles( + lambda _role, stream: stream.complete_local_round( + rollout_id=rollout_id, trajectories=trajectories, hold_submission=hold + ) + ) + wait_ts = unix_now() + if hold: + await self._wait_for_sweep(self._pending_sweep) + self._pending_sweep = None + if not pairing.any_finalizing: + for stream in pairing.streams.values(): + stream.finish_rebase() + if pairing.any_finalizing: + await pairing.finish_without_more_training() + return True + before = pairing.fragment_versions() + if rollout_id + 1 >= pairing.profile.total_rounds: + # Terminal round: nothing trains after it, so wait for the cut. + pairing.submit_ready() + submit_ts = unix_now() + await pairing.commit_staged_broadcasts() + await self._wait_for_sweep(before) + self._timing = SyncTiming(submit_ts, unix_now()) + await pairing.finish_without_more_training() + return True + await self._push_whole_sweep(before) + self._pending_sweep = before + for stream in pairing.streams.values(): + stream.mark_round_base() + # This boundary blocked on the previous sweep (wait + rebase) and on + # handing every fragment of this sweep to the wire. + self._timing = SyncTiming(wait_ts, unix_now()) + return False + + async def _push_whole_sweep(self, before: tuple[int, ...]) -> None: + """Overlap: submit every fragment of the sweep after ``before`` before + training on. A PULL that arrives later would otherwise wait for the + next boundary, a whole local round, and stall the sweep. The syncer + sends all P PULLs at once with ``--pipeline P`` (required by the + launcher), so this waits only for their arrival.""" + + pairing = self.pairing + fragments = pairing.profile.fragments + steps = {version + fragments if version else fragment_id + 1 + for fragment_id, version in enumerate(before)} + deadline = time.monotonic() + pairing.profile.wait_timeout + + def pushed() -> bool: + return steps <= {key[0] for key, pair in pairing.inflight.items() if pair.complete} + + pairing.submit_ready() + while not pushed(): + if time.monotonic() >= deadline: + error = TimeoutError("overlap could not submit the whole sweep before the wait timeout") + pairing.fail_stop(error) + raise error + await asyncio.sleep(pairing.profile.poll_seconds) + pairing.submit_ready() + async def _wait_for_sweep(self, before: tuple[int, ...]) -> None: """Strict sweep: block until all P fragments of this local step are applied, i.e. every fragment version moved past ``before``. An empty diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index b4de9148..2af06a5c 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -841,6 +841,62 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. "sao-streaming-full: full-parameter SAO, actor and critic streams " "(ports engine, --rl-placement fixed-partition, --rl-sao-contracts)", ), + cli_flag( + "--rl-sync-fragment-pipeline", + type=int, + default=1, + help="full-parameter only: fragment rounds the syncer gathers at once " + "(1..--fragments). Above 1, fragment k+1 uploads while fragment k merges " + "and downloads; merge results are unchanged (default 1 = serial)", + ), + cli_flag( + "--rl-sync-push-dtype", + choices=("f32", "bf16", "int8"), + default="f32", + help="full-parameter only: wire dtype of the islands' pushed outer deltas " + "(pseudo-gradients). bf16 halves upload bytes, int8 (block absmax scale, see " + "--rl-sync-push-int8-block) quarters them; the syncer keeps f32 masters, " + "momentum and f32 FINAL (default f32 = exact)", + ), + cli_flag( + "--rl-sync-overlap", + action="store_true", + help="full-parameter only (E2 option): train the next round while the " + "sweep syncs, then rebase the trainer onto the new global weights " + "(theta <- G_k + local progress). The local weights lag the global policy " + "by one outer step; rollouts stay on-policy, so --rl-max-policy-age stays 0. " + "Needs --rl-sync-fragment-pipeline equal to --fragments", + ), + cli_flag( + "--rl-sync-push-int8-block", + type=int, + default=128, + help="full-parameter only, with --rl-sync-push-dtype int8: values per f32 " + "scale (32..1048576; error <= block absmax / 254 per value)", + ), + cli_flag( + "--rl-sync-bcast-delta-full-every", + type=int, + default=0, + help="full-parameter only: above 0, the syncer sends each new fragment " + "version as a bf16 delta against the islands' shared f32 anchor, with an " + "exact f32 refresh every N-th version (N >= 2; digest-checked, a mismatch " + "reconnects for f32). 0 = always f32", + ), + cli_flag( + "--rl-sync-checkpoint-every", + type=int, + default=1, + help="full-parameter only: syncer checkpoint interval in fragment steps " + "(1, or a multiple of --fragments). Above 1, only checkpointed steps are " + "persisted before their broadcast; a crash resumes at the newest checkpoint", + ), + cli_flag( + "--rl-sync-async-checkpoint", + action="store_true", + help="full-parameter only: the syncer writes each per-round checkpoint off " + "its scheduler and broadcasts that round only once the checkpoint is durable", + ), cli_flag( "--rl-initial-adapter", default=None, diff --git a/yeto/rl/engine/sync_overlap.py b/yeto/rl/engine/sync_overlap.py new file mode 100644 index 00000000..3a7f17ec --- /dev/null +++ b/yeto/rl/engine/sync_overlap.py @@ -0,0 +1,59 @@ +"""Sync/train overlap for the full-parameter sweep (S26 E3 item 3, E2 option). + +Round k's outer delta is pushed at boundary k and the island trains round +k+1 immediately on its own weights. At boundary k+1 it waits for the global +result of sweep k (``G_k``) and rebases in place:: + + theta <- G_k + (theta_local_{k+1} - base_k) + +``base_k`` is the cut round k+1 started from. ``G_k`` becomes the anchor, so +the next push ``theta - G_k`` equals round k+1's progress exactly once: round +k's delta is already inside ``G_k`` and is not counted again. + +Staleness is configured apart from ``--rl-max-policy-age``. The driver +publishes the boundary's cut to the island's engines before the next +generation (boundary, then publish), so every rollout is still generated by +the weights the trainer then trains: the rollout policy age stays 0 and the +sweep keeps refusing any other value. What overlap adds is a lag of the local +weights behind the global policy of :data:`OUTER_STEP_LAG` outer step (like +one extra DiLoCo local step); it is bound into the role-stream session +contract (``sync_train_overlap``), not into the rollout policy age, which in +Miles would switch on partial rollout (carry-over) or fully_async generation. +""" + +from __future__ import annotations + +import numpy as np + +OUTER_STEP_LAG = 1 + + +def overlap_option_problems(*, max_policy_age: int, fragment_pipeline: int, fragments: int) -> list[str]: + """Launcher requirements of ``--rl-sync-overlap`` (empty when satisfied). + + * rollout policy age 0: see the module doc (overlap is not rollout staleness); + * syncer fragment pipeline = P: the syncer then PULLs every fragment of a + sweep at once, so the island pushes the whole sweep at the boundary. With + a smaller pipeline the later PULLs arrive while the next round trains and + wait for the next boundary, which serializes the sweep again.""" + problems = [] + if int(max_policy_age or 0) != 0: + problems.append("--rl-max-policy-age 0 with --rl-sync-overlap (overlap lags the global " + "policy by one outer step; rollouts stay on-policy)") + if int(fragment_pipeline) != int(fragments): + problems.append("--rl-sync-fragment-pipeline equal to --fragments with --rl-sync-overlap") + return problems + + +def rebase_values(anchor: np.ndarray, current: np.ndarray, base: np.ndarray) -> np.ndarray: + """``anchor + (current - base)`` in f32 (the trainer keeps its progress).""" + if not (anchor.shape == current.shape == base.shape): + raise ValueError("rebase operands have different shapes") + for name, value in (("anchor", anchor), ("current", current), ("base", base)): + if value.dtype != np.float32: + raise ValueError(f"rebase {name} must be float32") + out = np.subtract(current, base, dtype=np.float32) + np.add(anchor, out, out=out) + if not np.isfinite(out).all(): + raise ValueError("rebased fragment has a non-finite value") + return out diff --git a/yeto/rl/engine/sync_timing.py b/yeto/rl/engine/sync_timing.py index eaa4ad9b..eae22025 100644 --- a/yeto/rl/engine/sync_timing.py +++ b/yeto/rl/engine/sync_timing.py @@ -44,3 +44,42 @@ def fields(self) -> dict[str, float | None]: def unix_now() -> float: return time.time() + + +def upload_fields( + window, + *, + drained_ts: float | None, + timing: SyncTiming, + stalls: int, +) -> dict[str, object]: + """Per-boundary upload diagnostics of one island (S26 E3b). + + ``window`` is a :class:`yeto.wan_streams.UploadWindow` covering this + boundary. ``sync/upload_stream_mbps`` lists each data stream's rate over + the time it was busy; ``sync/upload_mbps`` is the bytes over the longest + busy stream (the streams run in parallel). ``sync/wait_after_upload_seconds`` + is the time from this island's last byte leaving to the global policy + being ready: waiting for slower islands, the merge and the broadcast.""" + rates = window.stream_mbps() + busy = max(window.busy_seconds_per_stream, default=0.0) + measured = [rate for rate in rates if rate is not None] + wait_after = None + if ( + drained_ts is not None + and timing.submit_ts is not None + and timing.global_ready_ts is not None + and drained_ts <= timing.global_ready_ts + ): + wait_after = timing.global_ready_ts - max(drained_ts, timing.submit_ts) + return { + "sync/upload_bytes": window.total_bytes, + "sync/upload_busy_seconds": busy, + "sync/upload_mbps": (window.total_bytes / busy / 1e6) if busy > 0 else None, + "sync/upload_stream_mbps": [None if rate is None else round(rate, 3) for rate in rates], + "sync/upload_min_stream_mbps": min(measured) if measured else None, + "sync/upload_active_streams": len(measured), + "sync/upload_drained_ts": drained_ts, + "sync/wait_after_upload_seconds": wait_after, + "sync/stream_stalls": stalls, + } diff --git a/yeto/rl/event_echo.py b/yeto/rl/event_echo.py index b3f2a33e..b9abdadc 100644 --- a/yeto/rl/event_echo.py +++ b/yeto/rl/event_echo.py @@ -133,6 +133,7 @@ def __init__(self, path, *, fresh: bool = True) -> None: self.discarded = 0 self.finalized = False self.left_pool = False # 0.27: elastic_left_pool seen (re-JOIN exhausted) + self.wan_rejects = 0 # S26 E3b: rl_wan_probe records with accepted=false self.closed = False @property @@ -179,8 +180,11 @@ def _record_line(self, line: str) -> None: with self.path.open("a", encoding="utf-8") as handle: handle.write(raw + "\n") self.count += 1 - event = json.loads(raw).get("event") - if event == FINALIZED_EVENT: + record = json.loads(raw) + event = record.get("event") + if event == "rl_wan_probe" and record.get("accepted") is False: + self.wan_rejects += 1 + elif event == FINALIZED_EVENT: self.finalized = True elif event == "elastic_left_pool": self.left_pool = True diff --git a/yeto/rl/wandb_rl.py b/yeto/rl/wandb_rl.py index 3467c503..60a08c7b 100644 --- a/yeto/rl/wandb_rl.py +++ b/yeto/rl/wandb_rl.py @@ -28,6 +28,7 @@ "train/*": TRAIN_STEP_KEY, "rl/*": STEP_KEY, "eval/*": STEP_KEY, + "sync/*": STEP_KEY, } # Closed output surface. Input aliases are limited to namespaced scalar fields @@ -115,11 +116,40 @@ def _non_negative_step(event: dict[str, Any], aliases: tuple[str, ...]) -> int | return None +# S26 E3b: per-sweep sync diagnostics of a full-parameter island, on the +# global policy-version axis (rl/step). +_SYNC_ROUND_OUTPUTS = ( + "sync/sync_wait_seconds", + "sync/upload_mbps", + "sync/upload_min_stream_mbps", + "sync/upload_active_streams", + "sync/wait_after_upload_seconds", + "sync/stream_stalls", +) + + +def _sync_round_metrics(event: dict[str, Any]) -> dict[str, Any] | None: + step = _non_negative_step(event, ("policy_version",)) + if step is None: + return None + metrics = { + key: value + for key in _SYNC_ROUND_OUTPUTS + if (value := _finite_scalar(event, (key,))) is not None + } + if not metrics: + return None + metrics[STEP_KEY] = step + return metrics + + def event_metrics(event: dict[str, Any]) -> dict[str, Any] | None: """Project one event onto the explicit scalar-only W&B surface.""" global _derived_step_logged event_name = event.get("event") + if event_name == "rl_full_parameter_round": + return _sync_round_metrics(event) if event_name not in {"rl_local_round", "rl_eval_result", "rl_round_trained"}: return None diff --git a/yeto/syncer_profile.py b/yeto/syncer_profile.py index 390d2bc4..ce23fa99 100644 --- a/yeto/syncer_profile.py +++ b/yeto/syncer_profile.py @@ -259,8 +259,12 @@ def _validate_general(self) -> None: raise ValueError("syncer profile resume requires checkpointing") if self.policy_sweep_fragments is not None: fragments = self.policy_sweep_fragments - if self.pipeline != 1 or self.total_steps % fragments: + if not 1 <= self.pipeline <= fragments or self.total_steps % fragments: raise ValueError("syncer profile policy sweep schedule is invalid") + if self.checkpoint_every != 1 and ( + self.checkpoint_every < 1 or self.checkpoint_every % fragments): + raise ValueError("syncer profile policy sweep checkpoint interval must be " + "1 or a multiple of the fragment count") def validate_sao(self) -> None: """Enforce the paired actor/critic fail-closed server profile.""" diff --git a/yeto/wan_probe.py b/yeto/wan_probe.py new file mode 100644 index 00000000..4ccb91fe --- /dev/null +++ b/yeto/wan_probe.py @@ -0,0 +1,149 @@ +"""Island-to-syncer link-bandwidth probe (S26 E3b). + +E3/E3b showed that the sync share follows the island *container*: on a slow +container every TCP stream is slow at once (1-5 MB/s each, both directions), +so no stream scheduling can help. An island can measure its link to the head +before it loads a model and exit early when the link is too slow; the +launcher then requests a fresh container (``--wan-probe-relaunches``). + +Each probe connection sends one ``MSG_BW_PROBE`` frame (HMAC-sealed like +HELLO), uploads ``upload_bytes`` raw bytes, reads the syncer's 8-byte +acknowledgement, then reads ``download_bytes`` raw bytes. The syncer never +connects a probe to a learner session. +""" + +from __future__ import annotations + +import socket +import struct +import threading +import time +from collections.abc import Callable +from dataclasses import dataclass + +from .protocol import MAGIC, MSG_BW_PROBE, ProtocolError, seal_hello, write_frame + +# Exit code of an island that rejected its own link (the launcher may relaunch it). +EXIT_SLOW_LINK = 75 +PROBE_EVENT = "rl_wan_probe" +PROBE_MARKER = "[yeto-wan-probe]" +_BLOCK = memoryview(bytes(1 << 20)) + + +@dataclass(frozen=True) +class ProbeResult: + streams: int + upload_bytes: int + download_bytes: int + upload_seconds: float + download_seconds: float + error: str | None = None + + @property + def upload_mbps(self) -> float: + return self.upload_bytes / self.upload_seconds / 1e6 if self.upload_seconds > 0 else 0.0 + + @property + def download_mbps(self) -> float: + return self.download_bytes / self.download_seconds / 1e6 if self.download_seconds > 0 else 0.0 + + def accepted(self, min_upload_mbps: float, min_download_mbps: float) -> bool: + return ( + self.error is None + and self.upload_mbps >= min_upload_mbps + and self.download_mbps >= min_download_mbps + ) + + def fields(self) -> dict[str, object]: + return { + "wan_probe/streams": self.streams, + "wan_probe/upload_bytes": self.upload_bytes, + "wan_probe/download_bytes": self.download_bytes, + "wan_probe/upload_mbps": round(self.upload_mbps, 3), + "wan_probe/download_mbps": round(self.download_mbps, 3), + "wan_probe/error": self.error, + } + + +def encode_probe(upload_bytes: int, download_bytes: int, key: bytes | None) -> bytes: + return seal_hello(struct.pack(" None: + sock = socket.create_connection(addr, timeout=timeout) + try: + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + write_frame(sock, MSG_BW_PROBE, encode_probe(upload, download, key)) + start = time.monotonic() + left = upload + while left: + take = min(left, len(_BLOCK)) + sock.sendall(_BLOCK[:take]) + left -= take + ack = b"" + while len(ack) < 8: + part = sock.recv(8 - len(ack)) + if not part: + # The syncer refused the probe: it answers with an ERROR frame + # and closes; the ERROR body is the reason. + raise ConnectionError("syncer closed the probe before acknowledging the upload") + ack += part + if ack.startswith(struct.pack(" ProbeResult: + """Measure aggregate upload and download rates over ``streams`` parallel + connections; ``bytes_per_direction`` is split evenly over the streams. + A rate is total bytes over the slowest stream's time in that direction.""" + if streams < 1 or bytes_per_direction < streams: + raise ValueError("probe needs at least one stream and one byte per stream") + per = bytes_per_direction // streams + results = [dict() for _ in range(streams)] + errors: list[str] = [] + + def run(index: int) -> None: + try: + stream_fn(addr, per, per, key, timeout, results[index]) + except (OSError, ProtocolError) as error: + errors.append(f"{type(error).__name__}: {error}") + + threads = [threading.Thread(target=run, args=(i,), daemon=True) for i in range(streams)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout + 5) + done = [r for r in results if "up" in r and "down" in r] + if len(done) < streams and not errors: + errors.append(f"{streams - len(done)} probe stream(s) did not finish in {timeout:.0f}s") + return ProbeResult( + streams=streams, + upload_bytes=per * len(done), + download_bytes=per * len(done), + upload_seconds=max((r["up"] for r in done), default=0.0), + download_seconds=max((r["down"] for r in done), default=0.0), + error=errors[0] if errors else None, + ) + diff --git a/yeto/wan_streams.py b/yeto/wan_streams.py new file mode 100644 index 00000000..8d2cb6ab --- /dev/null +++ b/yeto/wan_streams.py @@ -0,0 +1,190 @@ +"""Per-stream accounting, stall detection and rate-aware striping for the +island side of a striped syncer connection (S26 E3b). + +A ``SyncerClient`` stripes one logical message over ``num_streams`` TCP data +sockets as 4 MiB CHUNK envelopes. The syncer reassembles a message by +``(msg_id, offset)``, so any chunk may travel on any data stream of the same +connection group. This module holds the parts of that transport that do not +need a socket: + +* :class:`StreamMeter` counts bytes, busy time and the last byte progress of + one stream. The send loop feeds it; diagnostics and the stall watchdog read + it. +* :func:`pick_stream` chooses the data stream for the next chunk. ``"rr"`` + is the historical round robin. ``"adaptive"`` sends each chunk to the + stream with the earliest estimated finish time (backlog / measured rate), + so a slow stream gets fewer chunks and a stalled stream gets none. The + number of open streams is fixed by HELLO; the number of streams that carry + data follows the measured rates. +* :class:`StreamStalledError` is the failure a watchdog reports when a + stream holds unsent bytes and makes no byte progress for the stall timeout. + The client handles it like any socket error: the whole group is torn down + and redialed, and the syncer replays its PULLs, so the island pushes the + interrupted fragment again at the same (global_step, round_attempt). +""" + +from __future__ import annotations + +import threading +import time +from collections.abc import Sequence +from dataclasses import dataclass + +STREAM_SCHEDULES = ("rr", "adaptive") + +# EWMA weight of the newest completed send in the per-stream rate estimate. +RATE_EWMA_ALPHA = 0.3 +# Sends shorter than this carry no useful rate signal (kernel buffer absorbs +# them), so they do not update the estimate. +MIN_RATE_SAMPLE_SECONDS = 0.01 + + +class StreamStalledError(ConnectionError): + """One stream held unsent (or partially received) bytes with no progress + for longer than the stall timeout.""" + + +@dataclass(frozen=True) +class StreamSnapshot: + """Cumulative counters of one stream at one instant.""" + + bytes_sent: int + busy_seconds: float + backlog_bytes: int + rate_bps: float | None + + +class StreamMeter: + """Counters of one outbound stream. Thread-safe; cheap enough to update + on every ``send`` call.""" + + def __init__(self, clock=time.monotonic) -> None: + self._clock = clock + self._lock = threading.Lock() + self.bytes_sent = 0 + self.busy_seconds = 0.0 + self.queued_bytes = 0 # enqueued, not yet handed to the kernel + self.rate_bps: float | None = None + self._item_started: float | None = None + self._item_bytes = 0 + self._last_progress = clock() + + def enqueued(self, nbytes: int) -> None: + with self._lock: + self.queued_bytes += nbytes + + def begin(self) -> None: + with self._lock: + now = self._clock() + self._item_started = now + self._item_bytes = 0 + self._last_progress = now + + def progress(self, nbytes: int) -> None: + with self._lock: + self.bytes_sent += nbytes + self.queued_bytes -= nbytes + self._item_bytes += nbytes + self._last_progress = self._clock() + + def end(self) -> None: + with self._lock: + if self._item_started is None: + return + elapsed = self._clock() - self._item_started + self.busy_seconds += elapsed + if elapsed >= MIN_RATE_SAMPLE_SECONDS and self._item_bytes: + sample = self._item_bytes / elapsed + self.rate_bps = ( + sample + if self.rate_bps is None + else (1 - RATE_EWMA_ALPHA) * self.rate_bps + RATE_EWMA_ALPHA * sample + ) + self._item_started = None + + def drop_pending(self) -> None: + """The connection group died: whatever was queued is gone.""" + with self._lock: + self.queued_bytes = 0 + self._item_started = None + self._item_bytes = 0 + + @property + def idle(self) -> bool: + with self._lock: + return self._item_started is None and self.queued_bytes <= 0 + + def stalled_for(self, now: float | None = None) -> float: + """Seconds without byte progress while bytes are waiting; 0 when idle.""" + with self._lock: + if self._item_started is None and self.queued_bytes <= 0: + return 0.0 + return max(0.0, (self._clock() if now is None else now) - self._last_progress) + + def snapshot(self) -> StreamSnapshot: + with self._lock: + busy = self.busy_seconds + if self._item_started is not None: + busy += self._clock() - self._item_started + return StreamSnapshot(self.bytes_sent, busy, max(0, self.queued_bytes), self.rate_bps) + + +def pick_stream( + schedule: str, + meters: Sequence[StreamMeter], + rr_ticket: int, +) -> int: + """Index into ``meters`` (the data streams) for the next chunk. + + ``adaptive`` minimizes ``backlog / rate``. A stream without a rate sample + yet uses the mean measured rate (or 1 B/s when nothing is measured, which + reduces the rule to least backlog). Ties go to the round-robin ticket so + an idle start still spreads chunks over every stream.""" + count = len(meters) + if count == 0: + raise ValueError("no data streams") + if schedule == "rr": + return rr_ticket % count + if schedule != "adaptive": + raise ValueError(f"unknown stream schedule {schedule!r}") + snaps = [meter.snapshot() for meter in meters] + known = [snap.rate_bps for snap in snaps if snap.rate_bps] + default_rate = sum(known) / len(known) if known else 1.0 + best = None + for step in range(count): + index = (rr_ticket + step) % count + snap = snaps[index] + finish = snap.backlog_bytes / (snap.rate_bps or default_rate) + if best is None or finish < best[0]: + best = (finish, index) + return best[1] + + +@dataclass(frozen=True) +class UploadWindow: + """Data-stream traffic between two snapshots (one sync boundary).""" + + bytes_per_stream: tuple[int, ...] + busy_seconds_per_stream: tuple[float, ...] + + @property + def total_bytes(self) -> int: + return sum(self.bytes_per_stream) + + def stream_mbps(self) -> list[float | None]: + """MB/s (1e6 B) of each stream over the time it was busy.""" + return [ + (b / s / 1e6) if s > 0 and b > 0 else None + for b, s in zip(self.bytes_per_stream, self.busy_seconds_per_stream, strict=True) + ] + + +def upload_window( + before: Sequence[StreamSnapshot], after: Sequence[StreamSnapshot] +) -> UploadWindow: + if len(before) != len(after): # the group was redialed with another width + before = [StreamSnapshot(0, 0.0, 0, None)] * len(after) + return UploadWindow( + tuple(max(0, a.bytes_sent - b.bytes_sent) for b, a in zip(before, after, strict=True)), + tuple(max(0.0, a.busy_seconds - b.busy_seconds) for b, a in zip(before, after, strict=True)), + ) diff --git a/yeto/wandb_tape.py b/yeto/wandb_tape.py index 125afa62..59c4e584 100644 --- a/yeto/wandb_tape.py +++ b/yeto/wandb_tape.py @@ -110,6 +110,8 @@ def tape_metrics(rec: dict) -> dict: ): if source in rec: metrics[metric] = int(rec[source]) + if "checkpoint_ms" in rec: # per-round checkpoint wall time (S26 E3) + metrics["sync/checkpoint_ms"] = float(rec["checkpoint_ms"]) if "sweep_complete" in rec: metrics["sync/sweep_complete"] = int(bool(rec["sweep_complete"])) if expected: