Callback System - Entrenar - Training & Optimization Library
Description: A comprehensive guide to building neural network training systems with autograd, optimizers, LoRA/QLoRA, and quantization using EXTREME TDD methodology

# Callback System

The callback system provides extensible hooks into the training loop, enabling behaviors like early stopping, checkpointing, progress logging, and real-time monitoring without modifying the core trainer.

## Overview

    #![allow(unused)]
    fn main() {
    use entrenar::train::{
        TrainerCallback, CallbackContext, CallbackAction, CallbackManager,
        EarlyStopping, CheckpointCallback, ProgressCallback, MonitorCallback,
    };

    // Add multiple callbacks
    trainer.add_callback(EarlyStopping::new(5, 0.001));
    trainer.add_callback(CheckpointCallback::new("./ckpt"));
    trainer.add_callback(ProgressCallback::new(10));
    trainer.add_callback(MonitorCallback::new());
    }

## Callback Lifecycle

Callbacks fire at six points during training:

    train()
      │
      ├─► on_train_begin
      │
      ├─► for epoch in 0..max_epochs:
      │     │
      │     ├─► on_epoch_begin
      │     │
      │     ├─► for batch in batches:
      │     │     ├─► on_step_begin
      │     │     ├─► train_step()
      │     │     └─► on_step_end
      │     │
      │     └─► on_epoch_end
      │
      └─► on_train_end

## CallbackAction

Callbacks return an action that controls training flow:

    #![allow(unused)]
    fn main() {
    pub enum CallbackAction {
        Continue,   // Continue training normally
        Stop,       // Stop training immediately
        SkipEpoch,  // Skip to next epoch (epoch_begin only)
    }
    }

Behavior:

- `Continue` - Training proceeds normally
- `Stop` - Training stops, `TrainResult.stopped_early = true`
- `SkipEpoch` - Skip remaining steps in current epoch (only valid in `on_epoch_begin`)

## CallbackContext

Every callback receives context with current training state:

    #![allow(unused)]
    fn main() {
    pub struct CallbackContext {
        pub epoch: usize,           // Current epoch (0-indexed)
        pub max_epochs: usize,      // Maximum epochs configured
        pub step: usize,            // Current step within epoch
        pub steps_per_epoch: usize, // Total steps in epoch
        pub global_step: usize,     // Total steps across all epochs
        pub loss: f32,              // Current/latest loss
        pub lr: f32,                // Current learning rate
        pub best_loss: Option<f32>, // Best loss achieved so far
        pub val_loss: Option<f32>,  // Validation loss (if provided)
        pub elapsed_secs: f64,      // Seconds since training started
    }
    }

## TrainerCallback Trait

    #![allow(unused)]
    fn main() {
    pub trait TrainerCallback: Send {
        /// Called before training begins
        fn on_train_begin(&mut self, ctx: &CallbackContext) -> CallbackAction {
            CallbackAction::Continue
        }

        /// Called after training ends
        fn on_train_end(&mut self, ctx: &CallbackContext) {}

        /// Called at the start of each epoch
        fn on_epoch_begin(&mut self, ctx: &CallbackContext) -> CallbackAction {
            CallbackAction::Continue
        }

        /// Called at the end of each epoch
        fn on_epoch_end(&mut self, ctx: &CallbackContext) -> CallbackAction {
            CallbackAction::Continue
        }

        /// Called before each training step
        fn on_step_begin(&mut self, ctx: &CallbackContext) -> CallbackAction {
            CallbackAction::Continue
        }

        /// Called after each training step
        fn on_step_end(&mut self, ctx: &CallbackContext) -> CallbackAction {
            CallbackAction::Continue
        }

        /// Callback name for logging
        fn name(&self) -> &str;
    }
    }

## Built-in Callbacks

### EarlyStopping

Stops training when loss stops improving:

    #![allow(unused)]
    fn main() {
    pub struct EarlyStopping {
        patience: usize,      // Epochs to wait before stopping
        min_delta: f32,       // Minimum improvement threshold
        best_loss: f32,       // Best loss seen
        epochs_without_improvement: usize,
    }

    // Usage
    let es = EarlyStopping::new(5, 0.001);
    // Stops if loss doesn't improve by at least 0.001 for 5 epochs
    }

