foxhunt/ENSEMBLE_TRAINING_QUICK_START.md

Ensemble Training Quick Start Guide

TL;DR: Ensemble training coordinator is ready. Use this guide to get started.

---

🚀 Quick Start (5 minutes)

1. Run Tests

# Run all ensemble training tests
cargo test -p ml_training_service --test ensemble_training_tests

# Run basic tests only
cargo test -p ml_training_service --test ensemble_training_basic_tests

# Run unit tests
cargo test -p ml_training_service ensemble_training_coordinator

2. Basic Usage

use ml_training_service::ensemble_training_coordinator::{
    EnsembleTrainingConfig, EnsembleTrainingCoordinator
};

// Create config (see example below)
let config = create_ensemble_config();

// Create coordinator
let mut coordinator = EnsembleTrainingCoordinator::new(config).await?;

// Start training
let job_id = coordinator.start_ensemble_training().await?;

// Check status
let status = coordinator.get_model_status("DQN").await?;
println!("DQN status: {:?}", status);

3. Integration with Inference

use ml::ensemble::EnsembleTrainingIntegration;

// Create integration
let integration = EnsembleTrainingIntegration::new();

// Load checkpoints
let checkpoints = hashmap! {
    "DQN" => "path/to/dqn.safetensors",
    "PPO" => "path/to/ppo.safetensors",
    "MAMBA2" => "path/to/mamba2.safetensors",
    "TFT" => "path/to/tft.safetensors",
};

integration.load_ensemble_checkpoints(checkpoints).await?;

// Validate ready for production
integration.validate_production_readiness().await?;

---

📦 What's Included

Core Components

Component	Location	Purpose
EnsembleTrainingCoordinator	services/ml_training_service/src/ensemble_training_coordinator.rs	Main orchestrator
EnsembleTrainingIntegration	ml/src/ensemble/training_integration.rs	Inference bridge
Tests	services/ml_training_service/tests/ensemble_training_*.rs	TDD test suite

Key Features

✅ Multi-Model Training - DQN, PPO, MAMBA-2, TFT ✅ Dynamic Weights - Performance-based optimization ✅ Checkpoint Sync - Unified epoch management ✅ Failure Recovery - Automatic retry logic ✅ Ensemble Metrics - Aggregated performance tracking

---

🔧 Configuration Template

use std::collections::HashMap;
use ml::training_pipeline::*;
use ml::safety::*;
use uuid::Uuid;
use chrono::Utc;

fn create_ensemble_config() -> EnsembleTrainingConfig {
    let mut model_configs = HashMap::new();
    let mut model_weights = HashMap::new();

    // DQN (33% weight)
    model_configs.insert("DQN".to_string(), ProductionTrainingConfig {
        model_config: ModelArchitectureConfig {
            input_dim: 64,
            hidden_dims: vec![256, 128],
            output_dim: 32,
            dropout_rate: 0.1,
            activation: "relu".to_string(),
            batch_norm: true,
            residual_connections: false,
        },
        training_params: TrainingHyperparameters {
            learning_rate: 0.001,
            batch_size: 64,
            max_epochs: 100,
            patience: 10,
            validation_split: 0.2,
            l2_regularization: 0.0001,
            lr_decay_factor: 0.5,
            lr_decay_patience: 5,
        },
        safety_config: MLSafetyConfig {
            max_loss_value: 1000.0,
            max_prediction_value: 100.0,
            nan_check_interval: 10,
            enable_loss_scaling: true,
            convergence_window: 20,
        },
        gradient_config: GradientSafetyConfig {
            max_gradient_norm: 1.0,
            min_gradient_norm: 1e-8,
            gradient_clip_threshold: 5.0,
            enable_gradient_monitoring: true,
            gradient_check_interval: 1,
        },
        financial_config: FinancialValidationConfig {
            max_prediction_multiple: 2.0,
            min_prediction_confidence: 0.6,
            validate_position_sizing: true,
            max_position_fraction: 0.2,
            min_sharpe_threshold: 0.5,
        },
        performance_config: PerformanceConfig {
            device_preference: "cpu".to_string(),
            max_memory_bytes: 4_000_000_000,
            mixed_precision: false,
            num_workers: 2,
            gradient_accumulation_steps: 1,
        },
    });
    model_weights.insert("DQN".to_string(), 0.33);

    // PPO (33% weight) - same config structure
    model_configs.insert("PPO".to_string(), /* same as DQN */);
    model_weights.insert("PPO".to_string(), 0.33);

    // MAMBA-2 (17% weight)
    model_configs.insert("MAMBA2".to_string(), /* similar config */);
    model_weights.insert("MAMBA2".to_string(), 0.17);

    // TFT (17% weight)
    model_configs.insert("TFT".to_string(), /* similar config */);
    model_weights.insert("TFT".to_string(), 0.17);

    EnsembleTrainingConfig {
        job_id: Uuid::new_v4(),
        model_configs,
        model_weights,
        enable_weight_optimization: true,
        weight_optimization_interval_epochs: 5,
        checkpoint_interval_epochs: 1,
        max_epochs: 100,
        parallel_training: false,
        created_at: Utc::now(),
    }
}

