foxhunt/ml/tests/oom_recovery_integration_test.rs

//! OOM Recovery Integration Tests
//!
//! Comprehensive test suite for Out-Of-Memory recovery mechanisms in ML training.
//! Tests OOM detection, batch size reduction, retry limits, and state preservation.
//!
//! ## Test Coverage
//!
//! 1. ✅ OOM Error Detection (mock errors)
//! 2. ✅ Batch Size Reduction Strategy (exponential backoff)
//! 3. ✅ Retry Limits (max 3 attempts)
//! 4. ✅ Model State Preservation across retries
//! 5. ✅ Calibration State Preservation (QAT)
//! 6. ✅ Logging Output validation
//! 7. ✅ Edge Cases (immediate OOM, multiple recoveries, non-OOM errors)
//!
//! ## Test Strategy
//!
//! - Uses mocks to simulate OOM errors (no GPU required)
//! - Tests compile but don't run on GPU (CPU-only testing)
//! - Validates retry logic without actual memory pressure
//! - Verifies state preservation across retry attempts

use candle_core::Device;
use candle_nn::{VarBuilder, VarMap};
use ml::benchmark::batch_size_finder::BatchSizeFinder;
use ml::memory_optimization::auto_batch_size::AutoBatchSizer;
use ml::{MLError, MLResult};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

// ============================================================================
// Test Helpers
// ============================================================================

/// Mock OOM error generator for testing
struct OOMErrorSimulator {
    oom_after_calls: AtomicUsize,
    current_calls: AtomicUsize,
}

impl OOMErrorSimulator {
    /// Create a simulator that triggers OOM after N calls
    fn new(oom_after_calls: usize) -> Self {
        Self {
            oom_after_calls: AtomicUsize::new(oom_after_calls),
            current_calls: AtomicUsize::new(0),
        }
    }

    /// Simulate a training step that may OOM
    fn simulate_training_step(&self) -> MLResult<()> {
        let calls = self.current_calls.fetch_add(1, Ordering::SeqCst);
        let oom_threshold = self.oom_after_calls.load(Ordering::SeqCst);

        if calls >= oom_threshold && oom_threshold > 0 {
            Err(MLError::ModelError("CUDA error: out of memory".to_string()))
        } else {
            Ok(())
        }
    }

    /// Reset the simulator (for testing multiple retries)
    fn reset(&self, new_threshold: usize) {
        self.current_calls.store(0, Ordering::SeqCst);
        self.oom_after_calls.store(new_threshold, Ordering::SeqCst);
    }
}

// Use the existing is_oom_error from BatchSizeFinder
use BatchSizeFinder as OOMDetector;

/// Mock model state for testing state preservation
#[derive(Clone, Debug, PartialEq)]
struct MockModelState {
    weights: Vec<f32>,
    epoch: usize,
    loss: f32,
}

impl MockModelState {
    fn new() -> Self {
        Self {
            weights: vec![1.0, 2.0, 3.0, 4.0],
            epoch: 0,
            loss: 1.0,
        }
    }

    fn update(&mut self, epoch: usize, loss: f32) {
        self.epoch = epoch;
        self.loss = loss;
        // Simulate weight updates
        for w in &mut self.weights {
            *w *= 0.99; // Decay
        }
    }
}

// ============================================================================
// TEST 1: OOM Error Detection (Mock Errors)
// ============================================================================

