Skip to main content

rusty_stack/bootstrap/
env_setup.rs

1//! enhanced_setup_environment.sh equivalent — GPU filtering, system deps, env file, ROCm optimization.
2//!
3//! Ports `scripts/enhanced_setup_environment.sh` (1411 lines) to native Rust.
4//! This is the most complex bootstrap script, handling:
5//! - GPU detection and filtering (discrete vs integrated)
6//! - System dependency checking (kernel, ROCm, packages, Python)
7//! - Environment file creation (`~/.mlstack_env`)
8//! - ROCm optimization settings (HSA_OVERRIDE_GFX_VERSION, MIOpen, channel selection)
9//!
10//! # Validation Assertions
11//!
12//! - **VAL-VBA-014**: GPU filtering detects AMD GPUs via ROCm, sets HIP_VISIBLE_DEVICES
13//! - **VAL-VBA-015**: System dependency checking covers kernel, ROCm, packages, Python
14//! - **VAL-VBA-016**: Environment file creation produces shell-sourceable ~/.mlstack_env
15//! - **VAL-VBA-017**: ROCm optimization: HSA_OVERRIDE_GFX_VERSION, ROCM_MATH_LIBS, MIOPEN_DEBUG, channel selection
16
17use crate::installers::common::utils::{
18    command_exists, print_error, print_step, print_success, print_warning,
19};
20use std::fmt;
21
22// ===========================================================================
23// Types
24// ===========================================================================
25
26/// GPU architecture information detected from the system.
27#[derive(Debug, Clone)]
28pub struct GpuArchInfo {
29    /// The GPU architecture string (e.g., "gfx1100").
30    pub gpu_arch: String,
31    /// The HSA override GFX version (e.g., "11.0.0").
32    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    /// Create GpuArchInfo from a GPU architecture string.
46    ///
47    /// Maps the architecture to the correct HSA override version.
48    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", // default fallback
59        };
60        Self {
61            gpu_arch: arch.to_string(),
62            hsa_override_gfx_version: hsa_version.to_string(),
63        }
64    }
65}
66
67/// Result of GPU detection and filtering.
68#[derive(Debug, Clone)]
69pub struct GpuFilterResult {
70    /// Indices of discrete GPUs (integrated GPUs filtered out).
71    pub discrete_gpu_indices: Vec<u32>,
72    /// Comma-separated string for HIP_VISIBLE_DEVICES.
73    pub hip_visible_devices: String,
74    /// Number of detected GPUs.
75    pub gpu_count: usize,
76}
77
78/// A system dependency to check.
79#[derive(Debug, Clone)]
80pub struct SystemDependency {
81    /// The package name (generic, will be mapped per distro).
82    pub package: &'static str,
83    /// Human-readable description.
84    pub description: &'static str,
85}
86
87impl SystemDependency {
88    /// Returns the list of required system packages.
89    ///
90    /// Mirrors the `required_packages` array from enhanced_setup_environment.sh.
91    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/// Result of checking a single system dependency.
113#[derive(Debug, Clone)]
114pub struct DependencyCheckResult {
115    /// The generic package name.
116    pub package: String,
117    /// Whether the package is installed.
118    pub installed: bool,
119    /// The distro-mapped package name.
120    pub mapped_name: String,
121}
122
123/// Result of environment file creation.
124#[derive(Debug, Clone)]
125pub struct EnvFileResult {
126    /// Whether the operation was successful.
127    pub success: bool,
128    /// Path to the created environment file.
129    pub path: Option<String>,
130    /// Any errors encountered.
131    pub errors: Vec<String>,
132}
133
134impl EnvFileResult {
135    /// Create a successful result.
136    pub fn success(path: String) -> Self {
137        Self {
138            success: true,
139            path: Some(path),
140            errors: Vec::new(),
141        }
142    }
143
144    /// Create a failure result.
145    pub fn failure(error: String) -> Self {
146        Self {
147            success: false,
148            path: None,
149            errors: vec![error],
150        }
151    }
152}
153
154// ===========================================================================
155// GPU Detection and Filtering (VAL-VBA-014)
156// ===========================================================================
157
158/// Detect discrete AMD GPUs, filtering out integrated GPUs.
159///
160/// Uses rocminfo to enumerate GPUs and filters out known integrated GPUs
161/// (Raphael, Ryzen Graphics, etc.). Falls back to lspci if rocminfo is
162/// unavailable.
163///
164/// # Validation
165///
166/// - **VAL-VBA-014**: GPU filtering detects AMD GPUs via ROCm, sets HIP_VISIBLE_DEVICES
167pub fn detect_discrete_gpus() -> GpuFilterResult {
168    let mut discrete_indices: Vec<u32> = Vec::new();
169
170    // Try rocminfo first
171    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    // Fallback to lspci
177    if discrete_indices.is_empty() {
178        discrete_indices = detect_gpus_from_lspci();
179    }
180
181    // Fallback to render nodes
182    if discrete_indices.is_empty() {
183        discrete_indices = detect_gpus_from_render_nodes();
184    }
185
186    // Default to GPU 0 if nothing detected
187    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
205/// Detect GPUs from rocminfo output, filtering out integrated GPUs.
206fn 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; // Fallback counter for systems without GPU ID lines
217
218    for line in output.lines() {
219        let trimmed = line.trim();
220
221        // Detect new Agent sections (lines like "Agent N")
222        if trimmed.starts_with("Agent") && trimmed.contains(char::is_numeric) {
223            // Don't increment here — only for GPU agents
224        }
225
226        if trimmed.starts_with("GPU ID:") {
227            // Save previous GPU if it was discrete
228            if let Some(id) = current_gpu_id {
229                if !is_igpu && !discrete_indices.contains(&id) {
230                    discrete_indices.push(id);
231                }
232            }
233
234            // Parse new GPU ID
235            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            // Check if this is an integrated GPU
246            if is_integrated_gpu_name(marketing_name) {
247                is_igpu = true;
248            }
249
250            // If no GPU ID lines exist, track by agent counter
251            if current_gpu_id.is_none() {
252                // Only count GPU agents (those with Device Type: GPU)
253                // We'll finalize on Device Type line
254            }
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            // CPU devices are also not discrete GPUs
262            if device_type.contains("CPU") {
263                is_igpu = true;
264            }
265
266            // Finalize this GPU
267            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                // No GPU ID lines in this rocminfo output — use counter
273                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    // Handle last GPU
282    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
293/// Detect GPUs from lspci output.
294fn 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        // Filter out integrated GPUs
315        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
328/// Detect GPUs from /dev/dri/render* nodes.
329/// Filters out iGPU render nodes by checking the associated device description.
330fn 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    // Try to filter iGPUs via sysfs device description
345    let discrete_count = render_entries
346        .iter()
347        .filter(|entry| {
348            let name = entry.file_name().to_string_lossy().to_string();
349            // Extract render node number
350            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; // Can't determine, include it
357            }
358
359            // Check sysfs for the associated device description
360            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                // Look for PCI_ID to identify the device
363                for line in uevent.lines() {
364                    if line.starts_with("PCI_ID=") {
365                        // Could check specific IDs here if needed
366                        return true;
367                    }
368                }
369            }
370
371            // Alternative: check via /sys/class/drm/card<N>/device/device
372            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                // 0x164e = Raphael iGPU (7800X3D)
376                if pci_id == "0x164e" {
377                    return false; // Skip iGPU
378                }
379            }
380
381            true // Include by default
382        })
383        .count();
384
385    if discrete_count > 0 {
386        (0..discrete_count as u32).collect()
387    } else if !render_entries.is_empty() {
388        // Fallback: use all render nodes (safer than returning empty)
389        (0..render_entries.len() as u32).collect()
390    } else {
391        Vec::new()
392    }
393}
394
395/// Check if a GPU name indicates an integrated GPU.
396///
397/// Delegates to the canonical classifier `crate::gpu::is_integrated_gpu_name`
398/// (Stage 1: single source of truth for dGPU/iGPU classification).
399pub fn is_integrated_gpu_name(name: &str) -> bool {
400    crate::gpu::is_integrated_gpu_name(name)
401}
402
403// ===========================================================================
404// GPU Architecture Detection (VAL-VBA-017)
405// ===========================================================================
406
407/// Detect the correct GPU architecture from marketing name.
408///
409/// This corrects rocminfo bugs where RDNA3 cards report gfx1030 instead of
410/// gfx1100. Maps marketing names to correct architectures.
411///
412/// # Validation
413///
414/// - **VAL-VBA-017**: ROCm optimization — correct GPU arch detection
415pub 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    // Skip integrated GPUs
422    if is_integrated_gpu_name(name) {
423        return GpuArchInfo::default();
424    }
425
426    // Delegate to canonical name→gfx lookup; default to gfx1100 for unknown
427    let arch = crate::gpu::gfx_from_marketing_name(name).unwrap_or("gfx1100");
428
429    GpuArchInfo::from_arch(arch)
430}
431
432// ===========================================================================
433// System Dependency Checking (VAL-VBA-015)
434// ===========================================================================
435
436/// Check a single system dependency.
437///
438/// Returns whether the package is installed and the mapped name.
439/// Uses distro detection to map generic package names to distro-specific names.
440///
441/// # Validation
442///
443/// - **VAL-VBA-015**: System dependency checking covers kernel, ROCm, packages, Python
444pub 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
460/// Check all required system dependencies.
461///
462/// Returns results for each dependency, indicating which are missing.
463pub fn check_all_dependencies() -> Vec<DependencyCheckResult> {
464    SystemDependency::required_packages()
465        .into_iter()
466        .map(check_system_dependency)
467        .collect()
468}
469
470/// Check if a package is installed (best-effort, distro-agnostic).
471fn check_package_installed(package: &str) -> bool {
472    // Try dpkg first (Debian/Ubuntu)
473    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    // Try rpm (Fedora/RHEL/SUSE)
486    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    // Try pacman (Arch)
498    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
512// ===========================================================================
513// Environment File Creation (VAL-VBA-016)
514// ===========================================================================
515
516/// Create the `~/.mlstack_env` file with all required variables.
517///
518/// Generates a shell-sourceable file with GPU settings, ROCm paths,
519/// performance optimizations, and PyTorch/MPI settings.
520///
521/// # Validation
522///
523/// - **VAL-VBA-016**: Environment file creation produces shell-sourceable ~/.mlstack_env
524pub 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    // Write the bash env file
548    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    // Write the fish-compatible env file
555    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    // Ensure fish conf.d directory exists
565    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        // Non-fatal — fish env is best-effort
570        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
579/// Generate the environment file content.
580///
581/// This produces a shell-sourceable file with all required environment
582/// variables. The content uses `if [ -z ... ]` guards to avoid overriding
583/// existing values, matching the original script behavior exactly.
584///
585/// # Validation
586///
587/// - **VAL-VBA-016**: Produces shell-sourceable content
588/// - **VAL-VBA-017**: Includes ROCm optimization settings
589pub 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    // RCCL overlay env is MANAGED here (the persistent-env generator is the
600    // single source of truth for env vars — components no longer write env files).
601    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
718/// Generate a fish-compatible environment file.
719///
720/// Produces `~/.config/fish/conf.d/mlstack_env.fish` that natively loads
721/// all ML Stack environment variables without needing to source the bash file.
722/// Fish doesn't support bash syntax (`if [ -z ... ]`, `${VAR:-}`, `fi`, etc.),
723/// so this uses idiomatic fish constructs.
724pub 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    // Build ONNX Runtime PYTHONPATH check
736    let onnx_path = format!("{home}/onnxruntime_build/onnxruntime/build/Linux/Release");
737    // RCCL overlay env is MANAGED here (single source of truth — no component env files).
738    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
847// ===========================================================================
848// Component env contributions — MANAGED by the persistent-env generator.
849// The generator is the SINGLE source of truth for env vars; components must
850// NOT write their own env files (the old rccl `active/env.sh` did an un-guarded
851// `export PYTHONPATH=X${PYTHONPATH:+:$PYTHONPATH}` that accumulated a duplicate
852// entry on every `source ~/.mlstack_env`). Components instead declare state via
853// the filesystem (the `active` symlink + manifest); the generator reads it and
854// emits an idempotent block here.
855// ===========================================================================
856
857/// Path to the sealed RCCL overlay's `active` dir for `home`.
858fn rccl_active_dir(home: &str) -> String {
859    format!("{home}/.mlstack/components/rccl/active")
860}
861
862/// Read the RCCL overlay sha256 from `<active>/manifest` (line `sha256=…`), if
863/// the overlay is installed and the manifest is readable.
864fn 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    // Defense-in-depth: the SHA is written by the installer (always a 64-char
872    // hex string), but this value is interpolated unquoted into a file that is
873    // sourced on EVERY shell login. A corrupted/tampered manifest containing
874    // shell metacharacters (`;`, backticks, `$()`, …) would execute when the
875    // env file is sourced. Accept ONLY a 64-char lowercase-hex SHA-256; any
876    // other shape is omitted (preserving the missing-manifest behavior).
877    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
885/// Emit the **idempotent** bash block for the sealed RCCL overlay, or an empty
886/// string when no overlay is installed.
887///
888/// Uses the `$HOME`-relative `active` symlink path in the emitted exports so
889/// the block tracks the current RCCL build without an env refresh. The
890/// PYTHONPATH prepend is guarded (`grep -qF`) so re-sourcing the env never
891/// accumulates duplicates — the generator, not a component-written file, owns
892/// these vars.
893fn 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        // Quoted for defense-in-depth (the value is already validated 64-hex
899        // upstream, so it carries no metacharacters — but quoting keeps the
900        // sourced line robust if the validation contract ever loosens).
901        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
923/// Emit the **idempotent** fish block for the sealed RCCL overlay, or empty.
924/// Fish analogue of [`rccl_overlay_bash_block`]; `contains` is the membership guard.
925fn 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
952// ===========================================================================
953// Full Setup Orchestration
954// ===========================================================================
955
956/// Run the complete enhanced setup environment.
957///
958/// This is the Rust equivalent of the `main()` function in
959/// `enhanced_setup_environment.sh`:
960/// 1. Check system dependencies
961/// 2. Detect AMD GPUs
962/// 3. Detect ROCm
963/// 4. Configure environment variables
964/// 5. Create environment file
965///
966/// # Returns
967///
968/// A tuple of (env_file_result, gpu_filter_result, gpu_arch_info).
969pub fn setup_environment() -> (EnvFileResult, GpuFilterResult, GpuArchInfo) {
970    print_step("Enhanced ML Stack Environment Setup");
971
972    // Step 1: Detect GPUs and filter
973    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    // Step 2: Detect ROCm
985    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    // Step 3: Detect GPU architecture
1004    let gpu_arch_info = detect_gpu_arch_from_system();
1005
1006    // Step 4: Resolve canonical Python and create environment file
1007    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
1026/// Detect GPU architecture from the system (rocminfo-based).
1027fn 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            // Skip integrated GPUs
1049            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// ===========================================================================
1068// Tests
1069// ===========================================================================
1070
1071#[cfg(test)]
1072mod tests {
1073    use super::*;
1074
1075    // --- GPU Arch Info (VAL-VBA-017) ---
1076
1077    #[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"), // unknown defaults
1096        ];
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    // --- Integrated GPU Detection (VAL-VBA-014) ---
1109
1110    #[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    // --- GPU Detection (VAL-VBA-014) ---
1135
1136    #[test]
1137    fn test_detect_discrete_gpus_returns_result() {
1138        let result = detect_discrete_gpus();
1139        // Should always return at least the default GPU 0
1140        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        // Just verify it doesn't panic
1147        let result = detect_gpus_from_lspci();
1148        // May be empty if no AMD GPUs, that's fine
1149        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    // --- Marketing Name to Architecture (VAL-VBA-017) ---
1159
1160    #[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"); // default
1189    }
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"); // default fallback
1195    }
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"); // default (skipped iGPU)
1201    }
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    // --- System Dependencies (VAL-VBA-015) ---
1210
1211    #[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        // git should be installed on any dev system
1225        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    // --- Environment File Content (VAL-VBA-016, VAL-VBA-017) ---
1242
1243    #[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        // Device filter is MANAGED/unconditional (stale override must never
1423        // re-expose an iGPU): ROCR/HIP/CUDA_VISIBLE_DEVICES are bare exports.
1424        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        // Tuning vars still use "if [ -z ... ]" guards (respect user overrides).
1428        assert!(content.contains("if [ -z \"${ROCM_HOME:-}\" ]"));
1429    }
1430
1431    #[test]
1432    fn test_generate_env_file_gpu_arch_not_guarded() {
1433        // GPU_ARCH should be exported unconditionally (not guarded)
1434        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    // --- create_env_file (VAL-VBA-016) ---
1466
1467    #[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        // Temporarily set HOME to the temp dir
1474        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        // Restore HOME
1488        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        // Verify the bash file exists and has content
1498        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    // --- Fish env file generation ---
1507
1508    #[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        // Must NOT contain bash syntax
1520        assert!(!content.contains("if [ -z"));
1521        assert!(!content.contains("; fi"));
1522        assert!(!content.contains("export "));
1523        assert!(!content.contains("${"));
1524        // Must use fish syntax
1525        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        // Centralization: the generator must NOT source a component-written
1551        // env.fish — RCCL overlay env is emitted inline + idempotently by the
1552        // generator (rccl_overlay_fish_block), the single source of truth.
1553        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        // Should never contain bash-specific constructs
1569        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        // Centralization: the generator must NOT source a component-written
1590        // env.sh — RCCL overlay env is emitted inline + idempotently here.
1591        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        // No active RCCL overlay → generator emits nothing for it.
1597        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        // Active RCCL overlay → generator emits a guarded, idempotent block that
1604        // follows the `active` symlink (no component-written env file).
1605        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        // A valid SHA-256 (64 lowercase hex chars) — the installer always writes
1612        // this shape. The generator validates 64-hex before splicing (defense-
1613        // in-depth against a tampered manifest), so use a realistic value here.
1614        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        // SHA is quoted + validated (64-hex). An invalid SHA would be omitted.
1620        assert!(bash.contains(&format!("MLSTACK_RCCL_OVERLAY_SHA256=\"{valid_sha}\"")));
1621        // Idempotent guard present (no bare un-guarded prepend).
1622        assert!(bash.contains("grep -qF"));
1623        // Single source of truth: does NOT delegate to a component env file.
1624        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        // Invalid SHA (not 64-hex) → omitted, never spliced into the sourced
1633        // block. Guards against a corrupted/tampered manifest injecting shell.
1634        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    // --- setup_environment integration ---
1647
1648    #[test]
1649    fn test_setup_environment_returns_results() {
1650        let (env_result, gpu_result, gpu_arch) = setup_environment();
1651        // Should always return results without panicking
1652        let _ = env_result.success;
1653        assert!(!gpu_result.hip_visible_devices.is_empty());
1654        assert!(!gpu_arch.gpu_arch.is_empty());
1655    }
1656}