foxhunt/crates/ml-core/src/gpu/profile.rs

//! GPU configuration profiles loaded from TOML.
//!
//! Replaces hardcoded VRAM-based if/else chains with a declarative profile system.
//! Profiles are embedded at compile time via `include_str!` for zero-filesystem fallback
//! (CI, containers), with filesystem and env var overrides for flexibility.
//!
//! # Loading priority
//!
//! 1. `$FOXHUNT_GPU_PROFILE` env var (exact file path)
//! 2. `config/gpu/<profile>.toml` relative to workspace root
//! 3. Embedded defaults via `include_str!`

use serde::Deserialize;
use tracing::info;

use super::capabilities::GpuCapabilities;

// Embed all profile TOMLs at compile time so they work without the filesystem.
const DEFAULT_TOML: &str = include_str!("../../../../config/gpu/default.toml");
const RTX3050_TOML: &str = include_str!("../../../../config/gpu/rtx3050.toml");
const H100_TOML: &str = include_str!("../../../../config/gpu/h100.toml");
const A100_TOML: &str = include_str!("../../../../config/gpu/a100.toml");

/// GPU configuration profile with training, experience collection, and CUDA settings.
///
/// Load via `GpuProfile::load()` which auto-detects the GPU and selects the
/// appropriate profile. All values serve as defaults -- CLI args and hyperopt
/// results always override.
#[derive(Debug, Clone, Deserialize)]
pub struct GpuProfile {
    pub training: TrainingProfile,
    pub experience: ExperienceProfile,
    pub cuda: CudaProfile,
}

/// Training-related GPU parameters.
#[derive(Debug, Clone, Deserialize)]
pub struct TrainingProfile {
    pub batch_size: usize,
    pub num_atoms: usize,
    pub buffer_size: usize,
    pub hidden_dim_base: usize,
    pub replay_buffer_vram_fraction: f64,
}

/// GPU experience collection parameters.
#[derive(Debug, Clone, Deserialize)]
pub struct ExperienceProfile {
    pub gpu_n_episodes: usize,
    pub gpu_timesteps_per_episode: usize,
}

/// CUDA kernel configuration.
#[derive(Debug, Clone, Deserialize)]
pub struct CudaProfile {
    pub cuda_stack_bytes: usize,
}

impl GpuProfile {
    /// Load the GPU profile by auto-detecting the GPU hardware.
    ///
    /// Resolution order:
    /// 1. `$FOXHUNT_GPU_PROFILE` env var (exact file path, highest priority)
    /// 2. `config/gpu/<profile>.toml` relative to workspace root
    /// 3. Embedded compile-time defaults (always available)
    ///
    /// The detected GPU name determines which profile file to select:
    /// - Contains "3050" -> `rtx3050.toml`
    /// - Contains "H100" -> `h100.toml`
    /// - Contains "A100" -> `a100.toml`
    /// - Otherwise -> `default.toml`
    pub fn load() -> Self {
        let caps = GpuCapabilities::detect();
        Self::load_for_device(&caps.device_name)
    }

    /// Load a profile for a specific device name (useful for testing).
    pub fn load_for_device(device_name: &str) -> Self {
        let profile_name = Self::profile_name_for_device(device_name);

        // Priority 1: $FOXHUNT_GPU_PROFILE env var (exact file path)
        if let Ok(env_path) = std::env::var("FOXHUNT_GPU_PROFILE") {
            if let Ok(contents) = std::fs::read_to_string(&env_path) {
                match toml::from_str::<GpuProfile>(&contents) {
                    Ok(profile) => {
                        info!(
                            "GPU profile loaded from $FOXHUNT_GPU_PROFILE: {}",
                            env_path
                        );
                        return profile;
                    }
                    Err(e) => {
                        tracing::warn!(
                            "Failed to parse GPU profile from {}: {} -- falling back",
                            env_path,
                            e
                        );
                    }
                }
            }
        }

        // Priority 2: config/gpu/<profile>.toml relative to workspace root
        // Try CARGO_MANIFEST_DIR ancestors to find workspace root with config/gpu/
        if let Some(workspace_root) = Self::find_workspace_root() {
            let profile_path = workspace_root
                .join("config")
                .join("gpu")
                .join(format!("{}.toml", profile_name));
            if let Ok(contents) = std::fs::read_to_string(&profile_path) {
                match toml::from_str::<GpuProfile>(&contents) {
                    Ok(profile) => {
                        info!(
                            "GPU profile loaded from filesystem: {} (device: {})",
                            profile_path.display(),
                            device_name
                        );
                        return profile;
                    }
                    Err(e) => {
                        tracing::warn!(
                            "Failed to parse GPU profile {}: {} -- using embedded default",
                            profile_path.display(),
                            e
                        );
                    }
                }
            }
        }

        // Priority 3: embedded compile-time defaults (always available)
        let embedded_toml = Self::embedded_toml(profile_name);
        // SAFETY: embedded TOMLs are compile-time constants that we control.
        // If parsing fails, something is catastrophically wrong with the build.
        match toml::from_str::<GpuProfile>(embedded_toml) {
            Ok(profile) => {
                info!(
                    "GPU profile loaded from embedded defaults: {} (device: {})",
                    profile_name, device_name
                );
                profile
            }
            Err(e) => {
                tracing::error!(
                    "Failed to parse embedded GPU profile '{}': {} -- using hardcoded fallback",
                    profile_name,
                    e
                );
                Self::hardcoded_fallback()
            }
        }
    }

