foxhunt/services/ml_training_service/src/gpu_config.rs

//! GPU configuration management for ML Training Service
//!
//! Handles GPU resource configuration, validation, and optimization settings
//! for machine learning training workloads.

use anyhow::{Context, Result};
use config::manager::ConfigManager;
use config::TrainingConfig;
use serde::{Deserialize, Serialize};
use std::sync::Arc;

/// GPU configuration structure
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GpuConfig {
    /// GPU device ID to use (0-based)
    pub device_id: u32,
    /// Maximum GPU memory to use in GB
    pub max_memory_gb: f32,
    /// Enable mixed precision training
    pub enable_mixed_precision: bool,
    /// Enable GPU memory optimization
    pub enable_memory_optimization: bool,
    /// Batch size optimization factor
    pub batch_size_factor: f32,
    /// Enable CUDA graphs for optimization
    pub enable_cuda_graphs: bool,
    /// Enable tensor cores if available
    pub enable_tensor_cores: bool,
}

impl Default for GpuConfig {
    fn default() -> Self {
        Self {
            device_id: 0,
            max_memory_gb: 8.0,
            enable_mixed_precision: true,
            enable_memory_optimization: true,
            batch_size_factor: 1.0,
            enable_cuda_graphs: false,
            enable_tensor_cores: true,
        }
    }
}

/// GPU validation result
#[derive(Debug, Clone)]
pub struct GpuValidation {
    pub is_available: bool,
    pub memory_available_gb: f32,
    pub compute_capability: String,
    pub driver_version: String,
    pub issues: Vec<String>,
}

impl GpuValidation {
    /// Check if GPU is ready for training
    pub fn is_ready_for_training(&self) -> bool {
        self.is_available && self.issues.is_empty()
    }

    /// Get list of validation issues
    pub fn get_issues(&self) -> &[String] {
        &self.issues
    }
}

/// GPU configuration manager
pub struct GpuConfigManager {
    training_config: TrainingConfig,
    config_manager: Arc<ConfigManager>,
    gpu_config: Option<GpuConfig>,
}

impl GpuConfigManager {
    /// Create new GPU configuration manager
    pub fn new(training_config: TrainingConfig, config_manager: Arc<ConfigManager>) -> Self {
        Self {
            training_config,
            config_manager,
            gpu_config: None,
        }
    }

    /// Get a reference to the training configuration
    pub fn training_config(&self) -> &TrainingConfig {
        &self.training_config
    }

    /// Load GPU configuration from config manager
    pub async fn load_config(&mut self) -> Result<GpuConfig> {
        let gpu_config = self
            .load_gpu_config_from_manager()
            .await
            .unwrap_or_else(|_| self.create_default_config());

        self.gpu_config = Some(gpu_config.clone());
        Ok(gpu_config)
    }

    /// Load GPU configuration from the config manager
    async fn load_gpu_config_from_manager(&self) -> Result<GpuConfig> {
        // Get the service config which contains settings as JSON
        let service_config = self.config_manager.get_config();
        let settings = &service_config.settings;

        // Helper function to extract typed values from settings
        let get_value = |key: &str| -> Option<serde_json::Value> { settings.get(key).cloned() };

        let device_id = get_value("gpu_device_id")
            .and_then(|v| v.as_u64().map(|n| n as u32))
            .unwrap_or(0);

        let max_memory_gb = get_value("gpu_max_memory_gb")
            .and_then(|v| v.as_f64().map(|n| n as f32))
            .unwrap_or(8.0);

        let enable_mixed_precision = get_value("gpu_enable_mixed_precision")
            .and_then(|v| v.as_bool())
            .unwrap_or(true);

        let enable_memory_optimization = get_value("gpu_enable_memory_optimization")
            .and_then(|v| v.as_bool())
            .unwrap_or(true);

        let batch_size_factor = get_value("gpu_batch_size_factor")
            .and_then(|v| v.as_f64().map(|n| n as f32))
            .unwrap_or(1.0);

        let enable_cuda_graphs = get_value("gpu_enable_cuda_graphs")
            .and_then(|v| v.as_bool())
            .unwrap_or(false);

        let enable_tensor_cores = get_value("gpu_enable_tensor_cores")
            .and_then(|v| v.as_bool())
            .unwrap_or(true);

        Ok(GpuConfig {
            device_id,
            max_memory_gb,
            enable_mixed_precision,
            enable_memory_optimization,
            batch_size_factor,
            enable_cuda_graphs,
            enable_tensor_cores,
        })
    }

    /// Create default GPU configuration based on training config
    fn create_default_config(&self) -> GpuConfig {
        // TrainingConfig doesn't have GPU-specific fields, so use defaults
        // GPU configuration is managed separately through ConfigManager
        GpuConfig::default()
    }