---

📊 Common Operations

Check Training Status

// Get status for all models
for model in &["DQN", "PPO", "MAMBA2", "TFT"] {
    let status = coordinator.get_model_status(model).await?;
    println!("{}: {:?}", model, status);
}

Update Performance Metrics

// Set performance for each model
coordinator.set_model_performance("DQN", 0.85, 0.15).await?;
coordinator.set_model_performance("PPO", 0.80, 0.20).await?;
coordinator.set_model_performance("MAMBA2", 0.75, 0.25).await?;
coordinator.set_model_performance("TFT", 0.90, 0.10).await?;

// Trigger weight optimization
coordinator.optimize_weights().await?;

Get Ensemble Metrics

let metrics = coordinator.get_ensemble_metrics().await?;

println!("Ensemble train loss: {}", metrics["ensemble_train_loss"]);
println!("Ensemble val loss: {}", metrics["ensemble_val_loss"]);
println!("Ensemble accuracy: {}", metrics["ensemble_accuracy"]);
println!("Prediction diversity: {}", metrics["prediction_diversity"]);

Checkpoint Management

// Get all checkpoints
let checkpoints = coordinator.get_all_checkpoints().await?;
for (model, path) in checkpoints {
    println!("{}: {}", model, path);
}

// Load synchronized ensemble from epoch 50
coordinator.load_synchronized_ensemble(50).await?;

Handle Training Failures

// Check if model failed
if let ModelTrainingStatus::Failed = coordinator.get_model_status("PPO").await? {
    println!("PPO failed, attempting retry...");
    coordinator.retry_failed_model("PPO").await?;
}

---

🎯 Common Patterns

Pattern 1: Training Loop

let mut coordinator = EnsembleTrainingCoordinator::new(config).await?;
let job_id = coordinator.start_ensemble_training().await?;

for epoch in 1..=100 {
    // Simulate training (replace with actual training)
    coordinator.simulate_training_epochs(1).await?;

    // Check if optimization needed
    if epoch % 5 == 0 {
        coordinator.optimize_weights().await?;
        let weights = coordinator.get_current_weights().await?;
        println!("Epoch {}: Updated weights: {:?}", epoch, weights);
    }

    // Check for failures
    for model in &["DQN", "PPO", "MAMBA2", "TFT"] {
        if let ModelTrainingStatus::Failed = coordinator.get_model_status(model).await? {
            coordinator.retry_failed_model(model).await?;
        }
    }
}

Pattern 2: Checkpoint Loading

let integration = EnsembleTrainingIntegration::new();

// Load from latest epoch
let epoch = 100;
let checkpoints = hashmap! {
    "DQN" => format!("models/{}/dqn_epoch_{}.safetensors", job_id, epoch),
    "PPO" => format!("models/{}/ppo_epoch_{}.safetensors", job_id, epoch),
    "MAMBA2" => format!("models/{}/mamba2_epoch_{}.safetensors", job_id, epoch),
    "TFT" => format!("models/{}/tft_epoch_{}.safetensors", job_id, epoch),
};

integration.load_ensemble_checkpoints(checkpoints).await?;
integration.validate_production_readiness().await?;

Pattern 3: Performance Monitoring

// Track performance over time
let mut performance_history = Vec::new();

for epoch in 1..=100 {
    coordinator.simulate_training_epochs(1).await?;

    let metrics = coordinator.get_ensemble_metrics().await?;
    performance_history.push((
        epoch,
        metrics["ensemble_train_loss"],
        metrics["ensemble_val_loss"],
        metrics["ensemble_accuracy"],
    ));

    // Check for convergence
    if performance_history.len() > 10 {
        let recent_losses: Vec<_> = performance_history
            .iter()
            .rev()
            .take(10)
            .map(|(_, _, val_loss, _)| val_loss)
            .collect();

        let improving = recent_losses
            .windows(2)
            .all(|w| w[0] >= w[1]);

        if !improving {
            println!("Training plateaued at epoch {}", epoch);
            break;
        }
    }
}

