rusty-stack 0.4.0

//! Permanent ROCm environment setup — ports `scripts/setup_permanent_rocm_env.sh`.
//!
//! Creates venv at correct path, sets up symlinks and environment variables
//! in `~/.mlstack_env`.
//!
//! # Validation Assertions
//!
//! - **VAL-INSTALL-028**: Permanent env creates correct venv and symlinks

// ===========================================================================
// Types
// ===========================================================================

/// A constructed shell command.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ShellCommand {
    /// The program to run.
    pub program: String,
    /// Arguments to pass.
    pub args: Vec<String>,
    /// Environment variables to set.
    pub env: Vec<(String, String)>,
}

impl ShellCommand {
    /// Format as a shell command string.
    pub fn to_command_string(&self) -> String {
        let env_prefix = self
            .env
            .iter()
            .map(|(k, v)| format!("{k}={v}"))
            .collect::<Vec<_>>()
            .join(" ");
        let cmd = if self.args.is_empty() {
            self.program.clone()
        } else {
            format!("{} {}", self.program, self.args.join(" "))
        };
        if env_prefix.is_empty() {
            cmd
        } else {
            format!("{env_prefix} {cmd}")
        }
    }
}

/// Configuration for the permanent environment setup.
#[derive(Debug, Clone)]
pub struct PermanentEnvConfig {
    /// Python binary to use.
    pub python_bin: String,
    /// ROCm version (detected or default).
    pub rocm_version: String,
    /// GPU architecture (detected or default).
    pub gpu_arch: String,
    /// HSA override GFX version.
    pub hsa_override_gfx_version: String,
    /// Discrete GPU device list.
    pub discrete_gpu_list: String,
    /// Installation method.
    pub install_method: String,
    /// Whether to run in dry-run mode.
    pub dry_run: bool,
}

impl Default for PermanentEnvConfig {
    fn default() -> Self {
        Self {
            python_bin: "python3".to_string(),
            rocm_version: "7.2.0".to_string(),
            gpu_arch: "gfx1100".to_string(),
            hsa_override_gfx_version: "11.0.0".to_string(),
            discrete_gpu_list: "0".to_string(),
            install_method: "auto".to_string(),
            dry_run: false,
        }
    }
}

/// Default path for the environment file.
pub fn default_env_file_path() -> String {
    format!("{}/.mlstack_env", std::env::var("HOME").unwrap_or_default())
}

/// The permanent environment setup installer.
pub struct PermanentEnvInstaller {
    config: PermanentEnvConfig,
}

impl Default for PermanentEnvInstaller {
    fn default() -> Self {
        Self::with_defaults()
    }
}

impl PermanentEnvInstaller {
    /// Create a new permanent env installer with the given config.
    pub fn new(config: PermanentEnvConfig) -> Self {
        Self { config }
    }

    /// Create with default config.
    pub fn with_defaults() -> Self {
        Self::new(PermanentEnvConfig::default())
    }

    // -----------------------------------------------------------------------
    // Command construction (VAL-INSTALL-028)
    // -----------------------------------------------------------------------

    /// Get the environment file path.
    pub fn env_file_path(&self) -> String {
        default_env_file_path()
    }

    /// Generate the content of the `.mlstack_env` file.
    ///
    /// Matches the content generated by `setup_permanent_rocm_env.sh`.
    pub fn generate_env_file_content(&self) -> String {
        // Delegate to the single canonical generator (bootstrap/env_setup.rs) so
        // there is ONE source of truth for the ROCm env across the bootstrap,
        // installer, and permanent-env paths. The canonical generator emits the
        // full union (device filter incl. ROCR_VISIBLE_DEVICES, arch, paths,
        // perf/triton/vLLM/MPI) in bash form. (install metadata like
        // MLSTACK_INSTALL_METHOD is appended by ensure_mlstack_env where relevant.)
        crate::bootstrap::env_setup::generate_env_file_content(
            &self.config.discrete_gpu_list,
            "/opt/rocm",
            &self.config.rocm_version,
            "latest",
            &self.config.gpu_arch,
            &self.config.hsa_override_gfx_version,
            &self.config.python_bin,
        )
    }

    /// Generate HSA_TOOLS_LIB line based on profiler library presence.
    /// Uses the resolved path from `resolve_rocprofiler_tool_path()` to handle
    /// both the new ROCm 7.x layout (`lib/rocprofiler-sdk/`) and the old layout.
    pub fn generate_hsa_tools_lib_line(&self) -> String {
        if let Some(path) = crate::installers::common::utils::resolve_rocprofiler_tool_path() {
            format!("export HSA_TOOLS_LIB={}", path)
        } else {
            "# HSA_TOOLS_LIB not set (profiler not found)".to_string()
        }
    }

