foxhunt/crates/ml/src/benchmark/batch_size_finder.rs

use crate::MLError;
use candle_core::Device;
use tracing;

/// Configuration for batch size and gradient accumulation
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize, Default)]
pub struct BatchSizeConfig {
    /// Actual batch size that fits in GPU memory
    pub batch_size: usize,
    /// Number of gradient accumulation steps to simulate larger batches
    pub gradient_accumulation_steps: usize,
    /// Effective batch size (batch_size * gradient_accumulation_steps)
    pub effective_batch_size: usize,
}

impl BatchSizeConfig {
    /// Create a new batch size configuration
    pub fn new(batch_size: usize, gradient_accumulation_steps: usize) -> Self {
        Self {
            batch_size,
            gradient_accumulation_steps,
            effective_batch_size: batch_size * gradient_accumulation_steps,
        }
    }

    /// Create a configuration with gradient accumulation targeting effective batch size
    pub fn with_target_effective_size(batch_size: usize, target_effective: usize) -> Self {
        let gradient_accumulation_steps = (target_effective + batch_size - 1) / batch_size;
        Self::new(batch_size, gradient_accumulation_steps)
    }
}

/// Finds optimal batch size for GPU training without OOM crashes
#[derive(Debug)]
pub struct BatchSizeFinder {
    device: Device,
    min_batch: usize,
    max_batch: usize,
    target_effective_batch: usize,
    safety_margin: f64,
}

impl BatchSizeFinder {
    /// Create a new batch size finder
    ///
    /// # Arguments
    /// * `device` - CUDA device to test on
    pub fn new(device: Device) -> Self {
        Self {
            device,
            min_batch: 4,
            max_batch: 256,
            target_effective_batch: 64,
            safety_margin: 0.9, // Use 90% of max found for safety
        }
    }

    /// Create a new batch size finder with custom parameters
    ///
    /// # Arguments
    /// * `device` - CUDA device to test on
    /// * `min_batch` - Minimum batch size to test
    /// * `max_batch` - Maximum batch size to test
    /// * `target_effective_batch` - Target effective batch size for gradient accumulation
    /// * `safety_margin` - Fraction of max batch to use (0.0-1.0)
    pub fn with_params(
        device: Device,
        min_batch: usize,
        max_batch: usize,
        target_effective_batch: usize,
        safety_margin: f64,
    ) -> Self {
        Self {
            device,
            min_batch,
            max_batch,
            target_effective_batch,
            safety_margin: safety_margin.clamp(0.5, 1.0),
        }
    }

    /// Find optimal batch size using binary search
    ///
    /// # Arguments
    /// * `test_fn` - Function that tests if a batch size fits in memory.
    ///               Should return Ok(true) if successful, Ok(false) or Err if OOM.
    ///
    /// # Returns
    /// BatchSizeConfig with optimal batch size and gradient accumulation settings
    pub fn find_optimal_batch_size<F>(&self, test_fn: F) -> Result<BatchSizeConfig, MLError>
    where
        F: Fn(usize) -> Result<bool, MLError>,
    {
        let mut min = self.min_batch;
        let mut max = self.max_batch;
        let mut largest_successful = self.min_batch;
        let mut iterations = 0;
        const MAX_ITERATIONS: usize = 10;

        // Binary search for maximum viable batch size
        while min <= max && iterations < MAX_ITERATIONS {
            iterations += 1;
            let mid = (min + max) / 2;

            tracing::debug!(
                "Batch size finder iteration {}: testing batch_size={}",
                iterations,
                mid
            );

            match test_fn(mid) {
                Ok(true) => {
                    // Success - try larger batch
                    largest_successful = mid;
                    min = mid + 1;
                    tracing::debug!("  ✓ Batch size {} succeeded", mid);
                },
                Ok(false) | Err(_) => {
                    // OOM or failure - try smaller batch
                    max = mid.saturating_sub(1);
                    tracing::debug!("  ✗ Batch size {} failed (OOM)", mid);
                },
            }
        }

        // Apply safety margin to prevent crashes from memory fluctuations
        let safe_batch_size =
            ((largest_successful as f64 * self.safety_margin) as usize).max(self.min_batch);

        tracing::info!(
            "Batch size finder converged in {} iterations: max_viable={}, safe_batch={}",
            iterations,
            largest_successful,
            safe_batch_size
        );

        // Configure gradient accumulation if needed
        let config = if safe_batch_size < self.target_effective_batch {
            let grad_accum_steps = self.target_effective_batch / safe_batch_size;
            tracing::info!(
                "Configuring gradient accumulation: batch={}, grad_accum={}, effective={}",
                safe_batch_size,
                grad_accum_steps,
                safe_batch_size * grad_accum_steps
            );
            BatchSizeConfig::new(safe_batch_size, grad_accum_steps)
        } else {
            BatchSizeConfig::new(safe_batch_size, 1)
        };

        Ok(config)
    }

