gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

test(ml): add TFT + Mamba2 checkpoint roundtrip integration tests

Browse Source 
Extend checkpoint_roundtrip.rs with 4 new tests for sequence-buffered models:

- TFT adapter deterministic inference: verifies same adapter produces
  identical direction/confidence on repeated calls with stable buffer
- TFT quantile metadata: confirms quantiles are absent during buffering
  phase and present (with correct count) after buffer fills
- Mamba2 adapter deterministic inference: same pattern as TFT, verifies
  direction/confidence stability and correct model name ("MAMBA-2")
- Mamba2 CheckpointManager roundtrip: saves/loads via Checkpointable
  trait on Mamba2SSM, verifies metadata tags and hyperparameters

All 10 tests (6 existing DQN/PPO + 4 new TFT/Mamba2) pass consistently.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-02-22 23:55:23 +01:00
parent 2775d60c25
commit edd5a503da
1 changed files with 382 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

383
tests/integration/checkpoint_roundtrip.rs
View File

@@ -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<ml::ensemble::inference_adapter::EnsemblePrediction> {
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,
);
}