foxhunt/ml/src/checkpoint

Unified Model Weight Persistence System

A comprehensive checkpoint system for all 5 AI models in the Foxhunt HFT system: DQN, MAMBA, TFT, TGGN, and LNN (Liquid Neural Networks).

Overview

The unified checkpoint system provides a single, consistent interface for saving, loading, and managing trained model weights and states across all AI models. It includes versioning, metadata management, compression, validation, and lifecycle management capabilities.

Key Features

🔧 Unified Interface

• Single API for all 5 model types
• Consistent checkpoint format across models
• Async/await support for non-blocking operations

📦 Model Versioning

• Semantic versioning (major.minor.patch)
• Compatibility checking between versions
• Migration support for version upgrades

📊 Rich Metadata

• Training state (epoch, step, loss, accuracy)
• Hyperparameters and model configuration
• Performance metrics and statistics
• Custom tags for organization

🗜️ Compression Support

• Multiple algorithms: LZ4, Zstd, Gzip
• Automatic compression ratio optimization
• Configurable compression levels

✅ Validation & Integrity

• SHA-256 checksums for corruption detection
• Metadata consistency validation
• Model compatibility verification

🔄 Lifecycle Management

• Automatic cleanup of old checkpoints
• Configurable retention policies
• Search and filtering capabilities

Architecture

┌─────────────────────────────────────────────────────────────┐
│                   CheckpointManager                         │
├─────────────────┬─────────────────┬─────────────────────────┤
│   Versioning    │   Compression   │    Storage Backend      │
│                 │                 │                         │
│ • Semantic Ver  │ • LZ4/Zstd      │ • FileSystem            │
│ • Compatibility │ • Delta Saves   │ • Cloud Storage         │
│ • Migration     │ • Streaming     │ • Database              │
└─────────────────┴─────────────────┴─────────────────────────┘

Supported Models

Model	Type	Description
DQN	Deep Q-Learning Network	Reinforcement learning for trading decisions
MAMBA	State-Space Model	Mamba-2 with SSD layers for sequence modeling
TFT	Temporal Fusion Transformer	Multi-horizon forecasting with attention
TGGN	Temporal Graph Gated Network	Graph neural networks for market microstructure
LNN	Liquid Neural Network	Continuous-time RNNs for adaptive behavior

Quick Start

Basic Usage

use ml_models::checkpoint::{CheckpointManager, CheckpointConfig, CompressionType};
use ml_models::dqn::DQNAgent;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Create checkpoint manager
    let config = CheckpointConfig {
        base_dir: "./checkpoints".into(),
        compression: CompressionType::LZ4,
        max_checkpoints_per_model: 10,
        auto_cleanup: true,
        ..Default::default()
    };
    let manager = CheckpointManager::new(config)?;
    
    // Create and train a model
    let mut dqn_agent = DQNAgent::new(Default::default())?;
    // ... training code ...
    
    // Save checkpoint
    let checkpoint_id = manager.save_checkpoint(
        &dqn_agent,
        Some(vec!["production".to_string(), "validated".to_string()])
    ).await?;
    
    println!("Saved checkpoint: {}", checkpoint_id);
    
    // Load checkpoint
    let metadata = manager.load_checkpoint(&mut dqn_agent, &checkpoint_id).await?;
    println!("Loaded checkpoint from epoch {:?}", metadata.epoch);
    
    Ok(())
}

Advanced Features

// Load latest checkpoint for a model
let latest = manager.load_latest_checkpoint(&mut model).await?;

// Search checkpoints by tags
let production_checkpoints = manager.find_checkpoints_by_tags(&[
    "production".to_string()
]).await;

// List all checkpoints for a model type
let dqn_checkpoints = manager.list_checkpoints(ModelType::DQN, "my_model").await;

// Get checkpoint statistics
let stats = manager.get_stats();
println!("Total checkpoints saved: {}", stats.get("total_saved").unwrap_or(&0));

Configuration Options

CheckpointConfig

pub struct CheckpointConfig {
    /// Base directory for checkpoints
    pub base_dir: PathBuf,
    
    /// Default compression type
    pub compression: CompressionType,
    
    /// Default checkpoint format
    pub format: CheckpointFormat,
    