    /// Get the key environment variables that must be set.
    /// Returns variable name and their expected values.
    pub fn required_env_vars(&self) -> Vec<(&'static str, String)> {
        vec![
            ("ROCM_VERSION", self.config.rocm_version.clone()),
            ("ROCM_CHANNEL", "latest".to_string()),
            ("GPU_ARCH", self.config.gpu_arch.clone()),
            ("PYTORCH_ROCM_ARCH", self.config.gpu_arch.clone()),
            ("ROCM_HOME", "/opt/rocm".to_string()),
            ("ROCM_PATH", "/opt/rocm".to_string()),
            ("HIP_VISIBLE_DEVICES", self.config.discrete_gpu_list.clone()),
            (
                "CUDA_VISIBLE_DEVICES",
                self.config.discrete_gpu_list.clone(),
            ),
            ("MLSTACK_PYTHON_BIN", self.config.python_bin.clone()),
            (
                "HSA_OVERRIDE_GFX_VERSION",
                self.config.hsa_override_gfx_version.clone(),
            ),
            ("HSA_ENABLE_SDMA", "0".to_string()),
            ("GPU_MAX_HEAP_SIZE", "100".to_string()),
            ("GPU_MAX_ALLOC_PERCENT", "100".to_string()),
            ("MIOPEN_DEBUG_CONV_IMPLICIT_GEMM", "1".to_string()),
            ("MIOPEN_FIND_MODE", "3".to_string()),
            ("MIOPEN_FIND_ENFORCE", "3".to_string()),
            ("PIP_BREAK_SYSTEM_PACKAGES", "1".to_string()),
            ("UV_PIP_BREAK_SYSTEM_PACKAGES", "1".to_string()),
            ("UV_SYSTEM_PYTHON", "1".to_string()),
        ]
    }

    /// Construct the mkdir command for triton directories.
    pub fn build_mkdir_triton_dirs_command(&self) -> ShellCommand {
        let home = std::env::var("HOME").unwrap_or_default();
        ShellCommand {
            program: "mkdir".to_string(),
            args: vec![
                "-p".to_string(),
                format!("{home}/.cache/mlstack/triton/cache"),
                format!("{home}/.cache/mlstack/triton/dump"),
                format!("{home}/.cache/mlstack/triton/override"),
            ],
            env: vec![],
        }
    }

    /// Construct the command to detect GPU architecture via rocminfo.
    ///
    /// Uses the full path to rocminfo (e.g., `/opt/rocm/bin/rocminfo`) to avoid
    /// "command not found" errors when running via sudo, which strips the user's
    /// PATH environment.
    pub fn build_rocminfo_gpu_detect_command(&self) -> ShellCommand {
        let rocminfo_path = crate::installers::common::utils::resolve_rocminfo_path();
        ShellCommand {
            program: rocminfo_path,
            args: vec![],
            env: vec![],
        }
    }

    /// Construct the command to read ROCm version.
    pub fn build_rocm_version_detect_command(&self) -> ShellCommand {
        ShellCommand {
            program: "cat".to_string(),
            args: vec!["/opt/rocm/.info/version".to_string()],
            env: vec![],
        }
    }

    /// Get GPU arch from environment variable.
    pub fn gpu_arch_from_env() -> Option<String> {
        std::env::var("GPU_ARCH").ok().filter(|s| !s.is_empty())
    }

    /// Default GPU arch when detection fails.
    pub fn default_gpu_arch() -> &'static str {
        "gfx1100"
    }

    /// Convert gfx architecture string to HSA override version.
    /// E.g., "gfx1100" -> "11.0.0"
    pub fn gfx_to_hsa_override(gfx: &str) -> Option<String> {
        let raw = gfx.strip_prefix("gfx")?;
        match raw.len() {
            4 => {
                let major = &raw[0..2];
                let minor = &raw[2..3];
                let patch = &raw[3..4];
                Some(format!("{}.{}.{}", major, minor, patch))
            }
            3 => {
                let major = &raw[0..1];
                let minor = &raw[1..2];
                let patch = &raw[2..3];
                Some(format!("{}.{}.{}", major, minor, patch))
            }
            _ => None,
        }
    }
}

