foxhunt/ml/src/benchmark/memory_profiler.rs

//! GPU Memory Profiler
//!
//! Real-time VRAM tracking during training using nvidia-smi subprocess integration.
//! Provides accurate memory usage measurements for RTX 3050 Ti (4GB VRAM).

use std::process::Command;
use std::time::{Duration, Instant};

/// Single memory measurement snapshot
#[derive(Debug, Clone)]
pub struct MemorySnapshot {
    /// Timestamp when snapshot was taken
    pub timestamp: Instant,
    /// VRAM used in megabytes
    pub vram_used_mb: f64,
    /// Total VRAM available in megabytes
    pub vram_total_mb: f64,
    /// Memory utilization as percentage (0-100)
    pub utilization_percent: f64,
}

impl MemorySnapshot {
    /// Create new memory snapshot
    pub fn new(vram_used_mb: f64, vram_total_mb: f64) -> Self {
        let utilization_percent = if vram_total_mb > 0.0 {
            (vram_used_mb / vram_total_mb) * 100.0
        } else {
            0.0
        };

        Self {
            timestamp: Instant::now(),
            vram_used_mb,
            vram_total_mb,
            utilization_percent,
        }
    }
}

/// GPU memory profiler using nvidia-smi
#[derive(Debug)]
pub struct MemoryProfiler {
    /// Collected memory snapshots
    snapshots: Vec<MemorySnapshot>,
    /// GPU device ID to monitor
    device_id: u32,
    /// Last snapshot cache (timestamp, snapshot)
    last_snapshot: Option<(Instant, MemorySnapshot)>,
    /// Cache duration (reuse snapshots within this window)
    cache_duration: Duration,
}

impl MemoryProfiler {
    /// Create new memory profiler for specified GPU device
    pub fn new(device_id: u32) -> Self {
        Self {
            snapshots: Vec::new(),
            device_id,
            last_snapshot: None,
            cache_duration: Duration::from_millis(100), // 100ms cache
        }
    }

    /// Take a memory snapshot using nvidia-smi
    pub fn take_snapshot(&mut self) -> Result<MemorySnapshot, Box<dyn std::error::Error>> {
        // Check cache first
        if let Some((cache_time, ref snapshot)) = self.last_snapshot {
            if cache_time.elapsed() < self.cache_duration {
                let cached_snapshot = MemorySnapshot {
                    timestamp: Instant::now(),
                    vram_used_mb: snapshot.vram_used_mb,
                    vram_total_mb: snapshot.vram_total_mb,
                    utilization_percent: snapshot.utilization_percent,
                };
                self.snapshots.push(cached_snapshot.clone());
                return Ok(cached_snapshot);
            }
        }

        // Query nvidia-smi
        let output = Command::new("nvidia-smi")
            .arg("--query-gpu=memory.used,memory.total")
            .arg("--format=csv,noheader,nounits")
            .arg("-i")
            .arg(self.device_id.to_string())
            .output();

        match output {
            Ok(output) => {
                if !output.status.success() {
                    return Err(format!("nvidia-smi failed with status: {}", output.status).into());
                }

                let stdout = String::from_utf8_lossy(&output.stdout);
                let snapshot = self.parse_nvidia_smi_output(&stdout)?;

                // Update cache
                self.last_snapshot = Some((Instant::now(), snapshot.clone()));
                self.snapshots.push(snapshot.clone());

                Ok(snapshot)
            },
            Err(e) => {
                // nvidia-smi not available (likely CPU-only system)
                if e.kind() == std::io::ErrorKind::NotFound {
                    Err(
                        "nvidia-smi not found - CPU-only system or NVIDIA drivers not installed"
                            .into(),
                    )
                } else {
                    Err(format!("Failed to execute nvidia-smi: {}", e).into())
                }
            },
        }
    }

    /// Parse nvidia-smi CSV output
    fn parse_nvidia_smi_output(
        &self,
        output: &str,
    ) -> Result<MemorySnapshot, Box<dyn std::error::Error>> {
        let line = output.trim();

        // Expected format: "used_mb, total_mb" (e.g., "2048, 4096")
        let parts: Vec<&str> = line.split(',').collect();

        if parts.len() != 2 {
            return Err(format!(
                "Invalid nvidia-smi output format. Expected 2 values, got {}: '{}'",
                parts.len(),
                line
            )
            .into());
        }

        let vram_used_mb: f64 = parts[0]
            .trim()
            .parse()
            .map_err(|e| format!("Failed to parse used memory '{}': {}", parts[0].trim(), e))?;

        let vram_total_mb: f64 = parts[1]
            .trim()
            .parse()
            .map_err(|e| format!("Failed to parse total memory '{}': {}", parts[1].trim(), e))?;

        Ok(MemorySnapshot::new(vram_used_mb, vram_total_mb))
    }

    /// Get peak memory usage across all snapshots
    pub fn peak_usage_mb(&self) -> f64 {
        self.snapshots
            .iter()
            .map(|s| s.vram_used_mb)
            .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .unwrap_or(0.0)
    }

