Skip to main content

ipfrs_tensorlogic/gradient/
coordination.rs

1//! Gradient coordination protocol for distributed backward passes.
2//!
3//! This module provides content-addressed Arrow IPC gradient storage and a
4//! `BackwardPassCoordinator` that tracks multi-peer gradient contributions for
5//! federated backward pass rounds.
6
7use std::collections::HashMap;
8
9// ── BackwardPassId ─────────────────────────────────────────────────────────
10
11/// Identifies a backward pass round across distributed nodes.
12#[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// ── CoordinationStatus ─────────────────────────────────────────────────────
32
33/// Status of a coordinated backward pass.
34#[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// ── GradientContribution ───────────────────────────────────────────────────
60
61/// A gradient contribution from one participant peer.
62#[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// ── CoordinationError ──────────────────────────────────────────────────────
73
74#[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
95// ── BackwardPassCoordinator ────────────────────────────────────────────────
96
97/// Coordinates gradient collection and aggregation across peers.
98pub struct BackwardPassCoordinator {
99    /// pass_id_key → (status, contributions)
100    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    /// Start a new backward pass round. Returns `false` if `pass_id` already exists.
115    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    /// Submit a gradient contribution for a pass.
126    ///
127    /// Returns an error if the pass is not found or is already Complete/Failed.
128    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    /// Returns `true` when all expected peers have submitted contributions.
155    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    /// Mark aggregation as complete with result CID.
165    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    /// Mark pass as failed with a reason string.
183    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    /// Get current status of a pass.
197    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    /// Get all contributions for a pass.
204    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    /// Total active (non-complete, non-failed) passes.
214    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    /// Remove completed/failed passes with `pass_id.round < round`.
223    /// Returns the number of passes removed.
224    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            // Extract round from key: "{round}-{initiator_peer}"
232            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    /// Expose timeout setting (useful for callers).
242    #[inline]
243    pub fn timeout_ms(&self) -> u64 {
244        self.timeout_ms
245    }
246}
247
248// ── ArrowBlockError ────────────────────────────────────────────────────────
249
250#[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
276// ── GradientArrowBlock ─────────────────────────────────────────────────────
277
278const MAGIC: &[u8; 4] = b"GARW";
279
280/// A gradient stored as an Arrow IPC block with content-addressing.
281pub 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    /// Serialize gradient values to custom Arrow-flavoured IPC bytes.
291    ///
292    /// Layout:
293    /// ```text
294    /// 4  bytes  magic "GARW"
295    /// 4  bytes  shape_len  (u32 LE)
296    /// 4n bytes  shape      (u32 LE each)
297    /// 4  bytes  values_len (u32 LE)
298    /// 4m bytes  values     (f32 LE each)
299    /// 4  bytes  num_samples (u32 LE)
300    /// ```
301    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    /// Deserialize from bytes produced by [`Self::to_arrow_bytes`].
324    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        // Magic
332        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        // shape_len
341        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        // shape
352        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        // values_len
367        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        // values
378        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        // num_samples
393        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    /// Compute a deterministic CID string: `"grad-" + hex(FNV-1a hash of serialized bytes)`.
412    pub fn compute_cid(layer_name: &str, values: &[f32], num_samples: usize) -> String {
413        // Build a temporary block for hashing purposes
414        let tmp = Self {
415            cid: String::new(),
416            layer_name: layer_name.to_string(),
417            // Derive shape from values length: single-dimensional
418            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    /// FedAvg: weighted average of multiple gradient blocks.
428    ///
429    /// Weights are proportional to `num_samples`. All blocks must have the same
430    /// shape. Returns `ArrowBlockError::EmptyInput` for empty input and
431    /// `ArrowBlockError::ShapeMismatch` if shapes differ.
432    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            // Uniform average fallback when all num_samples == 0
453            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
473// ── FNV-1a hash (pure Rust, no external deps) ─────────────────────────────
474
475fn 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// ── Tests ──────────────────────────────────────────────────────────────────
487
488#[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    // ── BackwardPassCoordinator tests ──────────────────────────────────────
508
509    #[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        // Round 1 — complete
613        let p1 = make_pass_id(1);
614        coord.start_pass(p1.clone());
615        coord.mark_complete(&p1, "cid-1");
616
617        // Round 5 — complete
618        let p5 = make_pass_id(5);
619        coord.start_pass(p5.clone());
620        coord.mark_complete(&p5, "cid-5");
621
622        // Round 10 — active
623        let p10 = make_pass_id(10);
624        coord.start_pass(p10.clone());
625
626        // GC passes with round < 5 (i.e. round 1 should be removed)
627        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    // ── GradientArrowBlock tests ───────────────────────────────────────────
657
658    #[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        // Block A: values [0.0, 0.0], 100 samples
685        // Block B: values [2.0, 4.0], 300 samples
686        // Expected weighted avg: 0*0.25 + 2*0.75 = 1.5 and 0*0.25 + 4*0.75 = 3.0
687        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'; // corrupt magic
745        let result = GradientArrowBlock::from_arrow_bytes("x".to_string(), "l".to_string(), &bad);
746        assert!(matches!(result, Err(ArrowBlockError::InvalidMagic)));
747    }
748}