## Major Achievements ### 1. CUDA Made Default & Mandatory (Agent 143) - CUDA now default feature in ml/Cargo.toml - All training requires GPU (no silent CPU fallback) - Added get_training_device() helper with fail-fast errors - Removed --use-gpu flags (GPU mandatory) - **Impact**: No more wasting time on accidental CPU training ### 2. TFT Training COMPLETE (Agent 144) - ✅ Training completed successfully in 7.6 minutes - ✅ Early stopping at epoch 100/200 (best val loss: 0.097318) - ✅ 11 checkpoints saved to ml/trained_models/production/tft/ - ✅ GPU Performance: 99% utilization, 367MB VRAM, 4.4s/epoch - ✅ 10x speedup vs CPU (4.4s vs 43-55s per epoch) - **Status**: PRODUCTION READY ### 3. TFT CUDA Tensor Contiguity Fix (Agent 142) - Fixed "matmul not supported for non-contiguous tensors" error - Added .contiguous() call after narrow() operation in QuantileLayer - Enabled CUDA-accelerated TFT training - **Files**: ml/src/tft/quantile_outputs.rs ### 4. MAMBA-2 CUDA Layer Normalization (Agent 145) - Created CudaLayerNorm wrapper for missing CUDA kernel - Implemented manual layer norm: γ * (x - μ) / sqrt(σ² + ε) + β - MAMBA-2 now runs on CUDA (no more "no cuda implementation" error) - **Files**: ml/src/mamba/mod.rs ### 5. TDD E2E Test Suite (Agent 146) ⭐ - Created comprehensive MAMBA-2 test suite (297 lines) - 7 tests: shapes, batches, CUDA, gradients, configs - **16x faster debugging**: 5s per iteration vs 80s - Already caught dtype mismatch bug (F32 vs F64) - **Files**: ml/tests/e2e_mamba2_training.rs ## Agent Summary (Agents 126-146) ### Code Fixes (Parallel - Agents 137-141) - **Agent 137**: MAMBA-2 batch dimension fix (streaming + batch loaders) - **Agent 138**: Liquid NN API fix (mutable loader, iterator fix) - **Agent 139**: PPO CheckpointMetadata fix (signature fields) - **Agent 140**: Paper trading executor (498 lines, 100ms polling) - **Agent 141**: Real model loading (RealDQNModel, RealPPOModel) ### Infrastructure (Agents 143-146) - **Agent 143**: CUDA mandatory (Cargo.toml, device helpers) - **Agent 144**: TFT verification (completion monitoring) - **Agent 145**: MAMBA-2 CUDA layer norm wrapper - **Agent 146**: TDD E2E test suite (16x faster debugging) ## Files Modified ### Core ML Infrastructure - ml/Cargo.toml: Added default = ["minimal-inference", "cuda"] - ml/src/lib.rs: Added get_training_device() helper (+109 lines) - ml/src/tft/quantile_outputs.rs: Fixed tensor contiguity - ml/src/mamba/mod.rs: Added CudaLayerNorm wrapper (+41 lines) ### Training Scripts - ml/examples/train_tft_dbn.rs: Removed --use-gpu flag - ml/examples/train_ppo.rs: Removed --use-gpu flag - ml/examples/train_mamba2_dbn.rs: Forced CUDA-only mode - ml/examples/train_liquid_dbn.rs: Fixed API usage ### Data Loaders - ml/src/data_loaders/dbn_sequence_loader.rs: Fixed batch dimensions - ml/src/data_loaders/streaming_dbn_loader.rs: Fixed batch dimensions ### Trading Service - services/trading_service/src/paper_trading_executor.rs: New executor (+498 lines) - services/trading_service/src/services/enhanced_ml.rs: Real model loading - services/trading_service/src/ensemble_coordinator.rs: Integration ### Tests - ml/tests/e2e_mamba2_training.rs: New TDD test suite (+297 lines) ### Trainers - ml/src/trainers/tft.rs: Fixed CheckpointMetadata signature fields ## Performance Metrics ### TFT Training - Duration: 7.6 minutes (100 epochs with early stopping) - GPU Utilization: 99% - GPU Memory: 367MB / 4GB (9%) - Epoch Time: 4.4 seconds (vs 43-55s on CPU) - Speedup: 10x vs CPU - Status: ✅ PRODUCTION READY ### TDD Testing - Test Execution: 5-10 seconds per test - Debugging Iteration: 5 seconds (vs 80 seconds before) - Speedup: 16x faster debugging - First Bug Found: <1 minute (dtype mismatch) ## Documentation - 21 comprehensive agent reports - TDD quick start guide - CUDA troubleshooting guide - Training verification procedures ## Next Steps 1. Fix MAMBA-2 dtype mismatch (F32→F64) - 2 minutes 2. Run MAMBA-2 tests until passing - 5-10 minutes 3. Launch full MAMBA-2 training - 200 epochs 4. Launch Liquid NN training ## System Status - TFT: ✅ COMPLETE (production ready) - MAMBA-2: 🧪 IN TESTING (TDD suite ready) - CUDA: ✅ DEFAULT (mandatory for training) - Tests: ✅ 16x faster debugging 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
391 lines
11 KiB
Rust
391 lines
11 KiB
Rust
//! Comprehensive security integration tests
|
|
//!
|
|
//! Tests for:
|
|
//! - Checkpoint signature verification
|
|
//! - Prediction validation and model poisoning detection
|
|
//! - Ensemble anomaly detection
|
|
//! - End-to-end security workflows
|
|
|
|
use ml::checkpoint::{CheckpointManager, CheckpointMetadata, CheckpointSigner, CompressionType, FileSystemStorage};
|
|
use ml::ensemble::model::{EnsembleDecision, ModelVote, TradingAction};
|
|
use ml::security::{
|
|
AnomalyDetectorConfig, EnsembleAnomalyDetector, PredictionValidator, ValidationConfig,
|
|
};
|
|
use ml::ModelType;
|
|
use std::collections::HashMap;
|
|
use tempfile::TempDir;
|
|
|
|
#[tokio::test]
|
|
async fn test_checkpoint_signing_workflow() {
|
|
// Create temporary directory for checkpoints
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let storage = FileSystemStorage::new(temp_dir.path().to_path_buf());
|
|
let manager = CheckpointManager::new(storage);
|
|
|
|
// Create checkpoint metadata
|
|
let mut metadata = CheckpointMetadata::new(
|
|
ModelType::DQN,
|
|
"test_model".to_string(),
|
|
"1.0.0".to_string(),
|
|
);
|
|
|
|
// Create mock checkpoint data
|
|
let checkpoint_data = vec![1u8, 2, 3, 4, 5];
|
|
|
|
// Save checkpoint (should sign automatically)
|
|
let signer = CheckpointSigner::new(None);
|
|
let sig_info = signer
|
|
.sign_checkpoint(&checkpoint_data, ModelType::DQN)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Update metadata with signature
|
|
metadata.signature = Some(sig_info.signature.clone());
|
|
metadata.signature_algorithm = sig_info.algorithm.clone();
|
|
metadata.signing_key_id = sig_info.key_id.clone();
|
|
metadata.signed_at = Some(sig_info.signed_at);
|
|
|
|
// Verify signature
|
|
let result = signer
|
|
.verify_signature(
|
|
&checkpoint_data,
|
|
&sig_info.signature,
|
|
&sig_info.key_id,
|
|
ModelType::DQN,
|
|
)
|
|
.await;
|
|
|
|
assert!(result.is_ok(), "Signature verification should succeed");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_checkpoint_tampering_detection() {
|
|
let signer = CheckpointSigner::new(None);
|
|
let data = b"original checkpoint data";
|
|
|
|
// Sign original data
|
|
let sig_info = signer.sign_checkpoint(data, ModelType::DQN).await.unwrap();
|
|
|
|
// Tamper with data
|
|
let tampered_data = b"tampered checkpoint data";
|
|
|
|
// Verification should fail
|
|
let result = signer
|
|
.verify_signature(
|
|
tampered_data,
|
|
&sig_info.signature,
|
|
&sig_info.key_id,
|
|
ModelType::DQN,
|
|
)
|
|
.await;
|
|
|
|
assert!(
|
|
result.is_err(),
|
|
"Tampered checkpoint should fail verification"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_prediction_validation_normal() {
|
|
let validator = PredictionValidator::new();
|
|
|
|
// Add bootstrap samples
|
|
for i in 0..1000 {
|
|
let value = (i as f64 / 1000.0) - 0.5; // Range: -0.5 to 0.5
|
|
validator.update_statistics(value).await;
|
|
}
|
|
|
|
// Validate normal prediction
|
|
let result = validator.validate(0.3, 0.8, "DQN").await;
|
|
|
|
assert!(result.is_ok(), "Normal prediction should pass validation");
|
|
|
|
let validated = result.unwrap();
|
|
assert!(!validated.should_override, "Normal prediction should not be overridden");
|
|
assert!(validated.validation_flags.is_empty(), "No flags should be set");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_prediction_validation_outlier() {
|
|
let validator = PredictionValidator::new();
|
|
|
|
// Build statistics around 0.0
|
|
for _ in 0..1000 {
|
|
validator.update_statistics(0.0).await;
|
|
}
|
|
|
|
// Inject extreme outlier
|
|
let result = validator.validate(0.95, 0.8, "DQN").await;
|
|
|
|
assert!(result.is_ok(), "Outlier should be detected but not rejected");
|
|
|
|
let validated = result.unwrap();
|
|
assert!(validated.is_outlier, "Should be flagged as outlier");
|
|
assert!(validated.z_score.abs() > 3.0, "Z-score should exceed threshold");
|
|
assert!(validated.should_override, "Should recommend override");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_prediction_validation_out_of_bounds() {
|
|
let validator = PredictionValidator::new();
|
|
|
|
// Test upper bound
|
|
let result = validator.validate(1.5, 0.8, "DQN").await;
|
|
assert!(result.is_err(), "Out of bounds prediction should be rejected");
|
|
|
|
// Test lower bound
|
|
let result = validator.validate(-1.5, 0.8, "DQN").await;
|
|
assert!(result.is_err(), "Out of bounds prediction should be rejected");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_extreme_rate_limiting() {
|
|
let mut config = ValidationConfig::default();
|
|
config.max_extreme_rate = 0.05; // 5%
|
|
config.window_duration = std::time::Duration::from_secs(1);
|
|
|
|
let validator = PredictionValidator::with_config(config);
|
|
|
|
// Inject many extreme predictions quickly
|
|
let mut rejection_count = 0;
|
|
for i in 0..100 {
|
|
let result = validator.validate(0.95, 0.8, "DQN").await;
|
|
|
|
if result.is_err() {
|
|
rejection_count += 1;
|
|
// Should eventually hit rate limit
|
|
assert!(
|
|
result
|
|
.unwrap_err()
|
|
.to_string()
|
|
.contains("extreme predictions"),
|
|
"Should be rate limit error"
|
|
);
|
|
}
|
|
|
|
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
|
|
}
|
|
|
|
assert!(
|
|
rejection_count > 0,
|
|
"Rate limit should have been triggered"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ensemble_sudden_shift_detection() {
|
|
let detector = EnsembleAnomalyDetector::new();
|
|
|
|
// Build history with stable predictions
|
|
for _ in 0..20 {
|
|
let decision = create_test_decision(0.1, HashMap::new());
|
|
detector.update_history(&decision).await;
|
|
}
|
|
|
|
// Inject sudden shift
|
|
let decision = create_test_decision(0.8, HashMap::new());
|
|
let report = detector.detect_anomaly(&decision).await;
|
|
|
|
assert!(report.has_anomalies, "Sudden shift should be detected");
|
|
assert!(
|
|
report.severity >= ml::security::AnomalySeverity::Medium,
|
|
"Should have medium or higher severity"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ensemble_coordinated_attack_detection() {
|
|
let detector = EnsembleAnomalyDetector::new();
|
|
|
|
// Create decision with all models predicting extreme values
|
|
let mut model_votes = HashMap::new();
|
|
for i in 1..=4 {
|
|
model_votes.insert(
|
|
format!("model{}", i),
|
|
ModelVote {
|
|
signal: 0.95,
|
|
confidence: 0.9,
|
|
vote: TradingAction::Buy,
|
|
},
|
|
);
|
|
}
|
|
|
|
let decision = create_test_decision(0.95, model_votes);
|
|
let report = detector.detect_anomaly(&decision).await;
|
|
|
|
assert!(
|
|
report.has_anomalies,
|
|
"Coordinated attack should be detected"
|
|
);
|
|
assert_eq!(
|
|
report.severity,
|
|
ml::security::AnomalySeverity::Critical,
|
|
"Should have critical severity"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ensemble_model_drift_detection() {
|
|
let detector = EnsembleAnomalyDetector::new();
|
|
|
|
// Build history with stable model behavior
|
|
for _ in 0..30 {
|
|
let mut model_votes = HashMap::new();
|
|
model_votes.insert(
|
|
"model1".to_string(),
|
|
ModelVote {
|
|
signal: 0.1,
|
|
confidence: 0.8,
|
|
vote: TradingAction::Hold,
|
|
},
|
|
);
|
|
|
|
let decision = create_test_decision(0.1, model_votes);
|
|
detector.update_history(&decision).await;
|
|
}
|
|
|
|
// Inject drift
|
|
let mut model_votes = HashMap::new();
|
|
model_votes.insert(
|
|
"model1".to_string(),
|
|
ModelVote {
|
|
signal: 0.9,
|
|
confidence: 0.8,
|
|
vote: TradingAction::Buy,
|
|
},
|
|
);
|
|
|
|
let decision = create_test_decision(0.9, model_votes);
|
|
let report = detector.detect_anomaly(&decision).await;
|
|
|
|
assert!(report.has_anomalies, "Model drift should be detected");
|
|
|
|
// Check for drift anomaly
|
|
let has_drift = report.anomalies.iter().any(|a| {
|
|
matches!(
|
|
a,
|
|
ml::security::Anomaly::ModelDrift { .. }
|
|
)
|
|
});
|
|
assert!(has_drift, "Should contain model drift anomaly");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_end_to_end_security_workflow() {
|
|
// Test complete security workflow:
|
|
// 1. Sign checkpoint
|
|
// 2. Validate predictions
|
|
// 3. Detect ensemble anomalies
|
|
|
|
// 1. Checkpoint signing
|
|
let signer = CheckpointSigner::new(None);
|
|
let checkpoint_data = vec![1, 2, 3, 4, 5];
|
|
let sig_info = signer
|
|
.sign_checkpoint(&checkpoint_data, ModelType::DQN)
|
|
.await
|
|
.unwrap();
|
|
|
|
let verify_result = signer
|
|
.verify_signature(
|
|
&checkpoint_data,
|
|
&sig_info.signature,
|
|
&sig_info.key_id,
|
|
ModelType::DQN,
|
|
)
|
|
.await;
|
|
assert!(verify_result.is_ok(), "Checkpoint signature should verify");
|
|
|
|
// 2. Prediction validation
|
|
let validator = PredictionValidator::new();
|
|
for i in 0..100 {
|
|
validator.update_statistics(i as f64 / 100.0).await;
|
|
}
|
|
|
|
let pred_result = validator.validate(0.5, 0.8, "DQN").await;
|
|
assert!(pred_result.is_ok(), "Prediction should be valid");
|
|
|
|
// 3. Ensemble anomaly detection
|
|
let detector = EnsembleAnomalyDetector::new();
|
|
for i in 0..10 {
|
|
let decision = create_test_decision(0.5, HashMap::new());
|
|
detector.update_history(&decision).await;
|
|
}
|
|
|
|
let decision = create_test_decision(0.52, HashMap::new());
|
|
let anomaly_report = detector.detect_anomaly(&decision).await;
|
|
|
|
assert!(
|
|
!anomaly_report.has_anomalies,
|
|
"Normal decision should not trigger anomalies"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_adversarial_prediction_sequence() {
|
|
// Simulate adversarial attack with gradually increasing predictions
|
|
let validator = PredictionValidator::new();
|
|
|
|
// Build normal baseline
|
|
for _ in 0..1000 {
|
|
validator.update_statistics(0.0).await;
|
|
}
|
|
|
|
// Gradually increase predictions (mimicking adversarial poisoning)
|
|
let mut outlier_count = 0;
|
|
for i in 0..100 {
|
|
let value = 0.8 + (i as f64 / 1000.0); // 0.8 to 0.9
|
|
let result = validator.validate(value, 0.8, "adversarial_model").await;
|
|
|
|
if let Ok(validated) = result {
|
|
if validated.is_outlier {
|
|
outlier_count += 1;
|
|
}
|
|
}
|
|
}
|
|
|
|
assert!(
|
|
outlier_count > 50,
|
|
"Should detect many outliers in adversarial sequence (detected: {})",
|
|
outlier_count
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_statistics_bootstrap_phase() {
|
|
let mut config = ValidationConfig::default();
|
|
config.bootstrap_samples = 50;
|
|
|
|
let validator = PredictionValidator::with_config(config);
|
|
|
|
// Add bootstrap samples
|
|
for i in 0..50 {
|
|
validator.update_statistics(i as f64 / 50.0).await;
|
|
}
|
|
|
|
let stats = validator.get_statistics().await;
|
|
|
|
assert_eq!(stats.sample_count, 50, "Should have 50 samples");
|
|
assert!(stats.std_dev > 0.0, "Should have non-zero std dev");
|
|
assert!(
|
|
stats.mean > 0.0 && stats.mean < 1.0,
|
|
"Mean should be in valid range"
|
|
);
|
|
}
|
|
|
|
// Helper function to create test ensemble decision
|
|
fn create_test_decision(signal: f64, model_votes: HashMap<String, ModelVote>) -> EnsembleDecision {
|
|
EnsembleDecision {
|
|
signal,
|
|
confidence: 0.8,
|
|
action: if signal > 0.5 {
|
|
TradingAction::Buy
|
|
} else if signal < -0.5 {
|
|
TradingAction::Sell
|
|
} else {
|
|
TradingAction::Hold
|
|
},
|
|
disagreement_rate: 0.0,
|
|
model_votes,
|
|
}
|
|
}
|