#[test]
fn test_oom_error_detection() {
    println!("\n🧪 TEST 1: OOM Error Detection");

    // Test standard OOM error strings
    let oom1 = MLError::ModelError("CUDA error: out of memory".to_string());
    assert!(
        OOMDetector::is_oom_error(&oom1),
        "Should detect 'out of memory'"
    );

    let oom2 = MLError::TrainingError("OOM detected during forward pass".to_string());
    assert!(OOMDetector::is_oom_error(&oom2), "Should detect 'OOM'");

    let oom3 = MLError::ModelError("cuda error 2: allocation failed".to_string());
    assert!(
        OOMDetector::is_oom_error(&oom3),
        "Should detect 'cuda error 2'"
    );

    let oom4 = MLError::ModelError("Failed to allocate 500MB on GPU".to_string());
    assert!(
        OOMDetector::is_oom_error(&oom4),
        "Should detect 'failed to allocate'"
    );

    let oom5 = MLError::ModelError("allocation failed on device".to_string());
    assert!(
        OOMDetector::is_oom_error(&oom5),
        "Should detect 'allocation failed'"
    );

    // Test non-OOM errors
    let not_oom1 = MLError::ModelError("Invalid tensor shape".to_string());
    assert!(
        !OOMDetector::is_oom_error(&not_oom1),
        "Should not detect regular errors"
    );

    let not_oom2 = MLError::ConfigError {
        reason: "Missing parameter".to_string(),
    };
    assert!(
        !OOMDetector::is_oom_error(&not_oom2),
        "Should not detect config errors"
    );

    let not_oom3 = MLError::TrainingError("Gradient explosion detected".to_string());
    assert!(
        !OOMDetector::is_oom_error(&not_oom3),
        "Should not detect training errors"
    );

    println!("✅ TEST 1 PASSED: OOM error detection works correctly");
}

// ============================================================================
// TEST 2: Batch Size Reduction Strategy (Exponential Backoff)
// ============================================================================

#[test]
fn test_batch_size_reduction_strategy() {
    println!("\n🧪 TEST 2: Batch Size Reduction Strategy");

    // Test exponential backoff: 64 → 32 → 16 → 8
    let mut current_batch_size = 64;
    let mut batch_sizes = vec![current_batch_size];

    for _ in 0..3 {
        current_batch_size = AutoBatchSizer::reduce_batch_size(current_batch_size);
        batch_sizes.push(current_batch_size);
    }

    assert_eq!(batch_sizes, vec![64, 32, 16, 8]);
    println!("✓ Batch size reduction: {:?}", batch_sizes);

    // Test different starting sizes
    assert_eq!(AutoBatchSizer::reduce_batch_size(128), 64);
    assert_eq!(AutoBatchSizer::reduce_batch_size(32), 16);
    assert_eq!(AutoBatchSizer::reduce_batch_size(16), 8);
    assert_eq!(AutoBatchSizer::reduce_batch_size(8), 4);
    assert_eq!(AutoBatchSizer::reduce_batch_size(4), 2);
    println!("✓ Tested various starting sizes");

    // Test minimum batch size detection
    assert!(AutoBatchSizer::is_batch_size_too_small(1));
    assert!(AutoBatchSizer::is_batch_size_too_small(2));
    assert!(AutoBatchSizer::is_batch_size_too_small(3));
    assert!(!AutoBatchSizer::is_batch_size_too_small(4));
    assert!(!AutoBatchSizer::is_batch_size_too_small(8));
    println!("✓ Minimum batch size threshold: 4");

    println!("✅ TEST 2 PASSED: Batch size reduction strategy is correct");
}

// ============================================================================
// TEST 3: Retry Limits (Max 3 Attempts)
// ============================================================================

#[test]
fn test_retry_limits() {
    println!("\n🧪 TEST 3: Retry Limits");

    const MAX_OOM_RETRIES: usize = 3;
    let simulator = OOMErrorSimulator::new(0); // Always OOM

    let mut retry_count = 0;
    let mut current_batch_size = 64;
    let mut retry_succeeded = false;

    while retry_count < MAX_OOM_RETRIES {
        match simulator.simulate_training_step() {
            Ok(_) => {
                retry_succeeded = true;
                println!("✓ Training succeeded on retry {}", retry_count);
                break;
            },
            Err(e) if OOMDetector::is_oom_error(&e) => {
                retry_count += 1;
                current_batch_size = AutoBatchSizer::reduce_batch_size(current_batch_size);
                println!(
                    "✓ OOM detected, retry {}/{}, reducing batch size to {}",
                    retry_count, MAX_OOM_RETRIES, current_batch_size
                );

                // Stop if batch size becomes too small
                if AutoBatchSizer::is_batch_size_too_small(current_batch_size) {
                    println!("✗ Batch size {} too small, aborting", current_batch_size);
                    break;
                }
            },
            Err(e) => {
                // Non-OOM error, fail immediately
                println!("✗ Non-OOM error: {:?}", e);
                break;
            },
        }
    }

    // Should exhaust all retries
    assert_eq!(retry_count, MAX_OOM_RETRIES, "Should retry exactly 3 times");
    assert!(
        !retry_succeeded,
        "Should not succeed with always-OOM simulator"
    );
    assert_eq!(current_batch_size, 8, "Final batch size should be 8");

    println!("✅ TEST 3 PASSED: Retry limits enforced correctly");
}