    /// Maximum number of checkpoints to keep per model
    pub max_checkpoints_per_model: usize,
    
    /// Automatic cleanup of old checkpoints
    pub auto_cleanup: bool,
    
    /// Enable checksum validation
    pub validate_checksums: bool,
    
    /// Enable incremental checkpoints (delta saves)
    pub incremental_checkpoints: bool,
    
    /// Compression level (0-9, algorithm dependent)
    pub compression_level: u32,
    
    /// Enable async I/O operations
    pub async_io: bool,
    
    /// Buffer size for I/O operations
    pub buffer_size: usize,
}

Compression Types

• None: No compression (fastest)
• LZ4: Fast compression with good speed/ratio balance
• Zstd: Balanced compression for production use
• Gzip: High compression ratio for storage optimization

Checkpoint Formats

• Binary: Fastest serialization using bincode
• JSON: Human-readable for debugging
• MessagePack: Compact binary format
• Custom: Optimized format for specific models

Model Implementation

To make a model checkpointable, implement the Checkpointable trait:

use async_trait::async_trait;
use ml_models::checkpoint::{Checkpointable, ModelType};

#[async_trait]
impl Checkpointable for MyModel {
    fn model_type(&self) -> ModelType {
        ModelType::DQN
    }
    
    fn model_name(&self) -> &str {
        "my_custom_model"
    }
    
    fn model_version(&self) -> &str {
        "1.0.0"
    }
    
    async fn serialize_state(&self) -> Result<Vec<u8>, MLError> {
        // Serialize model weights and state
        let state = MyModelState {
            weights: self.get_weights(),
            config: self.config.clone(),
            // ... other state
        };
        Ok(bincode::serialize(&state)?)
    }
    
    async fn deserialize_state(&mut self, data: &[u8]) -> Result<(), MLError> {
        // Deserialize and restore model state
        let state: MyModelState = bincode::deserialize(data)?;
        self.set_weights(state.weights);
        self.config = state.config;
        Ok(())
    }
    
    fn get_training_state(&self) -> (Option<u64>, Option<u64>, Option<f64>, Option<f64>) {
        (Some(self.epoch), Some(self.step), Some(self.loss), Some(self.accuracy))
    }
    
    fn get_hyperparameters(&self) -> HashMap<String, serde_json::Value> {
        // Return model hyperparameters as JSON values
        let mut params = HashMap::new();
        params.insert("learning_rate".to_string(), json!(self.config.learning_rate));
        params.insert("batch_size".to_string(), json!(self.config.batch_size));
        params
    }
    
    fn get_metrics(&self) -> HashMap<String, f64> {
        // Return current model metrics
        let mut metrics = HashMap::new();
        metrics.insert("accuracy".to_string(), self.current_accuracy);
        metrics.insert("loss".to_string(), self.current_loss);
        metrics
    }
    
    fn get_architecture_info(&self) -> HashMap<String, serde_json::Value> {
        // Return architecture information
        let mut info = HashMap::new();
        info.insert("layers".to_string(), json!(self.num_layers));
        info.insert("hidden_size".to_string(), json!(self.hidden_size));
        info
    }
}

Version Management

The system supports semantic versioning with automatic compatibility checking:

use ml_models::checkpoint::{VersionManager, VersionChangeType};

let version_manager = VersionManager::new();

// Check compatibility between versions
let compat_info = version_manager.check_compatibility(
    "1.0.0",  // current version
    "1.1.0",  // checkpoint version
    ModelType::DQN
)?;

if compat_info.compatible {
    println!("Versions are compatible");
    if compat_info.risk == CompatibilityRisk::Medium {
        println!("Some features may behave differently");
    }
} else {
    println!("Versions are incompatible - migration required");
}

// Suggest next version
let next_version = version_manager.suggest_next_version(
    "1.2.3", 
    VersionChangeType::Minor
)?;
println!("Next version: {}", next_version); // "1.3.0"

Storage Backends

FileSystem Storage (Default)

Stores checkpoints on the local filesystem with organized directory structure:

checkpoints/
├── dqn_model_v1.0.0_e50_s1000_20230815_143022.dqn
├── mamba_model_v2.1.0_e100_20230815_143055.mamba
├── metadata/
│   ├── dqn_model_v1.0.0_e50_s1000_20230815_143022.metadata.json
│   └── mamba_model_v2.1.0_e100_20230815_143055.metadata.json
└── ...

Memory Storage (Testing)

In-memory storage backend for unit tests and development:

use ml_models::checkpoint::storage::MemoryStorage;

let storage = MemoryStorage::new();
// Use for testing without filesystem dependencies

Performance Characteristics

Save Performance

• DQN: ~1ms for typical model size
• MAMBA: ~5ms for large state-space models
• TFT: ~3ms for transformer weights
• TGGN: ~2ms for graph embeddings
• LNN: ~1ms for continuous-time parameters

Compression Ratios

• LZ4: ~2-3x reduction, fastest
• Zstd: ~3-5x reduction, balanced
• Gzip: ~4-6x reduction, smallest

Storage Requirements

• Metadata: ~1-5KB per checkpoint
• Model Weights: 100KB - 100MB depending on model
• Compressed: 30-80% size reduction typical

Error Handling

The system uses comprehensive error handling with detailed error messages:

use ml_models::checkpoint::MLError;

match manager.load_checkpoint(&mut model, &checkpoint_id).await {
    Ok(metadata) => {
        println!("Loaded successfully: {:?}", metadata);
    }
    Err(MLError::ModelError(msg)) => {
        eprintln!("Model error: {}", msg);
    }
    Err(e) => {
        eprintln!("Checkpoint error: {}", e);
    }
}

Monitoring and Statistics

Track checkpoint system performance and usage:

let stats = manager.get_stats();
println!("Checkpoint Statistics:");
println!("  Total saved: {}", stats.get("total_saved").unwrap_or(&0));
println!("  Total loaded: {}", stats.get("total_loaded").unwrap_or(&0));
println!("  Compression savings: {} KB", stats.get("compression_savings").unwrap_or(&0) / 1024);
println!("  Average save time: {}μs", stats.get("avg_save_time_us").unwrap_or(&0));
println!("  Failed operations: {}", stats.get("failed_operations").unwrap_or(&0));

Best Practices

1. Version Strategy

• Use semantic versioning consistently
• Increment major version for breaking changes
• Test version compatibility before production deployment

2. Compression Selection

• Use LZ4 for development and testing (speed)
• Use Zstd for production (balanced)
• Use Gzip for archival storage (size)

3. Metadata Management

• Include meaningful tags for organization
• Store training metrics for analysis
• Document model architecture changes

4. Lifecycle Management

• Set reasonable retention policies
• Use auto-cleanup in production
• Monitor storage usage regularly

5. Error Handling

• Always validate loaded checkpoints
• Implement fallback strategies for load failures
• Log checkpoint operations for debugging

Testing

The system includes comprehensive tests covering all functionality:

# Run all checkpoint tests
cargo test checkpoint

# Run integration tests
cargo test integration_tests

# Run specific model tests
cargo test test_dqn_checkpoint
cargo test test_mamba_checkpoint

Examples

See the examples/ directory for complete working examples:

• checkpoint_integration_demo.rs: Comprehensive demonstration
• version_management_demo.rs: Version compatibility examples
• compression_benchmark.rs: Performance comparisons

Performance Tuning

For Development

let config = CheckpointConfig {
    compression: CompressionType::None,
    validate_checksums: false,
    async_io: false,
    ..Default::default()
};

For Production

let config = CheckpointConfig {
    compression: CompressionType::Zstd,
    compression_level: 3,
    validate_checksums: true,
    auto_cleanup: true,
    max_checkpoints_per_model: 10,
    async_io: true,
    ..Default::default()
};

For Storage-Constrained Environments

let config = CheckpointConfig {
    compression: CompressionType::Gzip,
    compression_level: 9,
    max_checkpoints_per_model: 3,
    auto_cleanup: true,
    ..Default::default()
};

Contributing

When adding new models to the checkpoint system:

1. Implement the Checkpointable trait
2. Add comprehensive serialization/deserialization
3. Include metadata extraction methods
4. Add integration tests
5. Update documentation

License

This checkpoint system is part of the Foxhunt HFT trading system and follows the same licensing terms.