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 TORCH_CUDA_ARCH_LIST="7.0;8.0;9.0"
159export PYTORCH_ALLOC_CONF="max_split_size_mb:512"
160export PYTORCH_HIP_ALLOC_CONF="max_split_size_mb:512"
161
162# vLLM RDNA3 Support (v0.15.0+)
163export FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE
164export VLLM_ROCM_USE_AITER=0
165export VLLM_WORKER_MULTIPROC_METHOD=spawn
166export MLSTACK_TRITON_HOME=$HOME/.cache/mlstack/triton
167export TRITON_HOME=$MLSTACK_TRITON_HOME
168export TRITON_CACHE_DIR=$MLSTACK_TRITON_HOME/cache
169export TRITON_DUMP_DIR=$MLSTACK_TRITON_HOME/dump
170export TRITON_OVERRIDE_DIR=$MLSTACK_TRITON_HOME/override
171
172# Global flags for seamless installs
173export PIP_BREAK_SYSTEM_PACKAGES=1
174export UV_PIP_BREAK_SYSTEM_PACKAGES=1
175export UV_SYSTEM_PYTHON=1
176
177# MPI/UCX Settings (ROCm-aware)
178export OMPI_MCA_opal_cuda_support=true
179export OMPI_MCA_pml=ucx
180export OMPI_MCA_osc=ucx
181export OMPI_MCA_btl=^openib,uct
182export OMPI_MCA_btl_openib_allow_ib=true
183export OMPI_MCA_btl_openib_warn_no_device_params_found=0
184export OMPI_MCA_coll_hcoll_enable=0
185# Prevent UCX ROCm memory registration segfault on multi-GPU systems
186export UCX_MEMTYPE_CACHE=no
187export UCX_LOG_LEVEL=error"
188"#,
189            rocm_version = self.config.rocm_version,
190            gpu_arch = self.config.gpu_arch,
191            gpu_list = self.config.discrete_gpu_list,
192            first_gpu = self
193                .config
194                .discrete_gpu_list
195                .split(',')
196                .next()
197                .unwrap_or("0"),
198            python_bin = self.config.python_bin,
199            install_method = self.config.install_method,
200            hsa_override = self.config.hsa_override_gfx_version,
201            hsa_tools_lib_line = self.generate_hsa_tools_lib_line(),
202            uid = uid_output,
203        )
204    }
205
206    /// Generate HSA_TOOLS_LIB line based on profiler library presence.
207    /// Uses the resolved path from `resolve_rocprofiler_tool_path()` to handle
208    /// both the new ROCm 7.x layout (`lib/rocprofiler-sdk/`) and the old layout.
209    pub fn generate_hsa_tools_lib_line(&self) -> String {
210        if let Some(path) = crate::installers::common::utils::resolve_rocprofiler_tool_path() {
211            format!("export HSA_TOOLS_LIB={}", path)
212        } else {
213            "# HSA_TOOLS_LIB not set (profiler not found)".to_string()
214        }
215    }
216
217    /// Get the key environment variables that must be set.
218    /// Returns variable name and their expected values.
219    pub fn required_env_vars(&self) -> Vec<(&'static str, String)> {
220        vec![
221            ("ROCM_VERSION", self.config.rocm_version.clone()),
222            ("ROCM_CHANNEL", "latest".to_string()),
223            ("GPU_ARCH", self.config.gpu_arch.clone()),
224            ("PYTORCH_ROCM_ARCH", self.config.gpu_arch.clone()),
225            ("ROCM_HOME", "/opt/rocm".to_string()),
226            ("ROCM_PATH", "/opt/rocm".to_string()),
227            ("HIP_VISIBLE_DEVICES", self.config.discrete_gpu_list.clone()),
228            (
229                "CUDA_VISIBLE_DEVICES",
230                self.config.discrete_gpu_list.clone(),
231            ),
232            ("MLSTACK_PYTHON_BIN", self.config.python_bin.clone()),
233            (
234                "HSA_OVERRIDE_GFX_VERSION",
235                self.config.hsa_override_gfx_version.clone(),
236            ),
237            ("HSA_ENABLE_SDMA", "0".to_string()),
238            ("GPU_MAX_HEAP_SIZE", "100".to_string()),
239            ("GPU_MAX_ALLOC_PERCENT", "100".to_string()),
240            ("MIOPEN_DEBUG_CONV_IMPLICIT_GEMM", "1".to_string()),
241            ("MIOPEN_FIND_MODE", "3".to_string()),
242            ("MIOPEN_FIND_ENFORCE", "3".to_string()),
243            ("PIP_BREAK_SYSTEM_PACKAGES", "1".to_string()),
244            ("UV_PIP_BREAK_SYSTEM_PACKAGES", "1".to_string()),
245            ("UV_SYSTEM_PYTHON", "1".to_string()),
246        ]
247    }
248
249    /// Construct the mkdir command for triton directories.
250    pub fn build_mkdir_triton_dirs_command(&self) -> ShellCommand {
251        let home = std::env::var("HOME").unwrap_or_default();
252        ShellCommand {
253            program: "mkdir".to_string(),
254            args: vec![
255                "-p".to_string(),
256                format!("{home}/.cache/mlstack/triton/cache"),
257                format!("{home}/.cache/mlstack/triton/dump"),
258                format!("{home}/.cache/mlstack/triton/override"),
259            ],
260            env: vec![],
261        }
262    }
263
264    /// Construct the command to detect GPU architecture via rocminfo.
265    ///
266    /// Uses the full path to rocminfo (e.g., `/opt/rocm/bin/rocminfo`) to avoid
267    /// "command not found" errors when running via sudo, which strips the user's
268    /// PATH environment.
269    pub fn build_rocminfo_gpu_detect_command(&self) -> ShellCommand {
270        let rocminfo_path = crate::installers::common::utils::resolve_rocminfo_path();
271        ShellCommand {
272            program: rocminfo_path,
273            args: vec![],
274            env: vec![],
275        }
276    }
277
278    /// Construct the command to read ROCm version.
279    pub fn build_rocm_version_detect_command(&self) -> ShellCommand {
280        ShellCommand {
281            program: "cat".to_string(),
282            args: vec!["/opt/rocm/.info/version".to_string()],
283            env: vec![],
284        }
285    }
286
287    /// Get GPU arch from environment variable.
288    pub fn gpu_arch_from_env() -> Option<String> {
289        std::env::var("GPU_ARCH").ok().filter(|s| !s.is_empty())
290    }
291
292    /// Default GPU arch when detection fails.
293    pub fn default_gpu_arch() -> &'static str {
294        "gfx1100"
295    }
296
297    /// Convert gfx architecture string to HSA override version.
298    /// E.g., "gfx1100" -> "11.0.0"
299    pub fn gfx_to_hsa_override(gfx: &str) -> Option<String> {
300        let raw = gfx.strip_prefix("gfx")?;
301        match raw.len() {
302            4 => {
303                let major = &raw[0..2];
304                let minor = &raw[2..3];
305                let patch = &raw[3..4];
306                Some(format!("{}.{}.{}", major, minor, patch))
307            }
308            3 => {
309                let major = &raw[0..1];
310                let minor = &raw[1..2];
311                let patch = &raw[2..3];
312                Some(format!("{}.{}.{}", major, minor, patch))
313            }
314            _ => None,
315        }
316    }
317}
318
319// ===========================================================================
320// Tests
321// ===========================================================================
322
323#[cfg(test)]
324mod tests {
325    use super::*;
326
327    // -----------------------------------------------------------------------
328    // VAL-INSTALL-028: Permanent env creates correct venv and symlinks
329    // -----------------------------------------------------------------------
330
331    #[test]
332    fn test_env_file_path_is_home_mlstack_env() {
333        let path = default_env_file_path();
334        assert!(
335            path.ends_with("/.mlstack_env"),
336            "Env file must be at ~/.mlstack_env, got: {path}"
337        );
338    }
339
340    #[test]
341    fn test_generate_env_file_contains_required_vars() {
342        let installer = PermanentEnvInstaller::with_defaults();
343        let content = installer.generate_env_file_content();
344
345        // Check all required variables from the original script
346        assert!(content.contains("export ROCM_VERSION="));
347        assert!(content.contains("export ROCM_CHANNEL=latest"));
348        assert!(content.contains("export GPU_ARCH="));
349        assert!(content.contains("export PYTORCH_ROCM_ARCH="));
350        assert!(content.contains("export ROCM_HOME=/opt/rocm"));
351        assert!(content.contains("export ROCM_PATH=/opt/rocm"));
352        assert!(content.contains("export HIP_PATH=/opt/rocm"));
353        assert!(content.contains("export HIP_VISIBLE_DEVICES="));
354        assert!(content.contains("export CUDA_VISIBLE_DEVICES="));
355        assert!(content.contains("export MLSTACK_PYTHON_BIN="));
356        assert!(content.contains("export UV_PYTHON="));
357        assert!(content.contains("export HSA_OVERRIDE_GFX_VERSION="));
358        assert!(content.contains("export HSA_ENABLE_SDMA=0"));
359        assert!(content.contains("export GPU_MAX_HEAP_SIZE=100"));
360        assert!(content.contains("export GPU_MAX_ALLOC_PERCENT=100"));
361        assert!(content.contains("export MIOPEN_DEBUG_CONV_IMPLICIT_GEMM=1"));
362        assert!(content.contains("export MIOPEN_FIND_MODE=3"));
363        assert!(content.contains("export MIOPEN_FIND_ENFORCE=3"));
364        assert!(content.contains("export PIP_BREAK_SYSTEM_PACKAGES=1"));
365        assert!(content.contains("export UV_PIP_BREAK_SYSTEM_PACKAGES=1"));
366        assert!(content.contains("export UV_SYSTEM_PYTHON=1"));
367    }
368
369    #[test]
370    fn test_env_file_contains_path_settings() {
371        let installer = PermanentEnvInstaller::with_defaults();
372        let content = installer.generate_env_file_content();
373
374        assert!(content.contains("/opt/rocm/bin"));
375        assert!(content.contains("/opt/rocm/lib"));
376        assert!(content.contains("LD_LIBRARY_PATH"));
377    }
378
379    #[test]
380    fn test_env_file_contains_pytorch_optimization() {
381        let installer = PermanentEnvInstaller::with_defaults();
382        let content = installer.generate_env_file_content();
383
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