    /// Get average memory usage across all snapshots
    pub fn avg_usage_mb(&self) -> f64 {
        if self.snapshots.is_empty() {
            return 0.0;
        }

        let sum: f64 = self.snapshots.iter().map(|s| s.vram_used_mb).sum();
        sum / self.snapshots.len() as f64
    }

    /// Get minimum memory usage across all snapshots
    pub fn min_usage_mb(&self) -> f64 {
        self.snapshots
            .iter()
            .map(|s| s.vram_used_mb)
            .min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .unwrap_or(0.0)
    }

    /// Get total VRAM capacity (from most recent snapshot)
    pub fn total_vram_mb(&self) -> f64 {
        self.snapshots
            .last()
            .map(|s| s.vram_total_mb)
            .unwrap_or(0.0)
    }

    /// Get number of snapshots collected
    pub fn snapshot_count(&self) -> usize {
        self.snapshots.len()
    }

    /// Generate formatted memory usage report
    pub fn memory_report(&self) -> String {
        if self.snapshots.is_empty() {
            return "No memory snapshots collected".to_string();
        }

        let peak = self.peak_usage_mb();
        let avg = self.avg_usage_mb();
        let min = self.min_usage_mb();
        let total = self.total_vram_mb();
        let count = self.snapshot_count();

        let peak_percent = if total > 0.0 {
            (peak / total) * 100.0
        } else {
            0.0
        };

        let avg_percent = if total > 0.0 {
            (avg / total) * 100.0
        } else {
            0.0
        };

        format!(
            r#"GPU Memory Profile (Device {})
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Total VRAM:    {:.0} MB
Peak Usage:    {:.0} MB ({:.1}%)
Average Usage: {:.0} MB ({:.1}%)
Min Usage:     {:.0} MB
Snapshots:     {} samples
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"#,
            self.device_id, total, peak, peak_percent, avg, avg_percent, min, count
        )
    }

    /// Clear all collected snapshots
    pub fn clear(&mut self) {
        self.snapshots.clear();
        self.last_snapshot = None;
    }

