edd5a503da
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>
793 lines
27 KiB
Rust
793 lines
27 KiB
Rust
//! Checkpoint roundtrip integration tests
|
|
//!
|
|
//! Verifies that saving a model checkpoint and loading it into a fresh model
|
|
//! produces identical inference output. Tests cover:
|
|
//!
|
|
//! - DQN: raw Q-network weight save/load via safetensors
|
|
//! - 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, 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
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn small_dqn_config() -> DQNConfig {
|
|
DQNConfig {
|
|
state_dim: 51,
|
|
num_actions: 45,
|
|
hidden_dims: vec![32, 32],
|
|
// Disable all advanced network variants so the forward pass goes through
|
|
// the base q_network (whose VarMap we save/load).
|
|
use_dueling: false,
|
|
use_distributional: false,
|
|
use_noisy_nets: false,
|
|
use_iqn: false,
|
|
use_cql: false,
|
|
..Default::default()
|
|
}
|
|
}
|
|
|
|
fn small_ppo_config() -> PPOConfig {
|
|
PPOConfig {
|
|
state_dim: 64,
|
|
num_actions: 45,
|
|
policy_hidden_dims: vec![32, 32],
|
|
value_hidden_dims: vec![32, 32],
|
|
..Default::default()
|
|
}
|
|
}
|
|
|
|
fn fixed_feature_vector_51() -> FeatureVector {
|
|
FeatureVector {
|
|
values: vec![0.3; 51],
|
|
timestamp: 1_700_000_000,
|
|
}
|
|
}
|
|
|
|
fn fixed_feature_vector_64() -> FeatureVector {
|
|
FeatureVector {
|
|
values: vec![0.3; 64],
|
|
timestamp: 1_700_000_000,
|
|
}
|
|
}
|
|
|
|
/// Extract Q-values from a DQN model given an f64 input vector.
|
|
///
|
|
/// Returns the raw Q-value vector (f64) for comparison.
|
|
fn dqn_q_values(model: &DQN, input: &[f64], device: &Device) -> Result<Vec<f64>, ml::MLError> {
|
|
let f32_values: Vec<f32> = input.iter().map(|&v| v as f32).collect();
|
|
let len = f32_values.len();
|
|
let tensor = Tensor::from_vec(f32_values, (1, len), device)
|
|
.map_err(|e| ml::MLError::ModelError(format!("Failed to create tensor: {e}")))?;
|
|
let output = model.forward(&tensor)?;
|
|
let squeezed = output
|
|
.squeeze(0)
|
|
.map_err(|e| ml::MLError::ModelError(format!("Failed to squeeze: {e}")))?;
|
|
let f32_vec: Vec<f32> = squeezed
|
|
.to_vec1()
|
|
.map_err(|e| ml::MLError::ModelError(format!("Failed to extract vec: {e}")))?;
|
|
Ok(f32_vec.iter().map(|&v| v as f64).collect())
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// DQN: raw Q-network weight roundtrip
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[tokio::test]
|
|
async fn test_dqn_raw_weight_roundtrip_exact_output() {
|
|
let config = small_dqn_config();
|
|
let device = Device::Cpu;
|
|
|
|
// 1. Create DQN model with random weights
|
|
let model_a = DQN::new(config.clone()).expect("DQN::new should succeed");
|
|
|
|
// 2. Run inference to get Q-values before save
|
|
let input = vec![0.3f64; 51];
|
|
let q_before = dqn_q_values(&model_a, &input, &device).expect("forward should succeed");
|
|
assert!(!q_before.is_empty(), "Q-values should not be empty");
|
|
|
|
// 3. Save weights to temp directory
|
|
let tmp = tempfile::tempdir().expect("failed to create tempdir");
|
|
let ckpt_path = tmp.path().join("dqn_roundtrip.safetensors");
|
|
model_a
|
|
.get_q_network_vars()
|
|
.save(ckpt_path.as_path())
|
|
.expect("VarMap::save should succeed");
|
|
|
|
// 4. Create second DQN model (same config, different random init)
|
|
let mut model_b = DQN::new(config).expect("DQN::new should succeed for model_b");
|
|
|
|
// 5. Load weights from checkpoint
|
|
let path_str = ckpt_path
|
|
.to_str()
|
|
.expect("temp path should be valid UTF-8");
|
|
model_b
|
|
.load_from_safetensors(path_str)
|
|
.expect("load_from_safetensors should succeed");
|
|
|
|
// 6. Run inference on same input
|
|
let q_after = dqn_q_values(&model_b, &input, &device).expect("forward after load should succeed");
|
|
|
|
// 7. Assert exact element-wise equality
|
|
assert_eq!(
|
|
q_before.len(),
|
|
q_after.len(),
|
|
"Q-value vector length mismatch: {} vs {}",
|
|
q_before.len(),
|
|
q_after.len(),
|
|
);
|
|
|
|
for (i, (a, b)) in q_before.iter().zip(q_after.iter()).enumerate() {
|
|
assert!(
|
|
(a - b).abs() < 1e-6,
|
|
"Q-value mismatch at index {i}: {a} vs {b}",
|
|
);
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// DQN: adapter roundtrip via from_checkpoint
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[tokio::test]
|
|
async fn test_dqn_adapter_roundtrip_via_from_checkpoint() {
|
|
let config = small_dqn_config();
|
|
|
|
// 1. Create a DQN model, save its weights
|
|
let model = DQN::new(config.clone()).expect("DQN::new should succeed");
|
|
let tmp = tempfile::tempdir().expect("failed to create tempdir");
|
|
let ckpt_path = tmp.path().join("dqn_adapter_test.safetensors");
|
|
model
|
|
.get_q_network_vars()
|
|
.save(ckpt_path.as_path())
|
|
.expect("VarMap::save should succeed");
|
|
|
|
// 2. Get reference Q-values from original model
|
|
let device = Device::Cpu;
|
|
let input = vec![0.3f64; 51];
|
|
let q_original =
|
|
dqn_q_values(&model, &input, &device).expect("original forward should succeed");
|
|
|
|
// 3. Load into a DqnInferenceAdapter via from_checkpoint
|
|
let path_str = ckpt_path
|
|
.to_str()
|
|
.expect("temp path should be valid UTF-8");
|
|
let adapter = DqnInferenceAdapter::from_checkpoint(config, path_str)
|
|
.expect("from_checkpoint should succeed");
|
|
|
|
// 4. Run inference through the adapter
|
|
let fv = fixed_feature_vector_51();
|
|
let pred = adapter.predict(&fv).expect("adapter predict should succeed");
|
|
|
|
// 5. Verify prediction shape and ranges
|
|
assert!(
|
|
pred.direction >= -1.0 && pred.direction <= 1.0,
|
|
"direction {} out of [-1,1]",
|
|
pred.direction,
|
|
);
|
|
assert!(
|
|
pred.confidence >= 0.0 && pred.confidence <= 1.0,
|
|
"confidence {} out of [0,1]",
|
|
pred.confidence,
|
|
);
|
|
|
|
// 6. Verify Q-values from adapter match original model
|
|
let adapter_q = pred
|
|
.metadata
|
|
.q_values
|
|
.as_ref()
|
|
.expect("DQN adapter should emit q_values");
|
|
assert_eq!(
|
|
q_original.len(),
|
|
adapter_q.len(),
|
|
"Q-value vector length mismatch",
|
|
);
|
|
for (i, (orig, loaded)) in q_original.iter().zip(adapter_q.iter()).enumerate() {
|
|
assert!(
|
|
(orig - loaded).abs() < 1e-5,
|
|
"Q-value mismatch at index {i}: original={orig} vs loaded={loaded}",
|
|
);
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// DQN: CheckpointManager + Checkpointable trait roundtrip
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[tokio::test]
|
|
async fn test_dqn_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 DQNAgent (the type that implements Checkpointable)
|
|
let config = small_dqn_config();
|
|
let agent_a = DQNAgent::new(config.clone()).expect("DQNAgent::new should succeed");
|
|
|
|
// Save checkpoint
|
|
let ckpt_id = manager
|
|
.save_checkpoint(&agent_a, Some(vec!["roundtrip-test".to_string()]))
|
|
.await
|
|
.expect("save_checkpoint should succeed");
|
|
|
|
// Create second agent with same config
|
|
let mut agent_b = DQNAgent::new(config).expect("DQNAgent::new should succeed for agent_b");
|
|
|
|
// Load checkpoint into agent_b
|
|
let metadata = manager
|
|
.load_checkpoint(&mut agent_b, &ckpt_id)
|
|
.await
|
|
.expect("load_checkpoint should succeed");
|
|
|
|
assert_eq!(metadata.model_name, "dqn_agent");
|
|
assert!(metadata.tags.contains(&"roundtrip-test".to_string()));
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// PPO: actor/critic checkpoint roundtrip
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[tokio::test]
|
|
async fn test_ppo_checkpoint_roundtrip_exact_output() {
|
|
let config = small_ppo_config();
|
|
|
|
// 1. Create PPO model with random weights
|
|
let ppo_a = PPO::new(config.clone()).expect("PPO::new should succeed");
|
|
|
|
// 2. Run inference to get action probabilities before save
|
|
let device = Device::Cpu;
|
|
let f32_input: Vec<f32> = vec![0.3f32; 64];
|
|
let input_tensor = Tensor::from_vec(f32_input.clone(), (1, 64), &device)
|
|
.expect("Tensor creation should succeed");
|
|
let probs_before = ppo_a
|
|
.actor
|
|
.action_probabilities(&input_tensor)
|
|
.expect("actor forward should succeed");
|
|
let probs_before_vec: Vec<f32> = probs_before
|
|
.squeeze(0)
|
|
.expect("squeeze should succeed")
|
|
.to_vec1()
|
|
.expect("to_vec1 should succeed");
|
|
|
|
// 3. Save checkpoint (actor + critic + metadata)
|
|
let tmp = tempfile::tempdir().expect("failed to create tempdir");
|
|
let actor_path = tmp.path().join("ppo_actor.safetensors");
|
|
let critic_path = tmp.path().join("ppo_critic.safetensors");
|
|
let meta_path = tmp.path().join("ppo_metadata.json");
|
|
|
|
let actor_str = actor_path
|
|
.to_str()
|
|
.expect("actor path should be valid UTF-8");
|
|
let critic_str = critic_path
|
|
.to_str()
|
|
.expect("critic path should be valid UTF-8");
|
|
let meta_str = meta_path
|
|
.to_str()
|
|
.expect("meta path should be valid UTF-8");
|
|
|
|
ppo_a
|
|
.save_checkpoint(actor_str, critic_str, meta_str)
|
|
.expect("PPO save_checkpoint should succeed");
|
|
|
|
// 4. Load into a fresh PPO model
|
|
let ppo_b = PPO::load_checkpoint(actor_str, critic_str, config.clone(), device)
|
|
.expect("PPO::load_checkpoint should succeed");
|
|
|
|
// 5. Run inference on same input
|
|
let input_tensor_b = Tensor::from_vec(f32_input, (1, 64), &Device::Cpu)
|
|
.expect("Tensor creation should succeed");
|
|
let probs_after = ppo_b
|
|
.actor
|
|
.action_probabilities(&input_tensor_b)
|
|
.expect("actor forward after load should succeed");
|
|
let probs_after_vec: Vec<f32> = probs_after
|
|
.squeeze(0)
|
|
.expect("squeeze should succeed")
|
|
.to_vec1()
|
|
.expect("to_vec1 should succeed");
|
|
|
|
// 6. Assert element-wise equality
|
|
assert_eq!(
|
|
probs_before_vec.len(),
|
|
probs_after_vec.len(),
|
|
"Probability vector length mismatch: {} vs {}",
|
|
probs_before_vec.len(),
|
|
probs_after_vec.len(),
|
|
);
|
|
for (i, (a, b)) in probs_before_vec
|
|
.iter()
|
|
.zip(probs_after_vec.iter())
|
|
.enumerate()
|
|
{
|
|
assert!(
|
|
(a - b).abs() < 1e-5,
|
|
"PPO action probability mismatch at index {i}: {a} vs {b}",
|
|
);
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// PPO: adapter roundtrip (save PPO, load into PpoInferenceAdapter-style check)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[tokio::test]
|
|
async fn test_ppo_adapter_inference_after_checkpoint() {
|
|
let config = small_ppo_config();
|
|
|
|
// 1. Create PPO model and save actor checkpoint
|
|
let ppo = PPO::new(config.clone()).expect("PPO::new should succeed");
|
|
let tmp = tempfile::tempdir().expect("failed to create tempdir");
|
|
let actor_path = tmp.path().join("ppo_actor.safetensors");
|
|
let critic_path = tmp.path().join("ppo_critic.safetensors");
|
|
let meta_path = tmp.path().join("ppo_meta.json");
|
|
|
|
let actor_str = actor_path
|
|
.to_str()
|
|
.expect("actor path should be valid UTF-8");
|
|
let critic_str = critic_path
|
|
.to_str()
|
|
.expect("critic path should be valid UTF-8");
|
|
let meta_str = meta_path
|
|
.to_str()
|
|
.expect("meta path should be valid UTF-8");
|
|
|
|
ppo.save_checkpoint(actor_str, critic_str, meta_str)
|
|
.expect("PPO save_checkpoint should succeed");
|
|
|
|
// 2. Load into a fresh PPO and wrap in PpoInferenceAdapter-style inference
|
|
let device = Device::Cpu;
|
|
let ppo_loaded = PPO::load_checkpoint(actor_str, critic_str, config, device)
|
|
.expect("PPO::load_checkpoint should succeed");
|
|
|
|
// 3. Run inference through the loaded model's actor
|
|
let f32_input: Vec<f32> = vec![0.3f32; 64];
|
|
let input_tensor = Tensor::from_vec(f32_input, (1, 64), &Device::Cpu)
|
|
.expect("Tensor creation should succeed");
|
|
let probs = ppo_loaded
|
|
.actor
|
|
.action_probabilities(&input_tensor)
|
|
.expect("actor inference should succeed");
|
|
|
|
let probs_vec: Vec<f32> = probs
|
|
.squeeze(0)
|
|
.expect("squeeze should succeed")
|
|
.to_vec1()
|
|
.expect("to_vec1 should succeed");
|
|
|
|
// 4. Verify basic properties: probabilities sum to ~1.0, all non-negative
|
|
let sum: f32 = probs_vec.iter().sum();
|
|
assert!(
|
|
(sum - 1.0).abs() < 1e-4,
|
|
"PPO action probabilities should sum to ~1.0, got {sum}",
|
|
);
|
|
for (i, &p) in probs_vec.iter().enumerate() {
|
|
assert!(
|
|
p >= 0.0,
|
|
"PPO action probability at index {i} should be non-negative, got {p}",
|
|
);
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// PPO: PpoInferenceAdapter deterministic on same weights
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[tokio::test]
|
|
async fn test_ppo_inference_adapter_deterministic() {
|
|
let config = small_ppo_config();
|
|
let adapter =
|
|
PpoInferenceAdapter::new(config).expect("PpoInferenceAdapter::new should succeed");
|
|
|
|
let fv = fixed_feature_vector_64();
|
|
let pred1 = adapter.predict(&fv).expect("first predict should succeed");
|
|
let pred2 = adapter.predict(&fv).expect("second predict should succeed");
|
|
|
|
assert_eq!(
|
|
pred1.direction, pred2.direction,
|
|
"PPO adapter should be deterministic: {} vs {}",
|
|
pred1.direction, pred2.direction,
|
|
);
|
|
assert_eq!(
|
|
pred1.confidence, pred2.confidence,
|
|
"PPO adapter confidence should be deterministic: {} vs {}",
|
|
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,
|
|
);
|
|
}
|