1use std::collections::HashMap;
8
9#[derive(Debug, Clone, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
13pub struct BackwardPassId {
14 pub round: u64,
15 pub initiator_peer: String,
16}
17
18impl BackwardPassId {
19 pub fn new(round: u64, initiator_peer: impl Into<String>) -> Self {
20 Self {
21 round,
22 initiator_peer: initiator_peer.into(),
23 }
24 }
25
26 pub fn as_key(&self) -> String {
27 format!("{}-{}", self.round, self.initiator_peer)
28 }
29}
30
31#[derive(Debug, Clone, PartialEq)]
35pub enum CoordinationStatus {
36 Pending,
37 CollectingGradients { received: usize, expected: usize },
38 Aggregating,
39 Complete { aggregated_cid: String },
40 Failed { reason: String },
41}
42
43impl CoordinationStatus {
44 fn is_terminal(&self) -> bool {
45 matches!(self, Self::Complete { .. } | Self::Failed { .. })
46 }
47
48 fn display_name(&self) -> &'static str {
49 match self {
50 Self::Pending => "Pending",
51 Self::CollectingGradients { .. } => "CollectingGradients",
52 Self::Aggregating => "Aggregating",
53 Self::Complete { .. } => "Complete",
54 Self::Failed { .. } => "Failed",
55 }
56 }
57}
58
59#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
63pub struct GradientContribution {
64 pub pass_id: BackwardPassId,
65 pub peer_id: String,
66 pub gradient_cid: String,
67 pub layer_name: String,
68 pub num_samples: usize,
69 pub submitted_at_ms: u64,
70}
71
72#[derive(Debug)]
75pub enum CoordinationError {
76 PassNotFound { pass_id: String },
77 PassAlreadyTerminal { status: String },
78}
79
80impl std::fmt::Display for CoordinationError {
81 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
82 match self {
83 Self::PassNotFound { pass_id } => {
84 write!(f, "Backward pass not found: {pass_id}")
85 }
86 Self::PassAlreadyTerminal { status } => {
87 write!(f, "Backward pass is already terminal (status: {status})")
88 }
89 }
90 }
91}
92
93impl std::error::Error for CoordinationError {}
94
95pub struct BackwardPassCoordinator {
99 passes: parking_lot::RwLock<HashMap<String, (CoordinationStatus, Vec<GradientContribution>)>>,
101 expected_peers: usize,
102 timeout_ms: u64,
103}
104
105impl BackwardPassCoordinator {
106 pub fn new(expected_peers: usize, timeout_ms: u64) -> Self {
107 Self {
108 passes: parking_lot::RwLock::new(HashMap::new()),
109 expected_peers,
110 timeout_ms,
111 }
112 }
113
114 pub fn start_pass(&self, pass_id: BackwardPassId) -> bool {
116 let key = pass_id.as_key();
117 let mut passes = self.passes.write();
118 if passes.contains_key(&key) {
119 return false;
120 }
121 passes.insert(key, (CoordinationStatus::Pending, Vec::new()));
122 true
123 }
124
125 pub fn submit_contribution(
129 &self,
130 contribution: GradientContribution,
131 ) -> Result<CoordinationStatus, CoordinationError> {
132 let key = contribution.pass_id.as_key();
133 let mut passes = self.passes.write();
134 let entry = passes
135 .get_mut(&key)
136 .ok_or_else(|| CoordinationError::PassNotFound {
137 pass_id: key.clone(),
138 })?;
139
140 let (status, contributions) = entry;
141 if status.is_terminal() {
142 return Err(CoordinationError::PassAlreadyTerminal {
143 status: status.display_name().to_string(),
144 });
145 }
146
147 contributions.push(contribution);
148 let received = contributions.len();
149 let expected = self.expected_peers;
150 *status = CoordinationStatus::CollectingGradients { received, expected };
151 Ok(status.clone())
152 }
153
154 pub fn is_ready_to_aggregate(&self, pass_id: &BackwardPassId) -> bool {
156 let key = pass_id.as_key();
157 let passes = self.passes.read();
158 passes
159 .get(&key)
160 .map(|(_, contributions)| contributions.len() >= self.expected_peers)
161 .unwrap_or(false)
162 }
163
164 pub fn mark_complete(
166 &self,
167 pass_id: &BackwardPassId,
168 aggregated_cid: impl Into<String>,
169 ) -> bool {
170 let key = pass_id.as_key();
171 let mut passes = self.passes.write();
172 if let Some((status, _)) = passes.get_mut(&key) {
173 *status = CoordinationStatus::Complete {
174 aggregated_cid: aggregated_cid.into(),
175 };
176 true
177 } else {
178 false
179 }
180 }
181
182 pub fn mark_failed(&self, pass_id: &BackwardPassId, reason: impl Into<String>) -> bool {
184 let key = pass_id.as_key();
185 let mut passes = self.passes.write();
186 if let Some((status, _)) = passes.get_mut(&key) {
187 *status = CoordinationStatus::Failed {
188 reason: reason.into(),
189 };
190 true
191 } else {
192 false
193 }
194 }
195
196 pub fn status(&self, pass_id: &BackwardPassId) -> Option<CoordinationStatus> {
198 let key = pass_id.as_key();
199 let passes = self.passes.read();
200 passes.get(&key).map(|(status, _)| status.clone())
201 }
202
203 pub fn contributions(&self, pass_id: &BackwardPassId) -> Vec<GradientContribution> {
205 let key = pass_id.as_key();
206 let passes = self.passes.read();
207 passes
208 .get(&key)
209 .map(|(_, contributions)| contributions.clone())
210 .unwrap_or_default()
211 }
212
213 pub fn active_count(&self) -> usize {
215 let passes = self.passes.read();
216 passes
217 .values()
218 .filter(|(status, _)| !status.is_terminal())
219 .count()
220 }
221
222 pub fn gc_before_round(&self, round: u64) -> usize {
225 let mut passes = self.passes.write();
226 let before = passes.len();
227 passes.retain(|key, (status, _)| {
228 if !status.is_terminal() {
229 return true;
230 }
231 let round_num: Option<u64> = key.split_once('-').and_then(|(r, _)| r.parse().ok());
233 match round_num {
234 Some(r) => r >= round,
235 None => true,
236 }
237 });
238 before - passes.len()
239 }
240
241 #[inline]
243 pub fn timeout_ms(&self) -> u64 {
244 self.timeout_ms
245 }
246}
247
248#[derive(Debug)]
251pub enum ArrowBlockError {
252 EmptyInput,
253 ShapeMismatch {
254 expected: Vec<usize>,
255 got: Vec<usize>,
256 },
257 InvalidMagic,
258 TruncatedData,
259}
260
261impl std::fmt::Display for ArrowBlockError {
262 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
263 match self {
264 Self::EmptyInput => write!(f, "No gradient blocks provided"),
265 Self::ShapeMismatch { expected, got } => {
266 write!(f, "Shape mismatch: expected {expected:?}, got {got:?}")
267 }
268 Self::InvalidMagic => write!(f, "Invalid magic bytes in Arrow block"),
269 Self::TruncatedData => write!(f, "Truncated data in Arrow block"),
270 }
271 }
272}
273
274impl std::error::Error for ArrowBlockError {}
275
276const MAGIC: &[u8; 4] = b"GARW";
279
280pub struct GradientArrowBlock {
282 pub cid: String,
283 pub layer_name: String,
284 pub shape: Vec<usize>,
285 pub values: Vec<f32>,
286 pub num_samples: usize,
287}
288
289impl GradientArrowBlock {
290 pub fn to_arrow_bytes(&self) -> Vec<u8> {
302 let shape_len = self.shape.len() as u32;
303 let values_len = self.values.len() as u32;
304 let cap = 4 + 4 + 4 * self.shape.len() + 4 + 4 * self.values.len() + 4;
305 let mut buf = Vec::with_capacity(cap);
306
307 buf.extend_from_slice(MAGIC);
308
309 buf.extend_from_slice(&shape_len.to_le_bytes());
310 for &dim in &self.shape {
311 buf.extend_from_slice(&(dim as u32).to_le_bytes());
312 }
313
314 buf.extend_from_slice(&values_len.to_le_bytes());
315 for &v in &self.values {
316 buf.extend_from_slice(&v.to_le_bytes());
317 }
318
319 buf.extend_from_slice(&(self.num_samples as u32).to_le_bytes());
320 buf
321 }
322
323 pub fn from_arrow_bytes(
325 cid: String,
326 layer_name: String,
327 data: &[u8],
328 ) -> Result<Self, ArrowBlockError> {
329 let mut pos = 0;
330
331 if data.len() < 4 {
333 return Err(ArrowBlockError::TruncatedData);
334 }
335 if &data[pos..pos + 4] != MAGIC {
336 return Err(ArrowBlockError::InvalidMagic);
337 }
338 pos += 4;
339
340 if data.len() < pos + 4 {
342 return Err(ArrowBlockError::TruncatedData);
343 }
344 let shape_len = u32::from_le_bytes(
345 data[pos..pos + 4]
346 .try_into()
347 .map_err(|_| ArrowBlockError::TruncatedData)?,
348 ) as usize;
349 pos += 4;
350
351 if data.len() < pos + 4 * shape_len {
353 return Err(ArrowBlockError::TruncatedData);
354 }
355 let mut shape = Vec::with_capacity(shape_len);
356 for _ in 0..shape_len {
357 let dim = u32::from_le_bytes(
358 data[pos..pos + 4]
359 .try_into()
360 .map_err(|_| ArrowBlockError::TruncatedData)?,
361 ) as usize;
362 shape.push(dim);
363 pos += 4;
364 }
365
366 if data.len() < pos + 4 {
368 return Err(ArrowBlockError::TruncatedData);
369 }
370 let values_len = u32::from_le_bytes(
371 data[pos..pos + 4]
372 .try_into()
373 .map_err(|_| ArrowBlockError::TruncatedData)?,
374 ) as usize;
375 pos += 4;
376
377 if data.len() < pos + 4 * values_len {
379 return Err(ArrowBlockError::TruncatedData);
380 }
381 let mut values = Vec::with_capacity(values_len);
382 for _ in 0..values_len {
383 let v = f32::from_le_bytes(
384 data[pos..pos + 4]
385 .try_into()
386 .map_err(|_| ArrowBlockError::TruncatedData)?,
387 );
388 values.push(v);
389 pos += 4;
390 }
391
392 if data.len() < pos + 4 {
394 return Err(ArrowBlockError::TruncatedData);
395 }
396 let num_samples = u32::from_le_bytes(
397 data[pos..pos + 4]
398 .try_into()
399 .map_err(|_| ArrowBlockError::TruncatedData)?,
400 ) as usize;
401
402 Ok(Self {
403 cid,
404 layer_name,
405 shape,
406 values,
407 num_samples,
408 })
409 }
410
411 pub fn compute_cid(layer_name: &str, values: &[f32], num_samples: usize) -> String {
413 let tmp = Self {
415 cid: String::new(),
416 layer_name: layer_name.to_string(),
417 shape: vec![values.len()],
419 values: values.to_vec(),
420 num_samples,
421 };
422 let bytes = tmp.to_arrow_bytes();
423 let hash = fnv1a_hash(&bytes);
424 format!("grad-{hash:016x}")
425 }
426
427 pub fn fedavg(blocks: &[GradientArrowBlock]) -> Result<Vec<f32>, ArrowBlockError> {
433 if blocks.is_empty() {
434 return Err(ArrowBlockError::EmptyInput);
435 }
436
437 let expected_shape = &blocks[0].shape;
438 for block in blocks.iter().skip(1) {
439 if &block.shape != expected_shape {
440 return Err(ArrowBlockError::ShapeMismatch {
441 expected: expected_shape.clone(),
442 got: block.shape.clone(),
443 });
444 }
445 }
446
447 let total_samples: usize = blocks.iter().map(|b| b.num_samples).sum();
448 let dim = blocks[0].values.len();
449 let mut avg = vec![0.0f32; dim];
450
451 if total_samples == 0 {
452 let n = blocks.len() as f32;
454 for block in blocks {
455 for (a, &v) in avg.iter_mut().zip(block.values.iter()) {
456 *a += v / n;
457 }
458 }
459 } else {
460 let total = total_samples as f32;
461 for block in blocks {
462 let weight = block.num_samples as f32 / total;
463 for (a, &v) in avg.iter_mut().zip(block.values.iter()) {
464 *a += v * weight;
465 }
466 }
467 }
468
469 Ok(avg)
470 }
471}
472
473fn fnv1a_hash(data: &[u8]) -> u64 {
476 const FNV_OFFSET_BASIS: u64 = 14_695_981_039_346_656_037;
477 const FNV_PRIME: u64 = 1_099_511_628_211;
478 let mut hash = FNV_OFFSET_BASIS;
479 for &byte in data {
480 hash ^= u64::from(byte);
481 hash = hash.wrapping_mul(FNV_PRIME);
482 }
483 hash
484}
485
486#[cfg(test)]
489mod tests {
490 use super::*;
491
492 fn make_pass_id(round: u64) -> BackwardPassId {
493 BackwardPassId::new(round, format!("peer-{round}"))
494 }
495
496 fn make_contribution(pass_id: BackwardPassId, peer_id: &str) -> GradientContribution {
497 GradientContribution {
498 pass_id,
499 peer_id: peer_id.to_string(),
500 gradient_cid: format!("grad-cid-{peer_id}"),
501 layer_name: "layer0".to_string(),
502 num_samples: 100,
503 submitted_at_ms: 0,
504 }
505 }
506
507 #[test]
510 fn test_start_pass() {
511 let coord = BackwardPassCoordinator::new(2, 5000);
512 let pass_id = make_pass_id(1);
513 assert!(coord.start_pass(pass_id.clone()));
514 assert_eq!(coord.status(&pass_id), Some(CoordinationStatus::Pending));
515 }
516
517 #[test]
518 fn test_start_pass_duplicate_returns_false() {
519 let coord = BackwardPassCoordinator::new(2, 5000);
520 let pass_id = make_pass_id(1);
521 assert!(coord.start_pass(pass_id.clone()));
522 assert!(!coord.start_pass(pass_id));
523 }
524
525 #[test]
526 fn test_submit_contribution_updates_status() {
527 let coord = BackwardPassCoordinator::new(2, 5000);
528 let pass_id = make_pass_id(2);
529 coord.start_pass(pass_id.clone());
530
531 let contribution = make_contribution(pass_id.clone(), "peer-a");
532 let status = coord
533 .submit_contribution(contribution)
534 .expect("submit should succeed");
535
536 assert_eq!(
537 status,
538 CoordinationStatus::CollectingGradients {
539 received: 1,
540 expected: 2
541 }
542 );
543 }
544
545 #[test]
546 fn test_is_ready_to_aggregate_after_all_peers() {
547 let coord = BackwardPassCoordinator::new(2, 5000);
548 let pass_id = make_pass_id(3);
549 coord.start_pass(pass_id.clone());
550
551 assert!(!coord.is_ready_to_aggregate(&pass_id));
552
553 coord
554 .submit_contribution(make_contribution(pass_id.clone(), "peer-a"))
555 .expect("peer-a submit");
556 assert!(!coord.is_ready_to_aggregate(&pass_id));
557
558 coord
559 .submit_contribution(make_contribution(pass_id.clone(), "peer-b"))
560 .expect("peer-b submit");
561 assert!(coord.is_ready_to_aggregate(&pass_id));
562 }
563
564 #[test]
565 fn test_mark_complete() {
566 let coord = BackwardPassCoordinator::new(1, 5000);
567 let pass_id = make_pass_id(4);
568 coord.start_pass(pass_id.clone());
569
570 assert!(coord.mark_complete(&pass_id, "agg-cid-42"));
571 assert_eq!(
572 coord.status(&pass_id),
573 Some(CoordinationStatus::Complete {
574 aggregated_cid: "agg-cid-42".to_string()
575 })
576 );
577 }
578
579 #[test]
580 fn test_mark_failed() {
581 let coord = BackwardPassCoordinator::new(1, 5000);
582 let pass_id = make_pass_id(5);
583 coord.start_pass(pass_id.clone());
584
585 assert!(coord.mark_failed(&pass_id, "timeout"));
586 assert_eq!(
587 coord.status(&pass_id),
588 Some(CoordinationStatus::Failed {
589 reason: "timeout".to_string()
590 })
591 );
592 }
593
594 #[test]
595 fn test_submit_to_terminal_pass_returns_error() {
596 let coord = BackwardPassCoordinator::new(1, 5000);
597 let pass_id = make_pass_id(6);
598 coord.start_pass(pass_id.clone());
599 coord.mark_failed(&pass_id, "network error");
600
601 let result = coord.submit_contribution(make_contribution(pass_id, "late-peer"));
602 assert!(matches!(
603 result,
604 Err(CoordinationError::PassAlreadyTerminal { .. })
605 ));
606 }
607
608 #[test]
609 fn test_gc_before_round() {
610 let coord = BackwardPassCoordinator::new(1, 5000);
611
612 let p1 = make_pass_id(1);
614 coord.start_pass(p1.clone());
615 coord.mark_complete(&p1, "cid-1");
616
617 let p5 = make_pass_id(5);
619 coord.start_pass(p5.clone());
620 coord.mark_complete(&p5, "cid-5");
621
622 let p10 = make_pass_id(10);
624 coord.start_pass(p10.clone());
625
626 let removed = coord.gc_before_round(5);
628 assert_eq!(removed, 1, "only round-1 should be removed");
629
630 assert!(coord.status(&p1).is_none(), "round-1 gone");
631 assert!(coord.status(&p5).is_some(), "round-5 remains");
632 assert!(coord.status(&p10).is_some(), "round-10 remains (active)");
633 }
634
635 #[test]
636 fn test_active_count() {
637 let coord = BackwardPassCoordinator::new(1, 5000);
638
639 let p1 = make_pass_id(1);
640 let p2 = make_pass_id(2);
641 let p3 = make_pass_id(3);
642
643 coord.start_pass(p1.clone());
644 coord.start_pass(p2.clone());
645 coord.start_pass(p3.clone());
646
647 assert_eq!(coord.active_count(), 3);
648
649 coord.mark_complete(&p1, "cid");
650 assert_eq!(coord.active_count(), 2);
651
652 coord.mark_failed(&p2, "err");
653 assert_eq!(coord.active_count(), 1);
654 }
655
656 #[test]
659 fn test_gradient_arrow_block_roundtrip() {
660 let values = vec![1.0f32, 2.0, 3.0, 4.0];
661 let cid = GradientArrowBlock::compute_cid("layer1", &values, 50);
662
663 let block = GradientArrowBlock {
664 cid: cid.clone(),
665 layer_name: "layer1".to_string(),
666 shape: vec![4],
667 values: values.clone(),
668 num_samples: 50,
669 };
670
671 let bytes = block.to_arrow_bytes();
672 let decoded =
673 GradientArrowBlock::from_arrow_bytes(cid.clone(), "layer1".to_string(), &bytes)
674 .expect("roundtrip");
675
676 assert_eq!(decoded.shape, vec![4]);
677 assert_eq!(decoded.values, values);
678 assert_eq!(decoded.num_samples, 50);
679 assert_eq!(decoded.cid, cid);
680 }
681
682 #[test]
683 fn test_fedavg_weighted() {
684 let block_a = GradientArrowBlock {
688 cid: "a".to_string(),
689 layer_name: "l".to_string(),
690 shape: vec![2],
691 values: vec![0.0, 0.0],
692 num_samples: 100,
693 };
694 let block_b = GradientArrowBlock {
695 cid: "b".to_string(),
696 layer_name: "l".to_string(),
697 shape: vec![2],
698 values: vec![2.0, 4.0],
699 num_samples: 300,
700 };
701 let avg = GradientArrowBlock::fedavg(&[block_a, block_b]).expect("fedavg");
702 assert!((avg[0] - 1.5).abs() < 1e-5, "avg[0] = {}", avg[0]);
703 assert!((avg[1] - 3.0).abs() < 1e-5, "avg[1] = {}", avg[1]);
704 }
705
706 #[test]
707 fn test_fedavg_empty_returns_error() {
708 let result = GradientArrowBlock::fedavg(&[]);
709 assert!(matches!(result, Err(ArrowBlockError::EmptyInput)));
710 }
711
712 #[test]
713 fn test_fedavg_shape_mismatch() {
714 let block_a = GradientArrowBlock {
715 cid: "a".to_string(),
716 layer_name: "l".to_string(),
717 shape: vec![3],
718 values: vec![1.0, 2.0, 3.0],
719 num_samples: 10,
720 };
721 let block_b = GradientArrowBlock {
722 cid: "b".to_string(),
723 layer_name: "l".to_string(),
724 shape: vec![2],
725 values: vec![1.0, 2.0],
726 num_samples: 10,
727 };
728 let result = GradientArrowBlock::fedavg(&[block_a, block_b]);
729 assert!(matches!(result, Err(ArrowBlockError::ShapeMismatch { .. })));
730 }
731
732 #[test]
733 fn test_compute_cid_deterministic() {
734 let values = vec![0.5f32, 1.0, 1.5];
735 let cid1 = GradientArrowBlock::compute_cid("layer2", &values, 64);
736 let cid2 = GradientArrowBlock::compute_cid("layer2", &values, 64);
737 assert_eq!(cid1, cid2, "CIDs must be deterministic");
738 assert!(cid1.starts_with("grad-"), "CID prefix");
739 }
740
741 #[test]
742 fn test_from_arrow_bytes_bad_magic() {
743 let mut bad = vec![0u8; 20];
744 bad[0] = b'X'; let result = GradientArrowBlock::from_arrow_bytes("x".to_string(), "l".to_string(), &bad);
746 assert!(matches!(result, Err(ArrowBlockError::InvalidMagic)));
747 }
748}