foxhunt/crates/ml/tests/ensemble_hot_swap_test.rs

#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
)]
//! Integration tests for ensemble checkpoint hot-swapping
//!
//! Tests the complete hot-swap workflow:
//! 1. Load DQN checkpoint into active buffer
//! 2. Stage new checkpoint in shadow buffer
//! 3. Validate staged checkpoint (1000 predictions, latency checks)
//! 4. Commit atomic swap (<1μs)
//! 5. Verify active buffer now points to new checkpoint
//! 6. Test rollback mechanism

use ml::ensemble::{
    CheckpointModel, CheckpointValidator, EnsembleMetrics, HotSwapManager, RollbackPolicy,
};
use ml::{Features, MLResult, ModelPrediction};
use std::sync::Arc;
use std::time::Instant;
use tracing::info;

/// Mock prediction function for DQN epoch 30
fn create_dqn_epoch_30_predictor(
) -> Arc<dyn Fn(&Features) -> MLResult<ModelPrediction> + Send + Sync> {
    Arc::new(|features: &Features| {
        let value = (features.values.iter().sum::<f64>() / features.values.len() as f64) * 0.8;
        Ok(ModelPrediction::new(
            "DQN_epoch_30".to_string(),
            value.tanh(),
            0.78,
        ))
    })
}

/// Mock prediction function for DQN epoch 50 (improved)
fn create_dqn_epoch_50_predictor(
) -> Arc<dyn Fn(&Features) -> MLResult<ModelPrediction> + Send + Sync> {
    Arc::new(|features: &Features| {
        let value = (features.values.iter().sum::<f64>() / features.values.len() as f64) * 0.9;
        Ok(ModelPrediction::new(
            "DQN_epoch_50".to_string(),
            value.tanh(),
            0.85,
        ))
    })
}

#[tokio::test]
async fn test_hot_swap_workflow_complete() {
    info!("Hot-Swap Workflow Test");

    // Initialize hot-swap manager
    let validator = CheckpointValidator::new();
    let policy = RollbackPolicy::default();
    let manager = HotSwapManager::new(validator, policy);

    // Step 1: Load DQN epoch 30 into active buffer
    info!("Step 1: Loading DQN epoch 30 into active buffer");
    let checkpoint_v30 = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "ml/checkpoints/dqn/checkpoint_epoch_30.safetensors".to_string(),
        create_dqn_epoch_30_predictor(),
    ));

    manager
        .register_model("DQN".to_string(), checkpoint_v30)
        .await
        .expect("Failed to register model");

    let active = manager.get_active_checkpoint("DQN").await.unwrap();
    assert_eq!(
        active.checkpoint_path,
        "ml/checkpoints/dqn/checkpoint_epoch_30.safetensors"
    );
    info!(checkpoint_path = %active.checkpoint_path, "Active checkpoint loaded");

    // Step 2: Stage DQN epoch 50 in shadow buffer
    info!("Step 2: Staging DQN epoch 50 in shadow buffer");
    let checkpoint_v50 = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "ml/checkpoints/dqn/checkpoint_epoch_50.safetensors".to_string(),
        create_dqn_epoch_50_predictor(),
    ));

    manager
        .stage_checkpoint("DQN", checkpoint_v50)
        .await
        .expect("Failed to stage checkpoint");
    info!("Staged checkpoint in shadow buffer");

    // Step 3: Validate staged checkpoint
    info!("Step 3: Validating staged checkpoint (1000 predictions)");
    let validation_start = Instant::now();
    let validation = manager
        .validate_staged_checkpoint("DQN")
        .await
        .expect("Failed to validate checkpoint");
    let validation_duration = validation_start.elapsed();

    assert!(
        validation.passed,
        "Validation failed: {:?}",
        validation.failure_reason
    );
    assert!(
        validation.p99_latency_us < 50,
        "P99 latency {}μs exceeds 50μs",
        validation.p99_latency_us
    );
    assert!(
        validation.predictions_in_range >= 950,
        "Only {} predictions in range",
        validation.predictions_in_range
    );

    info!("Validation PASSED");
    info!(avg_latency_us = validation.avg_latency_us, "Avg latency");
    info!(p99_latency_us = validation.p99_latency_us, "P99 latency");
    info!(
        predictions_in_range = validation.predictions_in_range,
        predictions_validated = validation.predictions_validated,
        validation_duration_ms = validation_duration.as_millis(),
        "Validation predictions in range"
    );

    // Record metrics
    EnsembleMetrics::record_validation(
        "DQN",
        true,
        validation_duration.as_millis() as u64,
        validation.p99_latency_us,
    );

    // Step 4: Commit atomic swap
    info!("Step 4: Committing atomic swap");
    let swap_latency = manager
        .commit_swap("DQN")
        .await
        .expect("Failed to commit swap");

    assert!(
        swap_latency.as_micros() < 100,
        "Swap latency {}μs exceeds 100μs",
        swap_latency.as_micros()
    );

    info!(swap_latency_us = swap_latency.as_micros(), "Atomic swap completed");

    // Record metrics
    EnsembleMetrics::record_swap_success("DQN", swap_latency.as_micros() as u64);

    // Step 5: Verify active checkpoint is now epoch 50
    info!("Step 5: Verifying active checkpoint");
    let active = manager.get_active_checkpoint("DQN").await.unwrap();
    assert_eq!(
        active.checkpoint_path,
        "ml/checkpoints/dqn/checkpoint_epoch_50.safetensors"
    );
    info!(checkpoint_path = %active.checkpoint_path, "Active checkpoint loaded");

    // Test prediction with new checkpoint
    let features = Features::new(
        vec![0.5, 0.6, 0.7, 0.8, 0.9],
        vec![
            "f1".to_string(),
            "f2".to_string(),
            "f3".to_string(),
            "f4".to_string(),
            "f5".to_string(),
        ],
    );

    let prediction = active.predict(&features).unwrap();
    assert_eq!(prediction.model_id, "DQN_epoch_50");
    assert!(prediction.confidence >= 0.85);
    info!(
        value = prediction.value,
        confidence = prediction.confidence,
        "Prediction with new checkpoint"
    );

    info!("Hot-Swap Workflow Test PASSED");
}