// ============================================================================
// TEST 4: Model State Preservation Across Retries
// ============================================================================

#[test]
fn test_model_state_preservation() {
    println!("\n🧪 TEST 4: Model State Preservation");

    let mut model_state = MockModelState::new();
    let initial_state = model_state.clone();

    println!("✓ Initial state: {:?}", initial_state);

    // Simulate training that OOMs after first epoch
    model_state.update(1, 0.8);
    let state_after_epoch1 = model_state.clone();

    println!("✓ State after epoch 1: {:?}", state_after_epoch1);

    // Simulate OOM and retry
    println!("⚠ OOM detected, preserving state...");
    let preserved_state = model_state.clone();

    // Simulate retry with reduced batch size (state should be preserved)
    assert_eq!(preserved_state, state_after_epoch1);
    println!("✓ State preserved: {:?}", preserved_state);

    // Continue training from preserved state
    model_state.update(2, 0.6);
    let state_after_epoch2 = model_state.clone();

    println!("✓ State after epoch 2: {:?}", state_after_epoch2);

    // Verify state evolved correctly
    assert_eq!(state_after_epoch2.epoch, 2);
    assert_eq!(state_after_epoch2.loss, 0.6);
    assert_ne!(
        state_after_epoch2.weights, initial_state.weights,
        "Weights should have changed"
    );

    println!("✅ TEST 4 PASSED: Model state preserved across retries");
}

// ============================================================================
// TEST 5: Calibration State Preservation (QAT)
// ============================================================================

#[test]
fn test_calibration_state_preservation() {
    println!("\n🧪 TEST 5: Calibration State Preservation (QAT)");

    // Mock calibration state
    #[derive(Clone, Debug, PartialEq)]
    struct CalibrationState {
        observer_count: usize,
        min_vals: Vec<f32>,
        max_vals: Vec<f32>,
        num_observations: usize,
    }

    let mut calib_state = CalibrationState {
        observer_count: 10,
        min_vals: vec![-1.0, -2.0, -3.0],
        max_vals: vec![1.0, 2.0, 3.0],
        num_observations: 100,
    };

    println!("✓ Initial calibration state: {:?}", calib_state);

    // Simulate calibration progress
    calib_state.num_observations += 50;
    calib_state.min_vals[0] = -1.5; // Observed lower value
    calib_state.max_vals[1] = 2.5; // Observed higher value

    let state_before_oom = calib_state.clone();
    println!("✓ Calibration state before OOM: {:?}", state_before_oom);

    // Simulate OOM during calibration
    println!("⚠ OOM detected during calibration, preserving state...");
    let preserved_calib_state = calib_state.clone();

    // Verify state preserved
    assert_eq!(preserved_calib_state, state_before_oom);
    println!("✓ Calibration state preserved: {:?}", preserved_calib_state);

    // Continue calibration with reduced batch size
    calib_state.num_observations += 25; // Smaller batch
    let state_after_retry = calib_state.clone();

    println!("✓ Calibration state after retry: {:?}", state_after_retry);

    // Verify calibration continued from preserved state
    assert_eq!(
        state_after_retry.num_observations,
        state_before_oom.num_observations + 25
    );
    assert_eq!(
        state_after_retry.observer_count, 10,
        "Observer count unchanged"
    );

    println!("✅ TEST 5 PASSED: Calibration state preserved during OOM recovery");
}

// ============================================================================
// TEST 6: Logging Output Validation
// ============================================================================

