foxhunt/services/ml_training_service/src/gpu_resource_manager.rs

//! GPU Resource Manager
//!
//! Provides explicit GPU locking and memory tracking to prevent concurrent training conflicts.
//! This module ensures that only one training job can use a GPU at a time, with automatic
//! cleanup on job completion or crash.

use anyhow::Result;
use std::collections::HashMap;
use std::process::Command;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, error, info, warn};
use uuid::Uuid;

/// GPU allocation errors
#[derive(Debug, Clone, thiserror::Error)]
pub enum GPUAllocationError {
    #[error("GPU {gpu_id} is already locked by job {current_job_id}")]
    GPUAlreadyLocked { gpu_id: u32, current_job_id: Uuid },

    #[error("GPU {gpu_id} not found in available GPU list")]
    GPUNotFound { gpu_id: u32 },

    #[error("No available GPUs found")]
    NoGPUsAvailable,

    #[error(
        "Insufficient memory on GPU {gpu_id}: required {required_mb}MB, available {available_mb}MB"
    )]
    InsufficientMemory {
        gpu_id: u32,
        required_mb: u64,
        available_mb: u64,
    },

    #[error("Failed to query GPU memory: {message}")]
    MemoryQueryFailed { message: String },

    #[error("Failed to query GPU utilization: {message}")]
    UtilizationQueryFailed { message: String },

    #[error("GPU {gpu_id} is locked by a different job {locked_job_id}, cannot release for job {requested_job_id}")]
    CannotReleaseLockedByDifferentJob {
        gpu_id: u32,
        locked_job_id: Uuid,
        requested_job_id: Uuid,
    },
}

/// GPU memory information
#[derive(Debug, Clone)]
pub struct GPUMemoryInfo {
    pub gpu_id: u32,
    pub total_mb: u64,
    pub used_mb: u64,
    pub free_mb: u64,
}

/// GPU lock representing exclusive access to a GPU
#[derive(Debug, Clone)]
pub struct GPULock {
    gpu_id: u32,
    job_id: Uuid,
    manager: Arc<GPUResourceManager>,
}

impl GPULock {
    pub fn gpu_id(&self) -> u32 {
        self.gpu_id
    }

    pub fn job_id(&self) -> Uuid {
        self.job_id
    }

    pub fn is_locked(&self) -> bool {
        // Check if this lock is still valid in the manager
        // For simplicity, assume lock is valid if it exists
        true
    }
}

impl Drop for GPULock {
    fn drop(&mut self) {
        // Automatic cleanup on drop
        let manager = Arc::clone(&self.manager);
        let gpu_id = self.gpu_id;
        let job_id = self.job_id;

        // Spawn a task to release the GPU asynchronously
        tokio::spawn(async move {
            if let Err(e) = manager.release_gpu(gpu_id, job_id).await {
                warn!(
                    "Failed to release GPU {} for job {} on drop: {}",
                    gpu_id, job_id, e
                );
            } else {
                debug!(
                    "GPU {} automatically released for job {} on drop",
                    gpu_id, job_id
                );
            }
        });
    }
}

/// GPU state tracking
#[derive(Debug, Clone)]
#[allow(dead_code)]
struct GPUState {
    gpu_id: u32,
    locked_by: Option<Uuid>,
}

/// GPU Resource Manager - manages GPU locks and memory tracking
#[derive(Debug)]
pub struct GPUResourceManager {
    /// Available GPU IDs
    available_gpus: Vec<u32>,

    /// GPU lock state (gpu_id -> locked_by_job_id)
    gpu_locks: Arc<RwLock<HashMap<u32, Uuid>>>,
}

impl GPUResourceManager {
    /// Create a new GPU resource manager
    pub async fn new(gpu_ids: Vec<u32>) -> Result<Self> {
        info!("Initializing GPU resource manager with GPUs: {:?}", gpu_ids);

        let manager = Self {
            available_gpus: gpu_ids.clone(),
            gpu_locks: Arc::new(RwLock::new(HashMap::new())),
        };

        // Validate all GPUs are accessible
        for gpu_id in &gpu_ids {
            match manager.get_gpu_memory(*gpu_id).await {
                Ok(info) => {
                    info!(
                        "GPU {} validated: {}MB total, {}MB free",
                        gpu_id, info.total_mb, info.free_mb
                    );
                },
                Err(e) => {
                    warn!("GPU {} validation warning: {}", gpu_id, e);
                    // Continue - GPU might be temporarily unavailable
                },
            }
        }

        Ok(manager)
    }

