# 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

```bash
# 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

```rust
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

```rust
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

```rust
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

```rust
// 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

```rust
// 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

```rust
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

```rust
// 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

```rust
// 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

```rust
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

```rust
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

```rust
// 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

```rust
// 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

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

### Issue: "Checkpoint not found"

**Solution**: Check file paths exist

```rust
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

```rust
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

```rust
// 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

```rust
// 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`