#![allow(unused_crate_dependencies)] //! Integration test: 4-model inference ensemble (DQN + PPO + Mamba2 + TFT) //! //! Constructs all four adapters with small configs, pre-warms the //! sequence-based models (Mamba2, TFT) by feeding feature vectors into //! their buffers, then runs the full ensemble and verifies bounded //! predictions from all 4 models. use ml::dqn::DQNConfig; use ml::ensemble::adapters::{ DqnInferenceAdapter, Mamba2InferenceAdapter, PpoInferenceAdapter, TftInferenceAdapter, }; use ml::ensemble::inference_ensemble::InferenceEnsemble; use ml::ensemble::{FeatureVector, ModelInferenceAdapter}; use ml::mamba::Mamba2Config; use ml::ppo::PPOConfig; use ml::tft::TFTConfig; /// Small DQN config -- 51 input features, 45 actions, two small hidden layers. fn small_dqn_config() -> DQNConfig { DQNConfig { state_dim: 51, num_actions: 45, hidden_dims: vec![64, 64], ..Default::default() } } /// Small PPO config -- 64 input (zero-padded from 51), 45 actions, small MLPs. fn small_ppo_config() -> PPOConfig { PPOConfig { state_dim: 64, num_actions: 45, policy_hidden_dims: vec![64, 64], value_hidden_dims: vec![64, 64], ..Default::default() } } /// Small Mamba2 config -- d_model=64 (>51 features), 1 layer, seq_len=4. fn small_mamba2_config() -> Mamba2Config { Mamba2Config { d_model: 64, d_state: 16, d_head: 16, num_heads: 2, expand: 2, num_layers: 1, max_seq_len: 8, dropout: 0.0, ..Default::default() } } /// Small TFT config -- input_dim=20, 1 layer, seq_len=4. /// input_dim = num_static(6) + num_known(6) + num_unknown(8) = 20 fn small_tft_config() -> TFTConfig { TFTConfig { input_dim: 20, hidden_dim: 32, num_heads: 2, num_layers: 1, prediction_horizon: 5, sequence_length: 4, num_quantiles: 9, num_static_features: 6, num_known_features: 6, num_unknown_features: 8, dropout_rate: 0.0, ..Default::default() } } const SEQ_LEN: usize = 4; const NUM_FEATURE_VECTORS: usize = 5; /// Build a 51-dimensional feature vector with a distinguishing pattern. fn make_feature_vector(index: usize) -> FeatureVector { let base = 0.01 * (index as f64 + 1.0); FeatureVector { values: vec![base; 51], timestamp: 1_700_000_000_000_000 + index as i64, } } #[test] fn test_four_model_ensemble_integration() { // ---- construct all 4 adapters ---- let dqn = DqnInferenceAdapter::new(small_dqn_config()) .expect("DqnInferenceAdapter::new should succeed with small config"); let ppo = PpoInferenceAdapter::new(small_ppo_config()) .expect("PpoInferenceAdapter::new should succeed with small config"); let mamba2 = Mamba2InferenceAdapter::new(small_mamba2_config(), SEQ_LEN) .expect("Mamba2InferenceAdapter::new should succeed with small config"); let tft = TftInferenceAdapter::new(small_tft_config(), SEQ_LEN) .expect("TftInferenceAdapter::new should succeed with small config"); // ---- DQN + PPO are ready immediately; Mamba2 + TFT need buffer fill ---- assert!(dqn.is_ready(), "DQN should be ready immediately"); assert!(ppo.is_ready(), "PPO should be ready immediately"); assert!( !mamba2.is_ready(), "Mamba2 should NOT be ready (empty buffer)" ); assert!(!tft.is_ready(), "TFT should NOT be ready (empty buffer)"); // ---- pre-warm sequence-based adapters (Mamba2, TFT) ---- // The ensemble only calls predict on ready adapters, so we must fill // the internal sequence buffers before assembling the ensemble. for i in 0..SEQ_LEN { let fv = make_feature_vector(i); let mamba2_pred = mamba2 .predict(&fv) .unwrap_or_else(|e| panic!("Mamba2 warm-up {} should not error: {}", i, e)); let tft_pred = tft .predict(&fv) .unwrap_or_else(|e| panic!("TFT warm-up {} should not error: {}", i, e)); // Before the buffer is full, direction and confidence should be neutral (0.0) if i < SEQ_LEN - 1 { assert_eq!( mamba2_pred.direction, 0.0, "Mamba2 should return neutral direction during warm-up (step {})", i ); assert_eq!( tft_pred.direction, 0.0, "TFT should return neutral direction during warm-up (step {})", i ); } else { // Buffer is now full -- should produce real predictions assert!( mamba2_pred.direction >= -1.0 && mamba2_pred.direction <= 1.0, "Mamba2 warm-up final: direction {} out of [-1, 1]", mamba2_pred.direction ); assert!( tft_pred.direction >= -1.0 && tft_pred.direction <= 1.0, "TFT warm-up final: direction {} out of [-1, 1]", tft_pred.direction ); } println!( "Warm-up {}: mamba2 dir={:.4} conf={:.4}, tft dir={:.4} conf={:.4}", i, mamba2_pred.direction, mamba2_pred.confidence, tft_pred.direction, tft_pred.confidence ); } // All 4 adapters should now be ready assert!(mamba2.is_ready(), "Mamba2 should be ready after warm-up"); assert!(tft.is_ready(), "TFT should be ready after warm-up"); // ---- build the ensemble with all 4 pre-warmed adapters ---- let adapters: Vec> = vec![Box::new(dqn), Box::new(ppo), Box::new(mamba2), Box::new(tft)]; let ensemble = InferenceEnsemble::new(adapters); // All 4 models should be ready assert_eq!( ensemble.ready_count(), 4, "All 4 adapters should be ready after warm-up" ); // ---- feed 5 feature vectors through the full ensemble ---- for i in 0..NUM_FEATURE_VECTORS { let fv = make_feature_vector(SEQ_LEN + i); let pred = ensemble .predict(&fv) .unwrap_or_else(|e| panic!("Ensemble prediction {} should succeed: {}", i, e)); // Direction must be in [-1.0, 1.0] assert!( pred.direction >= -1.0 && pred.direction <= 1.0, "FV {}: direction {} out of [-1, 1]", i, pred.direction ); // Confidence must be in [0.0, 1.0] assert!( pred.confidence >= 0.0 && pred.confidence <= 1.0, "FV {}: confidence {} out of [0, 1]", i, pred.confidence ); // All 4 models should contribute (ensemble name includes all) assert!( pred.model_name.contains("DQN"), "FV {}: ensemble name should include DQN, got: {}", i, pred.model_name ); assert!( pred.model_name.contains("PPO"), "FV {}: ensemble name should include PPO, got: {}", i, pred.model_name ); assert!( pred.model_name.contains("MAMBA-2"), "FV {}: ensemble name should include MAMBA-2, got: {}", i, pred.model_name ); assert!( pred.model_name.contains("TFT"), "FV {}: ensemble name should include TFT, got: {}", i, pred.model_name ); println!( "Ensemble FV {}: direction={:.4}, confidence={:.4}, model={}", i, pred.direction, pred.confidence, pred.model_name ); } // ---- confirm ready_count is still 4 after all predictions ---- assert_eq!( ensemble.ready_count(), 4, "All 4 adapters should remain ready after {} ensemble predictions", NUM_FEATURE_VECTORS ); println!("All 4-model ensemble integration checks passed."); }