Behavior:

- Tracks best loss seen during training
- Counts epochs without improvement (loss not decreasing by `min_delta`)
- Returns `CallbackAction::Stop` when patience exhausted

### CheckpointCallback

Saves model checkpoints periodically and/or when best loss achieved:

    #![allow(unused)]
    fn main() {
    pub struct CheckpointCallback {
        save_dir: PathBuf,
        save_every: Option<usize>,  // Save every N epochs
        save_best: bool,            // Save when best loss achieved
        best_loss: f32,
    }

    // Usage
    let ckpt = CheckpointCallback::new("./checkpoints")
        .save_every(5)      // Save every 5 epochs
        .save_best(true);   // Also save best model

    // Creates files like:
    // ./checkpoints/checkpoint_epoch_5.json
    // ./checkpoints/checkpoint_epoch_10.json
    // ./checkpoints/checkpoint_best.json
    }

### ProgressCallback

Logs training progress to stdout:

    #![allow(unused)]
    fn main() {
    pub struct ProgressCallback {
        log_interval: usize,  // Log every N steps
    }

    // Usage
    let progress = ProgressCallback::new(10);
    // Logs: "Epoch 1/100 [========>  ] Step 50/500: loss: 0.1234"
    }

Output format:

    Epoch 1/100 [========>  ] loss: 0.2345, lr: 0.001, elapsed: 12.3s
      Step 10/100: loss: 0.2456
      Step 20/100: loss: 0.2234

### MonitorCallback

Real-time monitoring with NaN/Inf detection and metrics collection:

    #![allow(unused)]
    fn main() {
    pub struct MonitorCallback {
        collector: MetricsCollector,
        andon: AndonSystem,
    }

    // Usage
    let monitor = MonitorCallback::new();
    // Automatically:
    // - Records loss, learning rate metrics
    // - Detects NaN/Inf and triggers Stop
    // - Integrates with Andon alerting system
    }

Automatic detection:

- NaN loss → `CallbackAction::Stop`
- Inf loss → `CallbackAction::Stop`
- Triggers Andon alert for investigation

## Custom Callbacks

### Basic Example

    #![allow(unused)]
    fn main() {
    use entrenar::train::{TrainerCallback, CallbackContext, CallbackAction};

    struct LossLogger {
        losses: Vec<f32>,
    }

    impl LossLogger {
        fn new() -> Self {
            Self { losses: Vec::new() }
        }
    }

    impl TrainerCallback for LossLogger {
        fn on_epoch_end(&mut self, ctx: &CallbackContext) -> CallbackAction {
            self.losses.push(ctx.loss);
            println!("Epoch {}: loss = {:.6}", ctx.epoch, ctx.loss);
            CallbackAction::Continue
        }

        fn name(&self) -> &str {
            "LossLogger"
        }
    }
    }

### Learning Rate Warmup

    #![allow(unused)]
    fn main() {
    struct WarmupCallback {
        warmup_epochs: usize,
        target_lr: f32,
    }

    impl TrainerCallback for WarmupCallback {
        fn on_epoch_begin(&mut self, ctx: &CallbackContext) -> CallbackAction {
            if ctx.epoch < self.warmup_epochs {
                let warmup_lr = self.target_lr * (ctx.epoch + 1) as f32
                              / self.warmup_epochs as f32;
                // Would need trainer access to set LR
                println!("Warmup LR: {:.6}", warmup_lr);
            }
            CallbackAction::Continue
        }

        fn name(&self) -> &str {
            "WarmupCallback"
        }
    }
    }