#[test]
fn test_logging_output() {
    println!("\n🧪 TEST 6: Logging Output Validation");

    const MAX_OOM_RETRIES: usize = 3;
    let simulator = OOMErrorSimulator::new(0); // Always OOM

    let mut retry_count = 0;
    let mut current_batch_size = 64;
    let mut log_messages = Vec::new();

    while retry_count < MAX_OOM_RETRIES {
        match simulator.simulate_training_step() {
            Ok(_) => {
                log_messages.push(format!("Training succeeded on retry {}", retry_count));
                break;
            },
            Err(e) if OOMDetector::is_oom_error(&e) => {
                retry_count += 1;
                current_batch_size = AutoBatchSizer::reduce_batch_size(current_batch_size);

                let log_msg = format!(
                    "OOM retry {}/{}: Reducing batch size from {} to {}",
                    retry_count,
                    MAX_OOM_RETRIES,
                    current_batch_size * 2,
                    current_batch_size
                );
                log_messages.push(log_msg.clone());
                println!("📝 {}", log_msg);

                if AutoBatchSizer::is_batch_size_too_small(current_batch_size) {
                    let abort_msg = format!(
                        "Batch size {} too small, aborting after {} retries",
                        current_batch_size, retry_count
                    );
                    log_messages.push(abort_msg.clone());
                    println!("📝 {}", abort_msg);
                    break;
                }
            },
            Err(e) => {
                let err_msg = format!("Non-OOM error: {:?}", e);
                log_messages.push(err_msg.clone());
                println!("📝 {}", err_msg);
                break;
            },
        }
    }

    // Verify log messages
    assert_eq!(log_messages.len(), 3, "Should have 3 retry log messages");

    assert!(log_messages[0].contains("OOM retry 1/3"));
    assert!(log_messages[0].contains("64 to 32"));

    assert!(log_messages[1].contains("OOM retry 2/3"));
    assert!(log_messages[1].contains("32 to 16"));

    assert!(log_messages[2].contains("OOM retry 3/3"));
    assert!(log_messages[2].contains("16 to 8"));

    println!("✓ All log messages validated");
    println!("✅ TEST 6 PASSED: Logging output is correct");
}

// ============================================================================
// TEST 7: Edge Case - Immediate OOM (batch_size=1)
// ============================================================================

#[test]
fn test_immediate_oom_edge_case() {
    println!("\n🧪 TEST 7: Edge Case - Immediate OOM (batch_size=1)");

    let mut current_batch_size = 4; // Start at minimum viable size
    let simulator = OOMErrorSimulator::new(0); // Always OOM

    println!("✓ Starting with batch_size={}", current_batch_size);

    // Simulate first OOM
    match simulator.simulate_training_step() {
        Err(e) if OOMDetector::is_oom_error(&e) => {
            current_batch_size = AutoBatchSizer::reduce_batch_size(current_batch_size);
            println!("✓ First OOM, reduced to batch_size={}", current_batch_size);
        },
        _ => panic!("Expected OOM error"),
    }

    // Check if batch size is too small
    assert_eq!(current_batch_size, 2, "Should reduce to 2");
    assert!(
        AutoBatchSizer::is_batch_size_too_small(current_batch_size),
        "Batch size 2 is too small"
    );

    println!("✓ Detected batch size too small, aborting immediately");
    println!("✅ TEST 7 PASSED: Immediate OOM handled correctly");
}

// ============================================================================
// TEST 8: Edge Case - Multiple OOM Recoveries
// ============================================================================

