d95e205d4b
Eliminate the entire mixed_precision runtime indirection layer: - Delete crates/ml-core/src/mixed_precision.rs (training_dtype, ensure_training_dtype, align_dim_for_tensor_cores) - Inline ~100 call sites across 130 files to constants: training_dtype(&device) → candle_core::DType::BF16 ensure_training_dtype(x) → x.to_dtype(candle_core::DType::BF16) align_dim_for_tensor_cores(x, &device) → (x + 7) & !7 - Remove re-exports from ml-dqn, ml-supervised, ml lib.rs - Clean config/toml/json/shell references No CPU/Metal training path exists — BF16 is the only dtype. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
530 lines
17 KiB
Rust
530 lines
17 KiB
Rust
#![allow(
|
|
clippy::unwrap_used,
|
|
clippy::expect_used,
|
|
clippy::indexing_slicing,
|
|
clippy::manual_range_contains
|
|
)]
|
|
//! TDD Tests for Ensemble Training Coordination
|
|
//!
|
|
//! These tests define the behavior we expect from the ensemble training system
|
|
//! BEFORE implementing the actual functionality.
|
|
//!
|
|
//! Test Coverage:
|
|
//! 1. Ensemble training configuration
|
|
//! 2. Multi-model coordination during training
|
|
//! 3. Ensemble weight optimization
|
|
//! 4. Checkpoint synchronization (all 4 models)
|
|
//! 5. Integration with ML Training Service
|
|
|
|
use std::collections::HashMap;
|
|
|
|
use chrono::Utc;
|
|
use ml::safety::{GradientSafetyConfig, MLSafetyConfig};
|
|
use ml::training_pipeline::{
|
|
FinancialValidationConfig, ModelArchitectureConfig, PerformanceConfig,
|
|
ProductionTrainingConfig, TrainingHyperparameters,
|
|
};
|
|
use ml_training_service::ensemble_training_coordinator::{
|
|
EnsembleTrainingConfig, EnsembleTrainingCoordinator, ModelTrainingStatus,
|
|
};
|
|
use uuid::Uuid;
|
|
|
|
/// Test 1: Ensemble training configuration validation
|
|
#[tokio::test]
|
|
async fn test_ensemble_training_config_validation() {
|
|
// Test 1.1: Valid configuration should be accepted
|
|
let config = create_valid_ensemble_config();
|
|
assert!(
|
|
config.is_valid(),
|
|
"Valid ensemble config should pass validation"
|
|
);
|
|
|
|
// Test 1.2: Must have all 4 models (DQN, PPO, MAMBA-2, TFT)
|
|
let mut models = config.model_configs.keys().cloned().collect::<Vec<_>>();
|
|
models.sort();
|
|
assert_eq!(
|
|
models,
|
|
vec!["DQN", "MAMBA2", "PPO", "TFT"],
|
|
"Must configure all 4 models"
|
|
);
|
|
|
|
// Test 1.3: Weights must sum to 1.0
|
|
let weight_sum: f64 = config.model_weights.values().sum();
|
|
assert!(
|
|
(weight_sum - 1.0).abs() < 1e-6,
|
|
"Model weights must sum to 1.0, got {}",
|
|
weight_sum
|
|
);
|
|
|
|
// Test 1.4: Each model must have matching training and weight configuration
|
|
for model_name in &["DQN", "PPO", "MAMBA2", "TFT"] {
|
|
let model_key = model_name.to_string();
|
|
assert!(
|
|
config.model_configs.contains_key(&model_key),
|
|
"Missing config for {}",
|
|
model_name
|
|
);
|
|
assert!(
|
|
config.model_weights.contains_key(&model_key),
|
|
"Missing weight for {}",
|
|
model_name
|
|
);
|
|
}
|
|
}
|
|
|
|
/// Test 2: Multi-model coordination during training
|
|
#[tokio::test]
|
|
async fn test_multi_model_training_coordination() {
|
|
let config = create_valid_ensemble_config();
|
|
let mut coordinator = create_ensemble_coordinator(config).await;
|
|
|
|
// Test 2.1: All models should start in Pending state
|
|
for model_name in &["DQN", "PPO", "MAMBA2", "TFT"] {
|
|
let status = coordinator.get_model_status(model_name).await.unwrap();
|
|
assert_eq!(
|
|
status,
|
|
ModelTrainingStatus::Pending,
|
|
"{} should start Pending",
|
|
model_name
|
|
);
|
|
}
|
|
|
|
// Test 2.2: Can start training for all models
|
|
let job_id = coordinator.start_ensemble_training().await.unwrap();
|
|
assert_ne!(job_id, Uuid::nil(), "Should return valid job ID");
|
|
|
|
// Test 2.3: At least one model should be training after start
|
|
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
|
|
let mut any_training = false;
|
|
for model in &["DQN", "PPO", "MAMBA2", "TFT"] {
|
|
if matches!(
|
|
coordinator.get_model_status(model).await,
|
|
Ok(ModelTrainingStatus::Training)
|
|
) {
|
|
any_training = true;
|
|
break;
|
|
}
|
|
}
|
|
assert!(
|
|
any_training,
|
|
"At least one model should be training after start"
|
|
);
|
|
}
|
|
|
|
/// Test 3: Ensemble weight optimization
|
|
#[tokio::test]
|
|
async fn test_ensemble_weight_optimization() {
|
|
let mut config = create_valid_ensemble_config();
|
|
config.enable_weight_optimization = true;
|
|
config.weight_optimization_interval_epochs = 5;
|
|
|
|
let coordinator = create_ensemble_coordinator(config).await;
|
|
|
|
// Test 3.1: Initial weights should match configuration
|
|
let initial_weights = coordinator.get_current_weights().await.unwrap();
|
|
assert_eq!(initial_weights.len(), 4, "Should have 4 model weights");
|
|
|
|
// Test 3.2: Simulate training progress and weight updates
|
|
coordinator.simulate_training_epochs(10).await.unwrap();
|
|
|
|
// Test 3.3: Weights should be updated after optimization interval
|
|
let updated_weights = coordinator.get_current_weights().await.unwrap();
|
|
assert_ne!(
|
|
initial_weights, updated_weights,
|
|
"Weights should be updated after optimization"
|
|
);
|
|
|
|
// Test 3.4: Updated weights should still sum to 1.0
|
|
let weight_sum: f64 = updated_weights.values().sum();
|
|
assert!(
|
|
(weight_sum - 1.0).abs() < 1e-6,
|
|
"Optimized weights must sum to 1.0, got {}",
|
|
weight_sum
|
|
);
|
|
|
|
// Test 3.5: Better-performing models should get higher weights
|
|
// (This test assumes DQN performs better in simulation)
|
|
let dqn_initial = initial_weights.get("DQN").expect("INVARIANT: Key should exist in map");
|
|
let dqn_updated = updated_weights.get("DQN").expect("INVARIANT: Key should exist in map");
|
|
// Weight adjustment logic will determine if this increases or decreases
|
|
assert_ne!(
|
|
dqn_initial, dqn_updated,
|
|
"DQN weight should be adjusted based on performance"
|
|
);
|
|
}
|
|
|
|
/// Test 4: Checkpoint synchronization for all models
|
|
#[tokio::test]
|
|
async fn test_checkpoint_synchronization() {
|
|
let config = create_valid_ensemble_config();
|
|
let mut coordinator = create_ensemble_coordinator(config).await;
|
|
|
|
// Test 4.1: Start training and wait for first checkpoint
|
|
let _job_id = coordinator.start_ensemble_training().await.unwrap();
|
|
coordinator.simulate_training_epochs(1).await.unwrap();
|
|
|
|
// Test 4.2: All models should have checkpoint paths after first epoch
|
|
for model_name in &["DQN", "PPO", "MAMBA2", "TFT"] {
|
|
let checkpoint = coordinator.get_latest_checkpoint(model_name).await.unwrap();
|
|
assert!(
|
|
checkpoint.is_some(),
|
|
"{} should have checkpoint after epoch 1",
|
|
model_name
|
|
);
|
|
|
|
let checkpoint_path = checkpoint.unwrap();
|
|
assert!(
|
|
checkpoint_path.contains(model_name),
|
|
"Checkpoint path should contain model name"
|
|
);
|
|
assert!(
|
|
checkpoint_path.contains("epoch_1"),
|
|
"Checkpoint should be from epoch 1"
|
|
);
|
|
}
|
|
|
|
// Test 4.3: Checkpoints should be synchronized (all from same epoch)
|
|
let checkpoints = coordinator.get_all_checkpoints().await.unwrap();
|
|
let epochs: Vec<_> = checkpoints
|
|
.iter()
|
|
.map(|(_, cp)| {
|
|
cp.split("epoch_")
|
|
.last()
|
|
.unwrap()
|
|
.split('_')
|
|
.next()
|
|
.unwrap()
|
|
.parse::<u32>()
|
|
.unwrap()
|
|
})
|
|
.collect();
|
|
|
|
let first_epoch = epochs[0];
|
|
assert!(
|
|
epochs.iter().all(|&e| e == first_epoch),
|
|
"All checkpoints should be from same epoch"
|
|
);
|
|
|
|
// Test 4.4: Can load synchronized ensemble from checkpoints
|
|
let ensemble_restored = coordinator.load_synchronized_ensemble(first_epoch).await;
|
|
assert!(
|
|
ensemble_restored.is_ok(),
|
|
"Should be able to load synchronized ensemble"
|
|
);
|
|
}
|
|
|
|
/// Test 5: Performance-based weight adjustment
|
|
#[tokio::test]
|
|
async fn test_performance_based_weight_adjustment() {
|
|
let mut config = create_valid_ensemble_config();
|
|
config.enable_weight_optimization = true;
|
|
|
|
let coordinator = create_ensemble_coordinator(config).await;
|
|
|
|
// Test 5.1: Set different performance metrics for each model
|
|
coordinator
|
|
.set_model_performance("DQN", 0.85, 0.15)
|
|
.await
|
|
.unwrap(); // High accuracy, low loss
|
|
coordinator
|
|
.set_model_performance("PPO", 0.75, 0.25)
|
|
.await
|
|
.unwrap(); // Medium
|
|
coordinator
|
|
.set_model_performance("MAMBA2", 0.65, 0.35)
|
|
.await
|
|
.unwrap(); // Lower
|
|
coordinator
|
|
.set_model_performance("TFT", 0.90, 0.10)
|
|
.await
|
|
.unwrap(); // Highest
|
|
|
|
// Test 5.2: Trigger weight optimization
|
|
coordinator.optimize_weights().await.unwrap();
|
|
|
|
// Test 5.3: TFT should have highest weight (best performance)
|
|
let weights = coordinator.get_current_weights().await.unwrap();
|
|
let tft_weight = weights.get("TFT").expect("INVARIANT: Key should exist in map");
|
|
|
|
for (model, weight) in weights.iter() {
|
|
if model != "TFT" {
|
|
assert!(
|
|
tft_weight >= weight,
|
|
"TFT (best performer) should have highest or equal weight"
|
|
);
|
|
}
|
|
}
|
|
|
|
// Test 5.4: MAMBA2 should have lowest weight (worst performance)
|
|
let mamba2_weight = weights.get("MAMBA2").expect("INVARIANT: Key should exist in map");
|
|
for (model, weight) in weights.iter() {
|
|
if model != "MAMBA2" {
|
|
assert!(
|
|
mamba2_weight <= weight,
|
|
"MAMBA2 (worst performer) should have lowest or equal weight"
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Test 6: Training failure recovery
|
|
#[tokio::test]
|
|
async fn test_training_failure_recovery() {
|
|
let config = create_valid_ensemble_config();
|
|
let mut coordinator = create_ensemble_coordinator(config).await;
|
|
|
|
// Test 6.1: Start training
|
|
let _job_id = coordinator.start_ensemble_training().await.unwrap();
|
|
|
|
// Test 6.2: Simulate one model failing
|
|
coordinator.simulate_model_failure("PPO").await.unwrap();
|
|
|
|
// Test 6.3: PPO should be in Failed state
|
|
let ppo_status = coordinator.get_model_status("PPO").await.unwrap();
|
|
assert_eq!(
|
|
ppo_status,
|
|
ModelTrainingStatus::Failed,
|
|
"PPO should be in Failed state"
|
|
);
|
|
|
|
// Test 6.4: Other models should continue training
|
|
for model_name in &["DQN", "MAMBA2", "TFT"] {
|
|
let status = coordinator.get_model_status(model_name).await.unwrap();
|
|
assert_ne!(
|
|
status,
|
|
ModelTrainingStatus::Failed,
|
|
"{} should not fail due to PPO failure",
|
|
model_name
|
|
);
|
|
}
|
|
|
|
// Test 6.5: Can retry failed model
|
|
let retry_result = coordinator.retry_failed_model("PPO").await;
|
|
assert!(retry_result.is_ok(), "Should be able to retry failed model");
|
|
|
|
// Test 6.6: PPO should return to training after retry
|
|
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
|
|
let ppo_status_after_retry = coordinator.get_model_status("PPO").await.unwrap();
|
|
assert_ne!(
|
|
ppo_status_after_retry,
|
|
ModelTrainingStatus::Failed,
|
|
"PPO should not be Failed after retry"
|
|
);
|
|
}
|
|
|
|
/// Test 7: Ensemble validation metrics
|
|
#[tokio::test]
|
|
async fn test_ensemble_validation_metrics() {
|
|
let config = create_valid_ensemble_config();
|
|
let mut coordinator = create_ensemble_coordinator(config).await;
|
|
|
|
// Test 7.1: Start training
|
|
coordinator.start_ensemble_training().await.unwrap();
|
|
coordinator.simulate_training_epochs(5).await.unwrap();
|
|
|
|
// Test 7.2: Should have ensemble-level metrics
|
|
let metrics = coordinator.get_ensemble_metrics().await.unwrap();
|
|
assert!(
|
|
metrics.contains_key("ensemble_train_loss"),
|
|
"Should have ensemble train loss"
|
|
);
|
|
assert!(
|
|
metrics.contains_key("ensemble_val_loss"),
|
|
"Should have ensemble val loss"
|
|
);
|
|
assert!(
|
|
metrics.contains_key("ensemble_accuracy"),
|
|
"Should have ensemble accuracy"
|
|
);
|
|
|
|
// Test 7.3: Ensemble metrics should be aggregated from all models
|
|
let ensemble_loss = metrics.get("ensemble_train_loss").expect("INVARIANT: Key should exist in map");
|
|
assert!(ensemble_loss > &0.0, "Ensemble loss should be positive");
|
|
|
|
// Test 7.4: Should track diversity metrics
|
|
assert!(
|
|
metrics.contains_key("prediction_diversity"),
|
|
"Should track prediction diversity"
|
|
);
|
|
let diversity = metrics.get("prediction_diversity").expect("INVARIANT: Key should exist in map");
|
|
assert!(
|
|
diversity >= &0.0 && diversity <= &1.0,
|
|
"Diversity should be in [0, 1]"
|
|
);
|
|
}
|
|
|
|
/// Test 8: Integration with ML Training Service
|
|
#[tokio::test]
|
|
async fn test_integration_with_ml_training_service() {
|
|
// This test verifies the coordinator integrates with existing ML training infrastructure
|
|
|
|
let config = create_valid_ensemble_config();
|
|
let mut coordinator = create_ensemble_coordinator(config).await;
|
|
|
|
// Test 8.1: Should use existing ProductionTrainingConfig
|
|
for model_name in &["DQN", "PPO", "MAMBA2", "TFT"] {
|
|
let model_config = coordinator
|
|
.get_model_training_config(model_name)
|
|
.await
|
|
.unwrap();
|
|
assert!(
|
|
model_config.model_config.input_dim > 0,
|
|
"Should have valid input dimension"
|
|
);
|
|
assert!(
|
|
!model_config.model_config.hidden_dims.is_empty(),
|
|
"Should have hidden layers"
|
|
);
|
|
}
|
|
|
|
// Test 8.2: Should respect existing safety configurations
|
|
let dqn_config = coordinator.get_model_training_config("DQN").await.unwrap();
|
|
assert!(
|
|
dqn_config.safety_config.safety_enabled,
|
|
"Should have safety enabled"
|
|
);
|
|
assert!(
|
|
dqn_config.gradient_config.max_gradient_norm > 0.0,
|
|
"Should have gradient clipping"
|
|
);
|
|
|
|
// Test 8.3: Should integrate with checkpoint manager
|
|
coordinator.start_ensemble_training().await.unwrap();
|
|
coordinator.simulate_training_epochs(1).await.unwrap();
|
|
|
|
let checkpoints = coordinator.get_all_checkpoints().await.unwrap();
|
|
assert_eq!(
|
|
checkpoints.len(),
|
|
4,
|
|
"Should have checkpoints for all 4 models"
|
|
);
|
|
}
|
|
|
|
// Helper functions for tests
|
|
|
|
/// Create valid ensemble configuration
|
|
fn create_valid_ensemble_config() -> EnsembleTrainingConfig {
|
|
let mut model_configs = HashMap::new();
|
|
let mut model_weights = HashMap::new();
|
|
|
|
// DQN configuration (33% weight)
|
|
model_configs.insert(
|
|
"DQN".to_string(),
|
|
create_model_config("DQN", 64, vec![256, 128], 32),
|
|
);
|
|
model_weights.insert("DQN".to_string(), 0.33);
|
|
|
|
// PPO configuration (33% weight)
|
|
model_configs.insert(
|
|
"PPO".to_string(),
|
|
create_model_config("PPO", 64, vec![256, 128], 32),
|
|
);
|
|
model_weights.insert("PPO".to_string(), 0.33);
|
|
|
|
// MAMBA-2 configuration (17% weight)
|
|
model_configs.insert(
|
|
"MAMBA2".to_string(),
|
|
create_model_config("MAMBA2", 64, vec![512, 256], 32),
|
|
);
|
|
model_weights.insert("MAMBA2".to_string(), 0.17);
|
|
|
|
// TFT configuration (17% weight)
|
|
model_configs.insert(
|
|
"TFT".to_string(),
|
|
create_model_config("TFT", 64, vec![512, 256, 128], 32),
|
|
);
|
|
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: 10,
|
|
checkpoint_interval_epochs: 1,
|
|
max_epochs: 100,
|
|
parallel_training: true,
|
|
created_at: Utc::now(),
|
|
}
|
|
}
|
|
|
|
/// Create model-specific training configuration
|
|
fn create_model_config(
|
|
_model_type: &str,
|
|
input_dim: usize,
|
|
hidden_dims: Vec<usize>,
|
|
output_dim: usize,
|
|
) -> ProductionTrainingConfig {
|
|
ProductionTrainingConfig {
|
|
model_config: ModelArchitectureConfig {
|
|
input_dim,
|
|
hidden_dims,
|
|
output_dim,
|
|
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 {
|
|
safety_enabled: true,
|
|
max_tensor_elements: 10_000_000,
|
|
max_inference_timeout_ms: 5000,
|
|
max_gpu_memory_bytes: 4_000_000_000,
|
|
drift_sensitivity: 0.1,
|
|
financial_precision: 8,
|
|
nan_infinity_checks: true,
|
|
max_prediction_value: 100.0,
|
|
min_prediction_value: -100.0,
|
|
bounds_checking: true,
|
|
auto_fallback: true,
|
|
max_retries: 3,
|
|
},
|
|
gradient_config: GradientSafetyConfig {
|
|
max_gradient_norm: 1.0,
|
|
min_gradient_norm: 1e-8,
|
|
max_individual_gradient: 5.0,
|
|
enable_norm_clipping: true,
|
|
enable_value_clipping: true,
|
|
enable_nan_detection: true,
|
|
gradient_history_size: 100,
|
|
explosion_threshold: 10.0,
|
|
min_gradient_history: 10,
|
|
enable_adaptive_scaling: false,
|
|
lr_adjustment_factor: 0.5,
|
|
base_learning_rate: 0.001,
|
|
},
|
|
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,
|
|
num_workers: 2,
|
|
gradient_accumulation_steps: 1,
|
|
},
|
|
}
|
|
}
|
|
|
|
/// Create ensemble coordinator instance
|
|
async fn create_ensemble_coordinator(
|
|
config: EnsembleTrainingConfig,
|
|
) -> EnsembleTrainingCoordinator {
|
|
EnsembleTrainingCoordinator::new(config)
|
|
.await
|
|
.expect("Failed to create coordinator")
|
|
}
|