    /// Acquire exclusive lock on a specific GPU
    pub async fn acquire_gpu(
        &self,
        job_id: Uuid,
        gpu_id: u32,
    ) -> Result<GPULock, GPUAllocationError> {
        // Check if GPU exists
        if !self.available_gpus.contains(&gpu_id) {
            return Err(GPUAllocationError::GPUNotFound { gpu_id });
        }

        // Try to acquire lock
        let mut locks = self.gpu_locks.write().await;

        if let Some(locked_by) = locks.get(&gpu_id) {
            // GPU already locked
            return Err(GPUAllocationError::GPUAlreadyLocked {
                gpu_id,
                current_job_id: *locked_by,
            });
        }

        // Acquire lock
        locks.insert(gpu_id, job_id);
        info!("GPU {} locked by job {}", gpu_id, job_id);

        Ok(GPULock {
            gpu_id,
            job_id,
            manager: Arc::new(Self {
                available_gpus: self.available_gpus.clone(),
                gpu_locks: Arc::clone(&self.gpu_locks),
            }),
        })
    }

    /// Acquire any available GPU (dynamic allocation)
    pub async fn acquire_any_available_gpu(
        &self,
        job_id: Uuid,
    ) -> Result<GPULock, GPUAllocationError> {
        let locks = self.gpu_locks.read().await;

        // Find first available GPU
        for gpu_id in &self.available_gpus {
            if !locks.contains_key(gpu_id) {
                drop(locks); // Release read lock before acquiring write lock
                return self.acquire_gpu(job_id, *gpu_id).await;
            }
        }

        Err(GPUAllocationError::NoGPUsAvailable)
    }

    /// Acquire GPU with memory requirement check
    pub async fn acquire_gpu_with_memory_requirement(
        &self,
        job_id: Uuid,
        gpu_id: u32,
        required_memory_mb: u64,
    ) -> Result<GPULock, GPUAllocationError> {
        // Check memory availability first
        let memory_info = self.get_gpu_memory(gpu_id).await.map_err(|e| {
            GPUAllocationError::MemoryQueryFailed {
                message: e.to_string(),
            }
        })?;

        if memory_info.free_mb < required_memory_mb {
            return Err(GPUAllocationError::InsufficientMemory {
                gpu_id,
                required_mb: required_memory_mb,
                available_mb: memory_info.free_mb,
            });
        }

        // Memory check passed, acquire lock
        self.acquire_gpu(job_id, gpu_id).await
    }

    /// Release GPU lock
    pub async fn release_gpu(&self, gpu_id: u32, job_id: Uuid) -> Result<(), GPUAllocationError> {
        let mut locks = self.gpu_locks.write().await;

        match locks.get(&gpu_id) {
            Some(locked_by) if *locked_by == job_id => {
                locks.remove(&gpu_id);
                info!("GPU {} released by job {}", gpu_id, job_id);
                Ok(())
            },
            Some(locked_by) => Err(GPUAllocationError::CannotReleaseLockedByDifferentJob {
                gpu_id,
                locked_job_id: *locked_by,
                requested_job_id: job_id,
            }),
            None => {
                // GPU not locked, this is OK (idempotent release)
                debug!("GPU {} release called but GPU was not locked", gpu_id);
                Ok(())
            },
        }
    }

    /// Release all GPU locks (cleanup)
    pub async fn release_all(&self) -> Result<()> {
        let mut locks = self.gpu_locks.write().await;
        let count = locks.len();
        locks.clear();
        info!("Released all GPU locks ({} locks cleared)", count);
        Ok(())
    }

    /// Get GPU memory information via nvidia-smi
    pub async fn get_gpu_memory(&self, gpu_id: u32) -> Result<GPUMemoryInfo> {
        // Execute nvidia-smi to get memory info
        let output = tokio::task::spawn_blocking(move || {
            Command::new("nvidia-smi")
                .args([
                    "--query-gpu=memory.total,memory.used,memory.free",
                    "--format=csv,noheader,nounits",
                    &format!("--id={}", gpu_id),
                ])
                .output()
        })
        .await??;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            return Err(anyhow::anyhow!("nvidia-smi failed: {}", stderr));
        }

