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 PermanentEnvInstaller {
91    /// Create a new permanent env installer with the given config.
92    pub fn new(config: PermanentEnvConfig) -> Self {
93        Self { config }
94    }
95
96    /// Create with default config.
97    pub fn with_defaults() -> Self {
98        Self::new(PermanentEnvConfig::default())
99    }
100
101    // -----------------------------------------------------------------------
102    // Command construction (VAL-INSTALL-028)
103    // -----------------------------------------------------------------------
104
105    /// Get the environment file path.
106    pub fn env_file_path(&self) -> String {
107        default_env_file_path()
108    }
109
110    /// Generate the content of the `.mlstack_env` file.
111    ///
112    /// Matches the content generated by `setup_permanent_rocm_env.sh`.
113    pub fn generate_env_file_content(&self) -> String {
114        let _home = std::env::var("HOME").unwrap_or_default();
115        let uid_output = std::process::Command::new("id")
116            .arg("-u")
117            .output()
118            .map(|o| String::from_utf8_lossy(&o.stdout).trim().to_string())
119            .unwrap_or_else(|_| "0".to_string());
120
121        format!(
122            r#"# Permanent ROCm Environment Setup (Generated by rusty-stack)
123export ROCM_VERSION={rocm_version}
124export ROCM_CHANNEL=latest
125export GPU_ARCH={gpu_arch}
126export PYTORCH_ROCM_ARCH={gpu_arch}
127export GPU_ARCHS={gpu_arch}
128export ROCM_HOME=/opt/rocm
129export ROCM_PATH=/opt/rocm
130export HIP_PATH=/opt/rocm
131# Discrete GPUs only (iGPUs filtered): {gpu_list}
132export HIP_VISIBLE_DEVICES={gpu_list}
133export CUDA_VISIBLE_DEVICES={gpu_list}
134export PYTORCH_ROCM_DEVICE={first_gpu}
135export MLSTACK_PYTHON_BIN={python_bin}
136export UV_PYTHON={python_bin}
137export MLSTACK_INSTALL_METHOD={install_method}
138export INSTALL_METHOD={install_method}
139export PYTHONPATH=/opt/rocm/lib:$PYTHONPATH
140export AMDGPU_ASIC_ID_TABLE_PATH=/usr/share/libdrm/amdgpu.ids
141export AMDGPU_ASIC_ID_TABLE_PATHS=/usr/share/libdrm
142
143# Path Settings
144export PATH="/usr/local/bin:/usr/bin:/bin:/opt/rocm/bin:/opt/rocm/hip/bin:$PATH"
145export 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"
146
147# Performance & Compatibility
148export HSA_OVERRIDE_GFX_VERSION={hsa_override}
149export HSA_ENABLE_SDMA=0
150export GPU_MAX_HEAP_SIZE=100
151export GPU_MAX_ALLOC_PERCENT=100
152{hsa_tools_lib_line}
153export MIOPEN_DEBUG_CONV_IMPLICIT_GEMM=1
154export MIOPEN_FIND_MODE=3
155export MIOPEN_FIND_ENFORCE=3
156
157# PyTorch Optimization
158export PYTORCH_ALLOC_CONF="max_split_size_mb:512"
159export PYTORCH_HIP_ALLOC_CONF="max_split_size_mb:512"
160
161# vLLM RDNA3 Support (v0.15.0+)
162export FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE
163export VLLM_ROCM_USE_AITER=0
164export VLLM_WORKER_MULTIPROC_METHOD=spawn
165export MLSTACK_TRITON_HOME=$HOME/.cache/mlstack/triton
166export TRITON_HOME=$MLSTACK_TRITON_HOME
167export TRITON_CACHE_DIR=$MLSTACK_TRITON_HOME/cache
168export TRITON_DUMP_DIR=$MLSTACK_TRITON_HOME/dump
169export TRITON_OVERRIDE_DIR=$MLSTACK_TRITON_HOME/override
170
171# Global flags for seamless installs
172export PIP_BREAK_SYSTEM_PACKAGES=1
173export UV_PIP_BREAK_SYSTEM_PACKAGES=1
174export UV_SYSTEM_PYTHON=1
175
176# MPI/UCX Settings (ROCm-aware — CUDA opal support disabled on an AMD stack)
177export OMPI_MCA_pml=ucx
178export OMPI_MCA_osc=ucx
179export OMPI_MCA_btl=^openib,uct
180export OMPI_MCA_btl_openib_allow_ib=true
181export OMPI_MCA_btl_openib_warn_no_device_params_found=0
182export OMPI_MCA_coll_hcoll_enable=0
183# Prevent UCX ROCm memory registration segfault on multi-GPU systems
184export UCX_MEMTYPE_CACHE=no
185export UCX_LOG_LEVEL=error"
186"#,
187            rocm_version = self.config.rocm_version,
188            gpu_arch = self.config.gpu_arch,
189            gpu_list = self.config.discrete_gpu_list,
190            first_gpu = self
191                .config
192                .discrete_gpu_list
193                .split(',')
194                .next()
195                .unwrap_or("0"),
196            python_bin = self.config.python_bin,
197            install_method = self.config.install_method,
198            hsa_override = self.config.hsa_override_gfx_version,
199            hsa_tools_lib_line = self.generate_hsa_tools_lib_line(),
200            uid = uid_output,
201        )
202    }
203
204    /// Generate HSA_TOOLS_LIB line based on profiler library presence.
205    /// Uses the resolved path from `resolve_rocprofiler_tool_path()` to handle
206    /// both the new ROCm 7.x layout (`lib/rocprofiler-sdk/`) and the old layout.
207    pub fn generate_hsa_tools_lib_line(&self) -> String {
208        if let Some(path) = crate::installers::common::utils::resolve_rocprofiler_tool_path() {
209            format!("export HSA_TOOLS_LIB={}", path)
210        } else {
211            "# HSA_TOOLS_LIB not set (profiler not found)".to_string()
212        }
213    }
214
215    /// Get the key environment variables that must be set.
216    /// Returns variable name and their expected values.
217    pub fn required_env_vars(&self) -> Vec<(&'static str, String)> {
218        vec![
219            ("ROCM_VERSION", self.config.rocm_version.clone()),
220            ("ROCM_CHANNEL", "latest".to_string()),
221            ("GPU_ARCH", self.config.gpu_arch.clone()),
222            ("PYTORCH_ROCM_ARCH", self.config.gpu_arch.clone()),
223            ("ROCM_HOME", "/opt/rocm".to_string()),
224            ("ROCM_PATH", "/opt/rocm".to_string()),
225            ("HIP_VISIBLE_DEVICES", self.config.discrete_gpu_list.clone()),
226            (
227                "CUDA_VISIBLE_DEVICES",
228                self.config.discrete_gpu_list.clone(),
229            ),
230            ("MLSTACK_PYTHON_BIN", self.config.python_bin.clone()),
231            (
232                "HSA_OVERRIDE_GFX_VERSION",
233                self.config.hsa_override_gfx_version.clone(),
234            ),
235            ("HSA_ENABLE_SDMA", "0".to_string()),
236            ("GPU_MAX_HEAP_SIZE", "100".to_string()),
237            ("GPU_MAX_ALLOC_PERCENT", "100".to_string()),
238            ("MIOPEN_DEBUG_CONV_IMPLICIT_GEMM", "1".to_string()),
239            ("MIOPEN_FIND_MODE", "3".to_string()),
240            ("MIOPEN_FIND_ENFORCE", "3".to_string()),
241            ("PIP_BREAK_SYSTEM_PACKAGES", "1".to_string()),
242            ("UV_PIP_BREAK_SYSTEM_PACKAGES", "1".to_string()),
243            ("UV_SYSTEM_PYTHON", "1".to_string()),
244        ]
245    }
246
247    /// Construct the mkdir command for triton directories.
248    pub fn build_mkdir_triton_dirs_command(&self) -> ShellCommand {
249        let home = std::env::var("HOME").unwrap_or_default();
250        ShellCommand {
251            program: "mkdir".to_string(),
252            args: vec![
253                "-p".to_string(),
254                format!("{home}/.cache/mlstack/triton/cache"),
255                format!("{home}/.cache/mlstack/triton/dump"),
256                format!("{home}/.cache/mlstack/triton/override"),
257            ],
258            env: vec![],
259        }
260    }
261
262    /// Construct the command to detect GPU architecture via rocminfo.
263    ///
264    /// Uses the full path to rocminfo (e.g., `/opt/rocm/bin/rocminfo`) to avoid
265    /// "command not found" errors when running via sudo, which strips the user's
266    /// PATH environment.
267    pub fn build_rocminfo_gpu_detect_command(&self) -> ShellCommand {
268        let rocminfo_path = crate::installers::common::utils::resolve_rocminfo_path();
269        ShellCommand {
270            program: rocminfo_path,
271            args: vec![],
272            env: vec![],
273        }
274    }
275
276    /// Construct the command to read ROCm version.
277    pub fn build_rocm_version_detect_command(&self) -> ShellCommand {
278        ShellCommand {
279            program: "cat".to_string(),
280            args: vec!["/opt/rocm/.info/version".to_string()],
281            env: vec![],
282        }
283    }
284
285    /// Get GPU arch from environment variable.
286    pub fn gpu_arch_from_env() -> Option<String> {
287        std::env::var("GPU_ARCH").ok().filter(|s| !s.is_empty())
288    }
289
290    /// Default GPU arch when detection fails.
291    pub fn default_gpu_arch() -> &'static str {
292        "gfx1100"
293    }
294
295    /// Convert gfx architecture string to HSA override version.
296    /// E.g., "gfx1100" -> "11.0.0"
297    pub fn gfx_to_hsa_override(gfx: &str) -> Option<String> {
298        let raw = gfx.strip_prefix("gfx")?;
299        match raw.len() {
300            4 => {
301                let major = &raw[0..2];
302                let minor = &raw[2..3];
303                let patch = &raw[3..4];
304                Some(format!("{}.{}.{}", major, minor, patch))
305            }
306            3 => {
307                let major = &raw[0..1];
308                let minor = &raw[1..2];
309                let patch = &raw[2..3];
310                Some(format!("{}.{}.{}", major, minor, patch))
311            }
312            _ => None,
313        }
314    }
315}
316
317// ===========================================================================
318// Tests
319// ===========================================================================
320
321#[cfg(test)]
322mod tests {
323    use super::*;
324
325    // -----------------------------------------------------------------------
326    // VAL-INSTALL-028: Permanent env creates correct venv and symlinks
327    // -----------------------------------------------------------------------
328
329    #[test]
330    fn test_env_file_path_is_home_mlstack_env() {
331        let path = default_env_file_path();
332        assert!(
333            path.ends_with("/.mlstack_env"),
334            "Env file must be at ~/.mlstack_env, got: {path}"
335        );
336    }
337
338    #[test]
339    fn test_generate_env_file_contains_required_vars() {
340        let installer = PermanentEnvInstaller::with_defaults();
341        let content = installer.generate_env_file_content();
342
343        // Check all required variables from the original script
344        assert!(content.contains("export ROCM_VERSION="));
345        assert!(content.contains("export ROCM_CHANNEL=latest"));
346        assert!(content.contains("export GPU_ARCH="));
347        assert!(content.contains("export PYTORCH_ROCM_ARCH="));
348        assert!(content.contains("export ROCM_HOME=/opt/rocm"));
349        assert!(content.contains("export ROCM_PATH=/opt/rocm"));
350        assert!(content.contains("export HIP_PATH=/opt/rocm"));
351        assert!(content.contains("export HIP_VISIBLE_DEVICES="));
352        assert!(content.contains("export CUDA_VISIBLE_DEVICES="));
353        assert!(content.contains("export MLSTACK_PYTHON_BIN="));
354        assert!(content.contains("export UV_PYTHON="));
355        assert!(content.contains("export HSA_OVERRIDE_GFX_VERSION="));
356        assert!(content.contains("export HSA_ENABLE_SDMA=0"));
357        assert!(content.contains("export GPU_MAX_HEAP_SIZE=100"));
358        assert!(content.contains("export GPU_MAX_ALLOC_PERCENT=100"));
359        assert!(content.contains("export MIOPEN_DEBUG_CONV_IMPLICIT_GEMM=1"));
360        assert!(content.contains("export MIOPEN_FIND_MODE=3"));
361        assert!(content.contains("export MIOPEN_FIND_ENFORCE=3"));
362        assert!(content.contains("export PIP_BREAK_SYSTEM_PACKAGES=1"));
363        assert!(content.contains("export UV_PIP_BREAK_SYSTEM_PACKAGES=1"));
364        assert!(content.contains("export UV_SYSTEM_PYTHON=1"));
365    }
366
367    #[test]
368    fn test_env_file_contains_path_settings() {
369        let installer = PermanentEnvInstaller::with_defaults();
370        let content = installer.generate_env_file_content();
371
372        assert!(content.contains("/opt/rocm/bin"));
373        assert!(content.contains("/opt/rocm/lib"));
374        assert!(content.contains("LD_LIBRARY_PATH"));
375    }
376
377    #[test]
378    fn test_env_file_contains_pytorch_optimization() {
379        let installer = PermanentEnvInstaller::with_defaults();
380        let content = installer.generate_env_file_content();
381
382        // CUDA arch list is intentionally NOT exported on an AMD/ROCm stack
383        // (Stage 2: no CUDA env leakage); ROCm uses PYTORCH_ROCM_ARCH instead.
384        assert!(!content.contains("TORCH_CUDA_ARCH_LIST"));
385        assert!(content.contains("PYTORCH_ALLOC_CONF"));
386        assert!(content.contains("PYTORCH_HIP_ALLOC_CONF"));
387    }
388
389    #[test]
390    fn test_env_file_contains_vllm_settings() {
391        let installer = PermanentEnvInstaller::with_defaults();
392        let content = installer.generate_env_file_content();
393
394        assert!(content.contains("FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE"));
395        assert!(content.contains("VLLM_ROCM_USE_AITER=0"));
396        assert!(content.contains("VLLM_WORKER_MULTIPROC_METHOD=spawn"));
397    }
398
399    #[test]
400    fn test_env_file_contains_triton_dirs() {
401        let installer = PermanentEnvInstaller::with_defaults();
402        let content = installer.generate_env_file_content();
403
404        assert!(content.contains("MLSTACK_TRITON_HOME"));
405        assert!(content.contains("TRITON_CACHE_DIR"));
406        assert!(content.contains("TRITON_DUMP_DIR"));
407        assert!(content.contains("TRITON_OVERRIDE_DIR"));
408    }
409
410    #[test]
411    fn test_required_env_vars_count() {
412        let installer = PermanentEnvInstaller::with_defaults();
413        let vars = installer.required_env_vars();
414        // Must have at least the key variables
415        assert!(
416            vars.len() >= 19,
417            "Expected at least 19 env vars, got {}",
418            vars.len()
419        );
420    }
421
422    #[test]
423    fn test_required_env_vars_values() {
424        let config = PermanentEnvConfig {
425            rocm_version: "7.2.0".to_string(),
426            gpu_arch: "gfx1100".to_string(),
427            hsa_override_gfx_version: "11.0.0".to_string(),
428            discrete_gpu_list: "0,1".to_string(),
429            python_bin: "python3.12".to_string(),
430            install_method: "auto".to_string(),
431            ..Default::default()
432        };
433        let installer = PermanentEnvInstaller::new(config);
434        let vars = installer.required_env_vars();
435
436        let get = |name: &str| -> String {
437            vars.iter()
438                .find(|(n, _)| *n == name)
439                .map(|(_, v)| v.clone())
440                .unwrap_or_default()
441        };
442
443        assert_eq!(get("ROCM_VERSION"), "7.2.0");
444        assert_eq!(get("GPU_ARCH"), "gfx1100");
445        assert_eq!(get("HIP_VISIBLE_DEVICES"), "0,1");
446        assert_eq!(get("MLSTACK_PYTHON_BIN"), "python3.12");
447        assert_eq!(get("HSA_OVERRIDE_GFX_VERSION"), "11.0.0");
448    }
449
450    #[test]
451    fn test_gfx_to_hsa_override_conversion() {
452        assert_eq!(
453            PermanentEnvInstaller::gfx_to_hsa_override("gfx1100"),
454            Some("11.0.0".to_string())
455        );
456        assert_eq!(
457            PermanentEnvInstaller::gfx_to_hsa_override("gfx1030"),
458            Some("10.3.0".to_string())
459        );
460        assert_eq!(
461            PermanentEnvInstaller::gfx_to_hsa_override("gfx1201"),
462            Some("12.0.1".to_string())
463        );
464        assert_eq!(
465            PermanentEnvInstaller::gfx_to_hsa_override("gfx900"),
466            Some("9.0.0".to_string())
467        );
468        assert_eq!(PermanentEnvInstaller::gfx_to_hsa_override("invalid"), None);
469    }
470
471    #[test]
472    fn test_mkdir_triton_dirs_command() {
473        let installer = PermanentEnvInstaller::with_defaults();
474        let cmd = installer.build_mkdir_triton_dirs_command();
475        assert_eq!(cmd.program, "mkdir");
476        assert!(cmd.args.contains(&"-p".to_string()));
477        let cmd_str = cmd.to_command_string();
478        assert!(cmd_str.contains("mlstack/triton/cache"));
479        assert!(cmd_str.contains("mlstack/triton/dump"));
480        assert!(cmd_str.contains("mlstack/triton/override"));
481    }
482
483    #[test]
484    fn test_rocm_version_detect_command() {
485        let installer = PermanentEnvInstaller::with_defaults();
486        let cmd = installer.build_rocm_version_detect_command();
487        assert_eq!(cmd.program, "cat");
488        assert!(cmd.args.contains(&"/opt/rocm/.info/version".to_string()));
489    }
490
491    #[test]
492    fn test_rocminfo_gpu_detect_command_uses_full_path() {
493        let installer = PermanentEnvInstaller::with_defaults();
494        let cmd = installer.build_rocminfo_gpu_detect_command();
495        // Must NOT be bare "rocminfo" — should resolve to full path when available
496        assert!(
497            cmd.program.contains("rocminfo"),
498            "Program should contain 'rocminfo', got: {}",
499            cmd.program
500        );
501        // If /opt/rocm/bin/rocminfo exists, must use the full path
502        if std::path::Path::new("/opt/rocm/bin/rocminfo").exists() {
503            assert_eq!(
504                cmd.program, "/opt/rocm/bin/rocminfo",
505                "Must use full path when /opt/rocm/bin/rocminfo exists"
506            );
507        }
508    }
509
510    #[test]
511    fn test_generate_hsa_tools_lib_line() {
512        let installer = PermanentEnvInstaller::with_defaults();
513        let line = installer.generate_hsa_tools_lib_line();
514        // Either exports the path or comments it out
515        assert!(
516            line.contains("HSA_TOOLS_LIB") || line.contains("profiler not found"),
517            "Should reference HSA_TOOLS_LIB"
518        );
519    }
520
521    #[test]
522    fn test_custom_config_values_in_env_file() {
523        let config = PermanentEnvConfig {
524            rocm_version: "7.0.0".to_string(),
525            gpu_arch: "gfx1030".to_string(),
526            discrete_gpu_list: "0".to_string(),
527            python_bin: "python3.11".to_string(),
528            ..Default::default()
529        };
530        let installer = PermanentEnvInstaller::new(config);
531        let content = installer.generate_env_file_content();
532
533        assert!(content.contains("ROCM_VERSION=7.0.0"));
534        assert!(content.contains("GPU_ARCH=gfx1030"));
535        assert!(content.contains("HIP_VISIBLE_DEVICES=0"));
536        assert!(content.contains("MLSTACK_PYTHON_BIN=python3.11"));
537    }
538}
rusty_stack/installers/components/permanent_env.rs

rusty_stack/installers/components/
permanent_env.rs
