ipfrs_tensorlogic/
checkpointer.rs

1//! TensorCheckpointer — periodic checkpointing of tensor computation state with rollback.
2//!
3//! Provides a configurable checkpoint manager that snapshots tensor values at
4//! regular intervals (or on demand), maintains a bounded rolling window of
5//! checkpoints, and supports rollback to any previously saved state.
6
7use std::collections::{HashMap, VecDeque};
8
9// ---------------------------------------------------------------------------
10// Configuration
11// ---------------------------------------------------------------------------
12
13/// Configuration for the [`TensorCheckpointer`].
14#[derive(Debug, Clone)]
15pub struct CheckpointConfig {
16    /// Maximum number of checkpoints to retain (FIFO eviction).
17    pub max_checkpoints: usize,
18    /// Number of ticks between automatic checkpoint triggers.
19    pub auto_checkpoint_interval: u64,
20}
21
22impl Default for CheckpointConfig {
23    fn default() -> Self {
24        Self {
25            max_checkpoints: 5,
26            auto_checkpoint_interval: 100,
27        }
28    }
29}
30
31// ---------------------------------------------------------------------------
32// Checkpoint
33// ---------------------------------------------------------------------------
34
35/// A single snapshot of tensor computation state.
36#[derive(Debug, Clone)]
37pub struct Checkpoint {
38    /// Unique, monotonically increasing identifier.
39    pub id: u64,
40    /// Tick at which this checkpoint was created.
41    pub tick: u64,
42    /// Human-readable label.
43    pub label: String,
44    /// Snapshot of tensor values — each inner `Vec<f64>` is one tensor.
45    pub tensor_data: Vec<Vec<f64>>,
46    /// Arbitrary key-value metadata attached to the checkpoint.
47    pub metadata: HashMap<String, String>,
48}
49
50// ---------------------------------------------------------------------------
51// Stats
52// ---------------------------------------------------------------------------
53
54/// Aggregate statistics about the checkpointer state.
55#[derive(Debug, Clone)]
56pub struct CheckpointerStats {
57    /// Number of checkpoints currently stored.
58    pub total_checkpoints: usize,
59    /// Tick of the oldest checkpoint, if any.
60    pub oldest_tick: Option<u64>,
61    /// Tick of the newest checkpoint, if any.
62    pub newest_tick: Option<u64>,
63    /// Total number of rollbacks performed since creation.
64    pub rollbacks_performed: u64,
65}
66
67// ---------------------------------------------------------------------------
68// TensorCheckpointer
69// ---------------------------------------------------------------------------
70
71/// Manages a bounded, rolling window of tensor computation checkpoints with
72/// rollback support.
73pub struct TensorCheckpointer {
74    config: CheckpointConfig,
75    checkpoints: VecDeque<Checkpoint>,
76    next_id: u64,
77    current_tick: u64,
78    last_checkpoint_tick: u64,
79    rollbacks_performed: u64,
80}
81
82impl TensorCheckpointer {
83    /// Create a new checkpointer with the given configuration.
84    pub fn new(config: CheckpointConfig) -> Self {
85        Self {
86            config,
87            checkpoints: VecDeque::new(),
88            next_id: 0,
89            current_tick: 0,
90            last_checkpoint_tick: 0,
91            rollbacks_performed: 0,
92        }
93    }
94
95    /// Create a checkpoint with the given label, tensor data, and metadata.
96    ///
97    /// If the number of stored checkpoints exceeds `max_checkpoints`, the
98    /// oldest checkpoint is evicted (FIFO).
99    ///
100    /// Returns the unique id of the newly created checkpoint.
101    pub fn create_checkpoint(
102        &mut self,
103        label: &str,
104        tensor_data: Vec<Vec<f64>>,
105        metadata: HashMap<String, String>,
106    ) -> u64 {
107        let id = self.next_id;
108        self.next_id += 1;
109
110        let checkpoint = Checkpoint {
111            id,
112            tick: self.current_tick,
113            label: label.to_string(),
114            tensor_data,
115            metadata,
116        };
117
118        self.checkpoints.push_back(checkpoint);
119        self.last_checkpoint_tick = self.current_tick;
120
121        // Evict oldest if over capacity.
122        while self.checkpoints.len() > self.config.max_checkpoints {
123            self.checkpoints.pop_front();
124        }
125
126        id
127    }
128
129    /// Roll back to the checkpoint with the given `checkpoint_id`.
130    ///
131    /// Returns the checkpoint data and removes all checkpoints that were
132    /// created *after* the target checkpoint.  The target checkpoint itself is
133    /// also consumed (removed) so callers can restore from the returned data.
134    ///
135    /// Increments `rollbacks_performed`.
136    pub fn rollback(&mut self, checkpoint_id: u64) -> Result<Checkpoint, String> {
137        let pos = self
138            .checkpoints
139            .iter()
140            .position(|c| c.id == checkpoint_id)
141            .ok_or_else(|| format!("checkpoint id {} not found", checkpoint_id))?;
142
143        // Remove everything after the target checkpoint.
144        self.checkpoints.truncate(pos + 1);
145
146        // Pop the target checkpoint itself so we can return it by value.
147        let checkpoint = self
148            .checkpoints
149            .pop_back()
150            .ok_or_else(|| "internal error: checkpoint disappeared".to_string())?;
151
152        self.rollbacks_performed += 1;
153
154        Ok(checkpoint)
155    }
156
157    /// Return a reference to the most recent checkpoint, if any.
158    pub fn latest_checkpoint(&self) -> Option<&Checkpoint> {
159        self.checkpoints.back()
160    }
161
162    /// Look up a checkpoint by its id.
163    pub fn get_checkpoint(&self, id: u64) -> Option<&Checkpoint> {
164        self.checkpoints.iter().find(|c| c.id == id)
165    }
166
167    /// Number of checkpoints currently stored.
168    pub fn checkpoint_count(&self) -> usize {
169        self.checkpoints.len()
170    }
171
172    /// Returns `true` if the number of ticks since the last checkpoint is
173    /// greater than or equal to `auto_checkpoint_interval`.
174    pub fn should_auto_checkpoint(&self) -> bool {
175        self.current_tick.saturating_sub(self.last_checkpoint_tick)
176            >= self.config.auto_checkpoint_interval
177    }
178
179    /// Advance the internal tick clock by one.
180    pub fn tick(&mut self) {
181        self.current_tick += 1;
182    }
183
184    /// List all checkpoints as `(id, label, tick)` triples, ordered oldest to
185    /// newest.
186    pub fn list_checkpoints(&self) -> Vec<(u64, String, u64)> {
187        self.checkpoints
188            .iter()
189            .map(|c| (c.id, c.label.clone(), c.tick))
190            .collect()
191    }
192
193    /// Remove all stored checkpoints.
194    pub fn clear_all(&mut self) {
195        self.checkpoints.clear();
196    }
197
198    /// Compute aggregate statistics.
199    pub fn stats(&self) -> CheckpointerStats {
200        CheckpointerStats {
201            total_checkpoints: self.checkpoints.len(),
202            oldest_tick: self.checkpoints.front().map(|c| c.tick),
203            newest_tick: self.checkpoints.back().map(|c| c.tick),
204            rollbacks_performed: self.rollbacks_performed,
205        }
206    }
207}
208
209// ---------------------------------------------------------------------------
210// Tests
211// ---------------------------------------------------------------------------
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216
217    fn default_config() -> CheckpointConfig {
218        CheckpointConfig::default()
219    }
220
221    fn config_with_max(max: usize) -> CheckpointConfig {
222        CheckpointConfig {
223            max_checkpoints: max,
224            ..Default::default()
225        }
226    }
227
228    fn empty_meta() -> HashMap<String, String> {
229        HashMap::new()
230    }
231
232    fn sample_tensors() -> Vec<Vec<f64>> {
233        vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0]]
234    }
235
236    // -- basic creation -----------------------------------------------------
237
238    #[test]
239    fn test_create_checkpoint_returns_unique_ids() {
240        let mut cp = TensorCheckpointer::new(default_config());
241        let id0 = cp.create_checkpoint("a", vec![], empty_meta());
242        let id1 = cp.create_checkpoint("b", vec![], empty_meta());
243        assert_ne!(id0, id1);
244        assert_eq!(cp.checkpoint_count(), 2);
245    }
246
247    #[test]
248    fn test_create_checkpoint_stores_label() {
249        let mut cp = TensorCheckpointer::new(default_config());
250        let id = cp.create_checkpoint("my_label", vec![], empty_meta());
251        let c = cp.get_checkpoint(id);
252        assert!(c.is_some());
253        assert_eq!(c.map(|c| c.label.as_str()), Some("my_label"));
254    }
255
256    #[test]
257    fn test_create_checkpoint_stores_tensor_data() {
258        let mut cp = TensorCheckpointer::new(default_config());
259        let data = sample_tensors();
260        let id = cp.create_checkpoint("t", data.clone(), empty_meta());
261        let c = cp.get_checkpoint(id).expect("checkpoint should exist");
262        assert_eq!(c.tensor_data, data);
263    }
264
265    #[test]
266    fn test_create_checkpoint_records_tick() {
267        let mut cp = TensorCheckpointer::new(default_config());
268        cp.tick();
269        cp.tick();
270        let id = cp.create_checkpoint("at_tick_2", vec![], empty_meta());
271        let c = cp.get_checkpoint(id).expect("checkpoint should exist");
272        assert_eq!(c.tick, 2);
273    }
274
275    // -- eviction -----------------------------------------------------------
276
277    #[test]
278    fn test_max_checkpoints_eviction_fifo() {
279        let mut cp = TensorCheckpointer::new(config_with_max(3));
280        let id0 = cp.create_checkpoint("0", vec![], empty_meta());
281        let _id1 = cp.create_checkpoint("1", vec![], empty_meta());
282        let _id2 = cp.create_checkpoint("2", vec![], empty_meta());
283        assert_eq!(cp.checkpoint_count(), 3);
284
285        // Adding a 4th should evict the oldest (id0).
286        let _id3 = cp.create_checkpoint("3", vec![], empty_meta());
287        assert_eq!(cp.checkpoint_count(), 3);
288        assert!(cp.get_checkpoint(id0).is_none());
289    }
290
291    #[test]
292    fn test_eviction_preserves_newest() {
293        let mut cp = TensorCheckpointer::new(config_with_max(2));
294        let _id0 = cp.create_checkpoint("a", vec![], empty_meta());
295        let _id1 = cp.create_checkpoint("b", vec![], empty_meta());
296        let id2 = cp.create_checkpoint("c", vec![], empty_meta());
297
298        let latest = cp.latest_checkpoint().expect("should have latest");
299        assert_eq!(latest.id, id2);
300    }
301
302    #[test]
303    fn test_eviction_with_max_one() {
304        let mut cp = TensorCheckpointer::new(config_with_max(1));
305        let _id0 = cp.create_checkpoint("first", vec![], empty_meta());
306        let id1 = cp.create_checkpoint("second", vec![], empty_meta());
307        assert_eq!(cp.checkpoint_count(), 1);
308        assert!(cp.get_checkpoint(id1).is_some());
309    }
310
311    // -- rollback -----------------------------------------------------------
312
313    #[test]
314    fn test_rollback_returns_checkpoint_data() {
315        let mut cp = TensorCheckpointer::new(default_config());
316        let data = sample_tensors();
317        let id = cp.create_checkpoint("snap", data.clone(), empty_meta());
318        let restored = cp.rollback(id).expect("rollback should succeed");
319        assert_eq!(restored.tensor_data, data);
320        assert_eq!(restored.label, "snap");
321    }
322
323    #[test]
324    fn test_rollback_removes_newer_checkpoints() {
325        let mut cp = TensorCheckpointer::new(default_config());
326        let id0 = cp.create_checkpoint("0", vec![], empty_meta());
327        let id1 = cp.create_checkpoint("1", vec![], empty_meta());
328        let id2 = cp.create_checkpoint("2", vec![], empty_meta());
329
330        let _restored = cp.rollback(id1).expect("rollback should succeed");
331        // id1 itself is consumed, id2 removed; only id0 remains.
332        assert_eq!(cp.checkpoint_count(), 1);
333        assert!(cp.get_checkpoint(id0).is_some());
334        assert!(cp.get_checkpoint(id1).is_none());
335        assert!(cp.get_checkpoint(id2).is_none());
336    }
337
338    #[test]
339    fn test_rollback_unknown_id_errors() {
340        let mut cp = TensorCheckpointer::new(default_config());
341        let result = cp.rollback(999);
342        assert!(result.is_err());
343        let msg = result.expect_err("should be error");
344        assert!(msg.contains("999"));
345    }
346
347    #[test]
348    fn test_rollback_increments_counter() {
349        let mut cp = TensorCheckpointer::new(default_config());
350        let id = cp.create_checkpoint("x", vec![], empty_meta());
351        assert_eq!(cp.stats().rollbacks_performed, 0);
352        let _ = cp.rollback(id);
353        assert_eq!(cp.stats().rollbacks_performed, 1);
354    }
355
356    #[test]
357    fn test_rollback_to_oldest() {
358        let mut cp = TensorCheckpointer::new(default_config());
359        let id0 = cp.create_checkpoint("oldest", vec![vec![0.0]], empty_meta());
360        let _id1 = cp.create_checkpoint("mid", vec![], empty_meta());
361        let _id2 = cp.create_checkpoint("newest", vec![], empty_meta());
362
363        let restored = cp.rollback(id0).expect("rollback should succeed");
364        assert_eq!(restored.label, "oldest");
365        assert_eq!(cp.checkpoint_count(), 0);
366    }
367
368    #[test]
369    fn test_rollback_to_latest() {
370        let mut cp = TensorCheckpointer::new(default_config());
371        let _id0 = cp.create_checkpoint("a", vec![], empty_meta());
372        let id1 = cp.create_checkpoint("b", vec![], empty_meta());
373
374        let restored = cp.rollback(id1).expect("rollback should succeed");
375        assert_eq!(restored.label, "b");
376        assert_eq!(cp.checkpoint_count(), 1); // id0 remains
377    }
378
379    #[test]
380    fn test_multiple_rollbacks() {
381        let mut cp = TensorCheckpointer::new(default_config());
382        let _id0 = cp.create_checkpoint("a", vec![], empty_meta());
383        let id1 = cp.create_checkpoint("b", vec![], empty_meta());
384        let _ = cp.rollback(id1);
385        assert_eq!(cp.stats().rollbacks_performed, 1);
386
387        let id2 = cp.create_checkpoint("c", vec![], empty_meta());
388        let _ = cp.rollback(id2);
389        assert_eq!(cp.stats().rollbacks_performed, 2);
390    }
391
392    // -- latest / get -------------------------------------------------------
393
394    #[test]
395    fn test_latest_checkpoint_empty() {
396        let cp = TensorCheckpointer::new(default_config());
397        assert!(cp.latest_checkpoint().is_none());
398    }
399
400    #[test]
401    fn test_latest_checkpoint_returns_last_added() {
402        let mut cp = TensorCheckpointer::new(default_config());
403        let _id0 = cp.create_checkpoint("first", vec![], empty_meta());
404        let id1 = cp.create_checkpoint("second", vec![], empty_meta());
405        let latest = cp.latest_checkpoint().expect("should exist");
406        assert_eq!(latest.id, id1);
407    }
408
409    #[test]
410    fn test_get_checkpoint_nonexistent() {
411        let cp = TensorCheckpointer::new(default_config());
412        assert!(cp.get_checkpoint(42).is_none());
413    }
414
415    // -- auto checkpoint timing ---------------------------------------------
416
417    #[test]
418    fn test_should_auto_checkpoint_initially_true_when_interval_zero() {
419        let config = CheckpointConfig {
420            auto_checkpoint_interval: 0,
421            ..Default::default()
422        };
423        let cp = TensorCheckpointer::new(config);
424        assert!(cp.should_auto_checkpoint());
425    }
426
427    #[test]
428    fn test_should_auto_checkpoint_false_initially() {
429        let cp = TensorCheckpointer::new(default_config()); // interval = 100
430        assert!(!cp.should_auto_checkpoint());
431    }
432
433    #[test]
434    fn test_should_auto_checkpoint_after_enough_ticks() {
435        let config = CheckpointConfig {
436            auto_checkpoint_interval: 5,
437            ..Default::default()
438        };
439        let mut cp = TensorCheckpointer::new(config);
440        for _ in 0..4 {
441            cp.tick();
442            assert!(!cp.should_auto_checkpoint());
443        }
444        cp.tick(); // tick 5
445        assert!(cp.should_auto_checkpoint());
446    }
447
448    #[test]
449    fn test_auto_checkpoint_resets_after_create() {
450        let config = CheckpointConfig {
451            auto_checkpoint_interval: 3,
452            ..Default::default()
453        };
454        let mut cp = TensorCheckpointer::new(config);
455        for _ in 0..3 {
456            cp.tick();
457        }
458        assert!(cp.should_auto_checkpoint());
459        cp.create_checkpoint("auto", vec![], empty_meta());
460        assert!(!cp.should_auto_checkpoint());
461    }
462
463    // -- list checkpoints ---------------------------------------------------
464
465    #[test]
466    fn test_list_checkpoints_empty() {
467        let cp = TensorCheckpointer::new(default_config());
468        assert!(cp.list_checkpoints().is_empty());
469    }
470
471    #[test]
472    fn test_list_checkpoints_order() {
473        let mut cp = TensorCheckpointer::new(default_config());
474        cp.create_checkpoint("a", vec![], empty_meta());
475        cp.tick();
476        cp.create_checkpoint("b", vec![], empty_meta());
477        let list = cp.list_checkpoints();
478        assert_eq!(list.len(), 2);
479        assert_eq!(list[0].1, "a");
480        assert_eq!(list[1].1, "b");
481        assert!(list[0].2 < list[1].2);
482    }
483
484    #[test]
485    fn test_list_checkpoints_after_eviction() {
486        let mut cp = TensorCheckpointer::new(config_with_max(2));
487        cp.create_checkpoint("a", vec![], empty_meta());
488        cp.create_checkpoint("b", vec![], empty_meta());
489        cp.create_checkpoint("c", vec![], empty_meta());
490        let list = cp.list_checkpoints();
491        assert_eq!(list.len(), 2);
492        assert_eq!(list[0].1, "b");
493        assert_eq!(list[1].1, "c");
494    }
495
496    // -- clear_all ----------------------------------------------------------
497
498    #[test]
499    fn test_clear_all() {
500        let mut cp = TensorCheckpointer::new(default_config());
501        cp.create_checkpoint("a", vec![], empty_meta());
502        cp.create_checkpoint("b", vec![], empty_meta());
503        assert_eq!(cp.checkpoint_count(), 2);
504        cp.clear_all();
505        assert_eq!(cp.checkpoint_count(), 0);
506        assert!(cp.latest_checkpoint().is_none());
507    }
508
509    #[test]
510    fn test_clear_all_on_empty() {
511        let mut cp = TensorCheckpointer::new(default_config());
512        cp.clear_all(); // should not panic
513        assert_eq!(cp.checkpoint_count(), 0);
514    }
515
516    // -- stats --------------------------------------------------------------
517
518    #[test]
519    fn test_stats_empty() {
520        let cp = TensorCheckpointer::new(default_config());
521        let s = cp.stats();
522        assert_eq!(s.total_checkpoints, 0);
523        assert!(s.oldest_tick.is_none());
524        assert!(s.newest_tick.is_none());
525        assert_eq!(s.rollbacks_performed, 0);
526    }
527
528    #[test]
529    fn test_stats_with_checkpoints() {
530        let mut cp = TensorCheckpointer::new(default_config());
531        cp.create_checkpoint("a", vec![], empty_meta());
532        cp.tick();
533        cp.tick();
534        cp.create_checkpoint("b", vec![], empty_meta());
535        let s = cp.stats();
536        assert_eq!(s.total_checkpoints, 2);
537        assert_eq!(s.oldest_tick, Some(0));
538        assert_eq!(s.newest_tick, Some(2));
539    }
540
541    #[test]
542    fn test_stats_after_rollback() {
543        let mut cp = TensorCheckpointer::new(default_config());
544        let id0 = cp.create_checkpoint("a", vec![], empty_meta());
545        cp.create_checkpoint("b", vec![], empty_meta());
546        let _ = cp.rollback(id0);
547        let s = cp.stats();
548        assert_eq!(s.total_checkpoints, 0);
549        assert_eq!(s.rollbacks_performed, 1);
550    }
551
552    // -- metadata -----------------------------------------------------------
553
554    #[test]
555    fn test_metadata_preserved() {
556        let mut cp = TensorCheckpointer::new(default_config());
557        let mut meta = HashMap::new();
558        meta.insert("epoch".to_string(), "42".to_string());
559        meta.insert("loss".to_string(), "0.01".to_string());
560        let id = cp.create_checkpoint("meta_test", vec![], meta.clone());
561        let c = cp.get_checkpoint(id).expect("should exist");
562        assert_eq!(c.metadata, meta);
563    }
564
565    #[test]
566    fn test_metadata_preserved_after_rollback() {
567        let mut cp = TensorCheckpointer::new(default_config());
568        let mut meta = HashMap::new();
569        meta.insert("key".to_string(), "value".to_string());
570        let id = cp.create_checkpoint("m", vec![], meta.clone());
571        let restored = cp.rollback(id).expect("rollback should succeed");
572        assert_eq!(restored.metadata, meta);
573    }
574
575    // -- edge cases ---------------------------------------------------------
576
577    #[test]
578    fn test_empty_checkpointer_count() {
579        let cp = TensorCheckpointer::new(default_config());
580        assert_eq!(cp.checkpoint_count(), 0);
581    }
582
583    #[test]
584    fn test_tick_without_checkpoints() {
585        let mut cp = TensorCheckpointer::new(default_config());
586        for _ in 0..200 {
587            cp.tick();
588        }
589        // Should not panic, and auto-checkpoint should trigger.
590        assert!(cp.should_auto_checkpoint());
591    }
592
593    #[test]
594    fn test_create_after_clear() {
595        let mut cp = TensorCheckpointer::new(default_config());
596        cp.create_checkpoint("before", vec![], empty_meta());
597        cp.clear_all();
598        let id = cp.create_checkpoint("after", vec![], empty_meta());
599        assert_eq!(cp.checkpoint_count(), 1);
600        let c = cp.get_checkpoint(id).expect("should exist");
601        assert_eq!(c.label, "after");
602    }
603
604    #[test]
605    fn test_rollback_then_create() {
606        let mut cp = TensorCheckpointer::new(default_config());
607        let id0 = cp.create_checkpoint("a", vec![], empty_meta());
608        cp.create_checkpoint("b", vec![], empty_meta());
609        let _ = cp.rollback(id0);
610
611        // After rollback, we can still create new checkpoints.
612        let id_new = cp.create_checkpoint("c", vec![], empty_meta());
613        assert_eq!(cp.checkpoint_count(), 1);
614        let c = cp.get_checkpoint(id_new).expect("should exist");
615        assert_eq!(c.label, "c");
616    }
617
618    #[test]
619    fn test_large_tensor_data() {
620        let mut cp = TensorCheckpointer::new(default_config());
621        let big = vec![vec![0.5; 1000]; 10];
622        let id = cp.create_checkpoint("large", big.clone(), empty_meta());
623        let c = cp.get_checkpoint(id).expect("should exist");
624        assert_eq!(c.tensor_data.len(), 10);
625        assert_eq!(c.tensor_data[0].len(), 1000);
626        assert!((c.tensor_data[0][500] - 0.5).abs() < f64::EPSILON);
627    }
628
629    #[test]
630    fn test_checkpoint_ids_are_monotonic() {
631        let mut cp = TensorCheckpointer::new(default_config());
632        let mut prev = cp.create_checkpoint("0", vec![], empty_meta());
633        for i in 1..10 {
634            let id = cp.create_checkpoint(&i.to_string(), vec![], empty_meta());
635            assert!(id > prev);
636            prev = id;
637        }
638    }
639
640    #[test]
641    fn test_default_config_values() {
642        let cfg = CheckpointConfig::default();
643        assert_eq!(cfg.max_checkpoints, 5);
644        assert_eq!(cfg.auto_checkpoint_interval, 100);
645    }
646}
ipfrs_tensorlogic/checkpointer.rs

ipfrs_tensorlogic/
checkpointer.rs