    /// Determine which profile name to use based on the GPU device name.
    fn profile_name_for_device(device_name: &str) -> &'static str {
        let upper = device_name.to_uppercase();
        if upper.contains("3050") {
            "rtx3050"
        } else if upper.contains("H100") || upper.contains("H200") || upper.contains("GH200") {
            "h100"
        } else if upper.contains("A100") || upper.contains("L40") {
            "a100"
        } else {
            "default"
        }
    }

    /// Get the embedded TOML string for a profile name.
    fn embedded_toml(profile_name: &str) -> &'static str {
        match profile_name {
            "rtx3050" => RTX3050_TOML,
            "h100" => H100_TOML,
            "a100" => A100_TOML,
            _ => DEFAULT_TOML,
        }
    }

    /// Try to find the workspace root by looking for `config/gpu/` in parent directories.
    fn find_workspace_root() -> Option<std::path::PathBuf> {
        // Start from the current working directory
        if let Ok(cwd) = std::env::current_dir() {
            let mut dir = cwd.as_path();
            loop {
                if dir.join("config").join("gpu").is_dir() {
                    return Some(dir.to_path_buf());
                }
                match dir.parent() {
                    Some(parent) => dir = parent,
                    None => break,
                }
            }
        }
        None
    }

    /// Last-resort fallback if even embedded TOML parsing fails.
    /// This should never happen in practice.
    fn hardcoded_fallback() -> Self {
        Self {
            training: TrainingProfile {
                batch_size: 256,
                num_atoms: 21,
                buffer_size: 50_000,
                hidden_dim_base: 256,
                replay_buffer_vram_fraction: 0.50,
            },
            experience: ExperienceProfile {
                gpu_n_episodes: 64,
                gpu_timesteps_per_episode: 200,
            },
            cuda: CudaProfile {
                cuda_stack_bytes: 32768,
            },
        }
    }
}

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

    #[test]
    fn test_profile_name_for_rtx3050() {
        assert_eq!(
            GpuProfile::profile_name_for_device("NVIDIA GeForce RTX 3050 Ti Laptop GPU"),
            "rtx3050"
        );
    }

    #[test]
    fn test_profile_name_for_h100() {
        assert_eq!(
            GpuProfile::profile_name_for_device("NVIDIA H100 80GB HBM3"),
            "h100"
        );
    }

    #[test]
    fn test_profile_name_for_h200() {
        assert_eq!(
            GpuProfile::profile_name_for_device("NVIDIA H200"),
            "h100"
        );
    }

    #[test]
    fn test_profile_name_for_a100() {
        assert_eq!(
            GpuProfile::profile_name_for_device("NVIDIA A100-SXM4-80GB"),
            "a100"
        );
    }

    #[test]
    fn test_profile_name_for_l40s() {
        assert_eq!(
            GpuProfile::profile_name_for_device("NVIDIA L40S"),
            "a100"
        );
    }

    #[test]
    fn test_profile_name_for_unknown() {
        assert_eq!(
            GpuProfile::profile_name_for_device("Tesla V100"),
            "default"
        );
    }

    #[test]
    fn test_profile_name_for_cpu() {
        assert_eq!(GpuProfile::profile_name_for_device("CPU"), "default");
    }

    #[test]
    fn test_embedded_rtx3050_parses() {
        let profile: GpuProfile = toml::from_str(RTX3050_TOML).unwrap();
        assert_eq!(profile.training.batch_size, 64);
        assert_eq!(profile.training.num_atoms, 11);
        assert_eq!(profile.training.buffer_size, 5_000);
        assert_eq!(profile.experience.gpu_n_episodes, 16);
        assert_eq!(profile.experience.gpu_timesteps_per_episode, 100);
        assert_eq!(profile.cuda.cuda_stack_bytes, 16384);
    }

    #[test]
    fn test_embedded_h100_parses() {
        let profile: GpuProfile = toml::from_str(H100_TOML).unwrap();
        assert_eq!(profile.training.batch_size, 1024);
        assert_eq!(profile.training.num_atoms, 51);
        assert_eq!(profile.training.buffer_size, 500_000);
        assert_eq!(profile.experience.gpu_n_episodes, 256);
        assert_eq!(profile.experience.gpu_timesteps_per_episode, 500);
        assert_eq!(profile.cuda.cuda_stack_bytes, 65536);
    }

    #[test]
    fn test_embedded_a100_parses() {
        let profile: GpuProfile = toml::from_str(A100_TOML).unwrap();
        assert_eq!(profile.training.batch_size, 512);
        assert_eq!(profile.training.num_atoms, 51);
        assert_eq!(profile.training.buffer_size, 200_000);
        assert_eq!(profile.cuda.cuda_stack_bytes, 65536);
    }

    #[test]
    fn test_embedded_default_parses() {
        let profile: GpuProfile = toml::from_str(DEFAULT_TOML).unwrap();
        assert_eq!(profile.training.batch_size, 256);
        assert_eq!(profile.training.num_atoms, 21);
        assert_eq!(profile.training.buffer_size, 50_000);
        assert_eq!(profile.experience.gpu_n_episodes, 64);
        assert_eq!(profile.experience.gpu_timesteps_per_episode, 200);
        assert_eq!(profile.cuda.cuda_stack_bytes, 32768);
    }

    #[test]
    fn test_load_for_device_rtx3050() {
        let profile = GpuProfile::load_for_device("NVIDIA GeForce RTX 3050 Ti Laptop GPU");
        assert_eq!(profile.training.batch_size, 64);
        assert_eq!(profile.training.num_atoms, 11);
        assert_eq!(profile.cuda.cuda_stack_bytes, 16384);
    }

    #[test]
    fn test_load_for_device_h100() {
        let profile = GpuProfile::load_for_device("NVIDIA H100 80GB HBM3");
        assert_eq!(profile.training.batch_size, 1024);
        assert_eq!(profile.training.num_atoms, 51);
        assert_eq!(profile.cuda.cuda_stack_bytes, 65536);
    }

    #[test]
    fn test_load_for_device_unknown_uses_default() {
        let profile = GpuProfile::load_for_device("Unknown GPU");
        assert_eq!(profile.training.batch_size, 256);
        assert_eq!(profile.training.num_atoms, 21);
    }

    #[test]
    fn test_hardcoded_fallback_values() {
        let profile = GpuProfile::hardcoded_fallback();
        assert_eq!(profile.training.batch_size, 256);
        assert_eq!(profile.training.num_atoms, 21);
        assert_eq!(profile.training.buffer_size, 50_000);
        assert_eq!(profile.experience.gpu_n_episodes, 64);
        assert_eq!(profile.experience.gpu_timesteps_per_episode, 200);
        assert_eq!(profile.cuda.cuda_stack_bytes, 32768);
    }

    #[test]
    fn test_all_embedded_profiles_consistent_with_lookup() {
        // Verify that every profile name maps to a parseable embedded TOML
        for name in &["rtx3050", "h100", "a100", "default"] {
            let toml_str = GpuProfile::embedded_toml(name);
            let profile: Result<GpuProfile, _> = toml::from_str(toml_str);
            assert!(
                profile.is_ok(),
                "Embedded profile '{}' failed to parse: {:?}",
                name,
                profile.err()
            );
        }
    }
}