e8a68ee39f
- Created data/examples/download_ml_training_data.rs using reqwest + Databento HTTP API - Downloaded 90 days × 4 symbols (ES.FUT, NQ.FUT, ZN.FUT, 6E.FUT) - Files saved to test_data/real/databento/ml_training/ - Total: 360 files, 15 MB compressed DBN format - Used existing Rust pattern from download_nq_fut.rs - API key loaded from .env file - 100% success rate (360/360 files) - Ready for ML training benchmarks Next: Create simplified training benchmark for RTX 3050 Ti GPU measurements
346 lines
11 KiB
Rust
346 lines
11 KiB
Rust
#![allow(unused_crate_dependencies)]
|
|
use candle_core::{DType, Device, Tensor};
|
|
use ml::mamba::{Mamba2Config, Mamba2State, SelectiveStateSpace, StateImportance};
|
|
use tokio;
|
|
|
|
mod real_data_helpers;
|
|
use real_data_helpers::{load_tft_sequences, real_data_available};
|
|
|
|
/// Helper function to create test Mamba2Config with reasonable defaults
|
|
fn create_test_config(d_model: usize, d_state: usize, num_layers: usize) -> Mamba2Config {
|
|
Mamba2Config {
|
|
d_model,
|
|
d_state,
|
|
d_head: d_state,
|
|
num_heads: 4,
|
|
expand: 2,
|
|
num_layers,
|
|
dropout: 0.1,
|
|
use_ssd: true,
|
|
use_selective_state: true,
|
|
hardware_aware: false, // Disable for tests
|
|
target_latency_us: 100,
|
|
max_seq_len: 1024,
|
|
learning_rate: 1e-4,
|
|
weight_decay: 1e-5,
|
|
grad_clip: 1.0,
|
|
warmup_steps: 100,
|
|
batch_size: 1,
|
|
seq_len: 256,
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_mamba2_config_creation() {
|
|
let config = create_test_config(512, 64, 6);
|
|
assert_eq!(config.d_model, 512);
|
|
assert_eq!(config.d_state, 64);
|
|
assert_eq!(config.num_layers, 6);
|
|
assert_eq!(config.expand, 2);
|
|
assert!(config.dropout >= 0.0 && config.dropout <= 1.0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_mamba2_state_creation() {
|
|
let config = create_test_config(256, 32, 4);
|
|
let state_result = Mamba2State::zeros(&config);
|
|
|
|
assert!(state_result.is_ok(), "State creation should succeed");
|
|
let state = state_result.unwrap();
|
|
|
|
assert_eq!(state.hidden_states.len(), config.num_layers);
|
|
assert_eq!(state.ssm_states.len(), config.num_layers);
|
|
assert!(!state.selective_state.is_empty());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_selective_state_creation() {
|
|
let config = create_test_config(128, 16, 2);
|
|
let selective_state_result = SelectiveStateSpace::new(&config);
|
|
|
|
assert!(
|
|
selective_state_result.is_ok(),
|
|
"Selective state creation should succeed"
|
|
);
|
|
let selective_state = selective_state_result.unwrap();
|
|
|
|
// Verify initialization
|
|
assert_eq!(
|
|
selective_state.importance_tracker.len(),
|
|
config.d_model * config.expand
|
|
);
|
|
assert_eq!(selective_state.active_indices.len(), 0); // Initially empty
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_selective_state_importance_scoring() {
|
|
let config = create_test_config(128, 16, 2);
|
|
let mut selective_state = SelectiveStateSpace::new(&config).unwrap();
|
|
let mut test_state = Mamba2State::zeros(&config).unwrap();
|
|
|
|
let device = Device::Cpu;
|
|
let input = Tensor::randn(0.0, 1.0, &[1, 10, 128], &device).unwrap();
|
|
|
|
// Update importance scores
|
|
let result = selective_state.update_importance_scores(&input, &mut test_state);
|
|
if let Err(e) = &result {
|
|
eprintln!("Error in update_importance_scores: {:?}", e);
|
|
}
|
|
assert!(result.is_ok(), "Importance score update should succeed");
|
|
|
|
// Verify that importance tracking is working - active_indices should be populated
|
|
assert!(
|
|
selective_state.active_indices.len() > 0,
|
|
"Should have some active indices"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_selective_state_compression() {
|
|
let config = create_test_config(128, 16, 2);
|
|
let mut selective_state = SelectiveStateSpace::new(&config).unwrap();
|
|
let mut test_state = Mamba2State::zeros(&config).unwrap();
|
|
|
|
// Compress a state component
|
|
let result = selective_state.compress_state_component(0, &mut test_state);
|
|
assert!(result.is_ok(), "State compression should succeed");
|
|
|
|
// Verify compression occurred
|
|
assert!(
|
|
selective_state.compressed_states.len() > 0,
|
|
"Should have compressed states"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_selective_state_decompression() {
|
|
let config = create_test_config(128, 16, 2);
|
|
let mut selective_state = SelectiveStateSpace::new(&config).unwrap();
|
|
let mut test_state = Mamba2State::zeros(&config).unwrap();
|
|
|
|
// First compress
|
|
let _ = selective_state.compress_state_component(0, &mut test_state);
|
|
|
|
// Then decompress
|
|
let result = selective_state.decompress_state_component(0, &mut test_state);
|
|
assert!(result.is_ok(), "State decompression should succeed");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_mamba2_state_transitions() {
|
|
let config = create_test_config(64, 8, 2);
|
|
let state = Mamba2State::zeros(&config).unwrap();
|
|
|
|
// Test state initialization
|
|
assert_eq!(state.hidden_states.len(), config.num_layers);
|
|
assert_eq!(state.ssm_states.len(), config.num_layers);
|
|
assert!(!state.selective_state.is_empty());
|
|
|
|
// Test state structure
|
|
assert!(state.compression_indices.is_empty());
|
|
assert!(state.metrics.is_empty());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_state_importance_new() {
|
|
let importance = StateImportance::new();
|
|
|
|
assert_eq!(importance.score, 0.0);
|
|
assert_eq!(importance.usage_count, 0);
|
|
assert_eq!(importance.last_access, 0);
|
|
assert_eq!(importance.moving_average, 0.0);
|
|
assert_eq!(importance.variance, 0.0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_state_importance_update() {
|
|
let mut importance = StateImportance::new();
|
|
|
|
importance.update(1.0, 1, 0.9);
|
|
assert!(importance.score > 0.0);
|
|
assert_eq!(importance.usage_count, 1);
|
|
assert_eq!(importance.last_access, 1);
|
|
|
|
importance.update(0.5, 2, 0.9);
|
|
assert_eq!(importance.usage_count, 2);
|
|
assert_eq!(importance.last_access, 2);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_state_importance_effective_importance() {
|
|
let mut importance = StateImportance::new();
|
|
|
|
importance.update(1.0, 1, 0.9);
|
|
let effective = importance.effective_importance();
|
|
|
|
assert!(effective > 0.0);
|
|
assert!(effective <= 1.0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_tensor_creation() {
|
|
let device = Device::Cpu;
|
|
|
|
// Test various tensor shapes
|
|
let t1 = Tensor::zeros(&[1, 128], DType::F32, &device);
|
|
assert!(t1.is_ok());
|
|
|
|
let t2 = Tensor::randn(0.0, 1.0, &[1, 10, 256], &device);
|
|
assert!(t2.is_ok());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_config_validation() {
|
|
let config = create_test_config(256, 32, 4);
|
|
|
|
// Verify reasonable defaults
|
|
assert!(config.d_model > 0);
|
|
assert!(config.d_state > 0);
|
|
assert!(config.num_layers > 0);
|
|
assert!(config.expand > 0);
|
|
assert!(config.dropout >= 0.0 && config.dropout <= 1.0);
|
|
assert!(config.learning_rate > 0.0);
|
|
assert!(config.weight_decay >= 0.0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_multiple_state_operations() {
|
|
let config = create_test_config(64, 8, 2);
|
|
let mut selective_state = SelectiveStateSpace::new(&config).unwrap();
|
|
let mut test_state = Mamba2State::zeros(&config).unwrap();
|
|
|
|
// Perform multiple operations in sequence
|
|
let device = Device::Cpu;
|
|
let input = Tensor::randn(0.0, 1.0, &[1, 10, 64], &device).unwrap();
|
|
|
|
// Update importance
|
|
let _ = selective_state.update_importance_scores(&input, &mut test_state);
|
|
|
|
// Compress some states
|
|
if test_state.selective_state.len() > 2 {
|
|
let _ = selective_state.compress_state_component(0, &mut test_state);
|
|
let _ = selective_state.compress_state_component(1, &mut test_state);
|
|
}
|
|
|
|
// Verify operations completed
|
|
assert!(selective_state.active_indices.len() > 0);
|
|
}
|
|
|
|
// Basic integration-style test
|
|
#[tokio::test]
|
|
async fn test_selective_state_full_workflow() {
|
|
let config = create_test_config(128, 16, 2);
|
|
let mut selective_state = SelectiveStateSpace::new(&config).unwrap();
|
|
let mut state = Mamba2State::zeros(&config).unwrap();
|
|
|
|
let device = Device::Cpu;
|
|
|
|
// Simulate a few timesteps
|
|
for _i in 0..5 {
|
|
let input = Tensor::randn(0.0, 1.0, &[1, 10, 128], &device).unwrap();
|
|
|
|
// Update importance scores
|
|
let result = selective_state.update_importance_scores(&input, &mut state);
|
|
assert!(result.is_ok());
|
|
|
|
// Active indices should be maintained
|
|
assert!(selective_state.active_indices.len() > 0);
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Real Market Data Tests
|
|
// ============================================================================
|
|
|
|
/// Test selective state importance scoring with real market data
|
|
#[tokio::test]
|
|
async fn test_selective_state_importance_scoring_real_data() {
|
|
// Skip if no real data available
|
|
if !real_data_available().await {
|
|
eprintln!("Skipping test: real data not available");
|
|
return;
|
|
}
|
|
|
|
let config = create_test_config(128, 16, 2);
|
|
let mut selective_state = SelectiveStateSpace::new(&config).unwrap();
|
|
let mut test_state = Mamba2State::zeros(&config).unwrap();
|
|
|
|
// Load real market data sequences (10 sequences, 10 timesteps each, padded to 128 features)
|
|
let sequences = load_tft_sequences(10, 10, 128).await.unwrap();
|
|
if sequences.is_empty() {
|
|
eprintln!("Skipping test: no sequences loaded");
|
|
return;
|
|
}
|
|
|
|
let device = Device::Cpu;
|
|
|
|
// Convert first sequence to tensor [1, 10, 128]
|
|
let sequence = &sequences[0];
|
|
let input_data: Vec<f32> = sequence.clone();
|
|
let input = Tensor::from_vec(input_data, &[1, 10, 128], &device).unwrap();
|
|
|
|
// Update importance scores with real data
|
|
let result = selective_state.update_importance_scores(&input, &mut test_state);
|
|
assert!(
|
|
result.is_ok(),
|
|
"Importance score update should succeed with real data"
|
|
);
|
|
|
|
// Verify that importance tracking is working
|
|
assert!(
|
|
selective_state.active_indices.len() > 0,
|
|
"Should have some active indices with real data"
|
|
);
|
|
}
|
|
|
|
/// Test selective state full workflow with real market data
|
|
#[tokio::test]
|
|
async fn test_selective_state_full_workflow_real_data() {
|
|
// Skip if no real data available
|
|
if !real_data_available().await {
|
|
eprintln!("Skipping test: real data not available");
|
|
return;
|
|
}
|
|
|
|
let config = create_test_config(128, 16, 2);
|
|
let mut selective_state = SelectiveStateSpace::new(&config).unwrap();
|
|
let mut state = Mamba2State::zeros(&config).unwrap();
|
|
|
|
// Load real market data sequences (5 sequences for 5 timesteps)
|
|
let sequences = load_tft_sequences(5, 10, 128).await.unwrap();
|
|
if sequences.len() < 5 {
|
|
eprintln!("Skipping test: not enough sequences loaded");
|
|
return;
|
|
}
|
|
|
|
let device = Device::Cpu;
|
|
|
|
// Simulate timesteps with real data
|
|
for i in 0..5 {
|
|
let sequence = &sequences[i];
|
|
let input_data: Vec<f32> = sequence.clone();
|
|
let input = Tensor::from_vec(input_data, &[1, 10, 128], &device).unwrap();
|
|
|
|
// Update importance scores
|
|
let result = selective_state.update_importance_scores(&input, &mut state);
|
|
assert!(
|
|
result.is_ok(),
|
|
"Step {} should succeed with real data",
|
|
i
|
|
);
|
|
|
|
// Active indices should be maintained
|
|
assert!(
|
|
selective_state.active_indices.len() > 0,
|
|
"Step {} should have active indices",
|
|
i
|
|
);
|
|
}
|
|
|
|
// Verify state accumulated properly over timesteps
|
|
assert!(
|
|
state.selective_state.len() > 0,
|
|
"State should have accumulated over real data timesteps"
|
|
);
|
|
}
|