diff --git a/tests/integration/checkpoint_roundtrip.rs b/tests/integration/checkpoint_roundtrip.rs index a8e2059ce..4d2b91414 100644 --- a/tests/integration/checkpoint_roundtrip.rs +++ b/tests/integration/checkpoint_roundtrip.rs @@ -7,14 +7,22 @@ //! - DQN: adapter roundtrip via `DqnInferenceAdapter::from_checkpoint` //! - DQN: `CheckpointManager` + `Checkpointable` trait flow //! - PPO: actor/critic checkpoint save/load via `PPO::save_checkpoint` / `PPO::load_checkpoint` +//! - TFT: dual adapter deterministic inference with sequence buffer +//! - TFT: quantile metadata verification after buffer fill +//! - Mamba2: dual adapter deterministic inference with sequence buffer +//! - Mamba2: `CheckpointManager` + `Checkpointable` trait on `Mamba2SSM` use ml::checkpoint::{CheckpointConfig, CheckpointManager, CompressionType}; use ml::prelude::{Device, Tensor}; use ml::dqn::agent::DQNAgent; use ml::dqn::dqn::{DQNConfig, DQN}; -use ml::ensemble::adapters::{DqnInferenceAdapter, PpoInferenceAdapter}; +use ml::ensemble::adapters::{ + DqnInferenceAdapter, Mamba2InferenceAdapter, PpoInferenceAdapter, TftInferenceAdapter, +}; use ml::ensemble::inference_adapter::{FeatureVector, ModelInferenceAdapter}; +use ml::mamba::{Mamba2Config, Mamba2SSM}; use ml::ppo::ppo::{PPOConfig, PPO}; +use ml::tft::TFTConfig; // --------------------------------------------------------------------------- // Helpers @@ -409,3 +417,376 @@ async fn test_ppo_inference_adapter_deterministic() { pred1.confidence, pred2.confidence, ); } + +// --------------------------------------------------------------------------- +// TFT + Mamba2 Helpers +// --------------------------------------------------------------------------- + +/// Small TFT config for fast tests. +/// input_dim = num_static + num_known + num_unknown = 6 + 6 + 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() + } +} + +/// Small Mamba2 config for fast tests. +fn small_mamba2_config() -> Mamba2Config { + Mamba2Config { + d_model: 32, + d_state: 8, + d_head: 8, + num_heads: 2, + expand: 2, + num_layers: 1, + max_seq_len: 8, + dropout: 0.0, + ..Default::default() + } +} + +const TFT_SEQ_LEN: usize = 4; +const MAMBA2_SEQ_LEN: usize = 4; + +/// Build a feature vector with 51 values at a given timestamp. +/// The adapter extracts and zero-pads as needed for its model dimensions. +fn make_fv(ts: i64) -> FeatureVector { + FeatureVector { + values: vec![0.1; 51], + timestamp: ts, + } +} + +/// Feed `count` feature vectors into an adapter to fill its sequence buffer. +/// Returns the predictions from each call (including the neutral ones). +fn fill_adapter_buffer( + adapter: &dyn ModelInferenceAdapter, + count: usize, + base_ts: i64, +) -> Vec { + let mut preds = Vec::with_capacity(count); + for i in 0..count { + let fv = make_fv(base_ts + i as i64); + let pred = adapter.predict(&fv).expect("predict during buffer fill should succeed"); + preds.push(pred); + } + preds +} + +// --------------------------------------------------------------------------- +// TFT: single adapter deterministic inference (same input => same output) +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn test_tft_adapter_deterministic_same_buffer() { + let config = small_tft_config(); + + // 1. Create a single TFT adapter + let adapter = TftInferenceAdapter::new(config, TFT_SEQ_LEN) + .expect("TftInferenceAdapter creation should succeed"); + + // 2. Fill the buffer with identical feature vectors + let base_ts: i64 = 1_700_000_000; + for i in 0..TFT_SEQ_LEN { + let fv = make_fv(base_ts + i as i64); + let pred = adapter + .predict(&fv) + .expect("predict should succeed during fill"); + + // While buffering, should return neutral + if i < TFT_SEQ_LEN - 1 { + assert_eq!( + pred.direction, 0.0, + "Should return neutral direction while buffering (step {i})" + ); + assert_eq!( + pred.confidence, 0.0, + "Should return zero confidence while buffering (step {i})" + ); + } + } + + // 3. Adapter should now be ready + assert!( + adapter.is_ready(), + "Adapter should be ready after filling buffer" + ); + + // 4. Feed identical FVs and compare predictions (buffer is full, so each + // call runs a real forward pass with the same buffer content) + let fv = make_fv(base_ts); + let pred1 = adapter + .predict(&fv) + .expect("first full-buffer predict should succeed"); + let pred2 = adapter + .predict(&fv) + .expect("second full-buffer predict should succeed"); + + // 5. Assert identical outputs (same weights, same buffer content => same output) + assert!( + (pred1.direction - pred2.direction).abs() < 1e-6, + "TFT adapter should be deterministic: direction {} vs {}", + pred1.direction, + pred2.direction, + ); + assert!( + (pred1.confidence - pred2.confidence).abs() < 1e-6, + "TFT adapter should be deterministic: confidence {} vs {}", + pred1.confidence, + pred2.confidence, + ); + + // 6. Direction and confidence bounds check + assert!( + pred1.direction >= -1.0 && pred1.direction <= 1.0, + "TFT direction {} out of [-1,1]", + pred1.direction, + ); + assert!( + pred1.confidence >= 0.0 && pred1.confidence <= 1.0, + "TFT confidence {} out of [0,1]", + pred1.confidence, + ); + + // 7. Model name + assert_eq!(pred1.model_name, "TFT"); +} + +// --------------------------------------------------------------------------- +// TFT: quantile metadata verification +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn test_tft_quantile_metadata_present_after_buffer_fill() { + let config = small_tft_config(); + let adapter = TftInferenceAdapter::new(config.clone(), TFT_SEQ_LEN) + .expect("TftInferenceAdapter creation should succeed"); + + // Fill the buffer + let base_ts: i64 = 1_700_000_000; + let preds = fill_adapter_buffer(&adapter, TFT_SEQ_LEN, base_ts); + + // Predictions before buffer is full should NOT have quantile metadata + for (i, pred) in preds.iter().enumerate() { + if i < TFT_SEQ_LEN - 1 { + assert!( + pred.metadata.quantiles.is_none(), + "Neutral prediction at step {i} should not have quantile metadata" + ); + } + } + + // The final prediction (buffer just filled) should have quantile metadata + let last_pred = preds + .last() + .expect("should have at least one prediction"); + assert!( + last_pred.metadata.quantiles.is_some(), + "TFT prediction with full buffer should include quantile metadata" + ); + + // Verify quantile count matches config.num_quantiles + let quantiles = last_pred + .metadata + .quantiles + .as_ref() + .expect("quantiles should be present"); + assert_eq!( + quantiles.len(), + config.num_quantiles, + "Quantile count should match config.num_quantiles ({}), got {}", + config.num_quantiles, + quantiles.len(), + ); + + // Quantiles should be finite numbers + for (i, &q) in quantiles.iter().enumerate() { + assert!( + q.is_finite(), + "Quantile at index {i} should be finite, got {q}" + ); + } + + // Additional prediction should also have quantiles + let fv_extra = make_fv(base_ts + TFT_SEQ_LEN as i64); + let pred_extra = adapter + .predict(&fv_extra) + .expect("extra predict should succeed"); + assert!( + pred_extra.metadata.quantiles.is_some(), + "Subsequent TFT prediction should also include quantile metadata" + ); + let quantiles_extra = pred_extra + .metadata + .quantiles + .as_ref() + .expect("quantiles should be present in extra prediction"); + assert_eq!( + quantiles_extra.len(), + config.num_quantiles, + "Subsequent quantile count should match config" + ); + + // Model name should be "TFT" + assert_eq!(last_pred.model_name, "TFT"); +} + +// --------------------------------------------------------------------------- +// Mamba2: single adapter deterministic inference (same input => same output) +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn test_mamba2_adapter_deterministic_same_buffer() { + let config = small_mamba2_config(); + + // 1. Create a single Mamba2 adapter + let adapter = Mamba2InferenceAdapter::new(config, MAMBA2_SEQ_LEN) + .expect("Mamba2InferenceAdapter creation should succeed"); + + // 2. Fill the buffer with identical feature vectors + let base_ts: i64 = 1_700_000_000; + for i in 0..MAMBA2_SEQ_LEN { + let fv = make_fv(base_ts); + let pred = adapter + .predict(&fv) + .expect("predict should succeed during fill"); + + // While buffering, should return neutral + if i < MAMBA2_SEQ_LEN - 1 { + assert_eq!( + pred.direction, 0.0, + "Should return neutral direction while buffering (step {i})" + ); + assert_eq!( + pred.confidence, 0.0, + "Should return zero confidence while buffering (step {i})" + ); + } + } + + // 3. Adapter should now be ready + assert!( + adapter.is_ready(), + "Adapter should be ready after filling buffer" + ); + + // 4. Feed identical FVs and compare predictions (buffer is full, each + // call shifts the ring buffer with the same value, so content is stable) + let fv = make_fv(base_ts); + let pred1 = adapter + .predict(&fv) + .expect("first full-buffer predict should succeed"); + let pred2 = adapter + .predict(&fv) + .expect("second full-buffer predict should succeed"); + + // 5. Assert identical outputs (same weights, same buffer content => same output) + assert!( + (pred1.direction - pred2.direction).abs() < 1e-6, + "Mamba2 adapter should be deterministic: direction {} vs {}", + pred1.direction, + pred2.direction, + ); + assert!( + (pred1.confidence - pred2.confidence).abs() < 1e-6, + "Mamba2 adapter should be deterministic: confidence {} vs {}", + pred1.confidence, + pred2.confidence, + ); + + // 6. Direction and confidence bounds check + assert!( + pred1.direction >= -1.0 && pred1.direction <= 1.0, + "Mamba2 direction {} out of [-1,1]", + pred1.direction, + ); + assert!( + pred1.confidence >= 0.0 && pred1.confidence <= 1.0, + "Mamba2 confidence {} out of [0,1]", + pred1.confidence, + ); + + // 7. Model name should be "MAMBA-2" + assert_eq!(pred1.model_name, "MAMBA-2"); + + // 8. Mamba2 should NOT have quantile metadata (that is TFT-specific) + assert!( + pred1.metadata.quantiles.is_none(), + "Mamba2 predictions should not have quantile metadata" + ); +} + +// --------------------------------------------------------------------------- +// Mamba2: CheckpointManager + Checkpointable trait roundtrip on Mamba2SSM +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn test_mamba2_checkpoint_manager_roundtrip() { + let tmp = tempfile::tempdir().expect("failed to create tempdir"); + let ckpt_config = CheckpointConfig { + base_dir: tmp.path().to_path_buf(), + compression: CompressionType::None, + auto_cleanup: false, + ..Default::default() + }; + + let manager = + CheckpointManager::new(ckpt_config).expect("CheckpointManager::new should succeed"); + + // Create a Mamba2SSM (implements Checkpointable) + let config = small_mamba2_config(); + let device = Device::Cpu; + let model_a = Mamba2SSM::new(config.clone(), &device) + .expect("Mamba2SSM::new should succeed for model_a"); + + // Save checkpoint + let ckpt_id = manager + .save_checkpoint(&model_a, Some(vec!["mamba2-roundtrip".to_string()])) + .await + .expect("save_checkpoint should succeed for Mamba2SSM"); + + // Create second model with same config + let mut model_b = + Mamba2SSM::new(config, &device).expect("Mamba2SSM::new should succeed for model_b"); + + // Load checkpoint into model_b + let metadata = manager + .load_checkpoint(&mut model_b, &ckpt_id) + .await + .expect("load_checkpoint should succeed for Mamba2SSM"); + + // Verify metadata + assert!( + metadata.tags.contains(&"mamba2-roundtrip".to_string()), + "Checkpoint tags should contain 'mamba2-roundtrip', got {:?}", + metadata.tags, + ); + assert_eq!( + metadata.model_type, + ml::ModelType::MAMBA, + "Checkpoint model type should be MAMBA" + ); + + // Verify hyperparameters were preserved + let d_model_val = metadata + .hyperparameters + .get("d_model") + .and_then(|v| v.as_u64()); + assert_eq!( + d_model_val, + Some(32), + "Checkpoint should preserve d_model=32, got {:?}", + d_model_val, + ); +}