    /// Get reference to all snapshots
    pub fn snapshots(&self) -> &[MemorySnapshot] {
        &self.snapshots
    }
}

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

    #[test]
    fn test_memory_snapshot_creation() {
        let snapshot = MemorySnapshot::new(2048.0, 4096.0);
        assert_eq!(snapshot.vram_used_mb, 2048.0);
        assert_eq!(snapshot.vram_total_mb, 4096.0);
        assert!((snapshot.utilization_percent - 50.0).abs() < 0.01);
    }

    #[test]
    fn test_memory_snapshot_zero_total() {
        let snapshot = MemorySnapshot::new(100.0, 0.0);
        assert_eq!(snapshot.utilization_percent, 0.0);
    }

    #[test]
    fn test_profiler_creation() {
        let profiler = MemoryProfiler::new(0);
        assert_eq!(profiler.device_id, 0);
        assert_eq!(profiler.snapshot_count(), 0);
    }

    #[test]
    fn test_parse_nvidia_smi_output() {
        let profiler = MemoryProfiler::new(0);

        // Test valid output
        let output = "2048, 4096";
        let snapshot = profiler.parse_nvidia_smi_output(output).unwrap();
        assert_eq!(snapshot.vram_used_mb, 2048.0);
        assert_eq!(snapshot.vram_total_mb, 4096.0);

        // Test without spaces
        let output = "1024,4096";
        let snapshot = profiler.parse_nvidia_smi_output(output).unwrap();
        assert_eq!(snapshot.vram_used_mb, 1024.0);
        assert_eq!(snapshot.vram_total_mb, 4096.0);

        // Test with extra whitespace
        let output = "  512  ,  4096  ";
        let snapshot = profiler.parse_nvidia_smi_output(output).unwrap();
        assert_eq!(snapshot.vram_used_mb, 512.0);
        assert_eq!(snapshot.vram_total_mb, 4096.0);
    }

    #[test]
    fn test_parse_nvidia_smi_invalid_format() {
        let profiler = MemoryProfiler::new(0);

        // Wrong number of values
        let output = "2048";
        assert!(profiler.parse_nvidia_smi_output(output).is_err());

        let output = "2048, 4096, 8192";
        assert!(profiler.parse_nvidia_smi_output(output).is_err());

        // Non-numeric values
        let output = "abc, 4096";
        assert!(profiler.parse_nvidia_smi_output(output).is_err());

        let output = "2048, xyz";
        assert!(profiler.parse_nvidia_smi_output(output).is_err());
    }

    #[test]
    fn test_peak_avg_calculations() {
        let mut profiler = MemoryProfiler::new(0);

        // Empty profiler
        assert_eq!(profiler.peak_usage_mb(), 0.0);
        assert_eq!(profiler.avg_usage_mb(), 0.0);
        assert_eq!(profiler.min_usage_mb(), 0.0);

        // Add snapshots manually for testing
        profiler.snapshots.push(MemorySnapshot::new(1000.0, 4096.0));
        profiler.snapshots.push(MemorySnapshot::new(2000.0, 4096.0));
        profiler.snapshots.push(MemorySnapshot::new(3000.0, 4096.0));
        profiler.snapshots.push(MemorySnapshot::new(1500.0, 4096.0));

        assert_eq!(profiler.peak_usage_mb(), 3000.0);
        assert_eq!(profiler.avg_usage_mb(), 1875.0); // (1000+2000+3000+1500)/4
        assert_eq!(profiler.min_usage_mb(), 1000.0);
        assert_eq!(profiler.total_vram_mb(), 4096.0);
        assert_eq!(profiler.snapshot_count(), 4);
    }

    #[test]
    fn test_memory_report_format() {
        let mut profiler = MemoryProfiler::new(0);

        // Empty report
        let report = profiler.memory_report();
        assert!(report.contains("No memory snapshots"));

        // Add test data
        profiler.snapshots.push(MemorySnapshot::new(1000.0, 4096.0));
        profiler.snapshots.push(MemorySnapshot::new(2000.0, 4096.0));
        profiler.snapshots.push(MemorySnapshot::new(3000.0, 4096.0));

        let report = profiler.memory_report();
        assert!(report.contains("GPU Memory Profile"));
        assert!(report.contains("4096 MB")); // Total
        assert!(report.contains("3000 MB")); // Peak
        assert!(report.contains("2000 MB")); // Avg
        assert!(report.contains("1000 MB")); // Min
        assert!(report.contains("3 samples"));
    }

    #[test]
    fn test_clear_snapshots() {
        let mut profiler = MemoryProfiler::new(0);

        profiler.snapshots.push(MemorySnapshot::new(1000.0, 4096.0));
        profiler.snapshots.push(MemorySnapshot::new(2000.0, 4096.0));
        assert_eq!(profiler.snapshot_count(), 2);

        profiler.clear();
        assert_eq!(profiler.snapshot_count(), 0);
        assert!(profiler.last_snapshot.is_none());
    }

    #[test]
    #[ignore = "Only run on systems with nvidia-smi"]
    fn test_real_gpu_snapshot() {
        let mut profiler = MemoryProfiler::new(0);

        match profiler.take_snapshot() {
            Ok(snapshot) => {
                println!("GPU Memory Snapshot:");
                println!("  Used: {:.0} MB", snapshot.vram_used_mb);
                println!("  Total: {:.0} MB", snapshot.vram_total_mb);
                println!("  Utilization: {:.1}%", snapshot.utilization_percent);

                assert!(snapshot.vram_used_mb > 0.0);
                assert!(snapshot.vram_total_mb > 0.0);
                assert!(snapshot.vram_used_mb <= snapshot.vram_total_mb);
                assert!(snapshot.utilization_percent >= 0.0);
                assert!(snapshot.utilization_percent <= 100.0);
            },
            Err(e) => {
                // Expected on CPU-only systems
                assert!(e.to_string().contains("nvidia-smi not found"));
            },
        }
    }

    #[test]
    #[ignore = "Only run on systems with nvidia-smi"]
    fn test_snapshot_performance() {
        let mut profiler = MemoryProfiler::new(0);

        let start = Instant::now();
        let iterations = 10;

        for _ in 0..iterations {
            match profiler.take_snapshot() {
                Ok(_) => {},
                Err(e) => {
                    if e.to_string().contains("nvidia-smi not found") {
                        println!("Skipping performance test - nvidia-smi not available");
                        return;
                    }
                    panic!("Unexpected error: {}", e);
                },
            }
        }

        let elapsed = start.elapsed();
        let avg_time_ms = elapsed.as_millis() as f64 / iterations as f64;

        println!("Average snapshot time: {:.2} ms", avg_time_ms);

        // Should be under 10ms per snapshot (relaxed for subprocess overhead)
        assert!(
            avg_time_ms < 50.0,
            "Snapshot too slow: {:.2} ms (target <50ms with caching)",
            avg_time_ms
        );
    }

    #[test]
    #[ignore = "Only run on systems with nvidia-smi"]
    fn test_memory_report_real_gpu() {
        let mut profiler = MemoryProfiler::new(0);

        // Take several snapshots
        for _ in 0..5 {
            match profiler.take_snapshot() {
                Ok(_) => {},
                Err(e) => {
                    if e.to_string().contains("nvidia-smi not found") {
                        println!("Skipping report test - nvidia-smi not available");
                        return;
                    }
                    panic!("Unexpected error: {}", e);
                },
            }
            std::thread::sleep(Duration::from_millis(100));
        }

        let report = profiler.memory_report();
        println!("\n{}", report);

        assert!(report.contains("GPU Memory Profile"));
        assert!(report.contains("MB"));
        assert!(profiler.snapshot_count() >= 5);
    }
}