#[tokio::test]
async fn test_hot_swap_rollback() {
    info!("Hot-Swap Rollback Test");

    let validator = CheckpointValidator::new();
    let policy = RollbackPolicy::default();
    let manager = HotSwapManager::new(validator, policy);

    // Register initial checkpoint
    let checkpoint_v30 = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "checkpoint_epoch_30.safetensors".to_string(),
        create_dqn_epoch_30_predictor(),
    ));

    manager
        .register_model("DQN".to_string(), checkpoint_v30)
        .await
        .unwrap();

    // Stage and commit new checkpoint
    let checkpoint_v50 = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "checkpoint_epoch_50.safetensors".to_string(),
        create_dqn_epoch_50_predictor(),
    ));

    manager
        .stage_checkpoint("DQN", checkpoint_v50)
        .await
        .unwrap();
    manager.commit_swap("DQN").await.unwrap();

    // Verify new checkpoint is active
    let active = manager.get_active_checkpoint("DQN").await.unwrap();
    assert_eq!(active.checkpoint_path, "checkpoint_epoch_50.safetensors");
    info!("Swapped to epoch 50");

    // Simulate failure and rollback
    info!("Simulating failure and rolling back");
    manager.rollback("DQN").await.unwrap();

    // Verify rollback restored epoch 30
    let active = manager.get_active_checkpoint("DQN").await.unwrap();
    assert_eq!(active.checkpoint_path, "checkpoint_epoch_30.safetensors");
    info!("Rolled back to epoch 30");

    // Record metrics
    EnsembleMetrics::record_swap_rollback("DQN");
    EnsembleMetrics::record_rollback("DQN", "test_failure");

    info!("Rollback Test PASSED");
}

