rusty-stack 0.3.1

//! 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 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 {
        let _home = std::env::var("HOME").unwrap_or_default();
        let uid_output = std::process::Command::new("id")
            .arg("-u")
            .output()
            .map(|o| String::from_utf8_lossy(&o.stdout).trim().to_string())
            .unwrap_or_else(|_| "0".to_string());

        format!(
            r#"# Permanent ROCm Environment Setup (Generated by rusty-stack)
export ROCM_VERSION={rocm_version}
export ROCM_CHANNEL=latest
export GPU_ARCH={gpu_arch}
export PYTORCH_ROCM_ARCH={gpu_arch}
export GPU_ARCHS={gpu_arch}
export ROCM_HOME=/opt/rocm
export ROCM_PATH=/opt/rocm
export HIP_PATH=/opt/rocm
# Discrete GPUs only (iGPUs filtered): {gpu_list}
export HIP_VISIBLE_DEVICES={gpu_list}
export CUDA_VISIBLE_DEVICES={gpu_list}
export PYTORCH_ROCM_DEVICE={first_gpu}
export MLSTACK_PYTHON_BIN={python_bin}
export UV_PYTHON={python_bin}
export MLSTACK_INSTALL_METHOD={install_method}
export INSTALL_METHOD={install_method}
export PYTHONPATH=/opt/rocm/lib:$PYTHONPATH
export AMDGPU_ASIC_ID_TABLE_PATH=/usr/share/libdrm/amdgpu.ids
export AMDGPU_ASIC_ID_TABLE_PATHS=/usr/share/libdrm

# Path Settings
export PATH="/usr/local/bin:/usr/bin:/bin:/opt/rocm/bin:/opt/rocm/hip/bin:$PATH"
export LD_LIBRARY_PATH="$HOME/.mlstack/libmpi-compat:$HOME/.mlstack/libmpi-compat-user-{uid}:/opt/rocm/lib:/opt/rocm/hip/lib:/opt/rocm/opencl/lib:$LD_LIBRARY_PATH"

# Performance & Compatibility
export HSA_OVERRIDE_GFX_VERSION={hsa_override}
export HSA_ENABLE_SDMA=0
export GPU_MAX_HEAP_SIZE=100
export GPU_MAX_ALLOC_PERCENT=100
{hsa_tools_lib_line}
export MIOPEN_DEBUG_CONV_IMPLICIT_GEMM=1
export MIOPEN_FIND_MODE=3
export MIOPEN_FIND_ENFORCE=3

# PyTorch Optimization
export PYTORCH_ALLOC_CONF="max_split_size_mb:512"
export PYTORCH_HIP_ALLOC_CONF="max_split_size_mb:512"

# vLLM RDNA3 Support (v0.15.0+)
export FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE
export VLLM_ROCM_USE_AITER=0
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export MLSTACK_TRITON_HOME=$HOME/.cache/mlstack/triton
export TRITON_HOME=$MLSTACK_TRITON_HOME
export TRITON_CACHE_DIR=$MLSTACK_TRITON_HOME/cache
export TRITON_DUMP_DIR=$MLSTACK_TRITON_HOME/dump
export TRITON_OVERRIDE_DIR=$MLSTACK_TRITON_HOME/override

# Global flags for seamless installs
export PIP_BREAK_SYSTEM_PACKAGES=1
export UV_PIP_BREAK_SYSTEM_PACKAGES=1
export UV_SYSTEM_PYTHON=1

# MPI/UCX Settings (ROCm-aware — CUDA opal support disabled on an AMD stack)
export OMPI_MCA_pml=ucx
export OMPI_MCA_osc=ucx
export OMPI_MCA_btl=^openib,uct
export OMPI_MCA_btl_openib_allow_ib=true
export OMPI_MCA_btl_openib_warn_no_device_params_found=0
export OMPI_MCA_coll_hcoll_enable=0
# Prevent UCX ROCm memory registration segfault on multi-GPU systems
export UCX_MEMTYPE_CACHE=no
export UCX_LOG_LEVEL=error"
"#,
            rocm_version = self.config.rocm_version,
            gpu_arch = self.config.gpu_arch,
            gpu_list = self.config.discrete_gpu_list,
            first_gpu = self
                .config
                .discrete_gpu_list
                .split(',')
                .next()
                .unwrap_or("0"),
            python_bin = self.config.python_bin,
            install_method = self.config.install_method,
            hsa_override = self.config.hsa_override_gfx_version,
            hsa_tools_lib_line = self.generate_hsa_tools_lib_line(),
            uid = uid_output,
        )
    }

    /// 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_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"));
    }
}