        let stdout = String::from_utf8_lossy(&output.stdout);
        let parts: Vec<&str> = stdout.trim().split(',').collect();

        if parts.len() != 3 {
            return Err(anyhow::anyhow!(
                "Unexpected nvidia-smi output format: {}",
                stdout
            ));
        }

        let total_mb = parts[0].trim().parse::<u64>()?;
        let used_mb = parts[1].trim().parse::<u64>()?;
        let free_mb = parts[2].trim().parse::<u64>()?;

        Ok(GPUMemoryInfo {
            gpu_id,
            total_mb,
            used_mb,
            free_mb,
        })
    }

    /// Get GPU utilization percentage via nvidia-smi
    pub async fn get_gpu_utilization(&self, gpu_id: u32) -> Result<f32> {
        let output = tokio::task::spawn_blocking(move || {
            Command::new("nvidia-smi")
                .args([
                    "--query-gpu=utilization.gpu",
                    "--format=csv,noheader,nounits",
                    &format!("--id={}", gpu_id),
                ])
                .output()
        })
        .await??;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            return Err(anyhow::anyhow!("nvidia-smi failed: {}", stderr));
        }

        let stdout = String::from_utf8_lossy(&output.stdout);
        let utilization = stdout.trim().parse::<f32>()?;

        Ok(utilization)
    }

    /// List all active jobs and their assigned GPUs
    pub async fn list_active_jobs(&self) -> Result<Vec<(u32, Uuid)>> {
        let locks = self.gpu_locks.read().await;
        Ok(locks
            .iter()
            .map(|(gpu_id, job_id)| (*gpu_id, *job_id))
            .collect())
    }

    /// Check if a specific GPU is locked
    pub async fn is_gpu_locked(&self, gpu_id: u32) -> bool {
        let locks = self.gpu_locks.read().await;
        locks.contains_key(&gpu_id)
    }

    /// Get the job currently using a GPU (if any)
    pub async fn get_gpu_owner(&self, gpu_id: u32) -> Option<Uuid> {
        let locks = self.gpu_locks.read().await;
        locks.get(&gpu_id).copied()
    }

    /// Get statistics about GPU usage
    pub async fn get_statistics(&self) -> GPUStatistics {
        let locks = self.gpu_locks.read().await;
        GPUStatistics {
            total_gpus: self.available_gpus.len(),
            locked_gpus: locks.len(),
            available_gpus: self.available_gpus.len() - locks.len(),
            active_jobs: locks.len(),
        }
    }
}

/// GPU usage statistics
#[derive(Debug, Clone)]
pub struct GPUStatistics {
    pub total_gpus: usize,
    pub locked_gpus: usize,
    pub available_gpus: usize,
    pub active_jobs: usize,
}

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

    #[tokio::test]
    async fn test_gpu_manager_creation() {
        let manager = GPUResourceManager::new(vec![0, 1]).await;
        assert!(manager.is_ok());
    }

    #[tokio::test]
    async fn test_lock_state_tracking() {
        let manager = GPUResourceManager::new(vec![0]).await.unwrap();
        let job_id = Uuid::new_v4();

        assert!(!manager.is_gpu_locked(0).await);

        let _lock = manager.acquire_gpu(job_id, 0).await.unwrap();
        assert!(manager.is_gpu_locked(0).await);
        assert_eq!(manager.get_gpu_owner(0).await, Some(job_id));
    }

    #[tokio::test]
    async fn test_statistics() {
        let manager = GPUResourceManager::new(vec![0, 1, 2]).await.unwrap();
        let job_id = Uuid::new_v4();

        let stats = manager.get_statistics().await;
        assert_eq!(stats.total_gpus, 3);
        assert_eq!(stats.available_gpus, 3);

        let _lock = manager.acquire_gpu(job_id, 0).await.unwrap();

        let stats = manager.get_statistics().await;
        assert_eq!(stats.locked_gpus, 1);
        assert_eq!(stats.available_gpus, 2);
    }
}