// ===========================================================================
// Tests
// ===========================================================================

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

    // -----------------------------------------------------------------------
    // VAL-INSTALL-028: Permanent env creates correct venv and symlinks
    // -----------------------------------------------------------------------

    #[test]
    fn test_env_file_path_is_home_mlstack_env() {
        let path = default_env_file_path();
        assert!(
            path.ends_with("/.mlstack_env"),
            "Env file must be at ~/.mlstack_env, got: {path}"
        );
    }

    #[test]
    fn test_generate_env_file_contains_required_vars() {
        let installer = PermanentEnvInstaller::with_defaults();
        let content = installer.generate_env_file_content();

        // Check all required variables from the original script
        assert!(content.contains("export ROCM_VERSION="));
        assert!(content.contains("export ROCM_CHANNEL=latest"));
        assert!(content.contains("export GPU_ARCH="));
        assert!(content.contains("export PYTORCH_ROCM_ARCH="));
        assert!(content.contains("export ROCM_HOME=/opt/rocm"));
        assert!(content.contains("export ROCM_PATH=/opt/rocm"));
        assert!(content.contains("export HIP_PATH=/opt/rocm"));
        assert!(content.contains("export HIP_VISIBLE_DEVICES="));
        assert!(content.contains("export CUDA_VISIBLE_DEVICES="));
        assert!(content.contains("export MLSTACK_PYTHON_BIN="));
        assert!(content.contains("export UV_PYTHON="));
        assert!(content.contains("export HSA_OVERRIDE_GFX_VERSION="));
        assert!(content.contains("export HSA_ENABLE_SDMA=0"));
        assert!(content.contains("export GPU_MAX_HEAP_SIZE=100"));
        assert!(content.contains("export GPU_MAX_ALLOC_PERCENT=100"));
        assert!(content.contains("export MIOPEN_DEBUG_CONV_IMPLICIT_GEMM=1"));
        assert!(content.contains("export MIOPEN_FIND_MODE=3"));
        assert!(content.contains("export MIOPEN_FIND_ENFORCE=3"));
        assert!(content.contains("export PIP_BREAK_SYSTEM_PACKAGES=1"));
        assert!(content.contains("export UV_PIP_BREAK_SYSTEM_PACKAGES=1"));
        assert!(content.contains("export UV_SYSTEM_PYTHON=1"));
    }

    #[test]
    fn test_env_file_contains_path_settings() {
        let installer = PermanentEnvInstaller::with_defaults();
        let content = installer.generate_env_file_content();

        assert!(content.contains("/opt/rocm/bin"));
        assert!(content.contains("/opt/rocm/lib"));
        assert!(content.contains("LD_LIBRARY_PATH"));
    }

    #[test]
    fn test_env_file_contains_pytorch_optimization() {
        let installer = PermanentEnvInstaller::with_defaults();
        let content = installer.generate_env_file_content();

        // CUDA arch list is intentionally NOT exported on an AMD/ROCm stack
        // (Stage 2: no CUDA env leakage); ROCm uses PYTORCH_ROCM_ARCH instead.
        assert!(!content.contains("TORCH_CUDA_ARCH_LIST"));
        assert!(content.contains("PYTORCH_CUDA_ALLOC_CONF"));
        assert!(content.contains("PYTORCH_HIP_ALLOC_CONF"));
    }

    #[test]
    fn test_env_file_contains_vllm_settings() {
        let installer = PermanentEnvInstaller::with_defaults();
        let content = installer.generate_env_file_content();

        assert!(content.contains("FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE"));
        assert!(content.contains("VLLM_ROCM_USE_AITER=0"));
        assert!(content.contains("VLLM_WORKER_MULTIPROC_METHOD=spawn"));
    }

    #[test]
    fn test_env_file_contains_triton_dirs() {
        let installer = PermanentEnvInstaller::with_defaults();
        let content = installer.generate_env_file_content();

        assert!(content.contains("MLSTACK_TRITON_HOME"));
        assert!(content.contains("TRITON_CACHE_DIR"));
        assert!(content.contains("TRITON_DUMP_DIR"));
        assert!(content.contains("TRITON_OVERRIDE_DIR"));
    }

    #[test]
    fn test_required_env_vars_count() {
        let installer = PermanentEnvInstaller::with_defaults();
        let vars = installer.required_env_vars();
        // Must have at least the key variables
        assert!(
            vars.len() >= 19,
            "Expected at least 19 env vars, got {}",
            vars.len()
        );
    }

    #[test]
    fn test_required_env_vars_values() {
        let config = PermanentEnvConfig {
            rocm_version: "7.2.0".to_string(),
            gpu_arch: "gfx1100".to_string(),
            hsa_override_gfx_version: "11.0.0".to_string(),
            discrete_gpu_list: "0,1".to_string(),
            python_bin: "python3.12".to_string(),
            install_method: "auto".to_string(),
            ..Default::default()
        };
        let installer = PermanentEnvInstaller::new(config);
        let vars = installer.required_env_vars();

        let get = |name: &str| -> String {
            vars.iter()
                .find(|(n, _)| *n == name)
                .map(|(_, v)| v.clone())
                .unwrap_or_default()
        };

        assert_eq!(get("ROCM_VERSION"), "7.2.0");
        assert_eq!(get("GPU_ARCH"), "gfx1100");
        assert_eq!(get("HIP_VISIBLE_DEVICES"), "0,1");
        assert_eq!(get("MLSTACK_PYTHON_BIN"), "python3.12");
        assert_eq!(get("HSA_OVERRIDE_GFX_VERSION"), "11.0.0");
    }

    #[test]
    fn test_gfx_to_hsa_override_conversion() {
        assert_eq!(
            PermanentEnvInstaller::gfx_to_hsa_override("gfx1100"),
            Some("11.0.0".to_string())
        );
        assert_eq!(
            PermanentEnvInstaller::gfx_to_hsa_override("gfx1030"),
            Some("10.3.0".to_string())
        );
        assert_eq!(
            PermanentEnvInstaller::gfx_to_hsa_override("gfx1201"),
            Some("12.0.1".to_string())
        );
        assert_eq!(
            PermanentEnvInstaller::gfx_to_hsa_override("gfx900"),
            Some("9.0.0".to_string())
        );
        assert_eq!(PermanentEnvInstaller::gfx_to_hsa_override("invalid"), None);
    }

    #[test]
    fn test_mkdir_triton_dirs_command() {
        let installer = PermanentEnvInstaller::with_defaults();
        let cmd = installer.build_mkdir_triton_dirs_command();
        assert_eq!(cmd.program, "mkdir");
        assert!(cmd.args.contains(&"-p".to_string()));
        let cmd_str = cmd.to_command_string();
        assert!(cmd_str.contains("mlstack/triton/cache"));
        assert!(cmd_str.contains("mlstack/triton/dump"));
        assert!(cmd_str.contains("mlstack/triton/override"));
    }

    #[test]
    fn test_rocm_version_detect_command() {
        let installer = PermanentEnvInstaller::with_defaults();
        let cmd = installer.build_rocm_version_detect_command();
        assert_eq!(cmd.program, "cat");
        assert!(cmd.args.contains(&"/opt/rocm/.info/version".to_string()));
    }

    #[test]
    fn test_rocminfo_gpu_detect_command_uses_full_path() {
        let installer = PermanentEnvInstaller::with_defaults();
        let cmd = installer.build_rocminfo_gpu_detect_command();
        // Must NOT be bare "rocminfo" — should resolve to full path when available
        assert!(
            cmd.program.contains("rocminfo"),
            "Program should contain 'rocminfo', got: {}",
            cmd.program
        );
        // If /opt/rocm/bin/rocminfo exists, must use the full path
        if std::path::Path::new("/opt/rocm/bin/rocminfo").exists() {
            assert_eq!(
                cmd.program, "/opt/rocm/bin/rocminfo",
                "Must use full path when /opt/rocm/bin/rocminfo exists"
            );
        }
    }

    #[test]
    fn test_generate_hsa_tools_lib_line() {
        let installer = PermanentEnvInstaller::with_defaults();
        let line = installer.generate_hsa_tools_lib_line();
        // Either exports the path or comments it out
        assert!(
            line.contains("HSA_TOOLS_LIB") || line.contains("profiler not found"),
            "Should reference HSA_TOOLS_LIB"
        );
    }

    #[test]
    fn test_custom_config_values_in_env_file() {
        let config = PermanentEnvConfig {
            rocm_version: "7.0.0".to_string(),
            gpu_arch: "gfx1030".to_string(),
            discrete_gpu_list: "0".to_string(),
            python_bin: "python3.11".to_string(),
            ..Default::default()
        };
        let installer = PermanentEnvInstaller::new(config);
        let content = installer.generate_env_file_content();

        assert!(content.contains("ROCM_VERSION=7.0.0"));
        assert!(content.contains("GPU_ARCH=gfx1030"));
        assert!(content.contains("HIP_VISIBLE_DEVICES=0"));
        assert!(content.contains("MLSTACK_PYTHON_BIN=python3.11"));
    }
}