    /// Validate GPU availability and configuration
    pub async fn validate_gpu_availability(&self) -> Result<GpuValidation> {
        let mut validation = GpuValidation {
            is_available: false,
            memory_available_gb: 0.0,
            compute_capability: "Unknown".to_string(),
            driver_version: "Unknown".to_string(),
            issues: Vec::new(),
        };

        // In a real implementation, this would use CUDA runtime API or similar
        // For now, we'll simulate basic validation
        if std::env::var("CUDA_VISIBLE_DEVICES").is_ok() {
            validation.is_available = true;
            validation.memory_available_gb = 8.0; // Simulated
            validation.compute_capability = "7.5".to_string(); // Simulated
            validation.driver_version = "11.8".to_string(); // Simulated
        } else {
            validation
                .issues
                .push("CUDA_VISIBLE_DEVICES not set".to_string());
        }

        // Check if requested device ID is valid
        if let Some(gpu_config) = &self.gpu_config {
            if gpu_config.device_id > 7 {
                validation.issues.push(format!(
                    "Device ID {} may be invalid (typically 0-7)",
                    gpu_config.device_id
                ));
            }

            if gpu_config.max_memory_gb > validation.memory_available_gb {
                validation.issues.push(format!(
                    "Requested memory {} GB exceeds available {} GB",
                    gpu_config.max_memory_gb, validation.memory_available_gb
                ));
            }
        }

        Ok(validation)
    }

    /// Get current GPU configuration
    pub fn get_config(&self) -> Option<&GpuConfig> {
        self.gpu_config.as_ref()
    }

    /// Update GPU configuration
    pub async fn update_config(&mut self, new_config: GpuConfig) -> Result<()> {
        // Validate the new configuration
        let temp_config = self.gpu_config.clone();
        self.gpu_config = Some(new_config.clone());

        match self.validate_gpu_availability().await {
            Ok(validation) if validation.is_ready_for_training() => {
                // Configuration is valid
                Ok(())
            },
            Ok(validation) => {
                // Restore previous configuration
                self.gpu_config = temp_config;
                Err(anyhow::anyhow!(
                    "GPU configuration validation failed: {:?}",
                    validation.issues
                ))
            },
            Err(e) => {
                // Restore previous configuration
                self.gpu_config = temp_config;
                Err(e).context("Failed to validate GPU configuration")
            },
        }
    }

    /// Get optimal batch size based on GPU configuration
    pub fn get_optimal_batch_size(&self, base_batch_size: u32) -> u32 {
        if let Some(config) = &self.gpu_config {
            (base_batch_size as f32 * config.batch_size_factor).round() as u32
        } else {
            base_batch_size
        }
    }

    /// Check if mixed precision is enabled
    pub fn is_mixed_precision_enabled(&self) -> bool {
        self.gpu_config
            .as_ref()
            .map(|c| c.enable_mixed_precision)
            .unwrap_or(false)
    }

    /// Check if memory optimization is enabled
    pub fn is_memory_optimization_enabled(&self) -> bool {
        self.gpu_config
            .as_ref()
            .map(|c| c.enable_memory_optimization)
            .unwrap_or(false)
    }

    /// Check if CUDA graphs are enabled
    pub fn is_cuda_graphs_enabled(&self) -> bool {
        self.gpu_config
            .as_ref()
            .map(|c| c.enable_cuda_graphs)
            .unwrap_or(false)
    }

    /// Check if tensor cores are enabled
    pub fn is_tensor_cores_enabled(&self) -> bool {
        self.gpu_config
            .as_ref()
            .map(|c| c.enable_tensor_cores)
            .unwrap_or(false)
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;

    #[test]
    fn test_gpu_config_default() {
        let config = GpuConfig::default();
        assert_eq!(config.device_id, 0);
        assert_eq!(config.max_memory_gb, 8.0);
        assert!(config.enable_mixed_precision);
        assert!(config.enable_memory_optimization);
        assert_eq!(config.batch_size_factor, 1.0);
        assert!(!config.enable_cuda_graphs);
        assert!(config.enable_tensor_cores);
    }

    #[test]
    fn test_gpu_validation() {
        let validation = GpuValidation {
            is_available: true,
            memory_available_gb: 8.0,
            compute_capability: "7.5".to_string(),
            driver_version: "11.8".to_string(),
            issues: vec![],
        };

        assert!(validation.is_ready_for_training());
        assert!(validation.get_issues().is_empty());
    }

    #[test]
    fn test_gpu_validation_with_issues() {
        let validation = GpuValidation {
            is_available: true,
            memory_available_gb: 8.0,
            compute_capability: "7.5".to_string(),
            driver_version: "11.8".to_string(),
            issues: vec!["Test issue".to_string()],
        };

        assert!(!validation.is_ready_for_training());
        assert_eq!(validation.get_issues().len(), 1);
    }
}