Spaces:
Building
Building
Download server/src/checkpoint.rs from wuhp/webgputest: direct link, hf CLI and curl.
- Browser
- Download file 18.4 kB
-
https://huggingface.co/spaces/wuhp/webgputest/resolve/main/server/src/checkpoint.rs
- Command line
-
hf download hf://spaces/wuhp/webgputest/server/src/checkpoint.rs
-
curl -L -o checkpoint.rs https://huggingface.co/spaces/wuhp/webgputest/resolve/main/server/src/checkpoint.rs
18.4 kB
| //! Checkpoint policy and peer placement primitives. | |
| //! | |
| //! This module plans metadata only. Checkpoint bytes stay with workers and move | |
| //! directly between peers; the coordinator should retain leases and manifests. | |
| use std::collections::{HashMap, HashSet}; | |
| pub struct CheckpointInput { | |
| pub full_bytes: u64, | |
| pub changed_bytes: u64, | |
| pub has_base_checkpoint: bool, | |
| pub observed_transfer_ms: f64, | |
| pub estimated_mttf_seconds: f64, | |
| pub work_per_second: f64, | |
| pub min_interval_seconds: f64, | |
| pub max_interval_seconds: f64, | |
| } | |
| pub enum StorageStrategy { | |
| Replication2x, | |
| Erasure2Plus1, | |
| Erasure2Plus2, | |
| } | |
| impl StorageStrategy { | |
| pub fn data_shards(self) -> usize { | |
| match self { Self::Replication2x => 1, Self::Erasure2Plus1 | Self::Erasure2Plus2 => 2 } | |
| } | |
| pub fn total_shards(self) -> usize { | |
| match self { Self::Replication2x => 2, Self::Erasure2Plus1 => 3, Self::Erasure2Plus2 => 4 } | |
| } | |
| pub fn tolerable_losses(self) -> usize { | |
| match self { Self::Replication2x | Self::Erasure2Plus1 => 1, Self::Erasure2Plus2 => 2 } | |
| } | |
| } | |
| pub struct WorkerLink { | |
| pub node_id: String, | |
| pub failure_domain: String, | |
| pub bandwidth_mbps: u32, | |
| pub rtt_ms: u32, | |
| } | |
| pub struct CheckpointPlan { | |
| pub interval_seconds: f64, | |
| pub interval_work_units: f64, | |
| pub payload_bytes: u64, | |
| pub is_delta: bool, | |
| pub strategy: StorageStrategy, | |
| } | |
| pub struct ArtifactShard { | |
| pub index: usize, | |
| pub node_id: String, | |
| pub failure_domain: String, | |
| pub bytes: u64, | |
| } | |
| pub struct ArtifactManifest { | |
| pub checkpoint_id: String, | |
| pub generation: u64, | |
| pub base_checkpoint_id: Option<String>, | |
| pub logical_bytes: u64, | |
| pub payload_bytes: u64, | |
| pub strategy: StorageStrategy, | |
| pub shards: Vec<ArtifactShard>, | |
| } | |
| pub struct ShardReceipt { | |
| pub checkpoint_id: String, | |
| pub shard_index: usize, | |
| pub node_id: String, | |
| pub sha256: String, | |
| } | |
| pub enum CheckpointError { | |
| InvalidManifest, | |
| StaleGeneration, | |
| MissingBase, | |
| UnknownTransaction, | |
| UnknownShard, | |
| WrongNode, | |
| HashMismatch, | |
| NotDurable, | |
| DuplicateCheckpoint, | |
| BrokenChain, | |
| } | |
| struct PendingCheckpoint { | |
| manifest: ArtifactManifest, | |
| expected_hashes: Vec<String>, | |
| receipts: HashSet<usize>, | |
| } | |
| /// Metadata-only two-phase checkpoint commit coordinator. | |
| /// | |
| /// Workers stage and hash their own shards, then submit receipts. The catalog | |
| /// publishes a new head only once enough distinct shards are verified to meet | |
| /// the strategy's reconstruction threshold. This type never handles artifact | |
| /// bytes and is intended to sit beside a direct peer transport. | |
| pub struct CheckpointCatalog { | |
| current_id: Option<String>, | |
| committed: HashMap<String, ArtifactManifest>, | |
| pending: HashMap<String, PendingCheckpoint>, | |
| } | |
| impl CheckpointCatalog { | |
| pub fn current(&self) -> Option<&ArtifactManifest> { | |
| self.current_id.as_ref().and_then(|id| self.committed.get(id)) | |
| } | |
| /// Start staging a manifest. `expected_hashes` has one SHA-256 digest per | |
| /// manifest shard, in shard-index order. A delta must extend the current | |
| /// head; a full checkpoint (no base ID) can compact the chain. | |
| pub fn begin( | |
| &mut self, | |
| manifest: ArtifactManifest, | |
| expected_hashes: Vec<String>, | |
| ) -> Result<(), CheckpointError> { | |
| if self.pending.contains_key(&manifest.checkpoint_id) || self.committed.contains_key(&manifest.checkpoint_id) { | |
| return Err(CheckpointError::DuplicateCheckpoint); | |
| } | |
| if manifest.shards.len() != manifest.strategy.total_shards() | |
| || expected_hashes.len() != manifest.shards.len() | |
| || expected_hashes.iter().any(|hash| !is_sha256(hash)) | |
| || manifest.shards.iter().enumerate().any(|(index, shard)| shard.index != index) | |
| || manifest.shards.iter().any(|shard| shard.node_id.is_empty() || shard.bytes == 0) | |
| || manifest.shards.iter().map(|shard| shard.node_id.as_str()).collect::<HashSet<_>>().len() != manifest.shards.len() | |
| { | |
| return Err(CheckpointError::InvalidManifest); | |
| } | |
| let expected_generation = self.current().map_or(1, |current| current.generation.saturating_add(1)); | |
| if manifest.generation != expected_generation { return Err(CheckpointError::StaleGeneration); } | |
| match (&self.current_id, &manifest.base_checkpoint_id) { | |
| (None, None) => {}, | |
| (Some(current), Some(base)) if current == base => {}, | |
| (Some(_), None) => {}, // full compaction | |
| _ => return Err(CheckpointError::MissingBase), | |
| } | |
| self.pending.insert(manifest.checkpoint_id.clone(), PendingCheckpoint { | |
| manifest, | |
| expected_hashes, | |
| receipts: HashSet::new(), | |
| }); | |
| Ok(()) | |
| } | |
| /// Accept a receipt only from the peer assigned to that shard, and only | |
| /// when its digest matches the expected digest from the job manifest. | |
| pub fn acknowledge(&mut self, receipt: ShardReceipt) -> Result<(), CheckpointError> { | |
| let pending = self.pending.get_mut(&receipt.checkpoint_id).ok_or(CheckpointError::UnknownTransaction)?; | |
| let shard = pending.manifest.shards.get(receipt.shard_index).ok_or(CheckpointError::UnknownShard)?; | |
| if shard.node_id != receipt.node_id { return Err(CheckpointError::WrongNode); } | |
| if !is_sha256(&receipt.sha256) || pending.expected_hashes[receipt.shard_index] != receipt.sha256 { | |
| return Err(CheckpointError::HashMismatch); | |
| } | |
| pending.receipts.insert(receipt.shard_index); | |
| Ok(()) | |
| } | |
| /// Publish the staged manifest when enough receipts exist to reconstruct | |
| /// it. Until then `current()` continues returning the previous good head. | |
| pub fn commit(&mut self, checkpoint_id: &str) -> Result<&ArtifactManifest, CheckpointError> { | |
| let pending = self.pending.get(checkpoint_id).ok_or(CheckpointError::UnknownTransaction)?; | |
| // Require the complete desired replica/parity set before publishing. | |
| // A later repair process may support degraded commits explicitly. | |
| if pending.receipts.len() < pending.manifest.strategy.total_shards() { | |
| return Err(CheckpointError::NotDurable); | |
| } | |
| let pending = self.pending.remove(checkpoint_id).expect("pending checkpoint was checked"); | |
| self.current_id = Some(checkpoint_id.to_owned()); | |
| self.committed.insert(checkpoint_id.to_owned(), pending.manifest); | |
| Ok(self.committed.get(checkpoint_id).expect("checkpoint was inserted")) | |
| } | |
| /// Drop an incomplete transaction without changing the published head. | |
| pub fn abort(&mut self, checkpoint_id: &str) -> bool { | |
| self.pending.remove(checkpoint_id).is_some() | |
| } | |
| /// Return the committed full base followed by each delta through the head. | |
| /// A missing parent means the artifact chain is incomplete and must not be | |
| /// presented as recoverable. | |
| pub fn restore_chain(&self) -> Result<Vec<&ArtifactManifest>, CheckpointError> { | |
| let mut chain = Vec::new(); | |
| let mut cursor = self.current().ok_or(CheckpointError::BrokenChain)?; | |
| loop { | |
| chain.push(cursor); | |
| match cursor.base_checkpoint_id.as_deref() { | |
| Some(parent_id) => cursor = self.committed.get(parent_id).ok_or(CheckpointError::BrokenChain)?, | |
| None => break, | |
| } | |
| } | |
| chain.reverse(); | |
| Ok(chain) | |
| } | |
| } | |
| fn is_sha256(value: &str) -> bool { | |
| value.len() == 64 && value.bytes().all(|byte| byte.is_ascii_hexdigit()) | |
| } | |
| /// Young/Daly-style interval estimate using observed checkpoint wall time and | |
| /// estimated mean time to failure, bounded for product-level safety limits. | |
| pub fn adaptive_interval(input: &CheckpointInput) -> f64 { | |
| let min = input.min_interval_seconds.max(0.1); | |
| let max = input.max_interval_seconds.max(min); | |
| let checkpoint_cost = (input.observed_transfer_ms / 1_000.0).max(0.001); | |
| let mttf = input.estimated_mttf_seconds.max(0.1); | |
| let estimate = (2.0 * checkpoint_cost * mttf).sqrt(); | |
| estimate.clamp(min, max) | |
| } | |
| /// Delta transfer is allowed only when a compatible base exists. Periodic full | |
| /// compaction can be expressed by passing `has_base_checkpoint = false`. | |
| pub fn payload_size(input: &CheckpointInput) -> (u64, bool) { | |
| if input.has_base_checkpoint && input.changed_bytes < input.full_bytes { | |
| (input.changed_bytes.max(1), true) | |
| } else { | |
| (input.full_bytes, false) | |
| } | |
| } | |
| /// Keep small/frequent artifacts simple; use 2+1 for large artifacts on slower | |
| /// links; use 2+2 only when the workload explicitly requires two-loss tolerance. | |
| pub fn choose_storage_strategy( | |
| payload_bytes: u64, | |
| interval_seconds: f64, | |
| bandwidth_mbps: u32, | |
| small_artifact_limit_bytes: u64, | |
| frequent_interval_seconds: f64, | |
| required_peer_loss_tolerance: usize, | |
| ) -> StorageStrategy { | |
| if required_peer_loss_tolerance >= 2 { return StorageStrategy::Erasure2Plus2; } | |
| if payload_bytes <= small_artifact_limit_bytes || interval_seconds <= frequent_interval_seconds { | |
| return StorageStrategy::Replication2x; | |
| } | |
| if bandwidth_mbps < 100 { StorageStrategy::Erasure2Plus1 } else { StorageStrategy::Replication2x } | |
| } | |
| /// Prefer high-throughput, low-latency peers while placing one shard per | |
| /// failure domain before reusing a domain. This helps avoid correlated loss. | |
| pub fn place_shards(workers: &[WorkerLink], count: usize) -> Vec<WorkerLink> { | |
| let mut ranked = workers.to_vec(); | |
| ranked.sort_by(|a, b| { | |
| let a_score = f64::from(a.bandwidth_mbps) / (1.0 + f64::from(a.rtt_ms)); | |
| let b_score = f64::from(b.bandwidth_mbps) / (1.0 + f64::from(b.rtt_ms)); | |
| b_score.total_cmp(&a_score).then_with(|| a.node_id.cmp(&b.node_id)) | |
| }); | |
| let mut selected = Vec::with_capacity(count.min(ranked.len())); | |
| for worker in &ranked { | |
| if selected.len() == count { return selected; } | |
| if !selected.iter().any(|picked: &WorkerLink| picked.failure_domain == worker.failure_domain) { | |
| selected.push(worker.clone()); | |
| } | |
| } | |
| for worker in ranked { | |
| if selected.len() == count { break; } | |
| if !selected.iter().any(|picked| picked.node_id == worker.node_id) { selected.push(worker); } | |
| } | |
| selected | |
| } | |
| pub fn build_manifest( | |
| checkpoint_id: String, | |
| generation: u64, | |
| base_checkpoint_id: Option<String>, | |
| logical_bytes: u64, | |
| payload_bytes: u64, | |
| strategy: StorageStrategy, | |
| workers: &[WorkerLink], | |
| ) -> Option<ArtifactManifest> { | |
| let selected = place_shards(workers, strategy.total_shards()); | |
| if selected.len() != strategy.total_shards() { return None; } | |
| let shard_bytes = payload_bytes.div_ceil(strategy.data_shards() as u64); | |
| let shards = selected.into_iter().enumerate().map(|(index, worker)| ArtifactShard { | |
| index, | |
| node_id: worker.node_id, | |
| failure_domain: worker.failure_domain, | |
| bytes: shard_bytes, | |
| }).collect(); | |
| Some(ArtifactManifest { checkpoint_id, generation, base_checkpoint_id, logical_bytes, payload_bytes, strategy, shards }) | |
| } | |
| mod tests { | |
| use super::*; | |
| fn input(mttf: f64, transfer_ms: f64) -> CheckpointInput { | |
| CheckpointInput { full_bytes: 100_000, changed_bytes: 10_000, has_base_checkpoint: true, | |
| observed_transfer_ms: transfer_ms, estimated_mttf_seconds: mttf, work_per_second: 1.0, | |
| min_interval_seconds: 1.0, max_interval_seconds: 120.0 } | |
| } | |
| fn unreliable_workers_get_shorter_intervals() { | |
| assert!(adaptive_interval(&input(30.0, 1_000.0)) < adaptive_interval(&input(3_600.0, 1_000.0))); | |
| } | |
| fn deltas_need_a_base_and_compaction_restores_full_snapshot() { | |
| assert_eq!(payload_size(&input(60.0, 100.0)), (10_000, true)); | |
| let mut no_base = input(60.0, 100.0); | |
| no_base.has_base_checkpoint = false; | |
| assert_eq!(payload_size(&no_base), (100_000, false)); | |
| } | |
| fn storage_policy_matches_size_frequency_and_tolerance() { | |
| assert_eq!(choose_storage_strategy(8_000, 60.0, 20, 16_000, 5.0, 0), StorageStrategy::Replication2x); | |
| assert_eq!(choose_storage_strategy(100_000, 60.0, 20, 16_000, 5.0, 0), StorageStrategy::Erasure2Plus1); | |
| assert_eq!(choose_storage_strategy(8_000, 1.0, 20, 16_000, 5.0, 2), StorageStrategy::Erasure2Plus2); | |
| } | |
| fn placement_prefers_speed_but_spreads_failure_domains() { | |
| let workers = vec![ | |
| WorkerLink { node_id: "fast-a".into(), failure_domain: "home-a".into(), bandwidth_mbps: 900, rtt_ms: 5 }, | |
| WorkerLink { node_id: "fast-b".into(), failure_domain: "home-a".into(), bandwidth_mbps: 800, rtt_ms: 6 }, | |
| WorkerLink { node_id: "medium-c".into(), failure_domain: "home-c".into(), bandwidth_mbps: 300, rtt_ms: 20 }, | |
| WorkerLink { node_id: "slow-d".into(), failure_domain: "home-d".into(), bandwidth_mbps: 100, rtt_ms: 40 }, | |
| ]; | |
| let picked = place_shards(&workers, 3); | |
| assert_eq!(picked[0].node_id, "fast-a"); | |
| assert_eq!(picked.iter().map(|w| &w.failure_domain).collect::<std::collections::HashSet<_>>().len(), 3); | |
| } | |
| fn manifest_contains_metadata_only_and_requires_enough_hosts() { | |
| let workers = vec![WorkerLink { node_id: "a".into(), failure_domain: "one".into(), bandwidth_mbps: 100, rtt_ms: 10 }]; | |
| assert!(build_manifest("cp-1".into(), 1, None, 100, 100, StorageStrategy::Erasure2Plus1, &workers).is_none()); | |
| } | |
| fn peers() -> Vec<WorkerLink> { | |
| vec![ | |
| WorkerLink { node_id: "node-a".into(), failure_domain: "home-a".into(), bandwidth_mbps: 500, rtt_ms: 10 }, | |
| WorkerLink { node_id: "node-b".into(), failure_domain: "home-b".into(), bandwidth_mbps: 400, rtt_ms: 20 }, | |
| WorkerLink { node_id: "node-c".into(), failure_domain: "home-c".into(), bandwidth_mbps: 300, rtt_ms: 30 }, | |
| WorkerLink { node_id: "node-d".into(), failure_domain: "home-d".into(), bandwidth_mbps: 200, rtt_ms: 40 }, | |
| ] | |
| } | |
| fn manifest(id: &str, generation: u64, base: Option<&str>) -> ArtifactManifest { | |
| build_manifest(id.into(), generation, base.map(str::to_owned), 100, 10, | |
| StorageStrategy::Erasure2Plus1, &peers()).unwrap() | |
| } | |
| fn acknowledge(catalog: &mut CheckpointCatalog, id: &str, index: usize, hash: &str) { | |
| let pending = catalog.pending.get(id).unwrap(); | |
| catalog.acknowledge(ShardReceipt { | |
| checkpoint_id: id.into(), | |
| shard_index: index, | |
| node_id: pending.manifest.shards[index].node_id.clone(), | |
| sha256: hash.into(), | |
| }).unwrap(); | |
| } | |
| fn checkpoint_head_moves_only_after_verified_durability_threshold() { | |
| let mut catalog = CheckpointCatalog::default(); | |
| let hashes = vec!["a".repeat(64), "b".repeat(64), "c".repeat(64)]; | |
| catalog.begin(manifest("base", 1, None), hashes.clone()).unwrap(); | |
| acknowledge(&mut catalog, "base", 0, &hashes[0]); | |
| acknowledge(&mut catalog, "base", 1, &hashes[1]); | |
| assert_eq!(catalog.commit("base").unwrap_err(), CheckpointError::NotDurable); | |
| acknowledge(&mut catalog, "base", 2, &hashes[2]); | |
| catalog.commit("base").unwrap(); | |
| catalog.begin(manifest("delta-1", 2, Some("base")), hashes.clone()).unwrap(); | |
| acknowledge(&mut catalog, "delta-1", 0, &hashes[0]); | |
| acknowledge(&mut catalog, "delta-1", 1, &hashes[1]); | |
| assert_eq!(catalog.commit("delta-1").unwrap_err(), CheckpointError::NotDurable); | |
| assert_eq!(catalog.current().unwrap().checkpoint_id, "base"); | |
| acknowledge(&mut catalog, "delta-1", 2, &hashes[2]); | |
| catalog.commit("delta-1").unwrap(); | |
| let chain = catalog.restore_chain().unwrap(); | |
| assert_eq!(chain.iter().map(|entry| entry.checkpoint_id.as_str()).collect::<Vec<_>>(), vec!["base", "delta-1"]); | |
| } | |
| fn rejects_unassigned_or_corrupt_shard_receipts() { | |
| let mut catalog = CheckpointCatalog::default(); | |
| let hashes = vec!["a".repeat(64), "b".repeat(64), "c".repeat(64)]; | |
| catalog.begin(manifest("base", 1, None), hashes.clone()).unwrap(); | |
| assert_eq!(catalog.acknowledge(ShardReceipt { | |
| checkpoint_id: "base".into(), shard_index: 0, node_id: "attacker".into(), sha256: hashes[0].clone(), | |
| }), Err(CheckpointError::WrongNode)); | |
| let assigned_node = catalog.pending["base"].manifest.shards[0].node_id.clone(); | |
| assert_eq!(catalog.acknowledge(ShardReceipt { | |
| checkpoint_id: "base".into(), shard_index: 0, node_id: assigned_node, sha256: "d".repeat(64), | |
| }), Err(CheckpointError::HashMismatch)); | |
| assert!(catalog.current().is_none()); | |
| } | |
| fn abort_and_full_compaction_preserve_a_recoverable_chain() { | |
| let mut catalog = CheckpointCatalog::default(); | |
| let hashes = vec!["a".repeat(64), "b".repeat(64), "c".repeat(64)]; | |
| catalog.begin(manifest("base", 1, None), hashes.clone()).unwrap(); | |
| acknowledge(&mut catalog, "base", 0, &hashes[0]); | |
| acknowledge(&mut catalog, "base", 1, &hashes[1]); | |
| acknowledge(&mut catalog, "base", 2, &hashes[2]); | |
| catalog.commit("base").unwrap(); | |
| catalog.begin(manifest("failed-delta", 2, Some("base")), hashes.clone()).unwrap(); | |
| assert!(catalog.abort("failed-delta")); | |
| assert_eq!(catalog.current().unwrap().checkpoint_id, "base"); | |
| catalog.begin(manifest("compacted", 2, None), hashes.clone()).unwrap(); | |
| acknowledge(&mut catalog, "compacted", 0, &hashes[0]); | |
| acknowledge(&mut catalog, "compacted", 1, &hashes[1]); | |
| acknowledge(&mut catalog, "compacted", 2, &hashes[2]); | |
| catalog.commit("compacted").unwrap(); | |
| let chain = catalog.restore_chain().unwrap(); | |
| assert_eq!(chain.iter().map(|entry| entry.checkpoint_id.as_str()).collect::<Vec<_>>(), vec!["compacted"]); | |
| } | |
| } | |