rusty_stack/installers/components/
permanent_env.rs

1//! Permanent ROCm environment setup — ports `scripts/setup_permanent_rocm_env.sh`.
2//!
3//! Creates venv at correct path, sets up symlinks and environment variables
4//! in `~/.mlstack_env`.
5//!
6//! # Validation Assertions
7//!
8//! - **VAL-INSTALL-028**: Permanent env creates correct venv and symlinks
9
10// ===========================================================================
11// Types
12// ===========================================================================
13
14/// A constructed shell command.
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct ShellCommand {
17    /// The program to run.
18    pub program: String,
19    /// Arguments to pass.
20    pub args: Vec<String>,
21    /// Environment variables to set.
22    pub env: Vec<(String, String)>,
23}
24
25impl ShellCommand {
26    /// Format as a shell command string.
27    pub fn to_command_string(&self) -> String {
28        let env_prefix = self
29            .env
30            .iter()
31            .map(|(k, v)| format!("{k}={v}"))
32            .collect::<Vec<_>>()
33            .join(" ");
34        let cmd = if self.args.is_empty() {
35            self.program.clone()
36        } else {
37            format!("{} {}", self.program, self.args.join(" "))
38        };
39        if env_prefix.is_empty() {
40            cmd
41        } else {
42            format!("{env_prefix} {cmd}")
43        }
44    }
45}
46
47/// Configuration for the permanent environment setup.
48#[derive(Debug, Clone)]
49pub struct PermanentEnvConfig {
50    /// Python binary to use.
51    pub python_bin: String,
52    /// ROCm version (detected or default).
53    pub rocm_version: String,
54    /// GPU architecture (detected or default).
55    pub gpu_arch: String,
56    /// HSA override GFX version.
57    pub hsa_override_gfx_version: String,
58    /// Discrete GPU device list.
59    pub discrete_gpu_list: String,
60    /// Installation method.
61    pub install_method: String,
62    /// Whether to run in dry-run mode.
63    pub dry_run: bool,
64}
65
66impl Default for PermanentEnvConfig {
67    fn default() -> Self {
68        Self {
69            python_bin: "python3".to_string(),
70            rocm_version: "7.2.0".to_string(),
71            gpu_arch: "gfx1100".to_string(),
72            hsa_override_gfx_version: "11.0.0".to_string(),
73            discrete_gpu_list: "0".to_string(),
74            install_method: "auto".to_string(),
75            dry_run: false,
76        }
77    }
78}
79
80/// Default path for the environment file.
81pub fn default_env_file_path() -> String {
82    format!("{}/.mlstack_env", std::env::var("HOME").unwrap_or_default())
83}
84
85/// The permanent environment setup installer.
86pub struct PermanentEnvInstaller {
87    config: PermanentEnvConfig,
88}
89
90impl Default for PermanentEnvInstaller {
91    fn default() -> Self {
92        Self::with_defaults()
93    }
94}
95
96impl PermanentEnvInstaller {
97    /// Create a new permanent env installer with the given config.
98    pub fn new(config: PermanentEnvConfig) -> Self {
99        Self { config }
100    }
101
102    /// Create with default config.
103    pub fn with_defaults() -> Self {
104        Self::new(PermanentEnvConfig::default())
105    }
106
107    // -----------------------------------------------------------------------
108    // Command construction (VAL-INSTALL-028)
109    // -----------------------------------------------------------------------
110
111    /// Get the environment file path.
112    pub fn env_file_path(&self) -> String {
113        default_env_file_path()
114    }
115
116    /// Generate the content of the `.mlstack_env` file.
117    ///
118    /// Matches the content generated by `setup_permanent_rocm_env.sh`.
119    pub fn generate_env_file_content(&self) -> String {
120        // Delegate to the single canonical generator (bootstrap/env_setup.rs) so
121        // there is ONE source of truth for the ROCm env across the bootstrap,
122        // installer, and permanent-env paths. The canonical generator emits the
123        // full union (device filter incl. ROCR_VISIBLE_DEVICES, arch, paths,
124        // perf/triton/vLLM/MPI) in bash form. (install metadata like
125        // MLSTACK_INSTALL_METHOD is appended by ensure_mlstack_env where relevant.)
126        crate::bootstrap::env_setup::generate_env_file_content(
127            &self.config.discrete_gpu_list,
128            "/opt/rocm",
129            &self.config.rocm_version,
130            "latest",
131            &self.config.gpu_arch,
132            &self.config.hsa_override_gfx_version,
133            &self.config.python_bin,
134        )
135    }
136
137    /// Generate HSA_TOOLS_LIB line based on profiler library presence.
138    /// Uses the resolved path from `resolve_rocprofiler_tool_path()` to handle
139    /// both the new ROCm 7.x layout (`lib/rocprofiler-sdk/`) and the old layout.
140    pub fn generate_hsa_tools_lib_line(&self) -> String {
141        if let Some(path) = crate::installers::common::utils::resolve_rocprofiler_tool_path() {
142            format!("export HSA_TOOLS_LIB={}", path)
143        } else {
144            "# HSA_TOOLS_LIB not set (profiler not found)".to_string()
145        }
146    }
147
148    /// Get the key environment variables that must be set.
149    /// Returns variable name and their expected values.
150    pub fn required_env_vars(&self) -> Vec<(&'static str, String)> {
151        vec![
152            ("ROCM_VERSION", self.config.rocm_version.clone()),
153            ("ROCM_CHANNEL", "latest".to_string()),
154            ("GPU_ARCH", self.config.gpu_arch.clone()),
155            ("PYTORCH_ROCM_ARCH", self.config.gpu_arch.clone()),
156            ("ROCM_HOME", "/opt/rocm".to_string()),
157            ("ROCM_PATH", "/opt/rocm".to_string()),
158            ("HIP_VISIBLE_DEVICES", self.config.discrete_gpu_list.clone()),
159            (
160                "CUDA_VISIBLE_DEVICES",
161                self.config.discrete_gpu_list.clone(),
162            ),
163            ("MLSTACK_PYTHON_BIN", self.config.python_bin.clone()),
164            (
165                "HSA_OVERRIDE_GFX_VERSION",
166                self.config.hsa_override_gfx_version.clone(),
167            ),
168            ("HSA_ENABLE_SDMA", "0".to_string()),
169            ("GPU_MAX_HEAP_SIZE", "100".to_string()),
170            ("GPU_MAX_ALLOC_PERCENT", "100".to_string()),
171            ("MIOPEN_DEBUG_CONV_IMPLICIT_GEMM", "1".to_string()),
172            ("MIOPEN_FIND_MODE", "3".to_string()),
173            ("MIOPEN_FIND_ENFORCE", "3".to_string()),
174            ("PIP_BREAK_SYSTEM_PACKAGES", "1".to_string()),
175            ("UV_PIP_BREAK_SYSTEM_PACKAGES", "1".to_string()),
176            ("UV_SYSTEM_PYTHON", "1".to_string()),
177        ]
178    }
179
180    /// Construct the mkdir command for triton directories.
181    pub fn build_mkdir_triton_dirs_command(&self) -> ShellCommand {
182        let home = std::env::var("HOME").unwrap_or_default();
183        ShellCommand {
184            program: "mkdir".to_string(),
185            args: vec![
186                "-p".to_string(),
187                format!("{home}/.cache/mlstack/triton/cache"),
188                format!("{home}/.cache/mlstack/triton/dump"),
189                format!("{home}/.cache/mlstack/triton/override"),
190            ],
191            env: vec![],
192        }
193    }
194
195    /// Construct the command to detect GPU architecture via rocminfo.
196    ///
197    /// Uses the full path to rocminfo (e.g., `/opt/rocm/bin/rocminfo`) to avoid
198    /// "command not found" errors when running via sudo, which strips the user's
199    /// PATH environment.
200    pub fn build_rocminfo_gpu_detect_command(&self) -> ShellCommand {
201        let rocminfo_path = crate::installers::common::utils::resolve_rocminfo_path();
202        ShellCommand {
203            program: rocminfo_path,
204            args: vec![],
205            env: vec![],
206        }
207    }
208
209    /// Construct the command to read ROCm version.
210    pub fn build_rocm_version_detect_command(&self) -> ShellCommand {
211        ShellCommand {
212            program: "cat".to_string(),
213            args: vec!["/opt/rocm/.info/version".to_string()],
214            env: vec![],
215        }
216    }
217
218    /// Get GPU arch from environment variable.
219    pub fn gpu_arch_from_env() -> Option<String> {
220        std::env::var("GPU_ARCH").ok().filter(|s| !s.is_empty())
221    }
222
223    /// Default GPU arch when detection fails.
224    pub fn default_gpu_arch() -> &'static str {
225        "gfx1100"
226    }
227
228    /// Convert gfx architecture string to HSA override version.
229    /// E.g., "gfx1100" -> "11.0.0"
230    pub fn gfx_to_hsa_override(gfx: &str) -> Option<String> {
231        let raw = gfx.strip_prefix("gfx")?;
232        match raw.len() {
233            4 => {
234                let major = &raw[0..2];
235                let minor = &raw[2..3];
236                let patch = &raw[3..4];
237                Some(format!("{}.{}.{}", major, minor, patch))
238            }
239            3 => {
240                let major = &raw[0..1];
241                let minor = &raw[1..2];
242                let patch = &raw[2..3];
243                Some(format!("{}.{}.{}", major, minor, patch))
244            }
245            _ => None,
246        }
247    }
248}
249
250// ===========================================================================
251// Tests
252// ===========================================================================
253
254#[cfg(test)]
255mod tests {
256    use super::*;
257
258    // -----------------------------------------------------------------------
259    // VAL-INSTALL-028: Permanent env creates correct venv and symlinks
260    // -----------------------------------------------------------------------
261
262    #[test]
263    fn test_env_file_path_is_home_mlstack_env() {
264        let path = default_env_file_path();
265        assert!(
266            path.ends_with("/.mlstack_env"),
267            "Env file must be at ~/.mlstack_env, got: {path}"
268        );
269    }
270
271    #[test]
272    fn test_generate_env_file_contains_required_vars() {
273        let installer = PermanentEnvInstaller::with_defaults();
274        let content = installer.generate_env_file_content();
275
276        // Check all required variables from the original script
277        assert!(content.contains("export ROCM_VERSION="));
278        assert!(content.contains("export ROCM_CHANNEL=latest"));
279        assert!(content.contains("export GPU_ARCH="));
280        assert!(content.contains("export PYTORCH_ROCM_ARCH="));
281        assert!(content.contains("export ROCM_HOME=/opt/rocm"));
282        assert!(content.contains("export ROCM_PATH=/opt/rocm"));
283        assert!(content.contains("export HIP_PATH=/opt/rocm"));
284        assert!(content.contains("export HIP_VISIBLE_DEVICES="));
285        assert!(content.contains("export CUDA_VISIBLE_DEVICES="));
286        assert!(content.contains("export MLSTACK_PYTHON_BIN="));
287        assert!(content.contains("export UV_PYTHON="));
288        assert!(content.contains("export HSA_OVERRIDE_GFX_VERSION="));
289        assert!(content.contains("export HSA_ENABLE_SDMA=0"));
290        assert!(content.contains("export GPU_MAX_HEAP_SIZE=100"));
291        assert!(content.contains("export GPU_MAX_ALLOC_PERCENT=100"));
292        assert!(content.contains("export MIOPEN_DEBUG_CONV_IMPLICIT_GEMM=1"));
293        assert!(content.contains("export MIOPEN_FIND_MODE=3"));
294        assert!(content.contains("export MIOPEN_FIND_ENFORCE=3"));
295        assert!(content.contains("export PIP_BREAK_SYSTEM_PACKAGES=1"));
296        assert!(content.contains("export UV_PIP_BREAK_SYSTEM_PACKAGES=1"));
297        assert!(content.contains("export UV_SYSTEM_PYTHON=1"));
298    }
299
300    #[test]
301    fn test_env_file_contains_path_settings() {
302        let installer = PermanentEnvInstaller::with_defaults();
303        let content = installer.generate_env_file_content();
304
305        assert!(content.contains("/opt/rocm/bin"));
306        assert!(content.contains("/opt/rocm/lib"));
307        assert!(content.contains("LD_LIBRARY_PATH"));
308    }
309
310    #[test]
311    fn test_env_file_contains_pytorch_optimization() {
312        let installer = PermanentEnvInstaller::with_defaults();
313        let content = installer.generate_env_file_content();
314
315        // CUDA arch list is intentionally NOT exported on an AMD/ROCm stack
316        // (Stage 2: no CUDA env leakage); ROCm uses PYTORCH_ROCM_ARCH instead.
317        assert!(!content.contains("TORCH_CUDA_ARCH_LIST"));
318        assert!(content.contains("PYTORCH_CUDA_ALLOC_CONF"));
319        assert!(content.contains("PYTORCH_HIP_ALLOC_CONF"));
320    }
321
322    #[test]
323    fn test_env_file_contains_vllm_settings() {
324        let installer = PermanentEnvInstaller::with_defaults();
325        let content = installer.generate_env_file_content();
326
327        assert!(content.contains("FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE"));
328        assert!(content.contains("VLLM_ROCM_USE_AITER=0"));
329        assert!(content.contains("VLLM_WORKER_MULTIPROC_METHOD=spawn"));
330    }
331
332    #[test]
333    fn test_env_file_contains_triton_dirs() {
334        let installer = PermanentEnvInstaller::with_defaults();
335        let content = installer.generate_env_file_content();
336
337        assert!(content.contains("MLSTACK_TRITON_HOME"));
338        assert!(content.contains("TRITON_CACHE_DIR"));
339        assert!(content.contains("TRITON_DUMP_DIR"));
340        assert!(content.contains("TRITON_OVERRIDE_DIR"));
341    }
342
343    #[test]
344    fn test_required_env_vars_count() {
345        let installer = PermanentEnvInstaller::with_defaults();
346        let vars = installer.required_env_vars();
347        // Must have at least the key variables
348        assert!(
349            vars.len() >= 19,
350            "Expected at least 19 env vars, got {}",
351            vars.len()
352        );
353    }
354
355    #[test]
356    fn test_required_env_vars_values() {
357        let config = PermanentEnvConfig {
358            rocm_version: "7.2.0".to_string(),
359            gpu_arch: "gfx1100".to_string(),
360            hsa_override_gfx_version: "11.0.0".to_string(),
361            discrete_gpu_list: "0,1".to_string(),
362            python_bin: "python3.12".to_string(),
363            install_method: "auto".to_string(),
364            ..Default::default()
365        };
366        let installer = PermanentEnvInstaller::new(config);
367        let vars = installer.required_env_vars();
368
369        let get = |name: &str| -> String {
370            vars.iter()
371                .find(|(n, _)| *n == name)
372                .map(|(_, v)| v.clone())
373                .unwrap_or_default()
374        };
375
376        assert_eq!(get("ROCM_VERSION"), "7.2.0");
377        assert_eq!(get("GPU_ARCH"), "gfx1100");
378        assert_eq!(get("HIP_VISIBLE_DEVICES"), "0,1");
379        assert_eq!(get("MLSTACK_PYTHON_BIN"), "python3.12");
380        assert_eq!(get("HSA_OVERRIDE_GFX_VERSION"), "11.0.0");
381    }
382
383    #[test]
384    fn test_gfx_to_hsa_override_conversion() {
385        assert_eq!(
386            PermanentEnvInstaller::gfx_to_hsa_override("gfx1100"),
387            Some("11.0.0".to_string())
388        );
389        assert_eq!(
390            PermanentEnvInstaller::gfx_to_hsa_override("gfx1030"),
391            Some("10.3.0".to_string())
392        );
393        assert_eq!(
394            PermanentEnvInstaller::gfx_to_hsa_override("gfx1201"),
395            Some("12.0.1".to_string())
396        );
397        assert_eq!(
398            PermanentEnvInstaller::gfx_to_hsa_override("gfx900"),
399            Some("9.0.0".to_string())
400        );
401        assert_eq!(PermanentEnvInstaller::gfx_to_hsa_override("invalid"), None);
402    }
403
404    #[test]
405    fn test_mkdir_triton_dirs_command() {
406        let installer = PermanentEnvInstaller::with_defaults();
407        let cmd = installer.build_mkdir_triton_dirs_command();
408        assert_eq!(cmd.program, "mkdir");
409        assert!(cmd.args.contains(&"-p".to_string()));
410        let cmd_str = cmd.to_command_string();
411        assert!(cmd_str.contains("mlstack/triton/cache"));
412        assert!(cmd_str.contains("mlstack/triton/dump"));
413        assert!(cmd_str.contains("mlstack/triton/override"));
414    }
415
416    #[test]
417    fn test_rocm_version_detect_command() {
418        let installer = PermanentEnvInstaller::with_defaults();
419        let cmd = installer.build_rocm_version_detect_command();
420        assert_eq!(cmd.program, "cat");
421        assert!(cmd.args.contains(&"/opt/rocm/.info/version".to_string()));
422    }
423
424    #[test]
425    fn test_rocminfo_gpu_detect_command_uses_full_path() {
426        let installer = PermanentEnvInstaller::with_defaults();
427        let cmd = installer.build_rocminfo_gpu_detect_command();
428        // Must NOT be bare "rocminfo" — should resolve to full path when available
429        assert!(
430            cmd.program.contains("rocminfo"),
431            "Program should contain 'rocminfo', got: {}",
432            cmd.program
433        );
434        // If /opt/rocm/bin/rocminfo exists, must use the full path
435        if std::path::Path::new("/opt/rocm/bin/rocminfo").exists() {
436            assert_eq!(
437                cmd.program, "/opt/rocm/bin/rocminfo",
438                "Must use full path when /opt/rocm/bin/rocminfo exists"
439            );
440        }
441    }
442
443    #[test]
444    fn test_generate_hsa_tools_lib_line() {
445        let installer = PermanentEnvInstaller::with_defaults();
446        let line = installer.generate_hsa_tools_lib_line();
447        // Either exports the path or comments it out
448        assert!(
449            line.contains("HSA_TOOLS_LIB") || line.contains("profiler not found"),
450            "Should reference HSA_TOOLS_LIB"
451        );
452    }
453
454    #[test]
455    fn test_custom_config_values_in_env_file() {
456        let config = PermanentEnvConfig {
457            rocm_version: "7.0.0".to_string(),
458            gpu_arch: "gfx1030".to_string(),
459            discrete_gpu_list: "0".to_string(),
460            python_bin: "python3.11".to_string(),
461            ..Default::default()
462        };
463        let installer = PermanentEnvInstaller::new(config);
464        let content = installer.generate_env_file_content();
465
466        assert!(content.contains("ROCM_VERSION=7.0.0"));
467        assert!(content.contains("GPU_ARCH=gfx1030"));
468        assert!(content.contains("HIP_VISIBLE_DEVICES=0"));
469        assert!(content.contains("MLSTACK_PYTHON_BIN=python3.11"));
470    }
471}
rusty_stack/installers/components/permanent_env.rs

rusty_stack/installers/components/
permanent_env.rs
