diff --git a/docs/PROTOCOL.md b/docs/PROTOCOL.md index f8c52ec2..bcec53f7 100644 --- a/docs/PROTOCOL.md +++ b/docs/PROTOCOL.md @@ -593,6 +593,18 @@ coordinator always rewrites `--checkpoint-path` at the final quiescent cut before terminal delivery. Thus a total-step count not divisible by the periodic interval still leaves an atomic final checkpoint. +SIGTERM or SIGINT stops the syncer gracefully: it launches no new round, +gives the rounds still gathering `--shutdown-timeout-s` (default 600) to +close (a round with no delta whose islands are all disconnected is abandoned +at once), then abandons the rest (they have not touched the state, so a resumed +run launches those steps again), commits every computed round, writes +`--checkpoint-path` synchronously (after any background checkpoint) and exits +0 with `stop: checkpoint written step=`. The checkpoint's global step is +then the last `outer step` on the event tape, and `--resume` continues the +tape at `t + 1`. A signal before training starts exits at once (nothing was +merged), and a second signal exits immediately without a checkpoint. The +terminal cut is not marked final. + `--checkpoint-at-steps T1,T2,...` (launcher: `--rl-sync-checkpoint-at-steps`, asynchronous full-parameter rounds) also writes the cut at each listed global step `t` (one fragment merge, the `step=` of the `outer step` log, the unit of diff --git a/syncer/src/main.rs b/syncer/src/main.rs index 9bb05cca..8b2d1417 100644 --- a/syncer/src/main.rs +++ b/syncer/src/main.rs @@ -19,6 +19,7 @@ mod protocol; mod sample_relay; mod server; mod sha256; +mod shutdown; mod state; mod subscriber; @@ -158,6 +159,13 @@ struct Args { /// during the write) host memory peaks at about twice params + momentum. #[arg(long)] background_checkpoint: bool, + /// Write --checkpoint-path once before the first round (after INIT or + /// --resume), so the rounds before the first --checkpoint-every one are + /// covered too. Skipped when --resume found that checkpoint (it is the + /// state just loaded). On a background thread with + /// --background-checkpoint. + #[arg(long)] + checkpoint_on_start: bool, /// 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)] @@ -167,6 +175,13 @@ struct Args { /// sessions are refused. #[arg(long)] bcast_delta_full_every: Option, + /// On SIGTERM / SIGINT: stop launching rounds, give the rounds still + /// gathering this many seconds to close (then drop them; they have not + /// touched the state), commit every computed round, write + /// --checkpoint-path synchronously and exit 0. A second signal exits at + /// once without a checkpoint. + #[arg(long, default_value_t = 600)] + shutdown_timeout_s: u64, /// Resume from --checkpoint-path if it exists. #[arg(long, default_value_t = false)] resume: bool, @@ -505,6 +520,7 @@ fn main() -> anyhow::Result<()> { checkpoint_every: args.checkpoint_every, checkpoint_at_steps: args.checkpoint_at_steps, background_checkpoint: args.background_checkpoint, + checkpoint_on_start: args.checkpoint_on_start, async_checkpoint: args.async_checkpoint, bcast_delta_full_every: args.bcast_delta_full_every, resume: args.resume, @@ -518,6 +534,9 @@ fn main() -> anyhow::Result<()> { island_hmac_key, expected_island_contract, syncer_epoch: args.syncer_epoch, + shutdown: Some(std::sync::Arc::new(shutdown::Shutdown::new(std::time::Duration::from_secs( + args.shutdown_timeout_s, + )))), membership, recovery, inline_resend: args.inline_resend, diff --git a/syncer/src/server.rs b/syncer/src/server.rs index f15d96b1..d944c457 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -104,6 +104,8 @@ pub struct Config { pub quorum_timeout_s: u64, /// Bounded wait for learners to apply and acknowledge the terminal cut. pub final_ack_timeout_s: u64, + /// Graceful stop (`crate::shutdown`); None: signals keep their default action. + pub shutdown: Option>, pub total_steps: u64, /// Strict dense-policy mode: one logical local step contributes every /// fragment in a serial sweep. None preserves the legacy independent @@ -125,6 +127,9 @@ pub struct Config { /// rounds keep merging (`GlobalState::save_checkpoint_in_background`) /// instead of on the scheduler. Every other checkpoint waits for it. pub background_checkpoint: bool, + /// Write the checkpoint once before the first round, unless `resume` + /// loaded it (`--checkpoint-on-start`). + pub checkpoint_on_start: bool, /// Global steps whose cut is also written to `.step-` /// off the scheduler (`crate::eval_checkpoint`); strictly increasing, /// below `total_steps`. Empty: none. @@ -1498,6 +1503,28 @@ fn round_grace( (grace, floor) } +/// `--checkpoint-on-start`: write the checkpoint before the first round. +/// Without it the first rounds (16 at ~860 s per write on the 120B run) have +/// no checkpoint, and a restart then loses the INIT and every merge so far. +/// A resume already has one (the state it loaded); a finished run needs none. +/// Returns whether a write was started. +fn save_start_checkpoint(cfg: &Config, st: &GlobalState, resumed: bool, last_step: u64) -> Result { + let Some(path) = cfg.checkpoint_path.as_ref().filter(|_| cfg.checkpoint_on_start) else { + return Ok(false); + }; + if resumed || st.global_step >= last_step { + return Ok(false); + } + if cfg.background_checkpoint { + st.save_checkpoint_in_background(path)?; + info!(step = st.global_step, path = %path.display(), "start checkpoint started in the background"); + } else { + st.save_checkpoint(path)?; + info!(step = st.global_step, path = %path.display(), "start checkpoint written"); + } + Ok(true) +} + /// Whether step `t` is the last round of its fragment within `total_steps`. fn is_terminal_round(schedule: &crate::shard::ShardSchedule, t: u64, total_steps: u64) -> bool { schedule.next_merge_of_fragment(t) > total_steps @@ -1829,6 +1856,9 @@ fn validate_config(cfg: &Config) -> Result<()> { without --policy-sweep-fragments (per-round checkpoints persist before their broadcast)" ); } + if cfg.checkpoint_on_start && cfg.checkpoint_path.is_none() { + bail!("--checkpoint-on-start needs --checkpoint-path"); + } if cfg.async_checkpoint { if cfg.policy_sweep_fragments.is_none() { bail!("--async-checkpoint requires --policy-sweep-fragments"); @@ -2080,6 +2110,9 @@ fn check_and_record_profile(cfg: &Config, profile: &[u8; 32]) -> Result<()> { pub async fn run(cfg: Config) -> Result<()> { validate_config(&cfg)?; + if let Some(stop) = &cfg.shutdown { + stop.install_signal_handlers()?; + } check_resume_checkpoint(&cfg)?; let semantic_profile_hash = cfg.semantic_profile_hash(); check_and_record_profile(&cfg, &semantic_profile_hash)?; @@ -3848,14 +3881,20 @@ async fn scheduler( break; } } - match recv_init_event( - &mut events, - &budget_cutoff, - budget_reports.len(), - cfg.learners, - ) - .await? - { + let init_event = recv_init_event(&mut events, &budget_cutoff, budget_reports.len(), cfg.learners); + let event = match cfg.shutdown.as_deref() { + Some(stop) => tokio::select! { + biased; + _ = stop.requested() => { + // Nothing merged yet: the checkpoint on disk (if any) is current. + info!("stop requested before training started; exiting"); + return Ok(()); + } + event = init_event => event?, + }, + None => init_event.await?, + }; + match event { Event::Fatal { metric, message } => { append_strict_failure(cfg.event_tape.as_deref(), metric, &message); bail!("RL strict failure {metric}: {message}"); @@ -4061,6 +4100,8 @@ async fn scheduler( info!(step = st.global_step, path = %path.display(), "checkpoint committed"); } } + } else { + save_start_checkpoint(&cfg, &st, resumed, last_step)?; } // Send everyone the initial (or resumed) global parameters so all // learners start bit-identical (also serves recovery for late joiners). @@ -4134,6 +4175,9 @@ async fn scheduler( if dynamic { write_pool_status(&cfg, ®istry, st.global_step, "running"); } + // A stop signal (`crate::shutdown`): set once every round has drained. + let mut stopped = false; + let mut stop_logged = false; 'outer: while next_t <= last_step || !inflight.is_empty() || !computing.is_empty() @@ -4144,6 +4188,55 @@ async fn scheduler( cutoff_established = true; break 'outer; } + let stop_deadline = cfg.shutdown.as_ref().and_then(|stop| stop.deadline()); + if let Some(deadline) = stop_deadline { + if !stop_logged { + stop_logged = true; + info!( + step = st.global_step, + gathering = inflight.len(), + computing = computing.len() + ready.len(), + "stopping: no new rounds; waiting for the rounds in flight" + ); + } + if Instant::now() >= deadline && !inflight.is_empty() { + // A gathering round has not touched the state; after --resume + // its step is launched again. + let dropped: Vec = inflight.iter().map(|round| round.t).collect(); + warn!(steps = ?dropped, "stop timeout: abandoning the rounds still gathering"); + for round in inflight.drain(..) { + busy_fragments.remove(&round.p); + } + } else if !inflight.is_empty() { + // A round that holds no delta and none of whose islands is + // still connected cannot close: do not wait the timeout out. + let orphaned: Vec = { + let registry = registry.lock().unwrap(); + inflight + .iter() + .enumerate() + .filter(|(_, round)| { + round.pushes.is_empty() + && round.carried.is_empty() + && !round.expected_members.iter().any(|m| registry.groups.contains_key(m)) + }) + .map(|(i, _)| i) + .collect() + }; + if !orphaned.is_empty() { + let steps: Vec = orphaned.iter().map(|&i| inflight[i].t).collect(); + warn!(?steps, "stopping: abandoning gathering rounds whose islands are all gone"); + for &i in orphaned.iter().rev() { + let round = inflight.remove(i); + busy_fragments.remove(&round.p); + } + } + } + if inflight.is_empty() && computing.is_empty() && ready.is_empty() && publishing.is_none() { + stopped = true; + break 'outer; + } + } let expired = expire_member_leases(®istry, &mut inflight, &mut step_rates).await; for member in &expired { lease_expired.insert(*member); @@ -4241,6 +4334,7 @@ async fn scheduler( // Keep the pipeline full (throttled by min_round_interval_ms). while inflight.len() + computing.len() + ready.len() < depth && next_t <= last_step + && stop_deadline.is_none() && Instant::now() >= next_launch { if !budget_cutoff.try_linearize_work() { @@ -4372,6 +4466,13 @@ async fn scheduler( cutoff_established = true; break 'outer; } + if stop_deadline.is_some() { + // Stopping: no new attempt; the round is launched again after --resume. + let round = inflight.remove(i); + busy_fragments.remove(&round.p); + warn!(step = round.t, "stop: quorum timeout, round abandoned"); + continue; + } let r = &mut inflight[i]; let groups = current_groups(®istry); if fixed_roster { @@ -4569,6 +4670,7 @@ async fn scheduler( .min(); let next_fragment_available = inflight.len() + computing.len() + ready.len() < depth && next_t <= last_step + && stop_deadline.is_none() && !busy_fragments.contains(&schedule.fragment_for_step(next_t)); if next_fragment_available { earliest = Some(earliest.map_or(next_launch, |d| d.min(next_launch))); @@ -4582,6 +4684,11 @@ async fn scheduler( let check = deadline_now + RESEND_CHECK_INTERVAL; earliest = Some(earliest.map_or(check, |d| d.min(check))); } + if cfg.shutdown.is_some() { + // Notice a stop signal (and its timeout) within a fraction of a second. + let check = deadline_now + SHUTDOWN_CHECK_INTERVAL; + earliest = Some(earliest.map_or(check, |d| d.min(check))); + } 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. @@ -5107,6 +5214,24 @@ async fn scheduler( resends.stop(®istry) }; + if stopped { + drop(forwarders); // a catch-up in progress is resent after the restart + st.drain_iso_backend().await?; + eval_checkpoints.finish().await?; + if let Some(path) = &cfg.checkpoint_path { + // Waits for a background checkpoint first, then writes this cut. + st.save_checkpoint(path)?; + info!(step = st.global_step, path = %path.display(), "stop: checkpoint written"); + } else { + warn!(step = st.global_step, "stop without --checkpoint-path: the merged state is not saved"); + } + if dynamic { + write_pool_status(&cfg, ®istry, st.global_step, "stopped"); + } + info!(step = st.global_step, "stopped; restart with --resume to continue"); + return Ok(()); + } + if cfg.learner_budget_steps.is_some() { wait_forwarders(std::mem::take(&mut forwarders)).await; cutoff_established |= budget_cutoff.deadline().is_some(); @@ -5878,21 +6003,25 @@ fn route_push( /// same fragment against its own base, weighted by its counters (and the /// carry discount), as long as its base lag stays within `--max-base-lag`. /// -/// A carried delta counts once: after the syncer carries a delta from an -/// anchor (fragment version), a later push of that island connection from -/// the same or an older anchor of that fragment is refused -/// (`anchor_reused`), and a late push from an anchor whose delta was -/// already merged is not carried. Pushes that never involve a carry keep -/// the existing stale-base rules. +/// A delta counts once: after the syncer carries a delta from an anchor +/// (fragment version), a later push of that island from the same or an +/// older anchor of that fragment is refused (`anchor_reused`) on any of its +/// connections; after it gathers one, the same holds for the island's later +/// connections. An island that redials (a WAN drop, a syncer it lost +/// for a while) releases the anchors its lost pushes held, so it may answer +/// the next pull from an anchor whose delta did reach this syncer; that +/// second delta repeats the first one's local steps and must not be merged +/// again. A round restart frees the anchors of the pushes it dropped +/// (`release_anchor`). Pushes from newer anchors keep the stale-base rules. struct LateCarry { limit: u64, /// Late pushes waiting for their fragment's next round. queued: HashMap>, - /// Newest base version of a carried delta per island connection and - /// fragment. - anchors: HashMap<(Member, usize), u64>, - /// Newest base version of a delta gathered by its own round. - gathered: HashMap<(Member, usize), u64>, + /// Newest base version of a carried delta per island and fragment. + anchors: HashMap<(u32, usize), u64>, + /// Newest base version of a delta gathered by its own round, per island + /// and fragment, with the connection generation that pushed it. + gathered: HashMap<(u32, usize), (u64, u64)>, } #[derive(Debug, Eq, PartialEq)] @@ -5920,32 +6049,42 @@ impl LateCarry { } } - /// True when `push` reuses an anchor this connection already spent. + /// True when `push` reuses an anchor whose delta this island already had + /// carried (on any connection) or gathered on an earlier connection. One + /// connection may answer consecutive rounds from the same anchor (an + /// asynchronous learner that has not installed the newer version yet); + /// a redialed one repeating the anchor resends a delta already taken. fn reuses_anchor(&self, member: Member, push: &Push) -> bool { - self.anchors - .get(&(member, push.fragment_id as usize)) - .is_some_and(|&spent| push.base_version <= spent) + let key = (member.learner_id, push.fragment_id as usize); + let carried = self.anchors.get(&key).is_some_and(|&taken| push.base_version <= taken); + let gathered = self + .gathered + .get(&key) + .is_some_and(|&(taken, generation)| push.base_version <= taken && member.generation != generation); + carried || gathered } fn spend_anchor(&mut self, member: Member, fragment: usize, base_version: u64) { - let spent = self.anchors.entry((member, fragment)).or_insert(0); + let spent = self.anchors.entry((member.learner_id, fragment)).or_insert(0); *spent = (*spent).max(base_version); } fn note_gathered(&mut self, member: Member, fragment: usize, base_version: u64) { - let taken = self.gathered.entry((member, fragment)).or_insert(0); - *taken = (*taken).max(base_version); + let taken = self.gathered.entry((member.learner_id, fragment)).or_insert((0, member.generation)); + if base_version >= taken.0 { + *taken = (base_version, member.generation); + } } /// A restarted round drops its gathered pushes; their islands answer the /// new attempt from the same anchor, so the anchor is free again. fn release_anchor(&mut self, member: Member, push: &Push) { - let key = (member, push.fragment_id as usize); - if self.gathered.get(&key) == Some(&push.base_version) { + let key = (member.learner_id, push.fragment_id as usize); + if let Some(taken) = self.gathered.get_mut(&key).filter(|taken| taken.0 == push.base_version) { if push.base_version == 0 { self.gathered.remove(&key); } else { - self.gathered.insert(key, push.base_version - 1); + taken.0 = push.base_version - 1; } } } @@ -5974,8 +6113,8 @@ impl LateCarry { } if self .gathered - .get(&(member, p)) - .is_some_and(|&taken| push.base_version <= taken) + .get(&(member.learner_id, p)) + .is_some_and(|&(taken, _)| push.base_version <= taken) { return CarryOutcome::Rejected(PushDisposition::Duplicate); } @@ -7696,6 +7835,8 @@ impl CatchUps { /// How often the scheduler checks background resends while any is active. const RESEND_CHECK_INTERVAL: Duration = Duration::from_millis(100); +/// How often the scheduler looks for a stop signal while it waits. +const SHUTDOWN_CHECK_INTERVAL: Duration = Duration::from_millis(200); /// At the end of the round loop, wait this long for a forwarder to hand /// over the fragment it is sending before it is aborted. const RESEND_STOP_WAIT: Duration = Duration::from_secs(30); @@ -9465,6 +9606,7 @@ mod tests { delta_correction: false, quorum_timeout_s: 1, final_ack_timeout_s: 1, + shutdown: None, total_steps, policy_sweep_fragments: None, outer_lr: 1.0, @@ -9476,6 +9618,7 @@ mod tests { checkpoint_every: 0, checkpoint_at_steps: Vec::new(), background_checkpoint: false, + checkpoint_on_start: false, async_checkpoint: false, bcast_delta_full_every: None, resume: false, @@ -11143,6 +11286,41 @@ mod tests { std::fs::remove_dir_all(&directory).ok(); } + #[test] + fn checkpoint_on_start_covers_the_rounds_before_the_first_periodic_one() { + let directory = unique_test_directory("start-checkpoint"); + for background in [false, true] { + let checkpoint = directory.join(format!("start-{background}.ckpt")); + let mut config = round_test_config(8); + config.checkpoint_path = Some(checkpoint.clone()); + config.checkpoint_every = 4; + config.background_checkpoint = background; + let mut state = two_fragment_sweep_state(); + state.policy_sweep_fragments = None; + // Off: nothing before the first periodic checkpoint. + assert!(!save_start_checkpoint(&config, &state, false, 8).unwrap()); + assert!(!checkpoint.exists()); + config.checkpoint_on_start = true; + validate_config(&config).unwrap(); + // A resume already holds the state on disk; a finished run needs none. + assert!(!save_start_checkpoint(&config, &state, true, 8).unwrap()); + state.global_step = 8; + assert!(!save_start_checkpoint(&config, &state, false, 8).unwrap()); + state.global_step = 0; + assert!(!checkpoint.exists()); + assert!(save_start_checkpoint(&config, &state, false, 8).unwrap()); + state.wait_for_background_checkpoint().unwrap(); + let mut on_disk = two_fragment_sweep_state(); + on_disk.policy_sweep_fragments = None; + on_disk.load_checkpoint(&checkpoint).unwrap(); + assert_eq!((on_disk.global_step, on_disk.versions.clone()), (0, state.versions.clone())); + } + let mut config = round_test_config(8); + config.checkpoint_on_start = true; + assert!(validate_config(&config).is_err(), "needs a checkpoint path"); + std::fs::remove_dir_all(&directory).ok(); + } + #[tokio::test] async fn background_checkpoint_keeps_rounds_going_and_writes_the_same_cuts() { let directory = unique_test_directory("background-checkpoint"); @@ -12457,6 +12635,32 @@ mod tests { assert!(carry.queued.is_empty()); } + #[test] + fn a_gathered_anchor_yields_one_delta_across_reconnects() { + // Island 1 answered round t=7 of fragment 1 from anchor 5; the delta + // is gathered. The link drops, the learner redials (generation 11), + // releases its spent anchor and answers the next pull from the same + // anchor 5 before the round's broadcast reached it: refused. + let mut carry = LateCarry::new(Some(4), false).unwrap(); + let (old, new) = (member(1, 10), member(1, 11)); + carry.note_gathered(old, 1, 5); + let again = test_push_for(1, 5); + assert!(carry.reuses_anchor(new, &again)); + // The connection that pushed it may answer the next round from the + // same anchor (it has not installed the newer version yet). + assert!(!carry.reuses_anchor(old, &again)); + // From the broadcast it then gets (version 7) it pushes again. + assert!(!carry.reuses_anchor(new, &test_push_for(1, 7))); + // A restart of that round drops the gathered delta: the anchor is + // free for the new attempt, from either connection. + carry.release_anchor(old, &test_push_for(1, 5)); + assert!(!carry.reuses_anchor(new, &again)); + // A restarted syncer starts with nothing taken: the delta it lost is + // merged once when the island answers from that anchor again. + let fresh = LateCarry::new(Some(4), false).unwrap(); + assert!(!fresh.reuses_anchor(new, &again)); + } + #[test] fn a_carried_anchor_yields_one_delta() { let mut carry = LateCarry::new(Some(2), false).unwrap(); @@ -12467,8 +12671,14 @@ mod tests { assert!(carry.reuses_anchor(island, &push)); assert!(carry.reuses_anchor(island, &test_push_for(1, 3))); assert!(!carry.reuses_anchor(island, &test_push_for(1, 7))); - // Another connection or fragment is unaffected. - assert!(!carry.reuses_anchor(member(1, 11), &push)); + // The island's next connection may not spend it again (a redialed + // learner answers from the same anchor); another island or + // fragment is unaffected. + assert!(carry.reuses_anchor(member(1, 11), &push)); + assert!(!carry.reuses_anchor(member(3, 10), &test_push_for(1, 5))); + let mut elsewhere = test_push_for(1, 5); + elsewhere.fragment_id += 1; + assert!(!carry.reuses_anchor(island, &elsewhere)); // A late push from an anchor already merged by its own round is not // carried; a restart clears the gathered push and frees the anchor. let other = member(2, 10); diff --git a/syncer/src/shutdown.rs b/syncer/src/shutdown.rs new file mode 100644 index 00000000..9caabd90 --- /dev/null +++ b/syncer/src/shutdown.rs @@ -0,0 +1,124 @@ +//! Graceful stop on SIGTERM / SIGINT. +//! +//! The first signal asks the scheduler to stop: no new round is launched, +//! the rounds already gathering get `--shutdown-timeout-s` to close, every +//! computed round commits, and the consistent cut is written synchronously +//! to `--checkpoint-path` before the process exits 0, so `--resume` starts +//! at the last outer step on the event tape. A second signal exits at once. + +use std::time::{Duration, Instant}; + +use tokio::sync::watch; +use tracing::warn; + +/// Exit status of a second signal while the first one is being served. +pub const FORCED_EXIT_CODE: i32 = 143; + +pub struct Shutdown { + tx: watch::Sender>, + timeout: Duration, +} + +impl Shutdown { + pub fn new(timeout: Duration) -> Self { + Self { + tx: watch::channel(None).0, + timeout, + } + } + + /// Ask for a stop; returns false if one was already requested. + pub fn request(&self) -> bool { + let mut first = false; + self.tx.send_if_modified(|at| { + first = at.is_none(); + if first { + *at = Some(Instant::now()); + } + first + }); + first + } + + pub fn requested_at(&self) -> Option { + *self.tx.borrow() + } + + /// When the rounds still gathering are abandoned. + pub fn deadline(&self) -> Option { + self.requested_at().map(|at| at + self.timeout) + } + + /// Resolves once a stop is requested. + pub async fn requested(&self) { + let mut rx = self.tx.subscribe(); + let _ = rx.wait_for(Option::is_some).await; + } + + /// Route SIGTERM and SIGINT here (inside a tokio runtime). + #[cfg(unix)] + pub fn install_signal_handlers(self: &std::sync::Arc) -> anyhow::Result<()> { + use tokio::signal::unix::{signal, SignalKind}; + + let mut term = signal(SignalKind::terminate())?; + let mut int = signal(SignalKind::interrupt())?; + let this = self.clone(); + tokio::spawn(async move { + loop { + let name = tokio::select! { + _ = term.recv() => "SIGTERM", + _ = int.recv() => "SIGINT", + }; + if this.request() { + warn!( + signal = name, + timeout_s = this.timeout.as_secs(), + "stop requested: no new rounds; draining in-flight rounds, then checkpoint and exit \ + (a second signal exits at once)" + ); + } else { + warn!(signal = name, "second stop signal: exiting without a checkpoint"); + std::process::exit(FORCED_EXIT_CODE); + } + } + }); + Ok(()) + } + + #[cfg(not(unix))] + pub fn install_signal_handlers(self: &std::sync::Arc) -> anyhow::Result<()> { + let this = self.clone(); + tokio::spawn(async move { + while tokio::signal::ctrl_c().await.is_ok() { + if !this.request() { + std::process::exit(FORCED_EXIT_CODE); + } + } + }); + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn first_request_wins_and_wakes_waiters() { + let stop = std::sync::Arc::new(Shutdown::new(Duration::from_secs(5))); + assert!(stop.requested_at().is_none() && stop.deadline().is_none()); + let waiter = { + let stop = stop.clone(); + tokio::spawn(async move { stop.requested().await }) + }; + tokio::task::yield_now().await; + assert!(stop.request()); + let at = stop.requested_at().unwrap(); + assert!(!stop.request()); + assert_eq!(stop.requested_at(), Some(at)); + assert_eq!(stop.deadline(), Some(at + Duration::from_secs(5))); + tokio::time::timeout(Duration::from_secs(1), waiter).await.unwrap().unwrap(); + // A waiter that starts after the request returns at once. + tokio::time::timeout(Duration::from_secs(1), stop.requested()).await.unwrap(); + } +} diff --git a/tests/test_diloco_sync.py b/tests/test_diloco_sync.py index bc086886..2ad902a4 100644 --- a/tests/test_diloco_sync.py +++ b/tests/test_diloco_sync.py @@ -364,3 +364,62 @@ def test_a_resent_pull_of_the_same_attempt_is_answered_once(): state.pending_pulls = [pull, PullRequest(fragment_id=0, global_step=3, round_attempt=1)] ready = take_ready_pulls(state, steps_total=5, units_total=50, anchored=lambda fid: True) assert [p for p, _, _ in ready] == [pull] and state.pending_pulls == [] + + +def test_resume_record_round_trips_and_picks_the_install_mode(): + import json + + from yeto.diloco_sync import ( + RESUME_ANCHOR, + RESUME_OVERWRITE, + RESUME_SHIFT, + resume_fragment_mode, + shifted_resume_value, + ) + + state = DiLoCoSyncState.create(3, track_anchors=False) + state.fragment_versions = [4, 7, 0] + state.steps_at_reset = [0, 12, 0] + state.units_at_reset = [0, 500, 0] + state.global_step = 7 + record = json.loads(json.dumps(state.resume_record( + anchored=[True, True, False], units_total=900, spent=[None, (7, 9), None]))) + assert record["spent"] == [None, [7, 9], None] and record["units_total"] == 900 + # Same version: the broadcast is the checkpoint's own anchor. + assert resume_fragment_mode(record, 0, 4, have_anchor=False) == RESUME_ANCHOR + # Newer, after a push from that anchor: the delta may be merged into it, + # so it is installed as is rather than counted twice. + assert resume_fragment_mode(record, 1, 9, have_anchor=True) == RESUME_OVERWRITE + # Newer, nothing pushed from the anchor: move the local progress, only + # when the old anchor is at hand. + unspent = dict(record, spent=[None, None, None]) + assert resume_fragment_mode(unspent, 1, 9, have_anchor=True) == RESUME_SHIFT + assert resume_fragment_mode(unspent, 1, 9, have_anchor=False) == RESUME_OVERWRITE + # A rewound global (a syncer resumed from an older checkpoint) still shifts. + assert resume_fragment_mode(record, 1, 5, have_anchor=True) == RESUME_SHIFT + # Never anchored before the checkpoint, or no record at all: a cold start. + assert resume_fragment_mode(record, 2, 0, have_anchor=True) == RESUME_OVERWRITE + assert resume_fragment_mode(None, 0, 4, have_anchor=True) == RESUME_OVERWRITE + fresh = DiLoCoSyncState.create(3, track_anchors=False) + fresh.steps_at_reset = [15, 15, 15] + fresh.restore_fragment(record, 1) + assert fresh.steps_at_reset == [15, 12, 15] and fresh.units_at_reset == [0, 500, 0] + new, local, old = torch.tensor([1.0, 2.0]), torch.tensor([0.5, 0.75]), torch.tensor([0.25, 1.0]) + assert torch.equal(shifted_resume_value(new, local, old), torch.tensor([1.25, 1.75])) + + +def test_live_install_mode_keeps_a_running_islands_progress_through_a_catch_up(): + from yeto.diloco_sync import LIVE_CONFIRM, LIVE_REBASE, live_install_mode + + def mode(version, catch_up, pushed_from=None, anchored=True): + return live_install_mode(anchored=anchored, current_version=5, version=version, + catch_up=catch_up, pushed_from=pushed_from) + + assert mode(9, False, anchored=False) is None and mode(9, True, anchored=False) is None + assert mode(5, False) == LIVE_CONFIRM # a replay of the anchor + assert mode(9, False) is None # an ordinary merge broadcast + assert mode(9, True) == LIVE_REBASE # a merge missed while away, not pushed from this anchor + assert mode(9, True, pushed_from=5) is None # its delta may be in it + assert mode(9, True, pushed_from=3) == LIVE_REBASE + assert mode(3, True, pushed_from=5) == LIVE_REBASE # a rewound syncer + assert mode(5, True) == LIVE_REBASE # same number, other values (the client compared the bytes) diff --git a/tests/test_head_mode.py b/tests/test_head_mode.py index 91f74e34..9392b08e 100644 --- a/tests/test_head_mode.py +++ b/tests/test_head_mode.py @@ -838,6 +838,8 @@ def test_local_syncer_command_matches_cluster_syncer_flags(): assert "--resume" in cmd assert "--mark-final-checkpoint" in cmd assert f"--port {launcher.SYNCER_PORT}" in cmd + # The SFT preset: periodic checkpoints off the scheduler, one before the first round. + assert "--background-checkpoint" in cmd and "--checkpoint-on-start" in cmd def test_rl_launch_records_the_run_under_its_own_name(fake_sky, monkeypatch, tmp_path): diff --git a/tests/test_integration.py b/tests/test_integration.py index a5aa75a5..f582925e 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -757,3 +757,75 @@ def test_single_learner_roundtrip(): if proc.poll() is None: proc.kill() print((proc.stdout.read() if proc.stdout else "")[-3000:]) + + +@pytest.mark.timeout(180) +def test_sigterm_checkpoints_the_last_outer_step_and_resume_continues_the_tape(tmp_path): + """SIGTERM stops a mid-run syncer cleanly: no new rounds, the round in + flight closes, the cut is written to --checkpoint-path (although no + --checkpoint-every boundary was reached) and the process exits 0; the + checkpoint's global step is the last outer step on the tape, and a + --resume run continues the tape from the next step.""" + import signal + + binary = build_syncer() + named = [("model.embed.weight", DIM // 4), ("model.body.weight", DIM)] + layout = build_layout(named, 2) + checkpoint = tmp_path / "state.ckpt" + tape = tmp_path / "tape.jsonl" + + def syncer(port, total_steps, *extra): + return group_popen( + [str(binary), "--port", str(port), "--learners", "1", "--quorum", "1", + "--grace-ms", "20", "--sync-interval-steps", "0", "--min-round-interval-ms", "50", + "--total-steps", str(total_steps), "--checkpoint-path", str(checkpoint), + "--checkpoint-every", "1000", "--event-tape", str(tape), + "--shutdown-timeout-s", "5", *extra], + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, + ) + + def steps(): + if not tape.exists(): + return [] + return [json.loads(line)["step"] for line in tape.read_text().splitlines() + if line.strip() and "step" in json.loads(line)] + + port = free_port() + first = syncer(port, 1000) + outputs = [] + try: + learner = ToyLearner(0, port, torch.ones(DIM + DIM // 4), layout, dtype=DTYPE_F32) + learner.start() + deadline = time.monotonic() + 60 + while len(steps()) < 4: + assert time.monotonic() < deadline, "no outer steps before the stop" + assert first.poll() is None, first.communicate()[0] + time.sleep(0.05) + first.send_signal(signal.SIGTERM) + output, _ = first.communicate(timeout=60) + outputs.append(output) + assert first.returncode == 0, output + assert "stop: checkpoint written" in output + stopped = steps() + assert stopped == list(range(1, len(stopped) + 1)) + assert read_checkpoint_global_step(checkpoint) == stopped[-1] + learner.join(timeout=30) + + port = free_port() + total = stopped[-1] + 4 + second = syncer(port, total, "--resume") + learner = ToyLearner(0, port, torch.ones(DIM + DIM // 4), layout, dtype=DTYPE_F32) + learner.start() + output, _ = second.communicate(timeout=90) + outputs.append(output) + assert second.returncode == 0, output + assert "resumed from checkpoint" in output + learner.join(timeout=30) + assert learner.final_manifest is not None and learner.final_manifest.global_step == total + assert steps() == list(range(1, total + 1)) + assert read_checkpoint_global_step(checkpoint) == total + finally: + for process in (first, locals().get("second")): + if process is not None and process.poll() is None: + process.kill() + print("\n".join(outputs)[-4000:]) diff --git a/tests/test_reconnect.py b/tests/test_reconnect.py index ce9710d2..5eb224a6 100644 --- a/tests/test_reconnect.py +++ b/tests/test_reconnect.py @@ -440,3 +440,125 @@ def test_lease_expired_error_redials_with_new_generation_even_without_reconnect_ finally: client.close() listener.close() + + +def _redial(client): + """What _connect_group does for a redial: every fragment awaits its catch-up.""" + with client._bcast_lock: + client._catch_up_pending = [True] * client.layout.num_fragments + + +def test_first_broadcast_after_a_redial_is_flagged_and_a_rewound_one_is_taken_on_request(): + layout = make_layout() + v2 = struct.pack("<4f", 2.0, 2.0, 2.0, 2.0) + v1 = struct.pack("<4f", 1.0, 1.0, 1.0, 1.0) + for accept in (False, True): + client = SyncerClient(("127.0.0.1", 1), LEARNER_ID, layout, dtype=DTYPE_F32, num_streams=0, + accept_rewound_catch_up=accept) + client._gen = 1 + client._dispatch(1, MSG_BCAST_FRAGMENT, struct.pack(" 0 and self.num_steps <= self.warmup: + return self.max_lr * self.num_steps / self.warmup + ratio = min(1.0, (self.num_steps - self.warmup) / (self.decay - self.warmup)) + return self.min_lr + (self.max_lr - self.min_lr) * 0.5 * (math.cos(math.pi * ratio) + 1.0) + + def step(self, increment): + self.num_steps += increment + for group in self.groups: + group["lr"] = self.get_lr() + + +def _lr_trainer(iteration, gbs=8, scheduler=None): + return types.SimpleNamespace(args=types.SimpleNamespace(global_batch_size=gbs), + state=types.SimpleNamespace(iteration=iteration), + opt_param_scheduler=scheduler or _Scheduler([{}])) + + +def test_a_resume_without_scheduler_state_continues_the_lr_schedule_in_place(monkeypatch): + """A params-only checkpoint (--no_save_optim) restores the iteration but + no scheduler: without the rebuild the resumed lr at iteration i is the + uninterrupted one at i - 5 (warmup again).""" + + monkeypatch.delenv("YETO_DILOCO_LR_RESUME", raising=False) + reference = _Scheduler([{}]) + want = {} + for it in range(1, 141): + reference.step(8) + want[it] = reference.groups[0]["lr"] + trainer = _lr_trainer(5) + assert sp.restore_lr_schedule(trainer) == 40 + resumed = trainer.opt_param_scheduler + assert resumed.groups[0]["lr"] == want[5] + for it in range(6, 141): + resumed.step(8) + assert resumed.groups[0]["lr"] == want[it], it + # Without it: shifted by the resumed iteration. + shifted = _Scheduler([{}]) + for it in range(6, 141): + shifted.step(8) + assert shifted.groups[0]["lr"] == want[it - 5] + + +def test_lr_schedule_is_left_alone_when_loaded_cold_or_off(monkeypatch): + loaded = _Scheduler([{}]) + loaded.step(5 * 8) # restored from a checkpoint with optimizer state + assert sp.restore_lr_schedule(_lr_trainer(5, scheduler=loaded)) is None and loaded.num_steps == 40 + cold = _lr_trainer(0) + assert sp.restore_lr_schedule(cold) is None and cold.opt_param_scheduler.num_steps == 0 + assert sp.restore_lr_schedule(types.SimpleNamespace(state=types.SimpleNamespace(iteration=5))) is None + monkeypatch.setenv("YETO_DILOCO_LR_RESUME", "off") + off = _lr_trainer(5) + assert sp.restore_lr_schedule(off) is None and off.opt_param_scheduler.num_steps == 0 + + +def test_callback_rebuilds_the_lr_schedule_on_attach(make): + trainer = FakeTrainer([([1], True)], start_iteration=15) + trainer.args.global_batch_size = 384 + trainer.opt_param_scheduler = _Scheduler([{}], warmup=10 * 384, decay=1000 * 384) + cb = sp.YetoDiLoCoCallback(trainer) + assert trainer.opt_param_scheduler.num_steps == 15 * 384 + cb.close() + + +def test_initial_cut_wait_is_on_by_default(monkeypatch): + for k, v in ENV.items(): + monkeypatch.setenv(k, v) + monkeypatch.delenv("YETO_DILOCO_INITIAL_CUT_TIMEOUT", raising=False) + assert sp.DiLoCoHookConfig.from_env().initial_cut_timeout == 3600.0 + monkeypatch.setenv("YETO_DILOCO_INITIAL_CUT_TIMEOUT", "0") + assert sp.DiLoCoHookConfig.from_env().initial_cut_timeout == 0.0 + + +def _logger_under(root_handler_level): + """A logger whose root mimics the Megatron SFT entry point's: one stdout + handler at ``root_handler_level`` (ERROR in the real process).""" + import io + import logging + + root = logging.RootLogger(logging.INFO) + sink = io.StringIO() + handler = logging.StreamHandler(sink) + handler.setLevel(root_handler_level) + root.addHandler(handler) + logger = logging.Logger("yeto-test-visible") + logger.parent = root + return logger, sink + + +def test_logs_are_made_visible_when_the_root_handler_drops_them(capsys): + import logging + + from yeto.swift_megatron.sync import ensure_visible_logs + + logger, root_sink = _logger_under(logging.ERROR) + ensure_visible_logs(logger) + logger.info("info line") + logger.warning("warning line") + logger.error("error line") + err = capsys.readouterr().err + assert "info line" in err and "warning line" in err and err.count("error line") == 1 + assert root_sink.getvalue() == "" # no longer propagated (no double ERROR) + ensure_visible_logs(logger) # idempotent + assert len(logger.handlers) == 1 + + +def test_logs_are_left_alone_when_the_process_prints_info(capsys): + import logging + + from yeto.swift_megatron.sync import ensure_visible_logs + + logger, root_sink = _logger_under(logging.INFO) + ensure_visible_logs(logger) + assert logger.handlers == [] and logger.propagate + logger.info("via root") + assert "via root" in root_sink.getvalue() and capsys.readouterr().err == "" + + +def test_the_lr_rebuild_warning_survives_an_error_level_root(make, monkeypatch, capsys): + import logging + + import yeto.swift_megatron.sync as sync_mod + + plugin_logger, _ = _logger_under(logging.ERROR) + sync_logger, _ = _logger_under(logging.ERROR) + monkeypatch.setattr(sp, "logger", plugin_logger) + monkeypatch.setattr(sync_mod, "log", sync_logger) + trainer = FakeTrainer([([1], True)], start_iteration=15) + trainer.args.global_batch_size = 384 + trainer.opt_param_scheduler = _Scheduler([{}], warmup=10 * 384, decay=1000 * 384) + cb = sp.YetoDiLoCoCallback(trainer) + err = capsys.readouterr().err + assert "restored no LR schedule; stepped it by 5760" in err + assert sync_logger.handlers and not sync_logger.propagate + cb.close() diff --git a/tests/test_swift_megatron_sync.py b/tests/test_swift_megatron_sync.py index 5b136d19..274ec284 100644 --- a/tests/test_swift_megatron_sync.py +++ b/tests/test_swift_megatron_sync.py @@ -74,6 +74,125 @@ def test_chunks_are_bounded_and_keep_order(): # ----------------------------------------------------------------------------- fake Megatron island +def test_heartbeat_thread_beats_between_steps_and_stops_on_a_stall(): + class Beats: + def __init__(self): + self.steps = [] + + def heartbeat(self, step): + self.steps.append(step) + + client, progress = Beats(), [time.monotonic()] + beat = sm._Heartbeat(client, 0.02, 0.3, lambda: 7, lambda: progress[0]) + deadline = time.monotonic() + 5 + while len(client.steps) < 3: + assert time.monotonic() < deadline + time.sleep(0.01) + assert set(client.steps) == {7} + time.sleep(0.5) # no progress for longer than the stall limit: silence + quiet = len(client.steps) + time.sleep(0.2) + assert len(client.steps) == quiet + progress[0] = time.monotonic() + while len(client.steps) == quiet: + assert time.monotonic() < deadline + time.sleep(0.01) + beat.stop() + assert not beat.thread.is_alive() + + +class _LinkClient: + """The client surface ``_PushSender`` and ``_ready_pushes`` use: a link + generation the test bumps, and a push result it chooses.""" + + def __init__(self, result=True): + self.connection_generation, self.result, self.pushes, self.pulls = 1, result, [], [] + + def push_fragment_parts(self, *head_and_parts): + if isinstance(self.result, BaseException): + raise self.result + self.pushes.append(head_and_parts[:-1]) + return self.result + + def drain_pulls(self): + pulls, self.pulls = self.pulls, [] + return pulls + + +@pytest.mark.parametrize("outcome", ["not-queued", "torso-lost", "socket-error"]) +def test_push_sender_treats_an_undelivered_push_as_a_drop(outcome): + """``push_fragment_parts`` returns False when the link is down (nothing + queued) and raises PartialMessageGenerationLost when the connection died + mid-message: neither reaches the syncer, so the anchor is released + instead of blocking that fragment's later pulls (or failing the island).""" + + from yeto.protocol import PartialMessageGenerationLost + + result = {"not-queued": False, + "torso-lost": PartialMessageGenerationLost(connection_generation=1, connection_epoch=1, + queued_chunks=1, operation="PUSH_FRAGMENT"), + "socket-error": ConnectionResetError("reset")}[outcome] + stats = {"pushed": 0, "push_bytes": 0, "dropped_pushes": 0} + dropped, sent = [], [] + sender = sm._PushSender(_LinkClient(result), stats, lambda *a: dropped.append(a), sent.append) + sender.submit((2, 9, 1, 7, 40, 3, 300), parts=[b"\0" * 8]) + assert sender.wait_idle(5) + sender.stop() + assert dropped == [(2, 7, 9)] and sent == [2] + assert stats["dropped_pushes"] == 1 and stats["pushed"] == 0 + sender.raise_error() # not an island failure + + +def _bare_sync(client, fragments=2): + from yeto.diloco_sync import DiLoCoSyncState + + sync = object.__new__(sm.MegatronDiLoCoSync) + sync.client, sync.sender, sync.interval = client, None, 1 + sync.state = DiLoCoSyncState.create(fragments, track_anchors=False) + sync.state.fragment_versions = [5] * fragments + sync.anchored = [True] * fragments + sync.spent, sync.spent_link, sync.pushed_from = [None] * fragments, [None] * fragments, [None] * fragments + sync.iteration, sync.units_total, sync.stats = 20, 100, {"deferred_pushes": 0} + return sync + + +def test_a_spent_anchor_is_released_once_the_island_redials(): + """A PUSH queued on a connection that later died may never have reached + the syncer, and a restarted syncer has lost the round it was carried + into; after the redial the next pull of that fragment is answered from + the same anchor (the syncer refuses it if it does hold that delta).""" + + from yeto.protocol import PullRequest + + client = _LinkClient() + sync = _bare_sync(client) + sync.spent[0] = (5, 7) # pushed from anchor 5 for round 7, on link 1 + sync._push_sent(0) + assert sync.spent_link[0] == 1 + client.pulls = [PullRequest(0, 9, 1)] # another round of the same anchor + assert sync._ready_pushes() == [] # same link: the syncer has the delta + assert sync.state.pending_pulls == [PullRequest(0, 9, 1)] + client.connection_generation = 2 # the link was replaced + assert [p[:3] for p in sync._ready_pushes()] == [(0, 9, 1)] + assert sync.spent == [None, None] and sync.stats["released_anchors"] == 1 + # A record from the process the island resumed from: released likewise. + sync.spent[1], sync.spent_link[1] = (5, 4), sm._PREVIOUS_PROCESS + client.pulls = [PullRequest(1, 10, 1)] + assert [p[:3] for p in sync._ready_pushes()] == [(1, 10, 1)] + + +def test_a_dropped_push_frees_its_anchor_on_the_same_link(): + from yeto.protocol import PullRequest + + client = _LinkClient() + sync = _bare_sync(client) + sync.spent[0] = (5, 7) + sync._push_sent(0) + sync._push_dropped(0, 5, 7) + client.pulls = [PullRequest(0, 9, 1)] + assert [p[:3] for p in sync._ready_pushes()] == [(0, 9, 1)] + + def _full(name: str, shape) -> torch.Tensor: """The initial full value of a global tensor (bf16-exact, so INIT is lossless).""" @@ -498,6 +617,375 @@ def test_initial_cut_is_installed_fragment_by_fragment_before_training(tmp_path) _join(procs, 200, tmp_path) +def _resume_worker(rank, world, port, out_dir): + _guard(out_dir, f"resume-{rank}", lambda: _resume_body(rank, world, port, out_dir)) + + +def _resume_body(rank, world, port, out_dir): + """An island checkpoints after three steps and restarts from it; the + syncer's rejoin catch-up must not undo the island's local progress.""" + + from yeto.diloco_sync import shifted_resume_value + from yeto.protocol import BcastFragment + from yeto.tensor_io import pack_tensor + + dist = _init_dist(rank, world, port) + tp, dp = rank % 2, rank // 2 + tp_groups = [dist.new_group([0, 1]), dist.new_group([2, 3])] + dp_groups = [dist.new_group([0, 2]), dist.new_group([1, 3])] + singles = [dist.new_group([r]) for r in range(world)] + ctx = sm.ParallelContext(rank, world, tp_groups[dp], tp, 2, singles[rank], 0, 1) + chunk = _chunk(tp=tp, pp=0, tp_size=2, pp_size=1) + trainable = [p for p in chunk.parameters() if p.requires_grad] + opt = _DistOpt(trainable, dp_rank=dp, dp_size=2, dp_group=dp_groups[tp]) + full_shapes = _global_shapes(2, 1) + old_g = {n: (_full(n, s) + 0.125).to(torch.bfloat16).float() for n, s in full_shapes.items()} + new_g = {n: (_full(n, s) - 0.25).to(torch.bfloat16).float() for n, s in full_shapes.items()} + cut = {} # fid -> (version, values) of the next sync's catch-up + + def factory(addrs, learner_id, layout, wire, streams, **kwargs): + bcasts = [BcastFragment(fid, v, pack_tensor(_flat(layout, fid, values), 2)) + for fid, (v, values) in sorted(cut.items())] + return _ScriptedClient(layout, {("bcast", 1): bcasts}) + + def make(resume_dir=None, resume_iteration=None): + return sm.MegatronDiLoCoSync([chunk], opt, syncer="tcp://127.0.0.1:1234", learner_id=1, num_learners=2, + fragments=3, sync_interval_steps=100, chunk_numel=40, parallel_context=ctx, + client_factory=factory, initial_cut_timeout=60, resume_dir=resume_dir, + resume_iteration=resume_iteration) + + def shard(full, p): + if getattr(p, "tensor_model_parallel", False): + size = p._full_shape[p.partition_dim] // 2 + full = full.narrow(p.partition_dim, tp * size, size) + return full + + ckpt, bare = Path(out_dir, "checkpoint-3"), Path(out_dir, "checkpoint-3-bare") + # The first run: the initial cut (version 4), three local steps, a checkpoint. + cut.update({fid: (4, old_g) for fid in range(3)}) + first = make() + for it in (1, 2, 3): + opt.step({id(p): _local_update(p, 1, it, tp, 2) for p in trainable}) + first.after_optimizer_step(it, 10) + assert first.save_state(str(ckpt), iteration=3, save_anchors=True) + assert first.save_state(str(bare), iteration=3, save_anchors=False) + fid_of = {name: fid for fid, frag in enumerate(first.layout.fragments) for name, _ in frag.tensors} + first.close() + if rank == 0: + record = json.loads((ckpt / sm.RESUME_STATE_FILE).read_text()) + assert record["versions"] == [4, 4, 4] and record["anchors"] and record["iteration"] == 3 + assert record["steps_at_reset"] == [0, 0, 0] and record["units_total"] == 30 + assert not json.loads((bare / sm.RESUME_STATE_FILE).read_text())["anchors"] + saved_master = {id(p): opt.master(p).clone() for p in trainable if id(p) in opt.owned} + saved_param = {id(p): p.detach().clone() for p in trainable} + + def restart(): # the trainer loads the checkpoint again + for p in trainable: + if id(p) in opt.owned: + opt.master(p).copy_(saved_master[id(p)]) + p.data.copy_(saved_param[id(p)]) + + def check(sync, expect): # expect(p) -> the full local FP32 value or None (= the checkpoint, bit for bit) + assert sync.anchored == [True] * 3 + for p in trainable: + want = expect(p) + if want is None: + if id(p) in opt.owned: + assert torch.equal(opt.master(p), saved_master[id(p)]), p._full_name + assert torch.equal(p.detach(), saved_param[id(p)]), p._full_name + continue + if id(p) in opt.owned: + s, e, _ = opt.owned[id(p)] + assert torch.equal(opt.master(p), want.reshape(-1)[s:e]), p._full_name + assert torch.equal(p.detach().reshape(-1)[s:e], want.reshape(-1)[s:e].to(torch.bfloat16)) + + def check_anchors(sync, values): + for p in trainable: + assert torch.equal(sync.anchors[id(p)].float(), shard(values[p._full_name], p)), p._full_name + + # (a) The syncer is still at the checkpoint's anchor version: anchor only. + restart() + cut.update({fid: (4, old_g) for fid in range(3)}) + same = make(str(ckpt), 3) + check(same, lambda p: None) + check_anchors(same, old_g) + assert same.resumed == {0: "anchor", 1: "anchor", 2: "anchor"} + same.after_optimizer_step(4, 10) + assert same.state.steps_at_reset == [0, 0, 0], "the H gate restarted with the island" + assert same.units_total == 40 + same.close() + + # (b) The syncer moved on: global_new + (local - anchor_old), exactly. + restart() + cut.update({fid: (9, new_g) for fid in range(3)}) + moved = make(str(ckpt), 3) + for p in trainable: + if id(p) in opt.owned: + s, e, _ = opt.owned[id(p)] + want = shifted_resume_value(shard(new_g[p._full_name], p).reshape(-1)[s:e], saved_master[id(p)], + shard(old_g[p._full_name], p).reshape(-1)[s:e]) + assert torch.equal(opt.master(p), want), p._full_name + assert torch.equal(p.detach().reshape(-1)[s:e], want.to(torch.bfloat16)), p._full_name + assert not torch.equal(want, shard(new_g[p._full_name], p).reshape(-1)[s:e]) + check_anchors(moved, new_g) + assert set(moved.resumed.values()) == {"shift"} and not moved._resume_anchors + moved.after_optimizer_step(4, 10) + assert moved.state.steps_at_reset == [0, 0, 0] and moved.state.fragment_versions == [9, 9, 9] + moved.close() + + # (c) Without saved anchors a moved fragment takes the global; an unmoved one is kept + # (a record without the contract, as written by hand for an older checkpoint, is accepted). + if rank == 0: + record = json.loads((bare / sm.RESUME_STATE_FILE).read_text()) + del record["contract"] + (bare / sm.RESUME_STATE_FILE).write_text(json.dumps(record)) + restart() + cut.update({0: (4, old_g), 1: (9, new_g), 2: (9, new_g)}) + mixed = make(str(bare), 3) + check(mixed, lambda p: None if fid_of[p._full_name] == 0 else shard(new_g[p._full_name], p)) + assert mixed.resumed == {0: "anchor", 1: "overwrite", 2: "overwrite"} + mixed.after_optimizer_step(4, 10) + assert mixed.state.steps_at_reset == [0, 3, 3] + mixed.close() + + # (d) A cold start (no record, or a record of another iteration) takes the global. + for resume in ((None, None), (str(ckpt), 7)): + restart() + cut.update({fid: (9, new_g) for fid in range(3)}) + cold = make(*resume) + check(cold, lambda p: shard(new_g[p._full_name], p)) + assert cold.resumed == {} and cold.resume_record is None + cold.close() + dist.destroy_process_group() + + +@pytest.mark.timeout(240) +def test_resumed_island_keeps_its_checkpoint_through_the_rejoin_catch_up(tmp_path): + procs = _run(_resume_worker, 4, free_port(), str(tmp_path)) + _join(procs, 200, tmp_path) + + +def _catch_up_worker(rank, world, port, out_dir): + _guard(out_dir, f"catch-up-{rank}", lambda: _catch_up_body(rank, world, port, out_dir)) + + +def _catch_up_body(rank, world, port, out_dir): + """A running island redials and gets the syncer's catch-up: a replay of + its anchor changes nothing (bit for bit), a rewound or newer catch-up of + a fragment it did not push from keeps its local progress (global_new + + (local - anchor)), a newer one of a fragment it pushed from its anchor is + installed as an ordinary merge.""" + + from yeto.diloco_sync import shifted_resume_value + from yeto.protocol import BcastFragment, PullRequest + from yeto.tensor_io import pack_tensor + + dist = _init_dist(rank, world, port) + tp, dp = rank % 2, rank // 2 + tp_groups = [dist.new_group([0, 1]), dist.new_group([2, 3])] + dp_groups = [dist.new_group([0, 2]), dist.new_group([1, 3])] + singles = [dist.new_group([r]) for r in range(world)] + ctx = sm.ParallelContext(rank, world, tp_groups[dp], tp, 2, singles[rank], 0, 1) + chunk = _chunk(tp=tp, pp=0, tp_size=2, pp_size=1) + trainable = [p for p in chunk.parameters() if p.requires_grad] + opt = _DistOpt(trainable, dp_rank=dp, dp_size=2, dp_group=dp_groups[tp]) + full_shapes = _global_shapes(2, 1) + old_g = {n: (_full(n, s) + 0.125).to(torch.bfloat16).float() for n, s in full_shapes.items()} + new_g = {n: (_full(n, s) - 0.25).to(torch.bfloat16).float() for n, s in full_shapes.items()} + + def factory(addrs, learner_id, layout, wire, streams, **kwargs): + assert kwargs["accept_rewound_catch_up"] is True + client = _ScriptedClient(layout, script) + script[("layout",)] = layout + return client + + def bcast(fid, version, values, catch_up): + layout = script[("layout",)] + return BcastFragment(fid, version, pack_tensor(_flat(layout, fid, values), 2), catch_up=catch_up) + + plan = {1: [(fid, 4, old_g, False) for fid in range(3)], # the initial cut + 3: "pull", # iteration 2: a pull of fragment 2 + 5: [(0, 4, old_g, False), (1, 2, new_g, True), (2, 9, new_g, True)], + 6: [(0, 9, new_g, True)]} + + class _Script(dict): # call -> scripted pulls / broadcasts (packed once the layout is known) + def get(self, key, default=None): + kind, call = key + entry = plan.get(call) + if kind == "pull": + return [PullRequest(2, 7, 1)] if entry == "pull" else [] + if isinstance(entry, list): + return [bcast(*row) for row in entry] + return [] + + script = _Script() + sync = sm.MegatronDiLoCoSync([chunk], opt, syncer="tcp://127.0.0.1:1234", learner_id=1, num_learners=2, + fragments=3, sync_interval_steps=1, chunk_numel=40, parallel_context=ctx, + client_factory=factory, initial_cut_timeout=60) + fid_of = {name: fid for fid, frag in enumerate(sync.layout.fragments) for name, _ in frag.tensors} + for it in (1, 2, 3): + opt.step({id(p): _local_update(p, 1, it, tp, 2) for p in trainable}) + sync.after_optimizer_step(it, 10) + if rank == 0: + assert sync.sender.wait_idle(10) and sync.stats["pushed"] == 1 and sync.pushed_from == [None, None, 4] + master = {id(p): opt.master(p).clone() for p in trainable if id(p) in opt.owned} + param = {id(p): p.detach().clone() for p in trainable} + + def shard(full, p): + if getattr(p, "tensor_model_parallel", False): + size = p._full_shape[p.partition_dim] // 2 + full = full.narrow(p.partition_dim, tp * size, size) + return full + + def owned(p, full): + s, e, _ = opt.owned[id(p)] + return shard(full, p).reshape(-1)[s:e] + + sync.after_optimizer_step(4, 10) # call 5: no local step since the snapshot + for p in trainable: + fid = fid_of[p._full_name] + if fid == 0: # a replay of its anchor: untouched + assert torch.equal(p.detach(), param[id(p)]), p._full_name + if id(p) in opt.owned: + assert torch.equal(opt.master(p), master[id(p)]) + assert torch.equal(sync.anchors[id(p)].float(), shard(old_g[p._full_name], p)) + continue + assert torch.equal(sync.anchors[id(p)].float(), shard(new_g[p._full_name], p)), p._full_name + if id(p) not in opt.owned: + continue + if fid == 1: # rewound to version 2: rebased exactly + want = shifted_resume_value(owned(p, new_g[p._full_name]), master[id(p)], owned(p, old_g[p._full_name])) + assert not torch.equal(want, owned(p, new_g[p._full_name])) + else: # newer, pushed from its anchor: the ordinary install + want = owned(p, new_g[p._full_name]) + assert torch.equal(opt.master(p), want), (fid, p._full_name) + assert torch.equal(p.detach().reshape(-1)[opt.owned[id(p)][0]:opt.owned[id(p)][1]], want.to(torch.bfloat16)) + assert sync.state.fragment_versions == [4, 2, 9] + assert sync.state.steps_at_reset[1] == 0 and sync.state.steps_at_reset[2] == 4, sync.state.steps_at_reset + before = {id(p): opt.master(p).clone() for p in trainable if id(p) in opt.owned and fid_of[p._full_name] == 0} + sync.after_optimizer_step(5, 10) # call 6: a newer catch-up of fragment 0, never pushed: rebased + for p in trainable: + if fid_of[p._full_name] == 0 and id(p) in opt.owned: + want = shifted_resume_value(owned(p, new_g[p._full_name]), before[id(p)], owned(p, old_g[p._full_name])) + assert torch.equal(opt.master(p), want), p._full_name + if rank == 0: + assert sync.stats["confirmed"] == 1 and sync.stats["rebased"] == 2, sync.stats + sync.close() + dist.destroy_process_group() + + +@pytest.mark.timeout(240) +def test_running_island_keeps_its_progress_through_a_catch_up(tmp_path): + procs = _run(_catch_up_worker, 4, free_port(), str(tmp_path)) + _join(procs, 200, tmp_path) + + +def _resumed_init_worker(rank, world, port, out_dir): + _guard(out_dir, f"resumed-init-{rank}", lambda: _resumed_init_body(rank, world, port, out_dir)) + + +def _resumed_init_body(rank, world, port, out_dir): + """Learner 0 resumed from a checkpoint with a record: an initialized + syncer answers its HELLO with the catch-up, so no INIT_PARAMS is sent; + a syncer that sends nothing within the probe gets INIT_PARAMS, streamed + chunk by chunk, byte-identical to a cold start's.""" + + from yeto.protocol import BcastFragment + from yeto.tensor_io import pack_tensor + + dist = _init_dist(rank, world, port) + tp, dp = rank % 2, rank // 2 + tp_groups = [dist.new_group([0, 1]), dist.new_group([2, 3])] + dp_groups = [dist.new_group([0, 2]), dist.new_group([1, 3])] + singles = [dist.new_group([r]) for r in range(world)] + ctx = sm.ParallelContext(rank, world, tp_groups[dp], tp, 2, singles[rank], 0, 1) + chunk = _chunk(tp=tp, pp=0, tp_size=2, pp_size=1) + trainable = [p for p in chunk.parameters() if p.requires_grad] + opt = _DistOpt(trainable, dp_rank=dp, dp_size=2, dp_group=dp_groups[tp]) + full = {n: (_full(n, s) + 0.125).to(torch.bfloat16).float() for n, s in _global_shapes(2, 1).items()} + clients, cut = [], {"calls": ()} + + def factory(addrs, learner_id, layout, wire, streams, **kwargs): + script = {("bcast", c): [BcastFragment(fid, 4, pack_tensor(_flat(layout, fid, full), 2)) for fid in range(3)] + for c in cut["calls"]} + clients.append(_ScriptedClient(layout, script)) + return clients[-1] + + def make(resume_dir=None, **kw): + return sm.MegatronDiLoCoSync([chunk], opt, syncer="tcp://127.0.0.1:1234", learner_id=0, num_learners=2, + fragments=3, sync_interval_steps=100, chunk_numel=40, parallel_context=ctx, + client_factory=factory, resume_dir=resume_dir, resume_iteration=3, **kw) + + cut["calls"] = (1,) + first = make(initial_cut_timeout=60) # a cold start: INIT, then the initial cut + ckpt = Path(out_dir, "checkpoint-3") + first.save_state(str(ckpt), iteration=3) + first.close() + cold_init = dict(clients[-1].inits) if rank == 0 else {} + if rank == 0: + assert sorted(cold_init) == [0, 1, 2] + resumed = make(str(ckpt), initial_cut_timeout=60, init_probe_s=30) # the catch-up arrives: no INIT + assert resumed.anchored == [True] * 3 and resumed.resumed == {0: "anchor", 1: "anchor", 2: "anchor"} + if rank == 0: + assert clients[-1].inits == {} + resumed.close() + cut["calls"] = () + started = time.monotonic() + fresh = make(str(ckpt), initial_cut_timeout=1, init_probe_s=0.5) # nothing arrives: INIT after the probe + assert time.monotonic() - started >= 0.5 + if rank == 0: # the island's current masters (the installed cut), streamed by chunk + layout = clients[-1].layout + assert clients[-1].inits == {fid: bytes(pack_tensor(_flat(layout, fid, full), fresh.bulk)) for fid in range(3)} + fresh.close() + dist.destroy_process_group() + + +@pytest.mark.timeout(240) +def test_a_resumed_learner_zero_sends_init_only_when_the_syncer_needs_it(tmp_path): + procs = _run(_resumed_init_worker, 4, free_port(), str(tmp_path)) + _join(procs, 200, tmp_path) + + +def _init_abort_worker(rank, world, port, out_dir): + _guard(out_dir, f"init-abort-{rank}", lambda: _init_abort_body(rank, world, port)) + + +def _init_abort_body(rank, world, port): + dist = _init_dist(rank, world, port) + single = dist.new_group([0]) + ctx = sm.ParallelContext(0, 1, single, 0, 1, single, 0, 1) + chunk = _chunk(tp=0, pp=0, tp_size=1, pp_size=1) + opt = _DistOpt([p for p in chunk.parameters() if p.requires_grad], dp_rank=0, dp_size=1, dp_group=single) + holder = {} + + def factory(addrs, learner_id, layout, wire, streams, **kwargs): + holder["client"] = _ScriptedClient(layout, {}) + return holder["client"] + + # learner 1: no INIT from the constructor; the test sends it by hand + sync = sm.MegatronDiLoCoSync([chunk], opt, syncer="127.0.0.1:1", learner_id=1, num_learners=2, + fragments=2, sync_interval_steps=1, chunk_numel=16, + parallel_context=ctx, client_factory=factory) + original = sync._export_fragment_chunks + + def failing(fid, delta): + for i, parts in enumerate(original(fid, delta=delta)): + if i == 1: + raise RuntimeError("collective timed out") + yield parts + + sync._export_fragment_chunks = failing + started = time.monotonic() + with pytest.raises(RuntimeError, match="collective timed out"): + sync._send_init() + assert time.monotonic() - started < 30 + + +def test_an_init_export_that_raises_fails_instead_of_hanging(tmp_path): + procs = _run(_init_abort_worker, 1, free_port(), str(tmp_path)) + _join(procs, 120, tmp_path) + + # ----------------------------------------------------------------------------- two islands, real syncer (TP2 x PP2) def build_syncer() -> Path: binary = ROOT / "syncer/target/debug/yeto-syncer" @@ -538,7 +1026,8 @@ def _island_body(rank, world, island, dist_port, syncer_port, out_dir, drop_at, syncer = sm.MegatronDiLoCoSync([chunk], opt, syncer=f"127.0.0.1:{syncer_port}", learner_id=island, num_learners=2, fragments=3, sync_interval_steps=2, wire_dtype=opts.get("wire", "bf16"), - chunk_numel=64, parallel_context=ctx, final_cut_payload="full") + chunk_numel=64, parallel_context=ctx, final_cut_payload="full", + heartbeat_s=opts.get("heartbeat_s", 30.0)) it = 0 dropped = False started = time.monotonic() @@ -571,7 +1060,7 @@ def _island_body(rank, world, island, dist_port, syncer_port, out_dir, drop_at, if rank == 0: torch.save({"names": names, "values": dict(zip(names, values, strict=True)), "events": list(syncer.events), "stats": dict(syncer.stats), "dropped": dropped, - "timeline": timeline, + "timeline": timeline, "lease_expiries": syncer.client.lease_expiries, "fragments": [[n for n, _ in f.tensors] for f in syncer.layout.fragments]}, Path(out_dir, f"island{island}.pt")) syncer.close() @@ -711,3 +1200,127 @@ def test_decoupled_quorum_one_keeps_the_fast_island_merging_while_the_other_stal # A late-carried delta counts towards its round's quorum, so the fast # island's rounds keep closing while the other island stalls. assert merged_during >= 3, merged_during + + +@pytest.mark.timeout(300) +@pytest.mark.parametrize("stop", ["sigterm", "sigkill"]) +def test_islands_keep_merging_across_a_syncer_restart(tmp_path, stop): + """The syncer is stopped while both islands keep training and answering + the pulls they already hold, then comes back with --resume on the same + port. ``sigterm``: it writes its checkpoint and exits. During the + outage a PUSH is not queued at all (``push_fragment_parts`` returns + False) and is a drop; a PUSH that went out on the old connection holds an + anchor whose round the restarted syncer no longer has, and is released + once the island redials. ``sigkill``: a crash; the syncer resumes from + its last periodic checkpoint, older than what the islands hold, and its + catch-up rewinds them (rebased onto it) instead of being dropped while + every push from the lost versions is refused as a future base. Both + islands must reach the restarted syncer's final cut.""" + + import signal + + binary = build_syncer() + from process_groups import group_popen + + port = free_port() + final_state = tmp_path / "final.bin" + ckpt = tmp_path / "syncer.ckpt" + flags = [str(binary), "--port", str(port), "--learners", "2", "--grace-ms", "200", + "--terminal-grace-ms", "2000", "--total-steps", "30", "--final-state", str(final_state), + "--checkpoint-path", str(ckpt), "--checkpoint-every", "4", "--shutdown-timeout-s", "5", + *DECOUPLED_SYNCER_FLAGS] + logs: list[str] = [] + + def start(extra): + proc = group_popen(flags + extra, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) + logs.append("") + return proc + + def read(i, proc): # the log file written so far + logs[i] += proc.stdout.read() + return logs[i] + + first = start([]) + second = None + procs = [] + try: + procs += _run(_island_worker, 4, 0, free_port(), port, str(tmp_path), None, {"step_s": 0.05}) + procs += _run(_island_worker, 4, 1, free_port(), port, str(tmp_path), None, {"step_s": 0.05}) + deadline = time.monotonic() + 120 + # sigkill: past a periodic checkpoint (every 4 rounds), not on one. + while read(0, first).count("outer step") < (8 if stop == "sigterm" else 10): + assert time.monotonic() < deadline, "the first syncer never merged enough rounds" + assert first.poll() is None, "the first syncer exited early" + time.sleep(0.05) + if stop == "sigterm": + first.send_signal(signal.SIGTERM) + assert first.wait(timeout=60) == 0, "SIGTERM did not end in a clean checkpoint" + else: + first.kill() + first.wait(timeout=60) + assert ckpt.exists() + time.sleep(3.0) # the outage: both islands keep stepping + second = start(["--resume"]) + _join(procs, 240, tmp_path) + assert second.wait(timeout=60) == 0, "the restarted syncer exited nonzero" + finally: + for p in procs: + if p.is_alive(): + p.kill() + for i, proc in enumerate((first, second)): + if proc is not None: + if proc.poll() is None: + proc.kill() + read(i, proc) + (tmp_path / f"syncer{i}.log").write_text(logs[i]) + print(logs[-1][-4000:]) + + islands = [torch.load(tmp_path / f"island{i}.pt") for i in range(2)] + for name in islands[0]["names"]: + assert torch.equal(islands[0]["values"][name], islands[1]["values"][name]), name + final = read_final_state(final_state) + for fid, frag_names in enumerate(islands[0]["fragments"]): + assert torch.equal(torch.cat([islands[0]["values"][n].reshape(-1) for n in frag_names]), final[fid]) + assert "outer step" in logs[1], "the restarted syncer merged nothing" + if stop == "sigkill": + assert sum(i["stats"].get("rebased", 0) for i in islands) > 0, "no island was rewound" + print(json.dumps([i["stats"] for i in islands])) + + +@pytest.mark.timeout(300) +@pytest.mark.parametrize("heartbeat_s", [0.3, 0.0]) +def test_heartbeats_keep_a_stalled_island_inside_its_lease(tmp_path, heartbeat_s): + """--island-lease-s 2 with an island that makes no boundary for 4 s: + rank 0's heartbeat thread keeps its lease; without heartbeats the syncer + expires it (the island redials as a re-JOIN).""" + + binary = build_syncer() + from process_groups import group_popen + + port = free_port() + final_state = tmp_path / "final.bin" + proc = group_popen([str(binary), "--port", str(port), "--learners", "2", "--grace-ms", "200", + "--terminal-grace-ms", "2000", "--total-steps", "30", "--final-state", str(final_state), + "--island-lease-s", "2", *DECOUPLED_SYNCER_FLAGS], + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) + procs = [] + try: + procs += _run(_island_worker, 4, 0, free_port(), port, str(tmp_path), None, + {"step_s": 0.05, "heartbeat_s": heartbeat_s}) + procs += _run(_island_worker, 4, 1, free_port(), port, str(tmp_path), None, + {"step_s": 0.05, "stall_at": 6, "stall_s": 4.0, "heartbeat_s": heartbeat_s}) + _join(procs, 240, tmp_path) + assert proc.wait(timeout=60) == 0, "syncer exited nonzero" + finally: + for p in procs: + if p.is_alive(): + p.kill() + if proc.poll() is None: + proc.kill() + log = proc.stdout.read() if proc.stdout else "" + (tmp_path / "syncer.log").write_text(log) + islands = [torch.load(tmp_path / f"island{i}.pt") for i in range(2)] + for name in islands[0]["names"]: + assert torch.equal(islands[0]["values"][name], islands[1]["values"][name]), name + expired = '"reason":"lease_expired"' in log or islands[1]["lease_expiries"] > 0 + assert expired == (heartbeat_s == 0), (islands[1]["lease_expiries"], heartbeat_s) diff --git a/yeto/diloco_sync.py b/yeto/diloco_sync.py index 7afbb042..035801b2 100644 --- a/yeto/diloco_sync.py +++ b/yeto/diloco_sync.py @@ -148,6 +148,94 @@ def create( anchors=[None] * num_fragments if track_anchors else None, ) + def resume_record(self, *, anchored: list[bool], units_total: int, + spent: list[tuple[int, int] | None] | None = None) -> dict: + """The per-fragment part of an island checkpoint (JSON-safe): which + version each fragment's anchor is, and the H-gate / token counters + measured from it. The anchor values themselves are not included.""" + n = len(self.fragment_versions) + spent = spent if spent is not None else [None] * n + return { + "versions": [int(v) for v in self.fragment_versions], + "anchored": [bool(a) for a in anchored], + "steps_at_reset": [int(s) for s in self.steps_at_reset], + "units_at_reset": [int(u) for u in self.units_at_reset], + "spent": [None if s is None else [int(s[0]), int(s[1])] for s in spent], + "units_total": int(units_total), + "global_step": int(self.global_step), + } + + def restore_fragment(self, record: dict, fid: int) -> None: + """Take fragment ``fid``'s counters from a ``resume_record`` (a resumed + island that kept its local progress keeps its H-gate clock too).""" + self.steps_at_reset[fid] = int(record["steps_at_reset"][fid]) + self.units_at_reset[fid] = int(record["units_at_reset"][fid]) + + +# What a resumed island does with the first broadcast of a fragment. +RESUME_OVERWRITE = "overwrite" # install the global (no usable record: a cold start) +RESUME_ANCHOR = "anchor" # same version as the checkpoint's anchor: anchor only +RESUME_SHIFT = "shift" # newer global: global_new + (local - anchor_old) + + +def resume_fragment_mode(record: dict | None, fid: int, version: int, *, have_anchor: bool) -> str: + """How a resumed island installs the syncer's first broadcast of ``fid``. + + The checkpoint holds ``local``, trained from the anchor of + ``record["versions"][fid]``. The same version again means the broadcast + *is* that anchor, so the local parameters stay as they are. Any other + version moves the island's local progress onto the new global, which + needs the old anchor value (``have_anchor``); without it the fragment + falls back to the cold-start overwrite. As in ``live_install_mode``, a + newer global of a fragment the island had pushed from that anchor + (``record["spent"]``) is installed as is: the pushed delta may already be + merged into it, and shifting would count it twice.""" + if record is None or not record["anchored"][fid]: + return RESUME_OVERWRITE + anchor = int(record["versions"][fid]) + if int(version) == anchor: + return RESUME_ANCHOR + spent = (record.get("spent") or [None] * len(record["versions"]))[fid] + if int(version) > anchor and spent is not None and int(spent[0]) == anchor: + return RESUME_OVERWRITE + return RESUME_SHIFT if have_anchor else RESUME_OVERWRITE + + +# What a running island does with a broadcast of a fragment it already +# holds an anchor of (``live_install_mode``); None = the ordinary install. +LIVE_CONFIRM = "confirm" # a replay of its anchor: nothing to install +LIVE_REBASE = "rebase" # a catch-up: global_new + (local - anchor) + + +def live_install_mode(*, anchored: bool, current_version: int, version: int, catch_up: bool, + pushed_from: int | None) -> str | None: + """How a running island installs a broadcast of a fragment. + + A broadcast of the version the island is anchored at is a replay of + that anchor (``LIVE_CONFIRM``): the local parameters stay bit for bit. + The first broadcast after a redial (``catch_up``) is the syncer's + current value, not the merge of a round the island answered: the + island's progress since its anchor is moved onto it (``LIVE_REBASE``), + also when the version went backwards (a syncer resumed from an older + checkpoint) or kept its number with other values. The exception is a + newer catch-up of a fragment the island pushed from its current anchor + (``pushed_from``): that delta may already be merged into it, so it is + installed like any merge broadcast. None: the ordinary install.""" + if not anchored: + return None + if not catch_up: + return LIVE_CONFIRM if int(version) == int(current_version) else None + if int(version) > int(current_version) and pushed_from is not None \ + and int(pushed_from) == int(current_version): + return None + return LIVE_REBASE + + +def shifted_resume_value(global_new: torch.Tensor, local: torch.Tensor, anchor_old: torch.Tensor) -> torch.Tensor: + """``global_new + (local - anchor_old)`` in FP32: the island's progress since + its checkpointed anchor, moved onto the newer global.""" + return global_new.to(torch.float32) + (local.to(torch.float32) - anchor_old.to(torch.float32)) + def sync_diloco_boundary( client: SyncerClient | None, diff --git a/yeto/launch/syncer.py b/yeto/launch/syncer.py index 99a2b344..77bbce57 100644 --- a/yeto/launch/syncer.py +++ b/yeto/launch/syncer.py @@ -880,6 +880,9 @@ def _sft_syncer_command(args, num_learners: int, binary: str, f" --total-steps {args.total_steps}" f"{_syncer_outer_optimizer(args)}" f" --checkpoint-path {checkpoint}" + # The SFT preset writes its periodic checkpoints (every 8 rounds by + # default) off the scheduler, and one before the first round. + f" --background-checkpoint --checkpoint-on-start" f" {_resume_if_exists(checkpoint)}" f" --mark-final-checkpoint" f" --event-tape {_shard_path(SYNCER_EVENT_TAPE, index)}" diff --git a/yeto/protocol.py b/yeto/protocol.py index 87d17989..e79bce91 100644 --- a/yeto/protocol.py +++ b/yeto/protocol.py @@ -575,7 +575,7 @@ def write_frame(sock: socket.socket, msg_type: int, payload: bytes) -> None: sock.sendall(_HEADER.pack(MAGIC, msg_type, len(payload)) + payload) -def _read_exact(sock: socket.socket, n: int) -> bytes: +def _read_exact(sock: socket.socket, n: int, *, copy: bool = True) -> bytes | bytearray: buf = bytearray(n) view = memoryview(buf) got = 0 @@ -584,14 +584,18 @@ def _read_exact(sock: socket.socket, n: int) -> bytes: if r == 0: raise ConnectionError("syncer closed connection") got += r - return bytes(buf) + return bytes(buf) if copy else buf def read_frame( sock: socket.socket, max_payload: int | Callable[[int], int] = MAX_HANDSHAKE_FRAME, + *, + raw_types: tuple[int, ...] = (), ) -> tuple[int, bytes]: - """Read one frame after checking its type-aware limit before allocation.""" + """Read one frame after checking its type-aware limit before allocation. + The payload of a type in ``raw_types`` is the receive buffer itself (a + bytearray, one copy less), every other payload is bytes.""" magic, msg_type, length = _HEADER.unpack(_read_exact(sock, _HEADER.size)) if magic != MAGIC: raise ValueError(f"bad magic {magic:#x}") @@ -600,7 +604,7 @@ def read_frame( raise ValueError( f"frame type {msg_type} has length {length}, exceeds limit {limit}" ) - return msg_type, _read_exact(sock, length) + return msg_type, _read_exact(sock, length, copy=msg_type not in raw_types) HELLO_MAC_DOMAIN = b"yeto-hello-mac-v1\0" @@ -1010,6 +1014,10 @@ class BcastFragment: payload_hash: bytes | None = field(default=None, compare=False) stored: bool = field(default=False, compare=False) discard: Callable[[], None] | None = field(default=None, repr=False, compare=False) + # The first broadcast of this fragment after a redial: the syncer's + # current value sent as the re-HELLO catch-up (or a newer one), not + # necessarily a merge of a round this island answered. + catch_up: bool = field(default=False, compare=False) @dataclass @@ -1048,6 +1056,113 @@ class _StagedInboundChunk: token: object | None +# Digest of a BCAST_FRAGMENT payload on the in-memory receive path: SHA-256 +# over the SHA-256 of every BCAST_DIGEST_BLOCK bytes, so the blocks of a +# large payload are hashed in parallel (hashlib releases the GIL) while its +# chunks still arrive. It only ever meets digests of this client's own +# receive path (``_bcast_seen``): a replay of the version held is one with +# the same digest, another payload under that version a conflict, as with a +# flat SHA-256. +BCAST_DIGEST_BLOCK = 32 << 20 +BCAST_DIGEST_THREADS = 4 +_BCAST_DIGEST_DOMAIN = b"yeto-bcast-blocks-sha256-v1\0" + + +def _block_digest_join(nbytes: int, blocks: Sequence[bytes]) -> bytes: + return hashlib.sha256(_BCAST_DIGEST_DOMAIN + struct.pack(" bytes: + """``_BlockDigest``'s digest of ``data``, computed in one call.""" + view = memoryview(data).cast("B") + blocks = [hashlib.sha256(view[i:i + BCAST_DIGEST_BLOCK]).digest() + for i in range(0, view.nbytes, BCAST_DIGEST_BLOCK)] + return _block_digest_join(view.nbytes, blocks) + + +class _BlockDigest: + """``bcast_payload_digest(buf[start:])``, computed while a chunked + message arrives: each block is hashed on the client's digest pool as + soon as all its bytes are in.""" + + def __init__(self, buf: bytearray, start: int, submit: Callable[[Callable[[], None]], None]) -> None: + self.view = memoryview(buf)[start:] + self.start = start + self.nbytes = self.view.nbytes + count = -(-self.nbytes // BCAST_DIGEST_BLOCK) + self.missing = [min(BCAST_DIGEST_BLOCK, self.nbytes - i * BCAST_DIGEST_BLOCK) for i in range(count)] + self.blocks: list[bytes | None] = [None] * count + self.submit = submit + self.lock = threading.Lock() + self.done = threading.Condition(self.lock) + self.pending = 0 # submitted, not yet hashed + self.cancelled = False + + def arrived(self, offset: int, end: int) -> None: + """``buf[offset:end]`` is in (called once per chunk, under the + reassembly lock): queue every block it completes.""" + lo, hi = max(offset - self.start, 0), max(end - self.start, 0) + ready = [] + with self.lock: + for block in range(lo // BCAST_DIGEST_BLOCK, -(-hi // BCAST_DIGEST_BLOCK)): + first = block * BCAST_DIGEST_BLOCK + self.missing[block] -= min(hi, first + BCAST_DIGEST_BLOCK) - max(lo, first) + if self.missing[block] == 0: + ready.append(block) + self.pending += len(ready) + for block in ready: + self.submit(lambda block=block: self._hash(block)) + + def _hash(self, block: int) -> None: + try: + if not self.cancelled: + at = block * BCAST_DIGEST_BLOCK + self.blocks[block] = hashlib.sha256(self.view[at:at + BCAST_DIGEST_BLOCK]).digest() + finally: + with self.lock: + self.pending -= 1 + self.done.notify_all() + + def finish(self) -> bytes: + """Every byte arrived: wait for the last blocks, return the digest.""" + with self.lock: + while self.pending: + self.done.wait() + for block, digest in enumerate(self.blocks): # a block the pool never got (e.g. closed) + if digest is None: + at = block * BCAST_DIGEST_BLOCK + self.blocks[block] = hashlib.sha256(self.view[at:at + BCAST_DIGEST_BLOCK]).digest() + self.view.release() + return _block_digest_join(self.nbytes, self.blocks) # type: ignore[arg-type] + + def cancel(self) -> None: + self.cancelled = True + + +class _Reassembly: + """One chunked inbound message being put together in memory.""" + + __slots__ = ("buf", "filled", "ranges", "digest") + + def __init__(self, total: int, digest_start: int | None, submit=None) -> None: + self.buf = bytearray(total) + self.filled = 0 + self.ranges: list[tuple[int, int]] = [] # sorted, disjoint + self.digest = None if digest_start is None else _BlockDigest(self.buf, digest_start, submit) + + def overlaps(self, offset: int, end: int) -> bool: + import bisect + + i = bisect.bisect_left(self.ranges, (offset, end)) + return (i > 0 and self.ranges[i - 1][1] > offset) or (i < len(self.ranges) and self.ranges[i][0] < end) + + def add(self, offset: int, end: int) -> None: + import bisect + + bisect.insort(self.ranges, (offset, end)) + self.filled += end - offset + + @dataclass class _StreamedReassembly: total: int @@ -1172,8 +1287,18 @@ def __init__( recover: bool = False, owned_fragments: range | None = None, expert_tokens: bool = False, + accept_rewound_catch_up: bool = False, ): - """``syncer_epoch``: the syncer incarnation this island was launched + """``accept_rewound_catch_up``: after a redial, take the first + broadcast of a fragment even when its version is older than (or equal + to, with other bytes) one already received: a syncer that crashed and + resumed from an older checkpoint has rewound, and its catch-up is the + only global it will merge pushes against (the island rebases onto it; + otherwise every push from the lost version is refused as a future + base). Off: such a broadcast is dropped (older) or a protocol error + (same version), the historical behaviour. + + ``syncer_epoch``: the syncer incarnation this island was launched for; a syncer with a newer epoch refuses it (not retried). ``recover``: this island lost its trainer state; every HELLO asks for peer recovery until the syncer sends RECOVERY_DONE (see ``yeto.rl.peer_recovery``). @@ -1301,10 +1426,16 @@ def __init__( self._threads: list[threading.Thread] = [] self._pulls: queue.Queue[PullRequest] = queue.Queue() self._bcasts: queue.Queue[BcastFragment] = queue.Queue() - self._bcast_seen: list[tuple[int, bytes] | None] = [ + # (version, digest, digest kind) of each fragment's newest BCAST: the + # in-memory path's block digest or a streamed sink's payload hash. + self._bcast_seen: list[tuple[int, bytes, str] | None] = [ None ] * layout.num_fragments self._bcast_lock = threading.Lock() + # Set for every fragment on each redial; cleared by its first BCAST. + self._catch_up_pending = [False] * layout.num_fragments + self.accept_rewound_catch_up = bool(accept_rewound_catch_up) + self.rewound_catch_ups = 0 # Raw globals are retained independently of the drain queue so a # final fragment that outruns its control-stream manifest can be # re-applied exactly after the manifest arrives. @@ -1314,11 +1445,9 @@ def __init__( self.finalizing = threading.Event() self.finalized = threading.Event() self._final_cond = threading.Condition() - self._reasm: dict[ - int, - tuple[bytearray, list[int], list[tuple[int, int]]] - | _StreamedReassembly, - ] = {} + self._reasm: dict[int, _Reassembly | _StreamedReassembly] = {} + # Hashes the blocks of large inbound BCAST payloads (_BlockDigest). + self._digest_pool: Any = None self._reasm_lock = threading.Lock() self._reasm_touched: dict[int, float] = {} # msg_id -> last chunk (monotonic) self._inbound_chunk_sink: object | None = None @@ -1438,6 +1567,9 @@ def _connect_group(self, patient: bool) -> None: _close_socket(s) raise with self._lock: + if self.connection_generation != 0: # a redial: the syncer catches us up + with self._bcast_lock: + self._catch_up_pending = [True] * self.layout.num_fragments self._gen += 1 gen = self._gen self.connection_generation = connection_generation @@ -1538,6 +1670,10 @@ def _supervise(self) -> None: def close(self) -> None: self._closed.set() self._failure.set() # unblock the supervisor so it exits instead of redialing + with self._reasm_lock: + pool, self._digest_pool = self._digest_pool, None + if pool is not None: + pool.shutdown(wait=False) # queued block hashes still run: a finish() may wait on them with self._lock: queues, socks = list(self._queues), list(self._socks) for q in queues: @@ -2497,7 +2633,7 @@ def _recv_loop(self, gen: int, sock: socket.socket) -> None: sock = _RateLimitedSocket(sock, test_rate) try: while True: - msg_type, payload = read_frame(sock, self._incoming_frame_limit) + msg_type, payload = read_frame(sock, self._incoming_frame_limit, raw_types=(MSG_CHUNK,)) if msg_type == MSG_CHUNK: inner = self._reassemble(gen, payload) if inner is not None: @@ -2505,6 +2641,7 @@ def _recv_loop(self, gen: int, sock: socket.socket) -> None: self._dispatch_streamed(gen, inner) else: self._dispatch(gen, *inner) + del inner, payload else: self._dispatch(gen, msg_type, payload) except (ValueError, ProtocolError) as e: @@ -2533,13 +2670,25 @@ def _incoming_frame_limit(self, msg_type: int) -> int: except KeyError: raise ValueError(f"unexpected syncer message type {msg_type}") from None + def _submit_digest(self, job: Callable[[], None]) -> None: + """Run one block hash of a ``_BlockDigest`` on the digest pool + (called under ``_reasm_lock``, which also guards the lazy start).""" + if self._digest_pool is None: + from concurrent.futures import ThreadPoolExecutor + + self._digest_pool = ThreadPoolExecutor(BCAST_DIGEST_THREADS, thread_name_prefix="yeto-digest") + try: + self._digest_pool.submit(job) + except RuntimeError: # the client closed its pool + job() + def _reassemble( - self, gen: int, payload: bytes - ) -> tuple[int, memoryview] | _StreamedInboundFragment | None: + self, gen: int, payload: bytes | bytearray + ) -> tuple[int, memoryview] | tuple[int, memoryview, bytes] | _StreamedInboundFragment | None: if len(payload) < _CHUNK_HEAD.size: raise ValueError("chunk header truncated") msg_id, total, offset = _CHUNK_HEAD.unpack_from(payload) - data = payload[_CHUNK_HEAD.size :] + data = memoryview(payload)[_CHUNK_HEAD.size :] if not data: raise ValueError("empty chunk") if total < _HEADER.size: @@ -2559,7 +2708,11 @@ def _reassemble( if len(self._reasm) >= MAX_PARTIAL_MESSAGES: raise ValueError("too many partial chunked messages") if self._inbound_chunk_sink is None: - self._reasm[msg_id] = (bytearray(total), [0], []) + # The digest of a large BCAST / FINAL payload (the bytes + # after its 12-byte head) runs while the chunks arrive. + start = _HEADER.size + 12 + big = total - start >= 2 * BCAST_DIGEST_BLOCK + self._reasm[msg_id] = _Reassembly(total, start if big else None, self._submit_digest) else: sink = self._inbound_chunk_sink self._reasm[msg_id] = _StreamedReassembly( @@ -2589,24 +2742,32 @@ def _reassemble( self._reasm.pop(msg_id, None) self._reasm_touched.pop(msg_id, None) return completed - buf, filled, ranges = partial + buf = partial.buf if len(buf) != total: raise ValueError("chunk total changed within one message") - if any(offset < old_end and old_start < end for old_start, old_end in ranges): + if partial.overlaps(offset, end): raise ValueError("overlapping chunk") buf[offset:end] = data - filled[0] += len(data) - ranges.append((offset, end)) - if filled[0] < total: - return None - if filled[0] != total: - raise ValueError("chunk byte count exceeds frame length") - del self._reasm[msg_id] - self._reasm_touched.pop(msg_id, None) + partial.add(offset, end) + digest = partial.digest + if digest is not None: + digest.arrived(offset, end) + complete = partial.filled >= total + if complete: + if partial.filled != total: + raise ValueError("chunk byte count exceeds frame length") + del self._reasm[msg_id] + self._reasm_touched.pop(msg_id, None) + if not complete: + return None magic, msg_type, length = _HEADER.unpack_from(buf) if magic != MAGIC or length != total - _HEADER.size: + if digest is not None: + digest.cancel() raise ValueError("corrupt reassembled frame") - return msg_type, memoryview(buf)[_HEADER.size :] + if digest is None: + return msg_type, memoryview(buf)[_HEADER.size :] + return msg_type, memoryview(buf)[_HEADER.size :], digest.finish() def _stream_chunk_locked( self, @@ -2772,6 +2933,8 @@ def _clear_partial_reassemblies(self) -> None: self._reasm.clear() self._reasm_touched.clear() for partial in partials: + if isinstance(partial, _Reassembly) and partial.digest is not None: + partial.digest.cancel() if isinstance(partial, _StreamedReassembly): try: self._abort_streamed_partial(partial) @@ -2781,8 +2944,10 @@ def _clear_partial_reassemblies(self) -> None: self._last_err = exc def _dispatch( - self, gen: int, msg_type: int, payload: bytes | memoryview + self, gen: int, msg_type: int, payload: bytes | memoryview, digest: bytes | None = None ) -> None: + """``digest``: SHA-256 of ``payload[12:]`` already computed while a + chunked BCAST payload arrived (``_BlockDigest``).""" with self._lock: if gen != self._gen: return # late message from a dead group @@ -2790,7 +2955,36 @@ def _dispatch( # atomic with respect to teardown/redial. Otherwise an obsolete # receiver could pass the check, lose the race to a reconnect, # then publish a stale terminal manifest or shutdown afterward. - self._dispatch_live(gen, msg_type, payload) + self._dispatch_live(gen, msg_type, payload, digest) + + def _admit_bcast_locked(self, fid: int, version: int, digest: bytes, kind: str = "blocks") -> bool | None: + """Under ``_bcast_lock``: whether a BCAST of ``fid`` is queued (its + ``catch_up`` flag) or dropped (None). Versions only move forward; + a replay of the version already held (same bytes) is dropped, other + bytes under that version are a protocol error, except for the first + BCAST after a redial with ``accept_rewound_catch_up`` (a rewound + syncer's catch-up, taken as the fragment's new global). Digests of + two kinds (a sink installed in between) cannot be compared: the same + version is then taken as a replay.""" + first = self._catch_up_pending[fid] + self._catch_up_pending[fid] = False + seen = self._bcast_seen[fid] + if seen is not None: + seen_version, seen_digest, seen_kind = seen + if version == seen_version and (digest == seen_digest or kind != seen_kind): + return None + if version <= seen_version: + if not (first and self.accept_rewound_catch_up and self._downlink is None): + if version < seen_version: + return None + raise ProtocolError( + f"conflicting BCAST_FRAGMENT payloads for fragment {fid} version {version}" + ) + self.rewound_catch_ups += 1 + _connect_log(f"fragment {fid}: the syncer rewound from version {seen_version} to " + f"{version} (resumed from an older checkpoint); taking its catch-up") + self._bcast_seen[fid] = (version, digest, kind) + return first def _dispatch_streamed( self, @@ -2811,25 +3005,16 @@ def _dispatch_streamed_live(self, inbound: _StreamedInboundFragment) -> None: digest = committed.payload_hash if inbound.msg_type == MSG_BCAST_FRAGMENT: with self._bcast_lock: - seen = self._bcast_seen[fragment_id] - if seen is not None: - seen_version, seen_digest = seen - if version < seen_version: - if committed.discard is not None: - committed.discard() - return - if version == seen_version: - if digest != seen_digest: - if committed.discard is not None: - committed.discard() - raise ProtocolError( - "conflicting BCAST_FRAGMENT payloads for fragment " - f"{fragment_id} version {version}" - ) - if committed.discard is not None: - committed.discard() - return - self._bcast_seen[fragment_id] = (version, digest) + try: + catch_up = self._admit_bcast_locked(fragment_id, version, digest, "sink") + except ProtocolError: + if committed.discard is not None: + committed.discard() + raise + if catch_up is None: + if committed.discard is not None: + committed.discard() + return self._bcasts.put( BcastFragment( fragment_id, @@ -2838,6 +3023,7 @@ def _dispatch_streamed_live(self, inbound: _StreamedInboundFragment) -> None: payload_hash=digest, stored=True, discard=committed.discard, + catch_up=catch_up, ) ) return @@ -2884,7 +3070,7 @@ def _dispatch_streamed_live(self, inbound: _StreamedInboundFragment) -> None: ) def _dispatch_live( - self, gen: int, msg_type: int, payload: bytes | memoryview + self, gen: int, msg_type: int, payload: bytes | memoryview, raw_digest: bytes | None = None ) -> None: if msg_type == MSG_ERROR: message = bytes(payload).decode("utf-8", errors="replace") @@ -2951,6 +3137,7 @@ def _dispatch_live( if fid >= self.layout.num_fragments: raise ValueError(f"BCAST_FRAGMENT for unknown fragment {fid}") if self._downlink is not None: + raw_digest = None # the digest is of the decoded values data = self._downlink.decode_whole( fid, version, payload[12:], self.layout.fragments[fid].numel) else: @@ -2963,26 +3150,16 @@ def _dispatch_live( f"expected {expected}" ) data = payload[12:] - digest = hashlib.sha256(data).digest() + digest = raw_digest if raw_digest is not None else bcast_payload_digest(data) with self._bcast_lock: - seen = self._bcast_seen[fid] - if seen is not None: - seen_version, seen_digest = seen - if version < seen_version: - return - if version == seen_version: - if digest != seen_digest: - raise ProtocolError( - f"conflicting BCAST_FRAGMENT payloads for fragment " - f"{fid} version {version}" - ) - return - self._bcast_seen[fid] = (version, digest) + catch_up = self._admit_bcast_locked(fid, version, digest) + if catch_up is None: + return # Receiver threads run once per striped socket. Queue the # accepted update under the same lock as the monotonicity # check so two versions cannot validate in order but enqueue # in reverse order after a thread reschedule. - self._bcasts.put(BcastFragment(fid, version, data)) + self._bcasts.put(BcastFragment(fid, version, data, catch_up=catch_up)) elif msg_type == MSG_FINAL_FRAGMENT: if len(payload) < 12: self._set_protocol_error("truncated FINAL_FRAGMENT") diff --git a/yeto/rl/spot_rollout/subscriber.py b/yeto/rl/spot_rollout/subscriber.py index c792a113..d1331d6a 100644 --- a/yeto/rl/spot_rollout/subscriber.py +++ b/yeto/rl/spot_rollout/subscriber.py @@ -106,7 +106,7 @@ def __init__(self, addr: tuple[str, int], subscriber_id: int, layout: FragmentLa self._beats = threading.Event() self._run_finished = threading.Event() - def _dispatch(self, gen: int, msg_type: int, payload) -> None: + def _dispatch(self, gen: int, msg_type: int, payload, digest: bytes | None = None) -> None: # The end-of-run notice counts on any connection generation. The # syncer closes the data streams right after it; when their end reaches # the client first, the client starts a new generation and the base class @@ -116,7 +116,7 @@ def _dispatch(self, gen: int, msg_type: int, payload) -> None: self._run_finished.set() self._closed.set() # no redial: the run is over self._failure.set() - super()._dispatch(gen, msg_type, payload) + super()._dispatch(gen, msg_type, payload, digest) def _hello_body_suffix(self) -> bytes: return self._trailer diff --git a/yeto/swift_megatron/swift_plugin.py b/yeto/swift_megatron/swift_plugin.py index 0ac2d695..0f4839e6 100644 --- a/yeto/swift_megatron/swift_plugin.py +++ b/yeto/swift_megatron/swift_plugin.py @@ -55,8 +55,41 @@ YETO_DILOCO_HMAC_KEY_PATH path to the HMAC key file (optional) YETO_DILOCO_DISABLE set to 1 to make the callback a no-op YETO_DILOCO_INITIAL_CUT_TIMEOUT seconds to install the syncer's initial cut - fragment by fragment before training (0 = - apply it at the boundaries; default 0) + (or a resumed island's catch-up) fragment by + fragment before training; default 3600. 0 = + apply it at the first boundaries instead: the + whole cut (4 bytes per parameter) then queues + on rank 0 first, and a resumed island takes it + only after one more step + YETO_DILOCO_RESUME_DIR checkpoint directory whose sync record to + resume from (default: ``args.mcore_model`` + when it holds one) + YETO_DILOCO_SAVE_ANCHORS 1 = also save every rank's anchor shards + with each checkpoint (~2 bytes per local + parameter per rank for bf16); default 0 + YETO_DILOCO_INIT_PROBE_S resumed learner 0 with an initial-cut wait: + send INIT_PARAMS only if the syncer sends no + catch-up within this many seconds (an + initialized syncer ignores INIT); default 1800 + YETO_DILOCO_HEARTBEAT_S seconds between rank 0's HEARTBEATs to the + syncer, sent from a thread (independent of + step ends) so its --island-lease-s can be on; + default 30, 0 = none + YETO_DILOCO_HEARTBEAT_STALL_S stop heartbeating once the training loop + made no progress for this long, so a hung + island loses its lease; default 0 (never) + YETO_DILOCO_LR_RESUME auto (default) | off: rebuild the learning + rate schedule of a resume that did not load + it (``restore_lr_schedule``) + +Resume: every checkpoint also gets the sync's record (``on_save`` -> +``MegatronDiLoCoSync.save_state``): each fragment's anchor version and +H-gate counters. On a restart from that checkpoint the syncer's catch-up +no longer overwrites the island: a fragment still at its checkpointed +version keeps its local values (anchor only), a newer one becomes +``global_new + (local - anchor_old)`` if the anchors were saved and the +global otherwise; without a record (a cold start) the island takes the +syncer's cut as before. """ from __future__ import annotations @@ -108,7 +141,12 @@ class DiLoCoHookConfig: hmac_key_path: Optional[str] = None final_cut_payload: str = "full" transport: str = "cpu" - initial_cut_timeout: float = 0.0 + initial_cut_timeout: float = 3600.0 + resume_dir: Optional[str] = None + save_anchors: bool = False + init_probe_s: float = 1800.0 + heartbeat_s: float = 30.0 + heartbeat_stall_s: float = 0.0 @classmethod def from_env(cls, env=None) -> "DiLoCoHookConfig": @@ -146,7 +184,12 @@ def get_float(name: str, default: float) -> float: hmac_key_path=get("HMAC_KEY_PATH"), final_cut_payload=get("FINAL_CUT_PAYLOAD", default="full"), # type: ignore[arg-type] transport=get("TRANSPORT", default="cpu"), # type: ignore[arg-type] - initial_cut_timeout=get_float("INITIAL_CUT_TIMEOUT", 0.0), + initial_cut_timeout=get_float("INITIAL_CUT_TIMEOUT", 3600.0), + resume_dir=get("RESUME_DIR"), + save_anchors=(get("SAVE_ANCHORS", default="0") or "").strip().lower() in {"1", "true", "yes", "on"}, + init_probe_s=get_float("INIT_PROBE_S", 1800.0), + heartbeat_s=get_float("HEARTBEAT_S", 30.0), + heartbeat_stall_s=get_float("HEARTBEAT_STALL_S", 0.0), ) if cfg.initial_cut_timeout < 0: raise ValueError(f"{ENV_PREFIX}INITIAL_CUT_TIMEOUT must be >= 0") @@ -223,7 +266,8 @@ def trainable_filter(name: str, param) -> bool: return bool(getattr(param, "requires_grad", False)) -def default_sync_factory(model_chunks, optimizer, cfg: DiLoCoHookConfig): +def default_sync_factory(model_chunks, optimizer, cfg: DiLoCoHookConfig, *, + resume_dir: Optional[str] = None, resume_iteration: Optional[int] = None): from yeto.swift_megatron.sync import MegatronDiLoCoSync return MegatronDiLoCoSync( @@ -240,9 +284,80 @@ def default_sync_factory(model_chunks, optimizer, cfg: DiLoCoHookConfig): final_cut_payload=cfg.final_cut_payload, transport=cfg.transport, initial_cut_timeout=cfg.initial_cut_timeout, + resume_dir=resume_dir, + resume_iteration=resume_iteration, + init_probe_s=cfg.init_probe_s, + heartbeat_s=cfg.heartbeat_s, + heartbeat_stall_s=cfg.heartbeat_stall_s, ) +def resume_dir_for(cfg: DiLoCoHookConfig, args) -> Optional[str]: + """The checkpoint directory whose sync record a resumed island reads: + ``YETO_DILOCO_RESUME_DIR``, else the checkpoint ms-swift loaded + (``--mcore_model``) when it holds a record. None: a cold start.""" + from yeto.swift_megatron.sync import RESUME_STATE_FILE + + if cfg.resume_dir: + return cfg.resume_dir + loaded = getattr(args, "mcore_model", None) + if isinstance(loaded, str) and os.path.isfile(os.path.join(loaded, RESUME_STATE_FILE)): + return loaded + return None + + +def save_dir_for(args, iteration: int, *call_args, **call_kwargs) -> Optional[str]: + """The directory of the checkpoint ``on_save`` reports: its + ``output_dir`` argument, else ``/checkpoint-``.""" + for value in (call_kwargs.get("output_dir"), call_kwargs.get("save_dir"), *call_args): + if isinstance(value, (str, os.PathLike)): + return os.fspath(value) + root = getattr(args, "output_dir", None) or getattr(args, "save", None) + return os.path.join(root, f"checkpoint-{iteration}") if root else None + + +def restore_lr_schedule(trainer) -> Optional[int]: + """Put a resumed run's LR schedule where an uninterrupted run would be. + + ms-swift loads the ``OptimizerParamScheduler`` state only together with + the optimizer (``load_mcore_checkpoint``: ``not finetune and not + no_load_optim``), and a checkpoint saved with ``--no_save_optim`` holds + no scheduler state at all. So an island that saves parameters only and + resumes with ``--no_load_optim true --finetune false`` gets back its + iteration and data position but restarts the schedule at step 0: warmup + again, the whole decay shifted by the resumed iteration, silently. + + The scheduler counts samples (ms-swift steps it by ``global_batch_size`` + per successful update), so a fresh scheduler (``num_steps == 0``) at a + resumed iteration ``i > 0`` is stepped by ``i * global_batch_size`` here, + which sets every parameter group's lr and weight decay the way ``i`` + uninterrupted updates would. A scheduler restored from the checkpoint + (``num_steps > 0``) and a cold start are left alone. Returns the samples + stepped, or None. Runs on every rank (each holds its own optimizer).""" + if (os.environ.get(ENV_PREFIX + "LR_RESUME", "auto").strip().lower() or "auto") == "off": + return None + scheduler = getattr(trainer, "opt_param_scheduler", None) + iteration = int(getattr(trainer.state, "iteration", 0) or 0) + if scheduler is None or iteration <= 0 or int(getattr(scheduler, "num_steps", 0) or 0) != 0: + return None + samples = iteration * int(trainer.args.global_batch_size) + scheduler.step(increment=samples) + logger.warning("yeto DiLoCo: the checkpoint of iteration %d restored no LR schedule; stepped it by %d " + "samples to resume it in place", iteration, samples) + return samples + + +def _visible_logs() -> None: + """This module's and the sync core's WARNING/INFO lines would otherwise + go to the root handler the Megatron entry point leaves at ERROR.""" + _ensure_repo_on_path() + try: + from yeto.swift_megatron.sync import ensure_visible_logs, log as sync_log + except ImportError: # pragma: no cover - the sync core needs torch + return + ensure_visible_logs(logger, sync_log) + + def _model_chunks(trainer) -> list: # Pass the DDP-wrapped chunks: the sync core needs the grad/param buffers # and the DistributedOptimizer's mapping, both keyed on the wrapped models. @@ -286,11 +401,16 @@ def __init__(self, trainer): logger.warning("yeto DiLoCo callback disabled via %sDISABLE", ENV_PREFIX) return self.config = DiLoCoHookConfig.from_env() + _visible_logs() + restore_lr_schedule(trainer) self._wrap_train_step() # Build here (not in on_train_begin) so a bad config fails before data # loading; the trainer has already loaded the resume checkpoint, which # is what learner 0's INIT_PARAMS carries. - self.sync = type(self).sync_factory(_model_chunks(trainer), trainer.optimizer, self.config) + start = int(self.state.iteration) + resume_dir = resume_dir_for(self.config, self.args) if start > 0 else None + self.sync = type(self).sync_factory(_model_chunks(trainer), trainer.optimizer, self.config, + resume_dir=resume_dir, resume_iteration=start) atexit.register(self.close) logger.info("yeto DiLoCo sync attached: learner %d/%d, interval %d, fragments %d, start iteration %d", self.config.learner_id, self.config.num_learners, self.config.sync_interval_steps, @@ -366,6 +486,17 @@ def _log_step(self, iteration: int, target_tokens: int, sync_s: float, except OSError as error: # never let bookkeeping kill a training step logger.warning("yeto DiLoCo: cannot write step log %s: %s", path, error) + def on_save(self, *args, **kwargs): + """Record the sync state next to the checkpoint just written (every rank).""" + if self.sync is None or self._closed or getattr(self.sync, "finished", False): + return + iteration = int(self.state.iteration) + output_dir = save_dir_for(self.args, iteration, *args, **kwargs) + if output_dir is None: + logger.warning("yeto DiLoCo: on_save without a checkpoint directory; sync state not recorded") + return + self.sync.save_state(output_dir, iteration=iteration, save_anchors=self.config.save_anchors) + def on_train_end(self): self.close() diff --git a/yeto/swift_megatron/sync.py b/yeto/swift_megatron/sync.py index d87585e8..3ed9d44e 100644 --- a/yeto/swift_megatron/sync.py +++ b/yeto/swift_megatron/sync.py @@ -76,6 +76,19 @@ param all-gather reads). The Adam moments are kept, exactly as the torch learner keeps its AdamW state when ``apply_fragment`` overwrites a fragment. +Redial (catch-up) + After a redial the syncer sends every fragment's current value. The + client drops a replay of a version already held (same bytes), so an + unchanged fragment stays bit for bit; a broadcast of the anchored + version that does reach the sync is only confirmed. The first broadcast + of a fragment after the redial (``BcastFragment.catch_up``) moves the + island's progress since its anchor onto it, ``global_new + (local - + anchor)``, unless the island pushed from that anchor (the delta may be in + the new version: the ordinary install). A syncer that crashed and + resumed from an older checkpoint rewinds; the client takes its catch-up + (``accept_rewound_catch_up``) and the island rebases onto it, instead of + keeping the lost versions as bases the syncer refuses forever. + Only global rank 0 talks to the syncer. Every boundary starts with one ``broadcast_object_list`` of rank 0's plan (fragments to apply, fragments to push, finalization, shutdown, error), so every rank runs the same collectives @@ -142,7 +155,9 @@ from __future__ import annotations import hashlib +import json import logging +import os import re import time from collections import deque @@ -152,12 +167,65 @@ log = logging.getLogger("swift-megatron-sync") + +def _emits_info(logger: logging.Logger) -> bool: + """Whether an INFO record of ``logger`` reaches some handler that prints it.""" + if not logger.isEnabledFor(logging.INFO): + return False + node: logging.Logger | None = logger + while node is not None: + if any(handler.level <= logging.INFO for handler in node.handlers): + return True + if not node.propagate: + return False + node = node.parent # type: ignore[assignment] + return False + + +def ensure_visible_logs(*loggers: logging.Logger) -> None: + """Make these loggers' INFO lines reach stderr. + + ms-swift configures only its own ``swift`` logger, and the Megatron SFT + entry point leaves the root logger with a single stdout handler at ERROR: + every INFO *and* WARNING of a yeto module (``resume: fragment ...``, + ``INIT_PARAMS skipped``, ``anchor ... released after a reconnect``, the LR + schedule rebuild) propagated into that handler and was dropped, and + logging's last resort never fires once root has a handler. A logger whose + records some handler on its propagation path already prints at INFO is + left alone; any other gets its own stderr handler at INFO and stops + propagating (so an ERROR is not printed twice).""" + for logger in loggers or (log,): + if logger.handlers or _emits_info(logger): + continue + handler = logging.StreamHandler() + handler.setFormatter(logging.Formatter(f"[{logger.name}] %(asctime)s %(levelname)s %(message)s")) + logger.addHandler(handler) + if logger.level == logging.NOTSET or logger.level > logging.INFO: + logger.setLevel(logging.INFO) + logger.propagate = False + + +_ensure_visible_logs = ensure_visible_logs + # Megatron's tied/untied output projection is an embedding-like tensor for the # direct-averaged fragment, as ``lm_head`` is for the torch learner. OUTPUT_LAYER_AVG_REGEX = r"(^|\.)output_layer\.weight$" DEFAULT_CHUNK_NUMEL = 1 << 28 CONTRACT_DOMAIN = b"yeto-swift-megatron-native-v1\0" FINAL_CUT_PAYLOADS = ("full", "manifest") +# Island checkpoint record of the sync (``save_state`` / ``resume_dir``), next +# to the trainer's checkpoint: the JSON on global rank 0, the optional anchor +# values one file per rank. +RESUME_STATE_FILE = "yeto_diloco_state.json" +RESUME_ANCHOR_DIR = "yeto_diloco_anchors" +RESUME_FORMAT = 1 +# ``yeto.diloco_sync.resume_fragment_mode`` results (that module imports torch). +RESUME_OVERWRITE, RESUME_ANCHOR, RESUME_SHIFT = "overwrite", "anchor", "shift" +# ``yeto.diloco_sync.live_install_mode`` results. +LIVE_CONFIRM, LIVE_REBASE = "confirm", "rebase" + +# ``spent_link`` of an anchor spent by the process the island resumed from. +_PREVIOUS_PROCESS = "previous-process" _LAYER_MODULE_RE = re.compile(r"(^|\.)decoder\.layers\.(\d+)$") _GROUPED_EXPERT_RE = re.compile(r"^(.*?)(weight|bias)(\d+)$") @@ -522,8 +590,10 @@ class _Plan: """Rank 0's decision for one boundary, broadcast to every rank.""" applies: list[tuple[int, int]] = field(default_factory=list) # (fragment, version) + modes: dict[int, str] = field(default_factory=dict) # fragment -> LIVE_* (else ordinary) finalizing: bool = False shutdown: bool = False + send_init: bool = False # initial-cut wait of a resumed learner 0: the syncer needs INIT error: str | None = None @@ -563,7 +633,9 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int, final_cut_payload: str = "full", chunk_numel: int = DEFAULT_CHUNK_NUMEL, transport: str = "cpu", parallel_context: ParallelContext | None = None, client_factory=None, observer: Callable[[dict], None] | None = None, - initial_cut_timeout: float = 0.0) -> None: + initial_cut_timeout: float = 0.0, resume_dir: str | None = None, + resume_iteration: int | None = None, init_probe_s: float = 1800.0, + heartbeat_s: float = 30.0, heartbeat_stall_s: float = 0.0) -> None: import torch import torch.distributed as dist @@ -585,6 +657,7 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int, raise ValueError("the ms-swift Megatron sync talks to one syncer; sharded syncers are not supported") if not dist.is_initialized(): raise SwiftMegatronSyncError("torch.distributed must be initialized (Megatron does this)") + _ensure_visible_logs() self.ctx = parallel_context or ParallelContext.from_megatron() self.rank, self.world = self.ctx.rank, self.ctx.world self.chunks = list(model_chunks) @@ -645,6 +718,15 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int, # for a newer broadcast instead of a gather that would be refused. A # restarted attempt of the same round may reuse it. self.spent: list[tuple[int, int] | None] = [None] * self.num_fragments + # The syncer link (connection generation) each spent PUSH went out on: + # once the link is replaced, an undelivered delta is gone with the old + # connection and a restarted syncer lost the rounds it was carried to, + # so the guard is released (``_release_spent_after_reconnect``). + self.spent_link: list[Any] = [None] * self.num_fragments + # The anchor version of each fragment's last PUSH that may have been + # delivered (a drop clears it; a redial does not): a newer catch-up + # of that fragment may already hold the delta (``live_install_mode``). + self.pushed_from: list[int | None] = [None] * self.num_fragments # The torch learner's per-fragment counters and versions, kept alike on # every rank (anchors live in ``self.anchors``, sharded, not here). self.state = DiLoCoSyncState.create(self.num_fragments, track_anchors=False) @@ -661,8 +743,25 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int, "deferred_pushes": 0, "apply_s": 0.0, "export_s": 0.0} self._observed_push_bytes = 0 # rank 0: push_bytes at the last observer report self.sender: _PushSender | None = None # rank 0: background PUSH streaming + # Rank 0: HEARTBEAT every ``heartbeat_s`` from a thread, so the syncer + # can lease the island (``--island-lease-s``) across ~10 min steps, + # INIT and checkpoint saves; ``_progress_at`` is the last sign of life + # of the training loop (``heartbeat_stall_s``: none for that long and + # the heartbeats stop, so a hung island loses its lease). + self.heartbeat: _Heartbeat | None = None + self._progress_at = time.monotonic() # (iteration, "apply" | "push", fragment, version | pull step), same on every rank. self.events: deque[tuple[int, str, int, int]] = deque(maxlen=4096) + # Resume (``resume_dir``): the record ``save_state`` wrote with the + # checkpoint the trainer resumed from, the old anchor of every local + # shard when it was saved too, and how each fragment's first broadcast + # after the restart was installed (fid -> mode, same on every rank). + self.resume_record: dict | None = None + self._resume_anchors: dict[int, Any] = {} # id(param) -> old anchor (host) + self._resume_has_anchors = False + self.resumed: dict[int, str] = {} + if resume_dir: + self._load_resume(resume_dir, resume_iteration) if self.rank == 0: log.info("swift-megatron sync: %d trainable tensors (%.2f B elements) -> %d fragments", @@ -671,20 +770,32 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int, if self.rank == 0: try: self.client = self._connect(syncer, hmac_key_path, wan_streams, client_factory) - self.sender = _PushSender(self.client, self.stats, self._push_dropped) + self.sender = _PushSender(self.client, self.stats, self._push_dropped, self._push_sent) + if heartbeat_s > 0: + self.heartbeat = _Heartbeat(self.client, float(heartbeat_s), float(heartbeat_stall_s), + lambda: self.iteration, lambda: self._progress_at) except Exception as exc: # raised on every rank below error = f"{type(exc).__name__}: {exc}" self._raise_if(self._bcast_object(error)) - if self.learner_id == 0: + # A resumed learner 0 (it holds a record, so the syncer of this session + # was initialized) first waits for the syncer's catch-up and sends + # INIT_PARAMS (~4 bytes per parameter, 480 GB for 120B) only if nothing + # arrives within ``init_probe_s``: an initialized syncer ignores INIT, + # and receiving the catch-up while gathering INIT doubled rank 0's RSS. + defer_init = self.learner_id == 0 and self.resume_record is not None and initial_cut_timeout > 0 + if self.learner_id == 0 and not defer_init: self._send_init() if initial_cut_timeout > 0: - self._install_initial_cut(float(initial_cut_timeout)) + self._install_initial_cut(float(initial_cut_timeout), + init_probe=float(init_probe_s) if defer_init else None) # ------------------------------------------------------------------ setup def _connect(self, syncer: str, hmac_key_path: str | None, wan_streams: int, client_factory): from yeto.diloco_sync import connect_syncer - kwargs = {"session_contract_hash": self.contract} + # A syncer resumed from an older checkpoint rewinds; its catch-up is + # taken (and the island rebased onto it) instead of being dropped. + kwargs = {"session_contract_hash": self.contract, "accept_rewound_catch_up": True} key = _read_hmac_key(hmac_key_path) if key is not None: kwargs["hmac_key"] = key @@ -696,53 +807,94 @@ def _connect(self, syncer: str, hmac_key_path: str | None, wan_streams: int, cli def _send_init(self) -> None: """INIT_PARAMS of every fragment from this island's current masters. - Rank 0 streams fragment k on a background thread while every rank - gathers fragment k + 1, so the island's gather and the link overlap; - at most two fragments' wire parts are on rank 0's host.""" + Streamed chunk by chunk: rank 0 hands each exported chunk to a + background sender, which writes it into the fragment's INIT_PARAMS + frame and drops it, while every rank gathers the next chunk. At most + a few chunks (``chunk_numel`` elements each) are on rank 0's host + instead of whole fragments.""" import queue import threading + end = object() jobs: queue.Queue = queue.Queue(maxsize=1) failure: list[str] = [] sender = None if self.rank == 0: + def chunks_of(pending: queue.Queue, ended: list): + while (item := pending.get()) is not end: + if isinstance(item, BaseException): + raise item + yield from item + ended.append(True) + def send() -> None: while (job := jobs.get()) is not None: - fid, parts = job - if failure: - continue # drain: the gathers go on until the error is broadcast + fid, pending = job + ended: list = [] try: - self.client.send_init_parts(fid, parts) + if not failure and self.client.send_init_parts(fid, chunks_of(pending, ended)) is False: + failure.append(f"INIT fragment {fid}: syncer link down") except Exception as exc: failure.append(f"INIT fragment {fid}: {type(exc).__name__}: {exc}") - del parts, job + # Drain what was not sent: the gathers go on until the error is broadcast. + while not ended and pending.get() is not end: + pass + del job, pending sender = threading.Thread(target=send, name="yeto-swift-init", daemon=True) sender.start() error = None + current = None # rank 0: the queue of the fragment being exported + aborted = False try: for fid in range(self.num_fragments): - parts = self._export_fragment(fid, delta=False) - if self.rank == 0 and error is None: - if self._export_error is not None: + pending = queue.Queue(maxsize=2) if self.rank == 0 else None + current = pending + if pending is not None: + jobs.put((fid, pending)) + for parts in self._export_fragment_chunks(fid, delta=False): + self._progress_at = time.monotonic() + if pending is None: + continue + if error is None and self._export_error is not None: error = f"INIT fragment {fid}: {self._export_error}" - elif failure: + failure.append(error) # the sender skips the fragments after it + pending.put(SwiftMegatronSyncError(error)) + elif error is None: + pending.put(parts) + del parts + if pending is not None: + pending.put(end) + current = None + if error is None and failure: error = failure[0] - else: - jobs.put((fid, parts)) - del parts + except BaseException: + # An export that raised (a collective timeout) leaves the sender + # waiting for the rest of the fragment: end it so the error, not + # a hang, reaches the caller. + aborted = True + if current is not None: + failure.append("INIT aborted") + try: + current.put(end, timeout=60) + except queue.Full: + pass + raise finally: if sender is not None: - jobs.put(None) - sender.join() + try: + jobs.put(None, timeout=60 if aborted else None) + except queue.Full: + pass + sender.join(timeout=60 if aborted else None) if self.rank == 0 and error is None and failure: error = failure[0] self._raise_if(self._bcast_object(error)) if self.rank == 0: log.info("sent INIT_PARAMS for %d fragments", self.num_fragments) - def _install_initial_cut(self, timeout: float) -> None: + def _install_initial_cut(self, timeout: float, *, init_probe: float | None = None) -> None: """Apply the syncer's first broadcast of every fragment as it arrives, before training starts (collective: every rank loops together). @@ -756,16 +908,34 @@ def _install_initial_cut(self, timeout: float) -> None: import time started = time.monotonic() + init_pending = init_probe is not None # resumed learner 0: INIT only if the syncer needs it while not all(self.anchored): + self._progress_at = time.monotonic() plan = self._plan() if self.rank == 0 else None - if self.rank == 0 and time.monotonic() - started > timeout and not plan.applies: - plan.shutdown = True # only ends the wait; the client's flag is unchanged + if self.rank == 0: + waited = time.monotonic() - started + if init_pending and not plan.applies and waited > init_probe: + plan.send_init = True + elif not init_pending and waited > timeout and not plan.applies: + plan.shutdown = True # only ends the wait; the client's flag is unchanged plan = self._bcast_object(plan) self._raise_if(plan.error) + if init_pending and (plan.applies or plan.send_init): + init_pending = False + if plan.send_init: + if self.rank == 0: + log.warning("no catch-up from the syncer within %.0f s: sending INIT_PARAMS", init_probe) + self._send_init() + elif self.rank == 0: + log.warning("the syncer is initialized (catch-up arriving): INIT_PARAMS skipped") + started = time.monotonic() updates = self._take_updates(plan) if self.rank == 0 else {} for fid, version in plan.applies: data = updates.pop(fid) if self.rank == 0 else None - self._apply_fragment(fid, data, self.bulk_torch, blend=True) + live = plan.modes.get(fid) + if live == LIVE_CONFIRM: + continue + self._apply_first_or_blend(fid, data, version, live=live) del data self.state.fragment_versions[fid] = version self.state.global_step = max(self.state.global_step, version) @@ -776,9 +946,191 @@ def _install_initial_cut(self, timeout: float) -> None: break if not plan.applies: time.sleep(0.2) + if init_pending: # the wait ended before the probe did: keep the old order's guarantee + self._send_init() + if self.rank == 0: + log.warning("initial cut: %d of %d fragments installed before training in %.0f s", + sum(self.anchored), self.num_fragments, time.monotonic() - started) + + # ------------------------------------------------------------------ resume + def _apply_first_or_blend(self, fid: int, data: Any, version: int, *, live: str | None = None) -> str | None: + """Apply a broadcast; the first one of a fragment after a resume by + ``resume_fragment_mode``, an anchored one by ``live`` (``_plan``: + ``LIVE_REBASE`` or the ordinary install). Returns the resume mode + (None: not a resumed first install) and records it in + ``self.resumed``.""" + + mode = None + if live == LIVE_REBASE and self.anchored[fid]: + if self.rank == 0: + log.info("fragment %d: catch-up version %d onto anchor %d -> rebase (local progress kept)", + fid, version, self.state.fragment_versions[fid]) + self._apply_fragment(fid, data, self.bulk_torch, blend=False, mode=LIVE_REBASE) + self.stats["rebased"] = self.stats.get("rebased", 0) + 1 + return None + if self.resume_record is not None and not self.anchored[fid] and fid not in self.resumed: + from yeto.diloco_sync import resume_fragment_mode + + mode = resume_fragment_mode(self.resume_record, fid, version, have_anchor=self._resume_has_anchors) + self.resumed[fid] = mode + if self.rank == 0: + log.info("resume: fragment %d version %d (checkpoint anchor %d) -> %s", fid, version, + self.resume_record["versions"][fid], mode) + self._apply_fragment(fid, data, self.bulk_torch, blend=True, mode=mode) + if mode is not None and len(self.resumed) == self.num_fragments: + self._resume_anchors.clear() # every old anchor is used or no longer needed + if self.rank == 0: + counts = {m: list(self.resumed.values()).count(m) for m in (RESUME_ANCHOR, RESUME_SHIFT, RESUME_OVERWRITE)} + log.warning("resume: all %d fragments installed: anchor %d, shift %d, overwrite %d", + self.num_fragments, counts[RESUME_ANCHOR], counts[RESUME_SHIFT], counts[RESUME_OVERWRITE]) + return mode + + def _restore_counters(self, fid: int, mode: str) -> None: + """A fragment whose checkpointed progress was kept (anchor / shift) + keeps its H-gate and token counters, and an unchanged anchor its + spent-anchor guard; an overwritten one starts afresh.""" + + if mode == RESUME_OVERWRITE: + return + self.state.restore_fragment(self.resume_record, fid) + if mode == RESUME_ANCHOR: + spent = self.resume_record["spent"][fid] + self.spent[fid] = None if spent is None else (int(spent[0]), int(spent[1])) + # Pushed by the previous process, on a link that is gone: released + # at the first boundary like any spent anchor after a reconnect. + self.spent_link[fid] = None if spent is None else _PREVIOUS_PROCESS + self.pushed_from[fid] = None if spent is None else int(spent[0]) + + def _resume_paths(self, resume_dir: str) -> tuple[str, str]: + return (os.path.join(resume_dir, RESUME_STATE_FILE), + os.path.join(resume_dir, RESUME_ANCHOR_DIR, f"rank{self.rank:05d}.pt")) + + def _shard_keys(self) -> list[tuple[str, int]]: + return [(s.name, s.tp_rank) for s in self.shards] + + def _load_resume(self, resume_dir: str, resume_iteration: int | None) -> None: + """Read the record ``save_state`` left in ``resume_dir`` (collective). + + A missing or foreign record (another learner, layout, parallel shape + or iteration) is a cold start: every fragment takes the global as + before. Old anchors are used only if every rank has its file.""" + + import torch + + record, reason = None, None + if self.rank == 0: + path = self._resume_paths(resume_dir)[0] + try: + with open(path, encoding="utf-8") as handle: + record = json.load(handle) + except FileNotFoundError: + reason = f"no {RESUME_STATE_FILE} in {resume_dir}" + except (OSError, ValueError) as exc: + reason = f"unreadable {path}: {exc}" + if record is not None: + checks = [ + (record.get("format") == RESUME_FORMAT, "format"), + (record.get("learner_id") == self.learner_id, "learner id"), + # A hand-written bridge record (an island checkpoint from before + # these records) may leave the contract out. + (record.get("contract", self.contract.hex()) == self.contract.hex(), + "layout / parallel shape"), + (len(record.get("versions", ())) == self.num_fragments, "fragment count"), + (resume_iteration is None or record.get("iteration") == int(resume_iteration), "iteration"), + ] + bad = [what for ok, what in checks if not ok] + if bad: + reason = f"{path} does not match this island ({', '.join(bad)})" + record = None + if record is None: + log.warning("resume: %s; the syncer's cut replaces every fragment (cold start)", reason) + record = self._bcast_object(record) + if record is None: + return + loaded: dict[int, Any] = {} + ok = bool(record.get("anchors")) + if ok: + try: + saved = torch.load(self._resume_paths(resume_dir)[1], map_location="cpu", weights_only=True) + if saved["keys"] != [list(k) for k in self._shard_keys()]: + raise ValueError("shard list differs") + for shard, value in zip(self.shards, saved["values"], strict=True): + if tuple(value.shape) != tuple(shard.local_shape): + raise ValueError(f"{shard.name} has shape {tuple(value.shape)}") + loaded[id(shard.param)] = value + except Exception as exc: # agreed on below; a missing anchor file is not fatal + log.warning("resume: rank %d cannot use its saved anchors: %s", self.rank, exc) + ok = False + ok = all(self._gather_object(ok)) + self.resume_record = record + self._resume_has_anchors = ok + self._resume_anchors = loaded if ok else {} + self.units_total = int(record["units_total"]) + self.state.global_step = int(record["global_step"]) if self.rank == 0: - log.info("initial cut: %d of %d fragments installed before training in %.0f s", - sum(self.anchored), self.num_fragments, time.monotonic() - started) + log.info("resume: island record of iteration %s (anchor versions %s..%s, old anchors %s)", + record.get("iteration"), min(record["versions"]), max(record["versions"]), + "loaded" if ok else "not saved: a newer fragment takes the global") + + def _gather_object(self, value: Any) -> list[Any]: + import torch.distributed as dist + + if self.world == 1: + return [value] + out: list[Any] = [None] * self.world + dist.all_gather_object(out, value, group=self.comm.world) + return out + + def save_state(self, output_dir: str, *, iteration: int, save_anchors: bool = False) -> bool: + """Record the sync's per-fragment state next to the trainer checkpoint + ``output_dir`` of ``iteration`` (call it on every rank after the + checkpoint is written; collective only with ``save_anchors``). + + Global rank 0 writes ``RESUME_STATE_FILE``: each fragment's anchor + version, the H-gate / token counters and the spent-anchor guard. + ``save_anchors`` also writes every rank's anchor shards (the + broadcast dtype, ~``2 * n_r`` bytes per rank for bf16) to + ``RESUME_ANCHOR_DIR``; only with them can a fragment that the syncer + moved past the checkpoint's anchor keep the island's local progress + (``global_new + (local - anchor_old)``). Never raises on an I/O + error (training goes on); returns whether rank 0 wrote the record + (True on the other ranks).""" + + import torch + + anchors_ok = False + if save_anchors: + path = self._resume_paths(output_dir)[1] + try: + os.makedirs(os.path.dirname(path), exist_ok=True) + values = [self.anchors.get(id(s.param)) for s in self.shards] + if any(v is None for v in values): + raise ValueError("a local shard has no anchor yet") + tmp = f"{path}.tmp" + torch.save({"keys": [list(k) for k in self._shard_keys()], "values": values}, tmp) + os.replace(tmp, path) + anchors_ok = True + except Exception as exc: + log.warning("save_state: rank %d cannot write its anchors: %s", self.rank, exc) + anchors_ok = all(self._gather_object(anchors_ok)) if save_anchors else False + written = True + if self.rank == 0: + record = {"format": RESUME_FORMAT, "learner_id": self.learner_id, "iteration": int(iteration), + "contract": self.contract.hex(), "anchors": anchors_ok, + **self.state.resume_record(anchored=self.anchored, units_total=self.units_total, + spent=self.spent)} + path = self._resume_paths(output_dir)[0] + try: + os.makedirs(output_dir, exist_ok=True) + tmp = f"{path}.tmp" + with open(tmp, "w", encoding="utf-8") as handle: + json.dump(record, handle) + os.replace(tmp, path) + log.info("save_state: %s (iteration %d, anchors %s)", path, iteration, anchors_ok) + except OSError as exc: + log.warning("save_state: cannot write %s: %s", path, exc) + written = False + return written # ------------------------------------------------------------------ collectives def _bcast_object(self, value: Any) -> Any: @@ -855,23 +1207,36 @@ def _export_fragment(self, fid: int, *, delta: bool) -> list[Any] | None: """The fragment's values on rank 0 as wire byte parts (INIT, or PUSH for bf16/f32 sessions) or, for a q4 push, one f32 tensor.""" + parts: list[Any] | None = [] if self.rank == 0 else None + for chunk in self._export_fragment_chunks(fid, delta=delta): + if parts is not None: + parts.extend(chunk) + return parts + + def _export_fragment_chunks(self, fid: int, *, delta: bool): + """``_export_fragment`` one chunk at a time (collective): yields each + chunk's parts on rank 0, None elsewhere; ``_export_error`` is set once + a non-finite value was seen.""" + from yeto.protocol import DTYPE_Q4 from yeto.tensor_io import pack_tensor - parts: list[Any] | None = [] if self.rank == 0 else None self._export_error = None q4 = delta and self.wire == DTYPE_Q4 for names in self.fragment_chunks[fid]: tensors = self._export_chunk(names, delta=delta) - if parts is None: + if tensors is None: + yield None continue + parts = [] for tensor in tensors: if self._export_error is None and not bool(tensor.isfinite().all()): # Raised on every rank by the caller's error broadcast. self._export_error = f"fragment {fid} holds a non-finite value" parts.append(tensor.reshape(-1) if q4 else pack_tensor(tensor.reshape(-1), self.bulk if not delta else self._push_part_dtype())) - return parts + del tensors + yield parts def _push_part_dtype(self) -> int: """The dtype of the PUSH parts handed to the client: f32 on the @@ -882,9 +1247,16 @@ def _push_part_dtype(self) -> int: return DTYPE_F32 if bulk_dtype(self.wire) == DTYPE_F32 else push_wire_dtype(self.wire) - def _apply_chunk(self, names: Sequence[str], flat: Any | None, dtype: Any, *, blend: bool) -> None: + def _apply_chunk(self, names: Sequence[str], flat: Any | None, dtype: Any, *, blend: bool, + mode: str | None = None) -> None: """Install the raw global values of ``names`` (``flat`` on rank 0, in the - broadcast dtype) as every rank's anchor, master and model copy.""" + broadcast dtype) as every rank's anchor, master and model copy. + + A resumed island's first install of a fragment (``mode``, + ``yeto.diloco_sync.resume_fragment_mode``): ``"anchor"`` only records + the anchor and leaves the master and model copy untouched, + ``"shift"`` writes ``global + (master - old anchor)``; ``"overwrite"`` + (or None) is the ordinary install.""" import torch import torch.distributed as dist @@ -932,12 +1304,36 @@ def piece(buffer, name, part_dim, tp_rank, size): n = int(shard.param.numel()) pieces.append(incoming[at:at + n].view(shard.local_shape)) at += n - targets = [] + targets, previous = [], [] for shard, whole in zip(local, pieces, strict=True): raw = whole.to(dtype=torch.float32).contiguous() + previous.append(self.anchors.get(id(shard.param))) self.anchors[id(shard.param)] = whole.to(device="cpu", dtype=self.anchor_dtype).contiguous().clone() targets.append(raw) - if blend and self.merge_alpha > 0: + if mode == RESUME_ANCHOR: + return # the broadcast is the checkpoint's own anchor: keep the local values + if mode == LIVE_REBASE: + # The island's progress since its anchor moves onto the catch-up; + # a shard whose global did not change keeps its values bit for bit. + from yeto.diloco_sync import shifted_resume_value + + moved = [i for i, (s, old) in enumerate(zip(local, previous, strict=True)) + if old is None or not torch.equal(old, self.anchors[id(s.param)])] + if not moved: + return + local = [local[i] for i in moved] + masters = self._masters(local) + targets = [targets[i] if previous[i] is None else shifted_resume_value(targets[i], m, previous[i]) + for i, m in zip(moved, masters, strict=True)] + del masters + elif mode == RESUME_SHIFT: + from yeto.diloco_sync import shifted_resume_value + + masters = self._masters(local) + targets = [shifted_resume_value(t, m, self._resume_anchors.pop(id(s.param))) + for s, m, t in zip(local, masters, targets, strict=True)] + del masters + elif blend and self.merge_alpha > 0: masters = self._masters(local) targets = [self.merge_alpha * m.to(torch.float32) + (1.0 - self.merge_alpha) * t for m, t in zip(masters, targets, strict=True)] @@ -961,8 +1357,10 @@ def _check_payload(self, fid: int, data: Any, dtype: Any) -> None: if nbytes != expected: raise SwiftMegatronSyncError(f"fragment {fid} payload has {nbytes} bytes, expected {expected}") - def _apply_fragment(self, fid: int, data: Any, dtype: Any, *, blend: bool) -> None: - """Apply one fragment's payload (``data`` on rank 0) chunk by chunk.""" + def _apply_fragment(self, fid: int, data: Any, dtype: Any, *, blend: bool, + mode: str | None = None) -> None: + """Apply one fragment's payload (``data`` on rank 0) chunk by chunk + (``mode``: a resumed island's first install, ``_apply_chunk``).""" import torch @@ -974,7 +1372,7 @@ def _apply_fragment(self, fid: int, data: Any, dtype: Any, *, blend: bool) -> No flat = None if view is not None: flat = torch.frombuffer(bytearray(view[offset:offset + nbytes]), dtype=torch.uint8).view(dtype) - self._apply_chunk(names, flat, dtype, blend=blend) + self._apply_chunk(names, flat, dtype, blend=blend, mode=mode) offset += nbytes # ------------------------------------------------------------------ boundary @@ -996,10 +1394,15 @@ def after_optimizer_step(self, iteration: int, target_tokens: int, *, if self.finished or self._closed: return + self._progress_at = time.monotonic() state = self.state iteration = int(iteration) if not self._counters_started: # count from the (possibly resumed) first step state.steps_at_reset = [iteration - 1] * self.num_fragments + # A fragment that kept its checkpointed progress keeps its clock: + # the H gate is not restarted by the island's own restart. + for fid, mode in self.resumed.items(): + self._restore_counters(fid, mode) self._counters_started = True self.iteration = iteration if update_successful: @@ -1016,11 +1419,19 @@ def after_optimizer_step(self, iteration: int, target_tokens: int, *, updates = self._take_updates(plan) if self.rank == 0 else {} for fid, version in plan.applies: data = updates.get(fid) if self.rank == 0 else None + live = plan.modes.get(fid) + if live == LIVE_CONFIRM: # a replay of the anchor: nothing changes + self.stats["confirmed"] = self.stats.get("confirmed", 0) + 1 + self.events.append((iteration, "confirm", fid, version)) + continue started = time.perf_counter() - self._apply_fragment(fid, data, self.bulk_torch, blend=True) + mode = self._apply_first_or_blend(fid, data, version, live=live) self.stats["apply_s"] += time.perf_counter() - started - state.steps_at_reset[fid] = iteration - state.units_at_reset[fid] = self.units_total + if live != LIVE_REBASE: # a rebased fragment keeps its local progress and clock + state.steps_at_reset[fid] = iteration + state.units_at_reset[fid] = self.units_total + if mode is not None: + self._restore_counters(fid, mode) state.fragment_versions[fid] = version self.anchored[fid] = True state.global_step = max(state.global_step, version) @@ -1071,19 +1482,35 @@ def _plan(self) -> _Plan: if plan.finalizing: return plan self._updates = {} + catch_up: set[int] = set() for update in client.drain_updates(): - # A later broadcast of the same fragment supersedes an earlier one. + # A later broadcast of the same fragment supersedes an earlier + # one; a catch-up among them is still the syncer's current value. self._updates[update.fragment_id] = update + if getattr(update, "catch_up", False): + catch_up.add(update.fragment_id) for fid, update in self._updates.items(): # Checked here, before the plan goes out: a bad payload found # while applying would fail rank 0 inside the chunk broadcasts. self._check_payload(fid, update.data, self.bulk_torch) plan.applies = sorted((fid, int(u.version)) for fid, u in self._updates.items()) + plan.modes = self._live_modes(plan.applies, catch_up) plan.shutdown = bool(client.shutdown.is_set()) except Exception as exc: plan.error = f"{type(exc).__name__}: {exc}" return plan + def _live_modes(self, applies: Sequence[tuple[int, int]], catch_up: set[int]) -> dict[int, str]: + from yeto.diloco_sync import live_install_mode + + modes = {} + for fid, version in applies: + mode = live_install_mode(anchored=self.anchored[fid], current_version=self.state.fragment_versions[fid], + version=version, catch_up=fid in catch_up, pushed_from=self.pushed_from[fid]) + if mode is not None: + modes[fid] = mode + return modes + def _ready_pushes(self) -> list[tuple[int, int, int, int, int]] | str: """Rank 0: the pulls answered now, by the torch learner's rule (``take_ready_pulls``) with the H-step gate; one push per fragment @@ -1102,6 +1529,7 @@ def _ready_pushes(self) -> list[tuple[int, int, int, int, int]] | str: if state.pending_pulls: self.stats["deferred_pushes"] += 1 return [] + self._release_spent_after_reconnect() ready = take_ready_pulls(state, steps_total=self.iteration, units_total=self.units_total, anchored=self.anchored.__getitem__, min_steps=self.interval) # One fragment per boundary: its PUSH streams in the background while @@ -1115,6 +1543,38 @@ def _ready_pushes(self) -> list[tuple[int, int, int, int, int]] | str: out.append((int(pull.fragment_id), int(pull.global_step), int(pull.round_attempt), c_steps, c_units)) return out + def link_of(self, fid: int) -> Any: + """Rank 0: the connection generation fragment ``fid`` is pushed on + (None for a client that does not expose it).""" + client = self.client + of = getattr(client, "connection_generation_of", None) + if of is not None: + return of(fid) + return getattr(client, "connection_generation", None) + + def _release_spent_after_reconnect(self) -> None: + """Rank 0, sender idle: free every spent anchor whose PUSH went out + on a connection that has since been replaced. + + Without this an island that pushed just before a syncer restart (or + whose PUSH was carried into a round the restarted syncer no longer + has) refuses every pull of that fragment until a newer broadcast, + which no syncer sends once every island waits like this. Answering + from the same anchor again cannot double-count it: a syncer that + still holds the delta refuses the second one per island + (``anchor_reused``, syncer ``LateCarry``), and one that lost it + merges it once.""" + for fid, spent in enumerate(self.spent): + if spent is None or self.spent_link[fid] is None: + continue + link = self.link_of(fid) + if link is not None and link != self.spent_link[fid]: + log.info("fragment %d: anchor %d released after a reconnect (push of step %d)", + fid, spent[0], spent[1]) + self.spent[fid] = None + self.spent_link[fid] = None + self.stats["released_anchors"] = self.stats.get("released_anchors", 0) + 1 + def _pushable(self, pull: Any) -> bool: fid = int(pull.fragment_id) spent = self.spent[fid] @@ -1124,6 +1584,13 @@ def _push_dropped(self, fid: int, base: int, step: int) -> None: """Sender thread: a dropped PUSH never reached the syncer, so its anchor is not spent.""" if self.spent[fid] == (base, step): self.spent[fid] = None + self.spent_link[fid] = None + if self.pushed_from[fid] == base: + self.pushed_from[fid] = None + + def _push_sent(self, fid: int) -> None: + """Sender thread, before the PUSH is queued: the link it goes out on.""" + self.spent_link[fid] = self.link_of(fid) def _take_updates(self, plan: _Plan) -> dict[int, Any]: updates, self._updates = self._updates, {} @@ -1143,6 +1610,8 @@ def _push(self, fid: int, step: int, attempt: int, c_steps: int, c_units: int, return f"PUSH fragment {fid}: {self._export_error}" head = (fid, step, attempt, self.state.fragment_versions[fid], self.iteration, c_steps, c_units) self.spent[fid] = (self.state.fragment_versions[fid], step) + self.spent_link[fid] = None # set by the sender once it goes out + self.pushed_from[fid] = self.state.fragment_versions[fid] if self.wire == DTYPE_Q4: # q4 blocks span tensors: the shared encoder takes the whole delta. self.sender.submit(head, payload=encode_push_delta(self.client, torch.cat(parts))) @@ -1211,6 +1680,8 @@ def close(self) -> None: if self._closed: return self._closed = True + if self.heartbeat is not None: + self.heartbeat.stop() if self.sender is not None: self.sender.stop() client, self.client = self.client, None @@ -1241,6 +1712,40 @@ def session_dtype(name: str) -> int: return code | DTYPE_FLAG_TENSOR_BLOCKS if key != "bf16" else code +class _Heartbeat: + """Rank 0's HEARTBEAT thread (``client.heartbeat(local_step)`` every + ``interval`` seconds): the syncer renews the island's lease on it, so a + lease can be on although a step (or an INIT, a checkpoint save) takes far + longer. ``stall`` > 0: no heartbeat once the training loop showed no + progress (``progress_at``) for that long. A send during an outage is + dropped by the client; the next one goes out after the redial.""" + + def __init__(self, client: Any, interval: float, stall: float, step: Callable[[], int], + progress_at: Callable[[], float]) -> None: + import threading + + self.client, self.interval, self.stall = client, interval, stall + self.step, self.progress_at = step, progress_at + self.sent = 0 + self._stop = threading.Event() + self.thread = threading.Thread(target=self._run, name="yeto-swift-heartbeat", daemon=True) + self.thread.start() + + def _run(self) -> None: + while not self._stop.wait(self.interval): + if self.stall > 0 and time.monotonic() - self.progress_at() > self.stall: + continue + try: + self.client.heartbeat(int(self.step())) + self.sent += 1 + except Exception as exc: # never take the island down for a heartbeat + log.debug("heartbeat not sent: %s", exc) + + def stop(self) -> None: + self._stop.set() + self.thread.join(timeout=5.0) + + class _PushSender: """Rank 0's background PUSH streaming: one PUSH in flight at a time, so a slow link or syncer never holds a training step (decoupled DiLoCo, as the @@ -1248,11 +1753,12 @@ class _PushSender: push (the syncer re-requests it after the redial); any other failure is re-raised by the next boundary on every rank.""" - def __init__(self, client: Any, stats: dict, on_drop: Callable[[int, int, int], None]) -> None: + def __init__(self, client: Any, stats: dict, on_drop: Callable[[int, int, int], None], + on_sent: Callable[[int], None] | None = None) -> None: import queue import threading - self.client, self.stats, self.on_drop = client, stats, on_drop + self.client, self.stats, self.on_drop, self.on_sent = client, stats, on_drop, on_sent self.jobs: queue.Queue = queue.Queue() self.idle = threading.Event() self.idle.set() @@ -1280,6 +1786,11 @@ def stop(self) -> None: self.thread.join(timeout=5.0) def _run(self) -> None: + from yeto.protocol import PartialMessageGenerationLost + + # A torso on a dead connection generation is dropped by the syncer, + # so all three mean the delta was never delivered. + undelivered = (ConnectionError, OSError, PartialMessageGenerationLost) while True: job = self.jobs.get() if job is None: @@ -1288,14 +1799,17 @@ def _run(self) -> None: fid = head[0] try: if payload is not None: - self.client.push_fragment(*head, payload) - nbytes = len(payload) - else: - self.client.push_fragment_parts(*head, parts) - nbytes = sum(len(part) for part in parts) + parts = [payload] # q4: already wire bytes, sent unchanged + nbytes = sum(len(part) for part in parts) + if self.on_sent is not None: + self.on_sent(fid) # before the send: the link it is queued on + if self.client.push_fragment_parts(*head, parts) is False: + # Nothing was queued: the link is down (an outage, a syncer + # restart). The syncer never sees this delta. + raise ConnectionError("syncer link down, nothing queued") self.stats["pushed"] += 1 self.stats["push_bytes"] += nbytes - except (ConnectionError, OSError) as exc: + except undelivered as exc: # the anchor is not spent self.stats["dropped_pushes"] += 1 self.on_drop(fid, head[3], head[1]) log.warning("push of fragment %d dropped (%s); the syncer re-requests it", fid, exc)