1use crate::installers::common::utils::{
18 command_exists, print_error, print_step, print_success, print_warning,
19};
20use std::fmt;
21
22#[derive(Debug, Clone)]
28pub struct GpuArchInfo {
29 pub gpu_arch: String,
31 pub hsa_override_gfx_version: String,
33}
34
35impl Default for GpuArchInfo {
36 fn default() -> Self {
37 Self {
38 gpu_arch: "gfx1100".to_string(),
39 hsa_override_gfx_version: "11.0.0".to_string(),
40 }
41 }
42}
43
44impl GpuArchInfo {
45 pub fn from_arch(arch: &str) -> Self {
49 let hsa_version = match arch {
50 "gfx1100" => "11.0.0",
51 "gfx1101" => "11.0.1",
52 "gfx1102" => "11.0.2",
53 "gfx1200" => "12.0.0",
54 "gfx1201" => "12.0.1",
55 "gfx1030" => "10.3.0",
56 "gfx1031" => "10.3.1",
57 "gfx1032" => "10.3.2",
58 _ => "11.0.0", };
60 Self {
61 gpu_arch: arch.to_string(),
62 hsa_override_gfx_version: hsa_version.to_string(),
63 }
64 }
65}
66
67#[derive(Debug, Clone)]
69pub struct GpuFilterResult {
70 pub discrete_gpu_indices: Vec<u32>,
72 pub hip_visible_devices: String,
74 pub gpu_count: usize,
76}
77
78#[derive(Debug, Clone)]
80pub struct SystemDependency {
81 pub package: &'static str,
83 pub description: &'static str,
85}
86
87impl SystemDependency {
88 pub fn required_packages() -> Vec<&'static str> {
92 vec![
93 "build-essential",
94 "cmake",
95 "git",
96 "python3-dev",
97 "python3-pip",
98 "libnuma-dev",
99 "pciutils",
100 "mesa-utils",
101 "clinfo",
102 ]
103 }
104}
105
106impl fmt::Display for SystemDependency {
107 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
108 write!(f, "{} ({})", self.package, self.description)
109 }
110}
111
112#[derive(Debug, Clone)]
114pub struct DependencyCheckResult {
115 pub package: String,
117 pub installed: bool,
119 pub mapped_name: String,
121}
122
123#[derive(Debug, Clone)]
125pub struct EnvFileResult {
126 pub success: bool,
128 pub path: Option<String>,
130 pub errors: Vec<String>,
132}
133
134impl EnvFileResult {
135 pub fn success(path: String) -> Self {
137 Self {
138 success: true,
139 path: Some(path),
140 errors: Vec::new(),
141 }
142 }
143
144 pub fn failure(error: String) -> Self {
146 Self {
147 success: false,
148 path: None,
149 errors: vec![error],
150 }
151 }
152}
153
154pub fn detect_discrete_gpus() -> GpuFilterResult {
168 let mut discrete_indices: Vec<u32> = Vec::new();
169
170 let rocminfo_path = crate::installers::common::utils::resolve_rocminfo_path();
172 if command_exists(&rocminfo_path) || command_exists("rocminfo") {
173 discrete_indices = detect_gpus_from_rocminfo();
174 }
175
176 if discrete_indices.is_empty() {
178 discrete_indices = detect_gpus_from_lspci();
179 }
180
181 if discrete_indices.is_empty() {
183 discrete_indices = detect_gpus_from_render_nodes();
184 }
185
186 if discrete_indices.is_empty() {
188 discrete_indices.push(0);
189 }
190
191 let gpu_count = discrete_indices.len();
192 let hip_visible_devices = discrete_indices
193 .iter()
194 .map(|i| i.to_string())
195 .collect::<Vec<_>>()
196 .join(",");
197
198 GpuFilterResult {
199 discrete_gpu_indices: discrete_indices,
200 hip_visible_devices,
201 gpu_count,
202 }
203}
204
205fn detect_gpus_from_rocminfo() -> Vec<u32> {
207 let rocminfo_path = crate::installers::common::utils::resolve_rocminfo_path();
208 let output = match std::process::Command::new(&rocminfo_path).output() {
209 Ok(o) if o.status.success() => String::from_utf8_lossy(&o.stdout).to_string(),
210 _ => return Vec::new(),
211 };
212
213 let mut discrete_indices: Vec<u32> = Vec::new();
214 let mut current_gpu_id: Option<u32> = None;
215 let mut is_igpu = false;
216 let mut gpu_agent_counter = 0u32; for line in output.lines() {
219 let trimmed = line.trim();
220
221 if trimmed.starts_with("Agent") && trimmed.contains(char::is_numeric) {
223 }
225
226 if trimmed.starts_with("GPU ID:") {
227 if let Some(id) = current_gpu_id {
229 if !is_igpu && !discrete_indices.contains(&id) {
230 discrete_indices.push(id);
231 }
232 }
233
234 current_gpu_id = trimmed
236 .strip_prefix("GPU ID:")
237 .and_then(|s| s.trim().parse::<u32>().ok());
238 is_igpu = false;
239 } else if trimmed.starts_with("Marketing Name:") {
240 let marketing_name = trimmed
241 .strip_prefix("Marketing Name:")
242 .map(|s| s.trim())
243 .unwrap_or("");
244
245 if is_integrated_gpu_name(marketing_name) {
247 is_igpu = true;
248 }
249
250 if current_gpu_id.is_none() {
252 }
255 } else if trimmed.starts_with("Device Type:") {
256 let device_type = trimmed
257 .strip_prefix("Device Type:")
258 .map(|s| s.trim())
259 .unwrap_or("");
260
261 if device_type.contains("CPU") {
263 is_igpu = true;
264 }
265
266 if let Some(id) = current_gpu_id {
268 if !is_igpu && !discrete_indices.contains(&id) {
269 discrete_indices.push(id);
270 }
271 } else if device_type.contains("GPU") {
272 if !is_igpu && !discrete_indices.contains(&gpu_agent_counter) {
274 discrete_indices.push(gpu_agent_counter);
275 }
276 gpu_agent_counter += 1;
277 }
278 }
279 }
280
281 if let Some(id) = current_gpu_id {
283 if !is_igpu && !discrete_indices.contains(&id) {
284 discrete_indices.push(id);
285 }
286 }
287
288 discrete_indices.sort();
289 discrete_indices.dedup();
290 discrete_indices
291}
292
293fn detect_gpus_from_lspci() -> Vec<u32> {
295 if !command_exists("lspci") {
296 return Vec::new();
297 }
298
299 let output = match std::process::Command::new("lspci").output() {
300 Ok(o) if o.status.success() => String::from_utf8_lossy(&o.stdout).to_string(),
301 _ => return Vec::new(),
302 };
303
304 let mut discrete_count = 0u32;
305 for line in output.lines() {
306 let lower = line.to_lowercase();
307 if !((lower.contains("amd")
308 || lower.contains("radeon")
309 || lower.contains("advanced micro devices"))
310 && (lower.contains("vga") || lower.contains("3d") || lower.contains("display")))
311 {
312 continue;
313 }
314 if is_integrated_gpu_name(line) {
316 continue;
317 }
318 discrete_count += 1;
319 }
320
321 if discrete_count > 0 {
322 (0..discrete_count).collect()
323 } else {
324 Vec::new()
325 }
326}
327
328fn detect_gpus_from_render_nodes() -> Vec<u32> {
331 let render_entries: Vec<_> = std::fs::read_dir("/dev/dri")
332 .map(|entries| {
333 entries
334 .filter_map(|e| e.ok())
335 .filter(|e| e.file_name().to_string_lossy().starts_with("render"))
336 .collect()
337 })
338 .unwrap_or_default();
339
340 if render_entries.is_empty() {
341 return Vec::new();
342 }
343
344 let discrete_count = render_entries
346 .iter()
347 .filter(|entry| {
348 let name = entry.file_name().to_string_lossy().to_string();
349 let node_num: String = name
351 .trim_start_matches("render")
352 .chars()
353 .take_while(|c| c.is_ascii_digit())
354 .collect();
355 if node_num.is_empty() {
356 return true; }
358
359 let sysfs_path = format!("/sys/class/drm/card{}-render/device/uevent", node_num);
361 if let Ok(uevent) = std::fs::read_to_string(&sysfs_path) {
362 for line in uevent.lines() {
364 if line.starts_with("PCI_ID=") {
365 return true;
367 }
368 }
369 }
370
371 let card_sysfs = format!("/sys/class/drm/card{}/device/device", node_num);
373 if let Ok(pci_id) = std::fs::read_to_string(&card_sysfs) {
374 let pci_id = pci_id.trim();
375 if pci_id == "0x164e" {
377 return false; }
379 }
380
381 true })
383 .count();
384
385 if discrete_count > 0 {
386 (0..discrete_count as u32).collect()
387 } else if !render_entries.is_empty() {
388 (0..render_entries.len() as u32).collect()
390 } else {
391 Vec::new()
392 }
393}
394
395pub fn is_integrated_gpu_name(name: &str) -> bool {
400 crate::gpu::is_integrated_gpu_name(name)
401}
402
403pub fn detect_correct_gpu_arch(marketing_name: Option<&str>) -> GpuArchInfo {
416 let name = match marketing_name {
417 Some(n) => n,
418 None => return GpuArchInfo::default(),
419 };
420
421 if is_integrated_gpu_name(name) {
423 return GpuArchInfo::default();
424 }
425
426 let arch = crate::gpu::gfx_from_marketing_name(name).unwrap_or("gfx1100");
428
429 GpuArchInfo::from_arch(arch)
430}
431
432pub fn check_system_dependency(package: &str) -> DependencyCheckResult {
445 use crate::installers::common::distro::DistroFacade;
446 use crate::installers::common::package_mappings::map_package_name;
447
448 let distro = DistroFacade::detect();
449 let pkg_mgr = distro.package_manager();
450 let mapped = map_package_name(package, pkg_mgr);
451 let installed = check_package_installed(&mapped);
452
453 DependencyCheckResult {
454 package: package.to_string(),
455 installed,
456 mapped_name: mapped,
457 }
458}
459
460pub fn check_all_dependencies() -> Vec<DependencyCheckResult> {
464 SystemDependency::required_packages()
465 .into_iter()
466 .map(check_system_dependency)
467 .collect()
468}
469
470fn check_package_installed(package: &str) -> bool {
472 if command_exists("dpkg-query") {
474 let output = std::process::Command::new("dpkg-query")
475 .args(["-W", "-f=${Status}", package])
476 .output();
477 if let Ok(o) = output {
478 let stdout = String::from_utf8_lossy(&o.stdout);
479 if stdout.contains("install ok installed") {
480 return true;
481 }
482 }
483 }
484
485 if command_exists("rpm") {
487 let output = std::process::Command::new("rpm")
488 .args(["-q", package])
489 .output();
490 if let Ok(o) = output {
491 if o.status.success() {
492 return true;
493 }
494 }
495 }
496
497 if command_exists("pacman") {
499 let output = std::process::Command::new("pacman")
500 .args(["-Q", package])
501 .output();
502 if let Ok(o) = output {
503 if o.status.success() {
504 return true;
505 }
506 }
507 }
508
509 false
510}
511
512pub fn create_env_file(
525 hip_visible_devices: &str,
526 rocm_path: &str,
527 rocm_version: &str,
528 rocm_channel: &str,
529 gpu_arch: &str,
530 hsa_override_gfx_version: &str,
531 python_bin: &str,
532) -> EnvFileResult {
533 let home = std::env::var("HOME").unwrap_or_else(|_| "/root".to_string());
534 let env_path = format!("{home}/.mlstack_env");
535 let fish_env_path = format!("{home}/.config/fish/conf.d/mlstack_env.fish");
536
537 let content = generate_env_file_content(
538 hip_visible_devices,
539 rocm_path,
540 rocm_version,
541 rocm_channel,
542 gpu_arch,
543 hsa_override_gfx_version,
544 python_bin,
545 );
546
547 if let Err(e) = std::fs::write(&env_path, &content) {
549 let msg = format!("Failed to write environment file: {e}");
550 print_error(&msg);
551 return EnvFileResult::failure(msg);
552 }
553
554 let fish_content = generate_fish_env_file_content(
556 hip_visible_devices,
557 rocm_path,
558 rocm_version,
559 rocm_channel,
560 gpu_arch,
561 hsa_override_gfx_version,
562 python_bin,
563 );
564 if let Some(parent) = std::path::Path::new(&fish_env_path).parent() {
566 let _ = std::fs::create_dir_all(parent);
567 }
568 if let Err(e) = std::fs::write(&fish_env_path, &fish_content) {
569 print_warning(&format!("Failed to write fish env file: {e}"));
571 } else {
572 print_success(&format!("Fish env file created: {fish_env_path}"));
573 }
574
575 print_success(&format!("Environment file created: {env_path}"));
576 EnvFileResult::success(env_path)
577}
578
579pub fn generate_env_file_content(
590 hip_visible_devices: &str,
591 rocm_path: &str,
592 rocm_version: &str,
593 rocm_channel: &str,
594 gpu_arch: &str,
595 hsa_override_gfx_version: &str,
596 python_bin: &str,
597) -> String {
598 let home = std::env::var("HOME").unwrap_or_else(|_| "$HOME".to_string());
599 let rccl_block = rccl_overlay_bash_block(&home);
602
603 format!(
604 r#"# ML Stack Environment File
605# Created by Enhanced ML Stack Environment Setup (Rust native)
606# Date: {date}
607
608# GPU Selection — device filter is MANAGED (unconditional). It must reflect the
609# detected discrete GPUs on every source so a stale override can never re-expose
610# an integrated GPU (the iGPU segfault). ROCR_VISIBLE_DEVICES is the authoritative
611# ROCr (HSA) runtime filter; HIP/CUDA alias it for the HIP + CUDA-compat layers.
612# Discrete GPUs only (iGPUs filtered): {hip_visible_devices}
613export ROCR_VISIBLE_DEVICES={hip_visible_devices}
614export HIP_VISIBLE_DEVICES={hip_visible_devices}
615export CUDA_VISIBLE_DEVICES={hip_visible_devices}
616export PYTORCH_ROCM_DEVICE={hip_visible_devices}
617
618# ROCm Settings
619# Only set if not already set
620if [ -z "${{ROCM_HOME:-}}" ]; then export ROCM_HOME={rocm_path}; fi
621if [ -z "${{ROCM_PATH:-}}" ]; then export ROCM_PATH={rocm_path}; fi
622if [ -z "${{HIP_PATH:-}}" ]; then export HIP_PATH={rocm_path}; fi
623if [ -z "${{CUDA_HOME:-}}" ]; then export CUDA_HOME={rocm_path}; fi
624if [ -z "${{ROCM_VERSION:-}}" ]; then export ROCM_VERSION={rocm_version}; fi
625if [ -z "${{ROCM_CHANNEL:-}}" ]; then export ROCM_CHANNEL={rocm_channel}; fi
626# GPU arch identity is MANAGED (unconditional): must match detected hardware.
627export GPU_ARCH={gpu_arch}
628export PYTORCH_ROCM_ARCH={gpu_arch}
629export GPU_ARCHS={gpu_arch}
630# libdrm amdgpu ASIC id table (ROCm's amdgpu.ids lookup)
631export AMDGPU_ASIC_ID_TABLE_PATH=/usr/share/libdrm/amdgpu.ids
632export AMDGPU_ASIC_ID_TABLE_PATHS=/usr/share/libdrm
633
634# Path Settings - Hardcoded safe paths to prevent "command not found" errors.
635# Idempotent: prepend the ROCm bin/lib blocks only if not already on the var, so
636# re-sourcing the env never grows PATH/LD_LIBRARY_PATH (the old un-guarded
637# `:$PATH` form accumulated a duplicate entry per source).
638case ":${{PATH:-}}:" in *:"{rocm_path}/bin":*) ;; *) export PATH="/usr/local/bin:/usr/bin:/bin:/usr/local/games:/usr/games:{rocm_path}/bin:{rocm_path}/hip/bin:${{PATH:-}}";; esac
639case ":${{LD_LIBRARY_PATH:-}}:" in *:"{rocm_path}/lib":*) ;; *) export LD_LIBRARY_PATH="$HOME/.mlstack/libmpi-compat:$HOME/.mlstack/libmpi-compat-user-$(id -u):{rocm_path}/lib:{rocm_path}/hip/lib:{rocm_path}/opencl/lib:${{LD_LIBRARY_PATH:-}}";; esac
640{rccl_block}
641
642# Performance Settings
643# HSA_OVERRIDE_GFX_VERSION is set based on detected GPU_ARCH
644export HSA_OVERRIDE_GFX_VERSION={hsa_override_gfx_version}
645export HSA_ENABLE_SDMA=0
646export GPU_MAX_HEAP_SIZE=100
647export GPU_MAX_ALLOC_PERCENT=100
648export HSA_TOOLS_LIB=/opt/rocm/lib/rocprofiler-sdk/librocprofiler-sdk-tool.so
649
650# MIOpen Settings
651# Only set if not already set
652if [ -z "${{MIOPEN_DEBUG_CONV_IMPLICIT_GEMM:-}}" ]; then export MIOPEN_DEBUG_CONV_IMPLICIT_GEMM=1; fi
653if [ -z "${{MIOPEN_FIND_MODE:-}}" ]; then export MIOPEN_FIND_MODE=3; fi
654if [ -z "${{MIOPEN_FIND_ENFORCE:-}}" ]; then export MIOPEN_FIND_ENFORCE=3; fi
655
656# PyTorch Settings
657# Only set if not already set
658# (Stage 2: the CUDA build-arch env var is intentionally NOT exported on a
659# ROCm stack — ROCm uses PYTORCH_ROCM_ARCH. No CUDA env leakage.)
660if [ -z "${{PYTORCH_CUDA_ALLOC_CONF:-}}" ]; then export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:512"; fi
661if [ -z "${{PYTORCH_HIP_ALLOC_CONF:-}}" ]; then export PYTORCH_HIP_ALLOC_CONF="max_split_size_mb:512"; fi
662if [ -z "${{VLLM_WORKER_MULTIPROC_METHOD:-}}" ]; then export VLLM_WORKER_MULTIPROC_METHOD=spawn; fi
663if [ -z "${{VLLM_ROCM_USE_AITER:-}}" ]; then export VLLM_ROCM_USE_AITER=0; fi
664if [ -z "${{MLSTACK_TRITON_HOME:-}}" ]; then export MLSTACK_TRITON_HOME="$HOME/.cache/mlstack/triton"; fi
665if [ -z "${{TRITON_HOME:-}}" ]; then export TRITON_HOME="$MLSTACK_TRITON_HOME"; fi
666if [ -z "${{TRITON_CACHE_DIR:-}}" ]; then export TRITON_CACHE_DIR="$TRITON_HOME/cache"; fi
667if [ -z "${{TRITON_DUMP_DIR:-}}" ]; then export TRITON_DUMP_DIR="$TRITON_HOME/dump"; fi
668if [ -z "${{TRITON_OVERRIDE_DIR:-}}" ]; then export TRITON_OVERRIDE_DIR="$TRITON_HOME/override"; fi
669mkdir -p "$TRITON_CACHE_DIR" "$TRITON_DUMP_DIR" "$TRITON_OVERRIDE_DIR" 2>/dev/null || true
670
671# MPI Settings
672# Only set if not already set
673# (Stage 2: the OMPI CUDA-opal support vars are intentionally NOT set — that is
674# the CUDA-built OpenMPI path; a ROCm stack uses the ROCm accelerator path.)
675if [ -z "${{OMPI_MCA_btl_openib_allow_ib:-}}" ]; then export OMPI_MCA_btl_openib_allow_ib=true; fi
676if [ -z "${{OMPI_MCA_btl_openib_warn_no_device_params_found:-}}" ]; then export OMPI_MCA_btl_openib_warn_no_device_params_found=0; fi
677if [ -z "${{OMPI_MCA_coll_hcoll_enable:-}}" ]; then export OMPI_MCA_coll_hcoll_enable=0; fi
678if [ -z "${{OMPI_MCA_pml:-}}" ]; then export OMPI_MCA_pml=ucx; fi
679if [ -z "${{OMPI_MCA_osc:-}}" ]; then export OMPI_MCA_osc=ucx; fi
680if [ -z "${{OMPI_MCA_btl:-}}" ]; then export OMPI_MCA_btl=^openib,uct; fi
681# Prevent UCX ROCm memory registration segfault on multi-GPU systems
682export UCX_MEMTYPE_CACHE=no
683export UCX_LOG_LEVEL=error
684
685# Global flags for seamless pip installs (PEP 668 override)
686export PIP_BREAK_SYSTEM_PACKAGES=1
687export UV_PIP_BREAK_SYSTEM_PACKAGES=1
688export UV_SYSTEM_PYTHON=1
689
690# Python Interpreter
691# Canonical Python binary — all ML components install to this interpreter.
692# Ensures a single unified environment regardless of how many Pythons exist.
693export MLSTACK_PYTHON_BIN={python_bin}
694export UV_PYTHON={python_bin}
695
696# Flash Attention AMD
697export FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE
698
699# ONNX Runtime Settings
700# Only add if not already in PYTHONPATH
701if ! echo "${{PYTHONPATH:-}}" | grep -q "{home}/onnxruntime_build/onnxruntime/build/Linux/Release"; then
702 export PYTHONPATH={home}/onnxruntime_build/onnxruntime/build/Linux/Release:${{PYTHONPATH:-}}
703fi
704"#,
705 date = chrono::Local::now().format("%Y-%m-%d %H:%M:%S"),
706 hip_visible_devices = hip_visible_devices,
707 rocm_path = rocm_path,
708 rocm_version = rocm_version,
709 rocm_channel = rocm_channel,
710 gpu_arch = gpu_arch,
711 hsa_override_gfx_version = hsa_override_gfx_version,
712 home = home,
713 python_bin = python_bin,
714 rccl_block = rccl_block,
715 )
716}
717
718pub fn generate_fish_env_file_content(
725 hip_visible_devices: &str,
726 rocm_path: &str,
727 rocm_version: &str,
728 rocm_channel: &str,
729 gpu_arch: &str,
730 hsa_override_gfx_version: &str,
731 python_bin: &str,
732) -> String {
733 let home = std::env::var("HOME").unwrap_or_else(|_| "$HOME".to_string());
734
735 let onnx_path = format!("{home}/onnxruntime_build/onnxruntime/build/Linux/Release");
737 let rccl_block = rccl_overlay_fish_block(&home);
739
740 format!(
741 r##"# ML Stack Environment File — Fish Shell Native
742# Generated by Enhanced ML Stack Environment Setup (Rust native)
743# Date: {date}
744#
745# This file is auto-generated. Edit via rusty-stack's env setup, not by hand.
746
747# --- GPU Selection — device filter is MANAGED (unconditional) so a stale
748# override can never re-expose an integrated GPU. ROCR_VISIBLE_DEVICES is the
749# authoritative ROCr (HSA) filter. ---
750# Discrete GPUs only (iGPUs filtered): {hip_visible_devices}
751set -gx ROCR_VISIBLE_DEVICES {hip_visible_devices}
752set -gx HIP_VISIBLE_DEVICES {hip_visible_devices}
753set -gx CUDA_VISIBLE_DEVICES {hip_visible_devices}
754set -gx PYTORCH_ROCM_DEVICE {hip_visible_devices}
755
756# --- ROCm Settings ---
757set -q ROCM_HOME; or set -gx ROCM_HOME {rocm_path}
758set -q ROCM_PATH; or set -gx ROCM_PATH {rocm_path}
759set -q HIP_PATH; or set -gx HIP_PATH {rocm_path}
760set -q CUDA_HOME; or set -gx CUDA_HOME {rocm_path}
761set -q ROCM_VERSION; or set -gx ROCM_VERSION {rocm_version}
762set -q ROCM_CHANNEL; or set -gx ROCM_CHANNEL {rocm_channel}
763# GPU arch identity is MANAGED (unconditional): must match detected hardware.
764set -gx GPU_ARCH {gpu_arch}
765set -gx PYTORCH_ROCM_ARCH {gpu_arch}
766set -gx GPU_ARCHS {gpu_arch}
767set -gx AMDGPU_ASIC_ID_TABLE_PATH /usr/share/libdrm/amdgpu.ids
768set -gx AMDGPU_ASIC_ID_TABLE_PATHS /usr/share/libdrm
769
770# --- Path Settings ---
771# Idempotent: prepend the ROCm bin/lib blocks only if not already on the var, so
772# re-sourcing never grows PATH/LD_LIBRARY_PATH (no per-source duplicate buildup).
773contains -- {rocm_path}/bin $PATH; or set -gx PATH /usr/local/bin /usr/bin /bin /usr/local/games /usr/games {rocm_path}/bin {rocm_path}/hip/bin $PATH
774contains -- {rocm_path}/lib $LD_LIBRARY_PATH; or set -gx LD_LIBRARY_PATH $HOME/.mlstack/libmpi-compat $HOME/.mlstack/libmpi-compat-user-(id -u) {rocm_path}/lib {rocm_path}/hip/lib {rocm_path}/opencl/lib $LD_LIBRARY_PATH
775{rccl_block}
776
777# --- Performance Settings ---
778set -gx HSA_OVERRIDE_GFX_VERSION {hsa_override_gfx_version}
779set -gx HSA_ENABLE_SDMA 0
780set -gx GPU_MAX_HEAP_SIZE 100
781set -gx GPU_MAX_ALLOC_PERCENT 100
782set -gx HSA_TOOLS_LIB /opt/rocm/lib/rocprofiler-sdk/librocprofiler-sdk-tool.so
783
784# --- MIOpen Settings ---
785set -q MIOPEN_DEBUG_CONV_IMPLICIT_GEMM; or set -gx MIOPEN_DEBUG_CONV_IMPLICIT_GEMM 1
786set -q MIOPEN_FIND_MODE; or set -gx MIOPEN_FIND_MODE 3
787set -q MIOPEN_FIND_ENFORCE; or set -gx MIOPEN_FIND_ENFORCE 3
788
789# --- PyTorch Settings ---
790# (Stage 2: the CUDA build-arch env var is intentionally NOT set on a ROCm
791# stack — ROCm uses PYTORCH_ROCM_ARCH. No CUDA env leakage.)
792set -q PYTORCH_CUDA_ALLOC_CONF; or set -gx PYTORCH_CUDA_ALLOC_CONF "max_split_size_mb:512"
793set -q PYTORCH_HIP_ALLOC_CONF; or set -gx PYTORCH_HIP_ALLOC_CONF "max_split_size_mb:512"
794set -q VLLM_WORKER_MULTIPROC_METHOD; or set -gx VLLM_WORKER_MULTIPROC_METHOD spawn
795set -q VLLM_ROCM_USE_AITER; or set -gx VLLM_ROCM_USE_AITER 0
796set -q MLSTACK_TRITON_HOME; or set -gx MLSTACK_TRITON_HOME "$HOME/.cache/mlstack/triton"
797set -q TRITON_HOME; or set -gx TRITON_HOME $MLSTACK_TRITON_HOME
798set -q TRITON_CACHE_DIR; or set -gx TRITON_CACHE_DIR "$TRITON_HOME/cache"
799set -q TRITON_DUMP_DIR; or set -gx TRITON_DUMP_DIR "$TRITON_HOME/dump"
800set -q TRITON_OVERRIDE_DIR; or set -gx TRITON_OVERRIDE_DIR "$TRITON_HOME/override"
801mkdir -p $TRITON_CACHE_DIR $TRITON_DUMP_DIR $TRITON_OVERRIDE_DIR 2>/dev/null; or true
802
803# --- Global flags for seamless pip installs (PEP 668 override) ---
804set -q PIP_BREAK_SYSTEM_PACKAGES; or set -gx PIP_BREAK_SYSTEM_PACKAGES 1
805set -q UV_PIP_BREAK_SYSTEM_PACKAGES; or set -gx UV_PIP_BREAK_SYSTEM_PACKAGES 1
806set -gx UV_SYSTEM_PYTHON 1
807
808# --- Python Interpreter ---
809# Canonical Python binary — all ML components install to this interpreter.
810# Ensures a single unified environment regardless of how many Pythons exist.
811set -gx MLSTACK_PYTHON_BIN {python_bin}
812set -gx UV_PYTHON {python_bin}
813
814# --- Flash Attention AMD ---
815set -gx FLASH_ATTENTION_TRITON_AMD_ENABLE TRUE
816
817# --- MPI/UCX Settings ---
818# (Stage 2: the OMPI CUDA-opal support vars are intentionally NOT set —
819# CUDA-built OpenMPI path; a ROCm stack uses the ROCm accelerator path.)
820set -q OMPI_MCA_btl_openib_allow_ib; or set -gx OMPI_MCA_btl_openib_allow_ib true
821set -q OMPI_MCA_btl_openib_warn_no_device_params_found; or set -gx OMPI_MCA_btl_openib_warn_no_device_params_found 0
822set -q OMPI_MCA_coll_hcoll_enable; or set -gx OMPI_MCA_coll_hcoll_enable 0
823set -q OMPI_MCA_pml; or set -gx OMPI_MCA_pml ucx
824set -q OMPI_MCA_osc; or set -gx OMPI_MCA_osc ucx
825set -q OMPI_MCA_btl; or set -gx OMPI_MCA_btl '^openib,uct'
826set -gx UCX_MEMTYPE_CACHE no
827set -gx UCX_LOG_LEVEL error
828
829# --- ONNX Runtime Settings ---
830if not string match -q '*{onnx_path}*' "$PYTHONPATH"
831 set -gx PYTHONPATH {onnx_path} $PYTHONPATH
832end
833"##,
834 date = chrono::Local::now().format("%Y-%m-%d %H:%M:%S"),
835 hip_visible_devices = hip_visible_devices,
836 rocm_path = rocm_path,
837 rocm_version = rocm_version,
838 rocm_channel = rocm_channel,
839 gpu_arch = gpu_arch,
840 hsa_override_gfx_version = hsa_override_gfx_version,
841 python_bin = python_bin,
842 onnx_path = onnx_path,
843 rccl_block = rccl_block,
844 )
845}
846
847fn rccl_active_dir(home: &str) -> String {
859 format!("{home}/.mlstack/components/rccl/active")
860}
861
862fn rccl_overlay_sha256(active: &str) -> Option<String> {
865 let manifest = std::fs::read_to_string(format!("{active}/manifest")).ok()?;
866 let raw = manifest.lines().find_map(|l| {
867 l.strip_prefix("sha256=")
868 .map(str::trim)
869 .filter(|s| !s.is_empty())
870 })?;
871 let is_hex64 = raw.len() == 64 && raw.chars().all(|c| c.is_ascii_hexdigit());
878 if is_hex64 {
879 Some(raw.to_string())
880 } else {
881 None
882 }
883}
884
885fn rccl_overlay_bash_block(home: &str) -> String {
894 if !std::path::Path::new(&rccl_active_dir(home)).exists() {
895 return String::new();
896 }
897 let sha_line = match rccl_overlay_sha256(&rccl_active_dir(home)) {
898 Some(sha) => format!(" export MLSTACK_RCCL_OVERLAY_SHA256=\"{sha}\""),
902 None => " # MLSTACK_RCCL_OVERLAY_SHA256 omitted (manifest missing or invalid)".to_string(),
903 };
904 format!(
905 r#"
906# Sealed RCCL overlay — MANAGED by the persistent-env generator (single source of
907# truth for env vars; the generator, not a component file, owns these). Follows
908# the `active` symlink; the PYTHONPATH prepend is guarded so re-sourcing the env
909# never accumulates duplicate entries.
910if [ -r "$HOME/.mlstack/components/rccl/active/lib/librccl.so.1.0" ]; then
911 export MLSTACK_RCCL_OVERLAY_LIB="$HOME/.mlstack/components/rccl/active/lib/librccl.so.1.0"
912{sha_line}
913 export NCCL_P2P_DISABLE=1
914 export RCCL_P2P_DISABLE=1
915 if ! echo "${{PYTHONPATH:-}}" | grep -qF "$HOME/.mlstack/components/rccl/active/python"; then
916 export PYTHONPATH="$HOME/.mlstack/components/rccl/active/python:${{PYTHONPATH:-}}"
917 fi
918fi
919"#
920 )
921}
922
923fn rccl_overlay_fish_block(home: &str) -> String {
926 if !std::path::Path::new(&rccl_active_dir(home)).exists() {
927 return String::new();
928 }
929 let sha_line = match rccl_overlay_sha256(&rccl_active_dir(home)) {
930 Some(sha) => format!(" set -gx MLSTACK_RCCL_OVERLAY_SHA256 \"{sha}\""),
931 None => {
932 " # MLSTACK_RCCL_OVERLAY_SHA256 omitted (manifest missing or invalid)".to_string()
933 }
934 };
935 format!(
936 r#"
937# --- Sealed RCCL overlay — MANAGED by the persistent-env generator (single
938# source of truth for env vars). Follows the `active` symlink; idempotent. ---
939if test -r $HOME/.mlstack/components/rccl/active/lib/librccl.so.1.0
940 set -gx MLSTACK_RCCL_OVERLAY_LIB $HOME/.mlstack/components/rccl/active/lib/librccl.so.1.0
941{sha_line}
942 set -gx NCCL_P2P_DISABLE 1
943 set -gx RCCL_P2P_DISABLE 1
944 if not contains -- $HOME/.mlstack/components/rccl/active/python $PYTHONPATH
945 set -gx PYTHONPATH $HOME/.mlstack/components/rccl/active/python $PYTHONPATH
946 end
947end
948"#
949 )
950}
951
952pub fn setup_environment() -> (EnvFileResult, GpuFilterResult, GpuArchInfo) {
970 print_step("Enhanced ML Stack Environment Setup");
971
972 print_step("Detecting AMD GPUs...");
974 let gpu_filter = detect_discrete_gpus();
975 if gpu_filter.gpu_count > 0 {
976 print_success(&format!(
977 "Detected {} discrete GPU(s): {}",
978 gpu_filter.gpu_count, gpu_filter.hip_visible_devices
979 ));
980 } else {
981 print_warning("No discrete GPUs detected, using defaults");
982 }
983
984 let rocm_env = crate::installers::common::rocm_env::RocmEnv::detect();
986 let rocm_path = rocm_env
987 .path()
988 .map(|p| p.to_string_lossy().to_string())
989 .unwrap_or_else(|| "/opt/rocm".to_string());
990 let rocm_version = if rocm_env.version().is_empty() {
991 "7.2.0".to_string()
992 } else {
993 rocm_env.version().to_string()
994 };
995 let rocm_channel = std::env::var("ROCM_CHANNEL").unwrap_or_else(|_| "latest".to_string());
996
997 if rocm_env.is_detected() {
998 print_success(&format!("ROCm {rocm_version} detected at {rocm_path}"));
999 } else {
1000 print_warning("ROCm not detected, using default paths");
1001 }
1002
1003 let gpu_arch_info = detect_gpu_arch_from_system();
1005
1006 print_step("Creating environment file...");
1008 let python_bin = crate::platform::environment::resolve_canonical_python_bin();
1009 let env_result = create_env_file(
1010 &gpu_filter.hip_visible_devices,
1011 &rocm_path,
1012 &rocm_version,
1013 &rocm_channel,
1014 &gpu_arch_info.gpu_arch,
1015 &gpu_arch_info.hsa_override_gfx_version,
1016 &python_bin,
1017 );
1018
1019 if env_result.success {
1020 print_success("Environment setup complete");
1021 }
1022
1023 (env_result, gpu_filter, gpu_arch_info)
1024}
1025
1026fn detect_gpu_arch_from_system() -> GpuArchInfo {
1028 let rocminfo_path = crate::installers::common::utils::resolve_rocminfo_path();
1029 if !command_exists(&rocminfo_path) && !command_exists("rocminfo") {
1030 return GpuArchInfo::default();
1031 }
1032
1033 let output = match std::process::Command::new(&rocminfo_path).output() {
1034 Ok(o) if o.status.success() => String::from_utf8_lossy(&o.stdout).to_string(),
1035 _ => return GpuArchInfo::default(),
1036 };
1037
1038 let mut marketing_name = String::new();
1039
1040 for line in output.lines() {
1041 let trimmed = line.trim();
1042 if trimmed.starts_with("Marketing Name:") {
1043 let name = trimmed
1044 .strip_prefix("Marketing Name:")
1045 .map(|s| s.trim())
1046 .unwrap_or("");
1047
1048 if is_integrated_gpu_name(name) {
1050 continue;
1051 }
1052
1053 if !name.is_empty() {
1054 marketing_name = name.to_string();
1055 break;
1056 }
1057 }
1058 }
1059
1060 detect_correct_gpu_arch(if marketing_name.is_empty() {
1061 None
1062 } else {
1063 Some(&marketing_name)
1064 })
1065}
1066
1067#[cfg(test)]
1072mod tests {
1073 use super::*;
1074
1075 #[test]
1078 fn test_gpu_arch_info_default() {
1079 let info = GpuArchInfo::default();
1080 assert_eq!(info.gpu_arch, "gfx1100");
1081 assert_eq!(info.hsa_override_gfx_version, "11.0.0");
1082 }
1083
1084 #[test]
1085 fn test_gpu_arch_info_from_arch_all_variants() {
1086 let cases = vec![
1087 ("gfx1100", "11.0.0"),
1088 ("gfx1101", "11.0.1"),
1089 ("gfx1102", "11.0.2"),
1090 ("gfx1200", "12.0.0"),
1091 ("gfx1201", "12.0.1"),
1092 ("gfx1030", "10.3.0"),
1093 ("gfx1031", "10.3.1"),
1094 ("gfx1032", "10.3.2"),
1095 ("gfx9999", "11.0.0"), ];
1097
1098 for (arch, expected_hsa) in cases {
1099 let info = GpuArchInfo::from_arch(arch);
1100 assert_eq!(info.gpu_arch, arch, "GPU arch mismatch for {arch}");
1101 assert_eq!(
1102 info.hsa_override_gfx_version, expected_hsa,
1103 "HSA version mismatch for {arch}"
1104 );
1105 }
1106 }
1107
1108 #[test]
1111 fn test_is_integrated_gpu_name_positive() {
1112 assert!(is_integrated_gpu_name("Raphael"));
1113 assert!(is_integrated_gpu_name("AMD Ryzen Graphics"));
1114 assert!(is_integrated_gpu_name("Integrated Graphics"));
1115 assert!(is_integrated_gpu_name("iGPU Device"));
1116 assert!(is_integrated_gpu_name("AMD Radeon Graphics"));
1117 }
1118
1119 #[test]
1120 fn test_is_integrated_gpu_name_negative() {
1121 assert!(!is_integrated_gpu_name("Radeon RX 7900 XTX"));
1122 assert!(!is_integrated_gpu_name("Radeon RX 7800 XT"));
1123 assert!(!is_integrated_gpu_name("Radeon RX 9070 XT"));
1124 assert!(!is_integrated_gpu_name("Unknown GPU"));
1125 }
1126
1127 #[test]
1128 fn test_is_integrated_gpu_name_case_insensitive() {
1129 assert!(is_integrated_gpu_name("raphael"));
1130 assert!(is_integrated_gpu_name("INTEGRATED"));
1131 assert!(is_integrated_gpu_name("AMD RADEON GRAPHICS"));
1132 }
1133
1134 #[test]
1137 fn test_detect_discrete_gpus_returns_result() {
1138 let result = detect_discrete_gpus();
1139 assert!(!result.hip_visible_devices.is_empty());
1141 assert!(result.gpu_count >= 1);
1142 }
1143
1144 #[test]
1145 fn test_detect_gpus_from_lspci_returns_vec() {
1146 let result = detect_gpus_from_lspci();
1148 let _ = result;
1150 }
1151
1152 #[test]
1153 fn test_detect_gpus_from_render_nodes_returns_vec() {
1154 let result = detect_gpus_from_render_nodes();
1155 let _ = result;
1156 }
1157
1158 #[test]
1161 fn test_detect_correct_gpu_arch_7900_xtx() {
1162 let info = detect_correct_gpu_arch(Some("Radeon RX 7900 XTX"));
1163 assert_eq!(info.gpu_arch, "gfx1100");
1164 assert_eq!(info.hsa_override_gfx_version, "11.0.0");
1165 }
1166
1167 #[test]
1168 fn test_detect_correct_gpu_arch_7800_xt() {
1169 let info = detect_correct_gpu_arch(Some("Radeon RX 7800 XT"));
1170 assert_eq!(info.gpu_arch, "gfx1101");
1171 }
1172
1173 #[test]
1174 fn test_detect_correct_gpu_arch_7600() {
1175 let info = detect_correct_gpu_arch(Some("Radeon RX 7600"));
1176 assert_eq!(info.gpu_arch, "gfx1102");
1177 }
1178
1179 #[test]
1180 fn test_detect_correct_gpu_arch_9070_xt() {
1181 let info = detect_correct_gpu_arch(Some("Radeon RX 9070 XT"));
1182 assert_eq!(info.gpu_arch, "gfx1201");
1183 }
1184
1185 #[test]
1186 fn test_detect_correct_gpu_arch_none() {
1187 let info = detect_correct_gpu_arch(None);
1188 assert_eq!(info.gpu_arch, "gfx1100"); }
1190
1191 #[test]
1192 fn test_detect_correct_gpu_arch_unknown() {
1193 let info = detect_correct_gpu_arch(Some("Unknown GPU Model"));
1194 assert_eq!(info.gpu_arch, "gfx1100"); }
1196
1197 #[test]
1198 fn test_detect_correct_gpu_arch_integrated_skipped() {
1199 let info = detect_correct_gpu_arch(Some("AMD Ryzen Graphics"));
1200 assert_eq!(info.gpu_arch, "gfx1100"); }
1202
1203 #[test]
1204 fn test_detect_correct_gpu_arch_case_insensitive() {
1205 let info = detect_correct_gpu_arch(Some("radeon rx 7900 xtx"));
1206 assert_eq!(info.gpu_arch, "gfx1100");
1207 }
1208
1209 #[test]
1212 fn test_system_dependency_required_packages() {
1213 let packages = SystemDependency::required_packages();
1214 assert!(!packages.is_empty());
1215 assert!(packages.contains(&"git"));
1216 assert!(packages.contains(&"cmake"));
1217 assert!(packages.contains(&"build-essential"));
1218 }
1219
1220 #[test]
1221 fn test_check_system_dependency_git() {
1222 let result = check_system_dependency("git");
1223 assert_eq!(result.package, "git");
1224 assert!(result.installed, "git should be installed");
1226 }
1227
1228 #[test]
1229 fn test_check_system_dependency_nonexistent() {
1230 let result = check_system_dependency("nonexistent-pkg-xyz-12345");
1231 assert!(!result.installed);
1232 }
1233
1234 #[test]
1235 fn test_check_all_dependencies_returns_results() {
1236 let results = check_all_dependencies();
1237 assert!(!results.is_empty(), "Should check at least one dependency");
1238 assert_eq!(results.len(), SystemDependency::required_packages().len());
1239 }
1240
1241 #[test]
1244 fn test_generate_env_file_has_header() {
1245 let content = generate_env_file_content(
1246 "0",
1247 "/opt/rocm",
1248 "7.2.0",
1249 "latest",
1250 "gfx1100",
1251 "11.0.0",
1252 "python3",
1253 );
1254 assert!(content.starts_with("# ML Stack Environment File"));
1255 }
1256
1257 #[test]
1258 fn test_generate_env_file_has_gpu_selection() {
1259 let content = generate_env_file_content(
1260 "0,1",
1261 "/opt/rocm",
1262 "7.2.0",
1263 "latest",
1264 "gfx1100",
1265 "11.0.0",
1266 "python3",
1267 );
1268 assert!(content.contains("HIP_VISIBLE_DEVICES"));
1269 assert!(content.contains("CUDA_VISIBLE_DEVICES"));
1270 assert!(content.contains("PYTORCH_ROCM_DEVICE"));
1271 assert!(content.contains("0,1"));
1272 }
1273
1274 #[test]
1275 fn test_generate_env_file_has_rocm_settings() {
1276 let content = generate_env_file_content(
1277 "0",
1278 "/opt/rocm",
1279 "7.2.0",
1280 "latest",
1281 "gfx1100",
1282 "11.0.0",
1283 "python3",
1284 );
1285 assert!(content.contains("ROCM_HOME"));
1286 assert!(content.contains("CUDA_HOME"));
1287 assert!(content.contains("ROCM_VERSION"));
1288 assert!(content.contains("ROCM_CHANNEL"));
1289 assert!(content.contains("GPU_ARCH"));
1290 assert!(content.contains("gfx1100"));
1291 }
1292
1293 #[test]
1294 fn test_generate_env_file_has_path_settings() {
1295 let content = generate_env_file_content(
1296 "0",
1297 "/opt/rocm",
1298 "7.2.0",
1299 "latest",
1300 "gfx1100",
1301 "11.0.0",
1302 "python3",
1303 );
1304 assert!(content.contains("PATH="));
1305 assert!(content.contains("LD_LIBRARY_PATH"));
1306 assert!(content.contains("/opt/rocm/bin"));
1307 }
1308
1309 #[test]
1310 fn test_generate_env_file_has_performance_settings() {
1311 let content = generate_env_file_content(
1312 "0",
1313 "/opt/rocm",
1314 "7.2.0",
1315 "latest",
1316 "gfx1100",
1317 "11.0.0",
1318 "python3",
1319 );
1320 assert!(content.contains("HSA_OVERRIDE_GFX_VERSION"));
1321 assert!(content.contains("HSA_ENABLE_SDMA"));
1322 assert!(content.contains("GPU_MAX_HEAP_SIZE"));
1323 assert!(content.contains("GPU_MAX_ALLOC_PERCENT"));
1324 assert!(content.contains("HSA_TOOLS_LIB"));
1325 }
1326
1327 #[test]
1328 fn test_generate_env_file_has_miopen_settings() {
1329 let content = generate_env_file_content(
1330 "0",
1331 "/opt/rocm",
1332 "7.2.0",
1333 "latest",
1334 "gfx1100",
1335 "11.0.0",
1336 "python3",
1337 );
1338 assert!(content.contains("MIOPEN_DEBUG_CONV_IMPLICIT_GEMM"));
1339 assert!(content.contains("MIOPEN_FIND_MODE"));
1340 assert!(content.contains("MIOPEN_FIND_ENFORCE"));
1341 }
1342
1343 #[test]
1344 fn test_generate_env_file_has_pytorch_settings() {
1345 let content = generate_env_file_content(
1346 "0",
1347 "/opt/rocm",
1348 "7.2.0",
1349 "latest",
1350 "gfx1100",
1351 "11.0.0",
1352 "python3",
1353 );
1354 assert!(!content.contains("TORCH_CUDA_ARCH_LIST"));
1355 assert!(content.contains("PYTORCH_CUDA_ALLOC_CONF"));
1356 assert!(content.contains("PYTORCH_HIP_ALLOC_CONF"));
1357 assert!(content.contains("VLLM_WORKER_MULTIPROC_METHOD"));
1358 }
1359
1360 #[test]
1361 fn test_generate_env_file_has_triton_settings() {
1362 let content = generate_env_file_content(
1363 "0",
1364 "/opt/rocm",
1365 "7.2.0",
1366 "latest",
1367 "gfx1100",
1368 "11.0.0",
1369 "python3",
1370 );
1371 assert!(content.contains("MLSTACK_TRITON_HOME"));
1372 assert!(content.contains("TRITON_HOME"));
1373 assert!(content.contains("TRITON_CACHE_DIR"));
1374 }
1375
1376 #[test]
1377 fn test_generate_env_file_has_mpi_settings() {
1378 let content = generate_env_file_content(
1379 "0",
1380 "/opt/rocm",
1381 "7.2.0",
1382 "latest",
1383 "gfx1100",
1384 "11.0.0",
1385 "python3",
1386 );
1387 assert!(
1388 !content.contains("opal_cuda_support"),
1389 "must NOT set any OMPI CUDA-opal MCA var (Stage 2: ROCm uses the ROCm accelerator path)"
1390 );
1391 assert!(content.contains("OMPI_MCA_pml"));
1392 assert!(content.contains("OMPI_MCA_osc"));
1393 assert!(content.contains("OMPI_MCA_btl"));
1394 }
1395
1396 #[test]
1397 fn test_generate_env_file_has_onnx_runtime() {
1398 let content = generate_env_file_content(
1399 "0",
1400 "/opt/rocm",
1401 "7.2.0",
1402 "latest",
1403 "gfx1100",
1404 "11.0.0",
1405 "python3",
1406 );
1407 assert!(content.contains("PYTHONPATH"));
1408 assert!(content.contains("onnxruntime"));
1409 }
1410
1411 #[test]
1412 fn test_generate_env_file_uses_guard_syntax() {
1413 let content = generate_env_file_content(
1414 "0",
1415 "/opt/rocm",
1416 "7.2.0",
1417 "latest",
1418 "gfx1100",
1419 "11.0.0",
1420 "python3",
1421 );
1422 assert!(content.contains("export ROCR_VISIBLE_DEVICES=0"));
1425 assert!(content.contains("export HIP_VISIBLE_DEVICES=0"));
1426 assert!(!content.contains("if [ -z \"${HIP_VISIBLE_DEVICES:-}\" ]"));
1427 assert!(content.contains("if [ -z \"${ROCM_HOME:-}\" ]"));
1429 }
1430
1431 #[test]
1432 fn test_generate_env_file_gpu_arch_not_guarded() {
1433 let content = generate_env_file_content(
1435 "0",
1436 "/opt/rocm",
1437 "7.2.0",
1438 "latest",
1439 "gfx1100",
1440 "11.0.0",
1441 "python3",
1442 );
1443 assert!(content.contains("export GPU_ARCH=gfx1100"));
1444 }
1445
1446 #[test]
1447 fn test_generate_env_file_custom_values() {
1448 let content = generate_env_file_content(
1449 "0,1",
1450 "/opt/rocm-7.0",
1451 "7.0.0",
1452 "legacy",
1453 "gfx1030",
1454 "10.3.0",
1455 "python3",
1456 );
1457 assert!(content.contains("0,1"));
1458 assert!(content.contains("/opt/rocm-7.0"));
1459 assert!(content.contains("7.0.0"));
1460 assert!(content.contains("legacy"));
1461 assert!(content.contains("gfx1030"));
1462 assert!(content.contains("10.3.0"));
1463 }
1464
1465 #[test]
1468 fn test_create_env_file_creates_file_and_fish() {
1469 let tmp_dir = tempfile::tempdir().expect("Failed to create temp dir");
1470 let _env_path = tmp_dir.path().join(".mlstack_env");
1471 let _fish_env_path = tmp_dir.path().join(".config/fish/conf.d/mlstack_env.fish");
1472
1473 let original_home = std::env::var("HOME").ok();
1475 std::env::set_var("HOME", tmp_dir.path());
1476
1477 let result = create_env_file(
1478 "0",
1479 "/opt/rocm",
1480 "7.2.0",
1481 "latest",
1482 "gfx1100",
1483 "11.0.0",
1484 "python3",
1485 );
1486
1487 if let Some(home) = original_home {
1489 std::env::set_var("HOME", home);
1490 } else {
1491 std::env::remove_var("HOME");
1492 }
1493
1494 assert!(result.success, "create_env_file should succeed");
1495 assert!(result.path.is_some());
1496
1497 let path = result.path.unwrap();
1499 let content = std::fs::read_to_string(&path).expect("Should be able to read env file");
1500 assert!(content.contains("HIP_VISIBLE_DEVICES"));
1501 assert!(content.contains("ROCM_HOME"));
1502 assert!(content.contains("PIP_BREAK_SYSTEM_PACKAGES=1"));
1503 assert!(content.contains("UV_SYSTEM_PYTHON=1"));
1504 }
1505
1506 #[test]
1509 fn test_generate_fish_env_file_has_fish_syntax() {
1510 let content = generate_fish_env_file_content(
1511 "0",
1512 "/opt/rocm",
1513 "7.2.0",
1514 "latest",
1515 "gfx1100",
1516 "11.0.0",
1517 "python3",
1518 );
1519 assert!(!content.contains("if [ -z"));
1521 assert!(!content.contains("; fi"));
1522 assert!(!content.contains("export "));
1523 assert!(!content.contains("${"));
1524 assert!(content.contains("set -gx"));
1526 assert!(content.contains("set -q"));
1527 assert!(content.contains("; or set -gx"));
1528 }
1529
1530 #[test]
1531 fn test_generate_fish_env_file_has_required_vars() {
1532 let content = generate_fish_env_file_content(
1533 "0,1",
1534 "/opt/rocm",
1535 "7.2.0",
1536 "latest",
1537 "gfx1100",
1538 "11.0.0",
1539 "python3",
1540 );
1541 assert!(content.contains("HIP_VISIBLE_DEVICES"));
1542 assert!(content.contains("ROCM_HOME"));
1543 assert!(content.contains("ROCM_VERSION"));
1544 assert!(content.contains("GPU_ARCH"));
1545 assert!(content.contains("HSA_OVERRIDE_GFX_VERSION"));
1546 assert!(content.contains("PIP_BREAK_SYSTEM_PACKAGES"));
1547 assert!(content.contains("UV_PIP_BREAK_SYSTEM_PACKAGES"));
1548 assert!(content.contains("UV_SYSTEM_PYTHON"));
1549 assert!(content.contains("FLASH_ATTENTION_TRITON_AMD_ENABLE"));
1550 assert!(!content.contains("source $HOME/.mlstack/components/rccl/active/env.fish"));
1554 assert!(content.contains("0,1"));
1555 }
1556
1557 #[test]
1558 fn test_generate_fish_env_file_has_no_bash_patterns() {
1559 let content = generate_fish_env_file_content(
1560 "0",
1561 "/opt/rocm",
1562 "7.2.0",
1563 "latest",
1564 "gfx1100",
1565 "11.0.0",
1566 "python3",
1567 );
1568 assert!(!content.contains("if ! "));
1570 assert!(!content.contains("fi\n"));
1571 assert!(!content.contains("; then"));
1572 }
1573
1574 #[test]
1575 fn test_generate_bash_env_file_has_pip_break_system_packages() {
1576 let content = generate_env_file_content(
1577 "0",
1578 "/opt/rocm",
1579 "7.2.0",
1580 "latest",
1581 "gfx1100",
1582 "11.0.0",
1583 "python3",
1584 );
1585 assert!(content.contains("export PIP_BREAK_SYSTEM_PACKAGES=1"));
1586 assert!(content.contains("export UV_PIP_BREAK_SYSTEM_PACKAGES=1"));
1587 assert!(content.contains("export UV_SYSTEM_PYTHON=1"));
1588 assert!(content.contains("export FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE"));
1589 assert!(!content.contains(". \"$HOME/.mlstack/components/rccl/active/env.sh\""));
1592 }
1593
1594 #[test]
1595 fn test_rccl_overlay_block_is_centralized_and_idempotent() {
1596 let tmp = std::env::temp_dir();
1598 let home_no_rccl = tmp.join("rusty_rccl_no_active_test");
1599 let _ = std::fs::remove_dir_all(&home_no_rccl);
1600 assert_eq!(rccl_overlay_bash_block(&home_no_rccl.to_string_lossy()), "");
1601 assert_eq!(rccl_overlay_fish_block(&home_no_rccl.to_string_lossy()), "");
1602
1603 let home = tmp.join("rusty_rccl_active_test");
1606 let active = home.join(".mlstack/components/rccl/active");
1607 let _ = std::fs::remove_dir_all(&home);
1608 std::fs::create_dir_all(active.join("python")).unwrap();
1609 std::fs::create_dir_all(active.join("lib")).unwrap();
1610 std::fs::write(active.join("lib/librccl.so.1.0"), b"").unwrap();
1611 let valid_sha = "a".repeat(64);
1615 std::fs::write(active.join("manifest"), format!("sha256={valid_sha}\n")).unwrap();
1616
1617 let bash = rccl_overlay_bash_block(&home.to_string_lossy());
1618 assert!(bash.contains("MLSTACK_RCCL_OVERLAY_LIB"));
1619 assert!(bash.contains(&format!("MLSTACK_RCCL_OVERLAY_SHA256=\"{valid_sha}\"")));
1621 assert!(bash.contains("grep -qF"));
1623 assert!(!bash.contains("env.sh"));
1625
1626 let fish = rccl_overlay_fish_block(&home.to_string_lossy());
1627 assert!(fish.contains("MLSTACK_RCCL_OVERLAY_LIB"));
1628 assert!(fish.contains(&format!("MLSTACK_RCCL_OVERLAY_SHA256 \"{valid_sha}\"")));
1629 assert!(fish.contains("contains --"));
1630 assert!(!fish.contains("env.fish"));
1631
1632 std::fs::write(active.join("manifest"), "sha256=NOT-HEX\n").unwrap();
1635 let bash_bad = rccl_overlay_bash_block(&home.to_string_lossy());
1636 assert!(
1637 !bash_bad.contains("MLSTACK_RCCL_OVERLAY_SHA256=\""),
1638 "invalid SHA must be omitted, not spliced: {bash_bad}"
1639 );
1640 assert!(bash_bad.contains("omitted"));
1641
1642 let _ = std::fs::remove_dir_all(&home);
1643 let _ = std::fs::remove_dir_all(&home_no_rccl);
1644 }
1645
1646 #[test]
1649 fn test_setup_environment_returns_results() {
1650 let (env_result, gpu_result, gpu_arch) = setup_environment();
1651 let _ = env_result.success;
1653 assert!(!gpu_result.hip_visible_devices.is_empty());
1654 assert!(!gpu_arch.gpu_arch.is_empty());
1655 }
1656}