#[tokio::test]
async fn test_swap_latency_benchmark() {
    info!("Swap Latency Benchmark");

    let validator = CheckpointValidator::new();
    let policy = RollbackPolicy::default();
    let manager = HotSwapManager::new(validator, policy);

    // Register initial checkpoint
    let checkpoint_v1 = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "checkpoint_v1.safetensors".to_string(),
        create_dqn_epoch_30_predictor(),
    ));

    manager
        .register_model("DQN".to_string(), checkpoint_v1)
        .await
        .unwrap();

    // Run 100 swap operations to benchmark
    let mut swap_latencies = Vec::new();

    for i in 0..100 {
        let checkpoint = Arc::new(CheckpointModel::new(
            "DQN".to_string(),
            format!("checkpoint_v{}.safetensors", i + 2),
            create_dqn_epoch_50_predictor(),
        ));

        manager.stage_checkpoint("DQN", checkpoint).await.unwrap();

        let swap_latency = manager.commit_swap("DQN").await.unwrap();
        swap_latencies.push(swap_latency.as_micros() as u64);
    }

    // Calculate statistics
    swap_latencies.sort_unstable();
    let avg_latency = swap_latencies.iter().sum::<u64>() / swap_latencies.len() as u64;
    let p50_latency = swap_latencies[50];
    let p99_latency = swap_latencies[99];

    info!("Swap latency statistics (100 swaps)");
    info!(avg_latency_us = avg_latency, "Average swap latency");
    info!(p50_latency_us = p50_latency, "P50 swap latency");
    info!(p99_latency_us = p99_latency, "P99 swap latency");
    info!(min_latency_us = swap_latencies[0], "Min swap latency");
    info!(max_latency_us = swap_latencies[99], "Max swap latency");

    // All swaps should be < 1μs (but we allow 100μs for CI/testing)
    assert!(
        p99_latency < 100,
        "P99 swap latency {}μs exceeds 100μs",
        p99_latency
    );

    info!("Swap Latency Benchmark PASSED");
}

#[tokio::test]
async fn test_zero_dropped_predictions() {
    info!("Zero Dropped Predictions Test");

    let validator = CheckpointValidator::new();
    let policy = RollbackPolicy::default();
    let manager = HotSwapManager::new(validator, policy);

    // Register initial checkpoint
    let checkpoint_v30 = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "checkpoint_epoch_30.safetensors".to_string(),
        create_dqn_epoch_30_predictor(),
    ));

    manager
        .register_model("DQN".to_string(), checkpoint_v30)
        .await
        .unwrap();

    // Spawn prediction workload (1000 predictions in background)
    let manager_clone = Arc::new(manager);
    let prediction_task = {
        let manager = manager_clone.clone();
        tokio::spawn(async move {
            let mut success_count = 0;
            let mut error_count = 0;

            for i in 0..1000 {
                let features = Features::new(
                    vec![
                        (i as f64 * 0.01).sin(),
                        (i as f64 * 0.02).cos(),
                        (i as f64 * 0.03).tanh(),
                        (i as f64 * 0.01 + 1.0).ln(),
                        (i as f64 * 0.001).exp().min(10.0),
                    ],
                    vec![
                        "f1".to_string(),
                        "f2".to_string(),
                        "f3".to_string(),
                        "f4".to_string(),
                        "f5".to_string(),
                    ],
                );

                match manager.get_active_checkpoint("DQN").await {
                    Ok(checkpoint) => match checkpoint.predict(&features) {
                        Ok(_) => success_count += 1,
                        Err(_) => error_count += 1,
                    },
                    Err(_) => error_count += 1,
                }

                // Small delay to simulate realistic prediction rate
                tokio::time::sleep(tokio::time::Duration::from_micros(100)).await;
            }

            (success_count, error_count)
        })
    };

    // Wait a bit for predictions to start
    tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;

    // Perform hot-swap while predictions are running
    info!("Performing hot-swap during active predictions");
    let checkpoint_v50 = Arc::new(CheckpointModel::new(
        "DQN".to_string(),
        "checkpoint_epoch_50.safetensors".to_string(),
        create_dqn_epoch_50_predictor(),
    ));

    manager_clone
        .stage_checkpoint("DQN", checkpoint_v50)
        .await
        .unwrap();
    let swap_latency = manager_clone.commit_swap("DQN").await.unwrap();
    info!(swap_latency_us = swap_latency.as_micros(), "Hot-swap completed during active predictions");

    // Wait for prediction task to complete
    let (success_count, error_count) = prediction_task.await.unwrap();

    info!("Prediction results during hot-swap");
    info!(success_count, "Successful predictions");
    info!(error_count, "Error predictions");
    info!(total = success_count + error_count, "Total predictions");

    // Verify zero dropped predictions (all predictions should succeed)
    assert_eq!(
        error_count, 0,
        "Found {} dropped predictions during hot-swap",
        error_count
    );
    assert_eq!(
        success_count, 1000,
        "Expected 1000 successful predictions, got {}",
        success_count
    );

    info!("Zero Dropped Predictions Test PASSED");
}