#[test]
fn test_multiple_oom_recoveries() {
    println!("\n🧪 TEST 8: Edge Case - Multiple OOM Recoveries");

    const MAX_OOM_RETRIES: usize = 3;
    let simulator = Arc::new(OOMErrorSimulator::new(1)); // OOM on second call

    let mut current_batch_size = 64;
    let mut total_oom_events = 0;
    let mut successful_batches = 0;

    // Simulate multiple epochs, each may OOM
    for epoch in 0..5 {
        println!("\n📊 Epoch {}", epoch);

        let mut retry_count = 0;
        let mut epoch_succeeded = false;

        while retry_count < MAX_OOM_RETRIES {
            match simulator.simulate_training_step() {
                Ok(_) => {
                    successful_batches += 1;
                    epoch_succeeded = true;
                    println!(
                        "✓ Epoch {} succeeded (batch_size={})",
                        epoch, current_batch_size
                    );
                    break;
                },
                Err(e) if OOMDetector::is_oom_error(&e) => {
                    retry_count += 1;
                    total_oom_events += 1;
                    current_batch_size = AutoBatchSizer::reduce_batch_size(current_batch_size);

                    println!(
                        "⚠ OOM event {} (retry {}/{}), batch_size reduced to {}",
                        total_oom_events, retry_count, MAX_OOM_RETRIES, current_batch_size
                    );

                    if AutoBatchSizer::is_batch_size_too_small(current_batch_size) {
                        println!("✗ Batch size too small, epoch failed");
                        break;
                    }

                    // Reset simulator for next retry (simulate success after batch size reduction)
                    simulator.reset(100); // Won't OOM next time
                },
                Err(e) => {
                    println!("✗ Non-OOM error: {:?}", e);
                    break;
                },
            }
        }

        if !epoch_succeeded {
            println!("✗ Epoch {} failed after {} retries", epoch, retry_count);
            break;
        }

        // Reset simulator for next epoch
        simulator.reset(1); // OOM on second call again
    }

    println!(
        "\n✓ Total OOM events: {}, Successful batches: {}",
        total_oom_events, successful_batches
    );
    assert!(total_oom_events > 0, "Should have encountered OOM events");
    assert!(
        successful_batches > 0,
        "Should have some successful batches"
    );

    println!("✅ TEST 8 PASSED: Multiple OOM recoveries handled correctly");
}

// ============================================================================
// TEST 9: Edge Case - Non-OOM Errors (Should Fail Fast)
// ============================================================================

#[test]
fn test_non_oom_errors_fail_fast() {
    println!("\n🧪 TEST 9: Edge Case - Non-OOM Errors (Fail Fast)");

    // Simulate non-OOM error (should fail immediately, no retries)
    let non_oom_error = MLError::ModelError("Invalid tensor shape".to_string());

    const MAX_OOM_RETRIES: usize = 3;
    let mut retry_count = 0;
    let mut failed_fast = false;

    // Attempt retry logic
    match &non_oom_error {
        e if OOMDetector::is_oom_error(e) => {
            retry_count += 1;
            println!("⚠ OOM detected, retrying...");
        },
        e => {
            println!("✗ Non-OOM error detected: {:?}", e);
            println!("✓ Failing fast without retry");
            failed_fast = true;
        },
    }

    assert_eq!(retry_count, 0, "Should not retry on non-OOM errors");
    assert!(failed_fast, "Should fail fast");

    println!("✅ TEST 9 PASSED: Non-OOM errors fail fast (no retries)");
}

// ============================================================================
// TEST 10: OOM Recovery with VarMap Preservation
// ============================================================================

#[test]
fn test_varmap_preservation_across_oom() {
    println!("\n🧪 TEST 10: VarMap Preservation Across OOM");

    let device = Device::Cpu;
    let varmap = VarMap::new();

    // Create some variables in VarMap (simulating model weights)
    let vb = VarBuilder::from_varmap(&varmap, candle_core::DType::F32, &device);

    // Simulate creating model parameters
    let _weight1 = vb
        .get((10, 20), "layer1.weight")
        .expect("Failed to create weight1");
    let _bias1 = vb.get(10, "layer1.bias").expect("Failed to create bias1");

    println!("✓ Created VarMap with 2 parameters");

    // Get initial parameter count
    let initial_var_count = varmap.all_vars().len();
    println!("✓ Initial variables: {} parameters", initial_var_count);

    // Simulate OOM during training (VarMap should be preserved)
    println!("⚠ Simulating OOM during training...");

    // Clone VarMap to preserve state
    let preserved_varmap = varmap.clone();

    println!("✓ VarMap preserved (cloned before retry)");

    // Verify preserved VarMap has same parameters
    let preserved_var_count = preserved_varmap.all_vars().len();

    assert_eq!(
        initial_var_count, preserved_var_count,
        "VarMap should have same number of variables after preservation"
    );

    println!(
        "✓ Preserved VarMap verified: {} parameters",
        preserved_var_count
    );

    // Simulate retry with reduced batch size (using preserved VarMap)
    let vb_retry = VarBuilder::from_varmap(&preserved_varmap, candle_core::DType::F32, &device);

    // Access preserved variables
    let _weight1_retry = vb_retry
        .get((10, 20), "layer1.weight")
        .expect("Failed to access weight1 after retry");
    let _bias1_retry = vb_retry
        .get(10, "layer1.bias")
        .expect("Failed to access bias1 after retry");

    println!("✓ Successfully accessed preserved variables after retry");

    println!("✅ TEST 10 PASSED: VarMap preserved across OOM recovery");
}