---

⚠️ Important Notes

Weight Constraints

• Weights MUST sum to 1.0
• Each model needs config AND weight
• Validation runs on coordinator creation

Model Requirements

Must include all 4 models:

• DQN - Deep Q-Network
• PPO - Proximal Policy Optimization
• MAMBA2 - MAMBA-2 architecture
• TFT - Temporal Fusion Transformer

Checkpoint Naming

Follow convention: {model}_epoch_{epoch}.safetensors

Example:

• dqn_epoch_50.safetensors
• ppo_epoch_50.safetensors
• mamba2_epoch_50.safetensors
• tft_epoch_50.safetensors

---

🐛 Troubleshooting

Issue: "Model not found"

Solution: Ensure all 4 models registered in config

// Check model count
assert_eq!(config.model_count(), 4);

// Check specific model
assert!(config.has_model("DQN"));

Issue: "Weights don't sum to 1.0"

Solution: Verify weight values

let total = config.total_weight();
assert!((total - 1.0).abs() < 1e-6);

Issue: "Checkpoint not found"

Solution: Check file paths exist

use std::path::Path;

for (model, path) in checkpoints {
    if !Path::new(&path).exists() {
        eprintln!("Checkpoint missing: {} at {}", model, path);
    }
}

Issue: "Training failed"

Solution: Check model status and retry

for model in &["DQN", "PPO", "MAMBA2", "TFT"] {
    match coordinator.get_model_status(model).await? {
        ModelTrainingStatus::Failed => {
            println!("Retrying {}", model);
            coordinator.retry_failed_model(model).await?;
        }
        status => println!("{}: {:?}", model, status),
    }
}

---

📚 API Reference

EnsembleTrainingCoordinator

// Creation
pub async fn new(config: EnsembleTrainingConfig) -> Result<Self>

// Training
pub async fn start_ensemble_training(&mut self) -> Result<Uuid>
pub async fn get_model_status(&self, model_name: &str) -> Result<ModelTrainingStatus>

// Weights
pub async fn optimize_weights(&self) -> Result<()>
pub async fn get_current_weights(&self) -> Result<HashMap<String, f64>>

// Checkpoints
pub async fn get_latest_checkpoint(&self, model_name: &str) -> Result<Option<String>>
pub async fn get_all_checkpoints(&self) -> Result<Vec<(String, String)>>
pub async fn load_synchronized_ensemble(&self, epoch: u32) -> Result<()>

// Performance
pub async fn set_model_performance(&self, model_name: &str, accuracy: f64, loss: f64) -> Result<()>
pub async fn get_ensemble_metrics(&self) -> Result<HashMap<String, f64>>

// Recovery
pub async fn retry_failed_model(&self, model_name: &str) -> Result<()>

// Configuration
pub async fn get_model_training_config(&self, model_name: &str) -> Result<ProductionTrainingConfig>

EnsembleTrainingIntegration

// Creation
pub fn new() -> Self

// Checkpoints
pub async fn load_ensemble_checkpoints(&self, checkpoints: HashMap<String, String>) -> Result<()>

// Weights
pub async fn update_weights_from_performance(&self, performance_metrics: HashMap<String, f64>) -> Result<()>

// Metrics
pub async fn aggregate_training_metrics(&self, model_metrics: HashMap<String, (f64, f64, f64)>) -> Result<(f64, f64, f64)>
pub fn calculate_diversity(predictions: &[ModelPrediction]) -> f64

// Validation
pub async fn validate_production_readiness(&self) -> Result<()>

---

✅ Verification Checklist

Before deploying to production:

• All tests pass: cargo test -p ml_training_service ensemble
• Configuration validated: config.is_valid() == true
• Weights sum to 1.0: config.total_weight() ≈ 1.0
• All 4 models present: config.model_count() == 4
• Checkpoints exist: Verify file paths
• Integration validated: validate_production_readiness() passes

---

🔗 Resources

• Full Documentation: ENSEMBLE_TRAINING_TDD_IMPLEMENTATION.md
• System Architecture: CLAUDE.md
• Training Pipeline: ml/src/training_pipeline.rs
• Ensemble Inference: ml/src/ensemble/coordinator.rs

---

Quick Start Complete! ✅

For detailed information, see ENSEMBLE_TRAINING_TDD_IMPLEMENTATION.md