### Gradient Explosion Detector

    #![allow(unused)]
    fn main() {
    struct GradientMonitor {
        max_loss: f32,
        loss_history: Vec<f32>,
    }

    impl TrainerCallback for GradientMonitor {
        fn on_step_end(&mut self, ctx: &CallbackContext) -> CallbackAction {
            self.loss_history.push(ctx.loss);

            // Detect sudden loss spike
            if self.loss_history.len() > 1 {
                let prev = self.loss_history[self.loss_history.len() - 2];
                if ctx.loss > prev * 10.0 {
                    eprintln!("WARNING: Loss spike detected! {} -> {}", prev, ctx.loss);
                    return CallbackAction::Stop;
                }
            }

            if ctx.loss > self.max_loss {
                eprintln!("ERROR: Loss exceeded threshold: {} > {}", ctx.loss, self.max_loss);
                return CallbackAction::Stop;
            }

            CallbackAction::Continue
        }

        fn name(&self) -> &str {
            "GradientMonitor"
        }
    }
    }

## CallbackManager

The `CallbackManager` orchestrates multiple callbacks:

    #![allow(unused)]
    fn main() {
    pub struct CallbackManager {
        callbacks: Vec<Box<dyn TrainerCallback>>,
    }

    impl CallbackManager {
        pub fn new() -> Self;
        pub fn add<C: TrainerCallback + 'static>(&mut self, callback: C);
        pub fn is_empty(&self) -> bool;
        pub fn len(&self) -> usize;

        // Event dispatchers (called by Trainer)
        pub fn on_train_begin(&mut self, ctx: &CallbackContext) -> CallbackAction;
        pub fn on_train_end(&mut self, ctx: &CallbackContext);
        pub fn on_epoch_begin(&mut self, ctx: &CallbackContext) -> CallbackAction;
        pub fn on_epoch_end(&mut self, ctx: &CallbackContext) -> CallbackAction;
        pub fn on_step_begin(&mut self, ctx: &CallbackContext) -> CallbackAction;
        pub fn on_step_end(&mut self, ctx: &CallbackContext) -> CallbackAction;
    }
    }

Dispatch behavior:

- Callbacks fire in order they were added
- If any callback returns `Stop`, remaining callbacks don't fire
- `on_train_end` always fires (even after early stop)

## Best Practices

### Callback Order

Add callbacks in order of priority:

    #![allow(unused)]
    fn main() {
    // Critical monitoring first
    trainer.add_callback(MonitorCallback::new());     // NaN detection
    trainer.add_callback(EarlyStopping::new(5, 0.001)); // Early stopping

    // Logging/checkpointing after
    trainer.add_callback(ProgressCallback::new(10));
    trainer.add_callback(CheckpointCallback::new("./ckpt"));
    }

### Stateful Callbacks

Callbacks can maintain state across training:

    #![allow(unused)]
    fn main() {
    struct StatefulCallback {
        epoch_losses: Vec<f32>,
        best_epoch: usize,
    }

    impl TrainerCallback for StatefulCallback {
        fn on_epoch_end(&mut self, ctx: &CallbackContext) -> CallbackAction {
            self.epoch_losses.push(ctx.loss);

            if ctx.best_loss == Some(ctx.loss) {
                self.best_epoch = ctx.epoch;
            }

            CallbackAction::Continue
        }

        fn on_train_end(&mut self, ctx: &CallbackContext) {
            println!("Best epoch: {} with loss {:.6}",
                self.best_epoch,
                self.epoch_losses[self.best_epoch]);
        }

        fn name(&self) -> &str {
            "StatefulCallback"
        }
    }
    }

### Thread Safety

Callbacks must be `Send` to support potential future parallelism:

    #![allow(unused)]
    fn main() {
    // Good: Uses Arc for shared state
    struct ThreadSafeCallback {
        counter: Arc<AtomicUsize>,
    }

    // Bad: Uses Rc (not Send)
    struct NotSendCallback {
        counter: Rc<RefCell<usize>>,  // Won't compile!
    }
    }

## See Also

- Trainer API - Main trainer documentation
- Early Stopping - Detailed early stopping guide
- Checkpointing - Checkpoint management
- Curriculum Learning - Progressive difficulty training (CITL)
- Explainability - Feature attribution callbacks
- Real-Time Monitoring - Monitor integration