// ============================================================================
// TEST 11: Comprehensive OOM Recovery Workflow
// ============================================================================

#[test]
fn test_comprehensive_oom_recovery_workflow() {
    println!("\n🧪 TEST 11: Comprehensive OOM Recovery Workflow");

    // Simulate full OOM recovery workflow
    const MAX_OOM_RETRIES: usize = 3;
    let simulator = Arc::new(OOMErrorSimulator::new(2)); // OOM after 2 calls

    let mut current_batch_size = 64;
    let mut model_state = MockModelState::new();
    let mut training_succeeded = false;
    let mut total_oom_events = 0;

    println!(
        "📊 Starting training with batch_size={}",
        current_batch_size
    );

    // Training loop with OOM recovery
    for epoch in 0..10 {
        let mut retry_count = 0;
        let mut epoch_succeeded = false;

        while retry_count < MAX_OOM_RETRIES {
            match simulator.simulate_training_step() {
                Ok(_) => {
                    // Training step succeeded
                    model_state.update(epoch, 1.0 / (epoch + 1) as f32);
                    epoch_succeeded = true;
                    println!(
                        "✓ Epoch {} completed (loss={:.4}, batch_size={})",
                        epoch, model_state.loss, current_batch_size
                    );
                    break;
                },
                Err(e) if OOMDetector::is_oom_error(&e) => {
                    // OOM detected
                    retry_count += 1;
                    total_oom_events += 1;

                    println!(
                        "⚠ OOM event {} (epoch {}, retry {}/{})",
                        total_oom_events, epoch, retry_count, MAX_OOM_RETRIES
                    );

                    // Preserve model state
                    let preserved_state = model_state.clone();
                    println!("✓ Model state preserved: epoch={}", preserved_state.epoch);

                    // Reduce batch size
                    let old_batch_size = current_batch_size;
                    current_batch_size = AutoBatchSizer::reduce_batch_size(current_batch_size);
                    println!(
                        "✓ Batch size reduced: {} → {}",
                        old_batch_size, current_batch_size
                    );

                    // Check if batch size too small
                    if AutoBatchSizer::is_batch_size_too_small(current_batch_size) {
                        println!(
                            "✗ Batch size {} too small, aborting training",
                            current_batch_size
                        );
                        training_succeeded = false;
                        break;
                    }

                    // Reset simulator to allow success on next attempt
                    simulator.reset(100);
                },
                Err(e) => {
                    // Non-OOM error, fail fast
                    println!("✗ Non-OOM error: {:?}", e);
                    training_succeeded = false;
                    break;
                },
            }
        }

        if !epoch_succeeded {
            println!("✗ Training failed at epoch {}", epoch);
            break;
        }

        // Check if training complete
        if epoch == 9 {
            training_succeeded = true;
            println!("✅ Training completed successfully!");
        }

        // Reset simulator for next epoch
        simulator.reset(2);
    }

    // Verify training outcome
    println!("\n📋 Training Summary:");
    println!("  - Final epoch: {}", model_state.epoch);
    println!("  - Final loss: {:.4}", model_state.loss);
    println!("  - Final batch size: {}", current_batch_size);
    println!("  - Total OOM events: {}", total_oom_events);
    println!("  - Training succeeded: {}", training_succeeded);

    assert!(total_oom_events > 0, "Should have encountered OOM events");
    assert!(training_succeeded, "Training should complete successfully");
    assert!(
        current_batch_size < 64,
        "Batch size should be reduced from initial 64"
    );

    println!("✅ TEST 11 PASSED: Comprehensive OOM recovery workflow successful");
}