    /// Check if an error is an OOM error
    pub fn is_oom_error(error: &MLError) -> bool {
        let error_msg = format!("{:?}", error).to_lowercase();
        error_msg.contains("out of memory")
            || error_msg.contains("oom")
            || error_msg.contains("cuda_error_out_of_memory")
            || error_msg.contains("memory allocation")
            || error_msg.contains("allocate")
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_batch_size_config_creation() {
        let config = BatchSizeConfig::new(16, 4);
        assert_eq!(config.batch_size, 16);
        assert_eq!(config.gradient_accumulation_steps, 4);
        assert_eq!(config.effective_batch_size, 64);
    }

    #[test]
    fn test_batch_size_config_with_target() {
        let config = BatchSizeConfig::with_target_effective_size(16, 64);
        assert_eq!(config.batch_size, 16);
        assert_eq!(config.gradient_accumulation_steps, 4);
        assert_eq!(config.effective_batch_size, 64);

        // Test rounding up
        let config2 = BatchSizeConfig::with_target_effective_size(16, 65);
        assert_eq!(config2.gradient_accumulation_steps, 5); // Rounds up to 5
        assert_eq!(config2.effective_batch_size, 80);
    }

    #[test]
    fn test_finder_creation() {
        let device = Device::Cpu; // Use CPU for testing
        let finder = BatchSizeFinder::new(device);
        assert_eq!(finder.min_batch, 4);
        assert_eq!(finder.max_batch, 256);
        assert_eq!(finder.target_effective_batch, 64);
        assert!((finder.safety_margin - 0.9).abs() < 1e-6);
    }

    #[test]
    fn test_finder_with_custom_params() {
        let device = Device::Cpu;
        let finder = BatchSizeFinder::with_params(device, 8, 128, 128, 0.85);
        assert_eq!(finder.min_batch, 8);
        assert_eq!(finder.max_batch, 128);
        assert_eq!(finder.target_effective_batch, 128);
        assert!((finder.safety_margin - 0.85).abs() < 1e-6);
    }

    #[test]
    fn test_binary_search_all_succeed() {
        let device = Device::Cpu;
        let finder = BatchSizeFinder::new(device);

        // Mock function that always succeeds
        let test_fn = |_batch_size: usize| Ok(true);

        let config = finder.find_optimal_batch_size(test_fn).unwrap();

        // Should find max batch with safety margin
        let expected_batch = (256.0 * 0.9) as usize;
        assert_eq!(config.batch_size, expected_batch);
        assert_eq!(config.gradient_accumulation_steps, 1);
    }

    #[test]
    fn test_binary_search_all_fail() {
        let device = Device::Cpu;
        let finder = BatchSizeFinder::new(device);

        // Mock function that always fails
        let test_fn = |_batch_size: usize| Ok(false);

        let config = finder.find_optimal_batch_size(test_fn).unwrap();

        // Should use minimum batch size
        assert_eq!(config.batch_size, 4);
        assert_eq!(config.gradient_accumulation_steps, 16); // 64 / 4
        assert_eq!(config.effective_batch_size, 64);
    }

    #[test]
    fn test_binary_search_threshold() {
        let device = Device::Cpu;
        let finder = BatchSizeFinder::new(device);

        // Mock function that fails above 32
        let test_fn = |batch_size: usize| {
            if batch_size <= 32 {
                Ok(true)
            } else {
                Ok(false)
            }
        };

        let config = finder.find_optimal_batch_size(test_fn).unwrap();

        // Should find 32 with safety margin
        let expected_batch = (32.0 * 0.9) as usize; // 28
        assert_eq!(config.batch_size, expected_batch);

        // Should configure gradient accumulation
        let expected_grad_accum = 64 / expected_batch;
        assert_eq!(config.gradient_accumulation_steps, expected_grad_accum);
    }

    #[test]
    fn test_binary_search_with_errors() {
        let device = Device::Cpu;
        let finder = BatchSizeFinder::new(device);

        // Mock function that returns errors for large batches
        let test_fn = |batch_size: usize| {
            if batch_size <= 64 {
                Ok(true)
            } else {
                Err(MLError::TrainingError("OOM: out of memory".to_owned()))
            }
        };

        let config = finder.find_optimal_batch_size(test_fn).unwrap();

        // Should handle errors as OOM
        let expected_batch = (64.0 * 0.9) as usize; // 57
        assert_eq!(config.batch_size, expected_batch);
    }

    #[test]
    fn test_oom_error_detection() {
        // Test various OOM error patterns
        let oom_errors = vec![
            MLError::TrainingError("CUDA out of memory".to_owned()),
            MLError::TrainingError("OOM".to_owned()),
            MLError::TrainingError("cuda_error_out_of_memory".to_owned()),
            MLError::TrainingError("Failed to allocate memory".to_owned()),
            MLError::TrainingError("Memory allocation failed".to_owned()),
        ];

        for error in oom_errors {
            assert!(
                BatchSizeFinder::is_oom_error(&error),
                "Failed to detect OOM error: {:?}",
                error
            );
        }

        // Test non-OOM errors
        let non_oom_errors = vec![
            MLError::TrainingError("Invalid input".to_owned()),
            MLError::TrainingError("Computation error".to_owned()),
        ];

        for error in non_oom_errors {
            assert!(
                !BatchSizeFinder::is_oom_error(&error),
                "False positive OOM detection: {:?}",
                error
            );
        }
    }

    #[test]
    fn test_convergence_iterations() {
        use std::sync::atomic::{AtomicUsize, Ordering};
        use std::sync::Arc;

        let device = Device::Cpu;
        let finder = BatchSizeFinder::new(device);

        let call_count = Arc::new(AtomicUsize::new(0));
        let call_count_clone = call_count.clone();

        let test_fn = move |batch_size: usize| {
            call_count_clone.fetch_add(1, Ordering::SeqCst);
            Ok(batch_size <= 48)
        };

        let _config = finder.find_optimal_batch_size(test_fn).unwrap();

        // Binary search should converge in <10 iterations
        let count = call_count.load(Ordering::SeqCst);
        assert!(count < 10, "Too many iterations: {}", count);
    }

    #[test]
    fn test_safety_margin_clamping() {
        let device = Device::Cpu;

        // Test margin too low
        let finder = BatchSizeFinder::with_params(device.clone(), 4, 256, 64, 0.1);
        assert_eq!(finder.safety_margin, 0.5); // Should clamp to 0.5

        // Test margin too high
        let finder = BatchSizeFinder::with_params(device.clone(), 4, 256, 64, 1.5);
        assert_eq!(finder.safety_margin, 1.0); // Should clamp to 1.0

        // Test valid margin
        let finder = BatchSizeFinder::with_params(device, 4, 256, 64, 0.8);
        assert_eq!(finder.safety_margin, 0.8);
    }
}