1use crate::benchmark_logs::{
2 benchmark_log_directories, extract_benchmark_json_value, find_latest_log_in_dirs,
3};
4use crate::component_status::{
5 component_verification_commands, is_component_installed_by_id, modules_available,
6 python_interpreters, python_search_paths, verification_commands, VerificationCommand,
7};
8use crate::config::InstallerConfig;
9use crate::state::{Category, Component};
10use anyhow::{bail, Context, Result};
11use chrono::Local;
12use std::collections::HashMap;
13use std::io::{BufRead, BufReader, Read, Write};
14use std::path::{Path, PathBuf};
15use std::process::{Command, Stdio};
16use std::sync::{
17 atomic::{AtomicBool, Ordering},
18 mpsc::Receiver,
19 mpsc::Sender,
20 Arc, Mutex,
21};
22use std::time::Duration;
23use std::{fs, thread};
24
25#[derive(Debug, Clone)]
26pub enum InstallerEvent {
27 Log(String, bool),
28 Progress {
29 component_id: String,
30 progress: f32,
31 message: String,
32 },
33 ComponentStart {
34 component_id: String,
35 name: String,
36 },
37 ComponentComplete {
38 component_id: String,
39 success: bool,
40 message: String,
41 },
42 VerificationReport {
43 component_id: String,
44 lines: Vec<String>,
45 },
46 Finished {
47 success: bool,
48 },
49}
50
51pub fn run_installation(
52 components: Vec<Component>,
53 config: InstallerConfig,
54 sudo_password: Option<String>,
55 sender: Sender<InstallerEvent>,
56 input_rx: Receiver<String>,
57) {
58 let scripts_dir = &config.scripts_dir;
59 let batch_mode = config.batch_mode;
60 let install_method = config.install_method.clone();
61 let total = components.len() as f32;
62 let mut index = 0usize;
63 let mut python_candidates = python_interpreters();
64 let user_home = resolve_mlstack_user_home();
65 let input_rx = Arc::new(Mutex::new(input_rx));
66
67 match ensure_mlstack_env(&user_home, &install_method) {
68 Ok(EnvUpdate::Created) => {
69 let _ = sender.send(InstallerEvent::Log(
70 format!("Created .mlstack_env in {}", user_home),
71 false,
72 ));
73 }
74 Ok(EnvUpdate::Updated) => {
75 let _ = sender.send(InstallerEvent::Log(
76 format!("Updated .mlstack_env in {}", user_home),
77 false,
78 ));
79 }
80 Ok(EnvUpdate::Unchanged) => {}
81 Err(err) => {
82 let _ = sender.send(InstallerEvent::Log(
83 format!("Failed to ensure .mlstack_env: {}", err),
84 false,
85 ));
86 }
87 }
88
89 let any_component_needs_sudo = components.iter().any(|c| c.needs_sudo);
91
92 if needs_sudo() && any_component_needs_sudo {
93 if let Some(password) = sudo_password.clone() {
94 if let Err(err) = validate_sudo(password) {
95 let _ = sender.send(InstallerEvent::Log(
96 format!("Sudo validation failed: {}", err),
97 false,
98 ));
99 let _ = sender.send(InstallerEvent::Finished { success: false });
100 return;
101 }
102 } else {
103 let _ = sender.send(InstallerEvent::Log(
104 "Sudo password missing; cannot continue".into(),
105 false,
106 ));
107 let _ = sender.send(InstallerEvent::Finished { success: false });
108 return;
109 }
110 }
111
112 let mut overall_success = true;
113 #[allow(clippy::needless_borrow)] let env_exports = load_mlstack_env_exports(&user_home);
115 let persistent_python = env_exports
116 .get("MLSTACK_PYTHON_BIN")
117 .map(|v| v.trim().to_string())
118 .filter(|v| !v.is_empty());
119 let python_bin = persistent_python.unwrap_or_else(resolve_python_bin);
120 python_candidates.retain(|candidate| candidate != &python_bin);
121 python_candidates.insert(0, python_bin.clone());
122 let _ = sender.send(InstallerEvent::Log(
123 format!("Selected python interpreter: {}", python_bin),
124 false,
125 ));
126
127 let components = sort_components_by_dependencies(components);
131
132 for component in components {
133 index += 1;
134 let _ = sender.send(InstallerEvent::ComponentStart {
135 component_id: component.id.clone(),
136 name: component.name.clone(),
137 });
138
139 if component.category == Category::Verification {
140 let outcome = run_verification(&component, &python_candidates, &sender, &user_home);
141 let _ = sender.send(InstallerEvent::VerificationReport {
142 component_id: component.id.clone(),
143 lines: outcome.report_lines.clone(),
144 });
145 if !outcome.success {
146 overall_success = false;
147 }
148 let _ = sender.send(InstallerEvent::ComponentComplete {
149 component_id: component.id.clone(),
150 success: outcome.success,
151 message: if outcome.success {
152 format!("{} verification complete", component.name)
153 } else {
154 format!("{} verification completed with issues", component.name)
155 },
156 });
157
158 let overall_progress = index as f32 / total;
159 let _ = sender.send(InstallerEvent::Progress {
160 component_id: "__overall__".into(),
161 progress: overall_progress,
162 message: format!(
163 "Finished {} ({:.0}%)",
164 component.name,
165 overall_progress * 100.0
166 ),
167 });
168 continue;
169 }
170
171 let is_native = crate::installers::components::is_native_component(&component.id);
174
175 let mut install_success = true;
176 let mut error_msg = String::new();
177
178 if is_native {
179 let ctx = NativeInstallerContext {
184 sudo_password: sudo_password.clone(),
185 batch_mode,
186 install_method: &install_method,
187 sender: &sender,
188 input_rx: Arc::clone(&input_rx),
189 user_home: &user_home,
190 env_exports: &env_exports,
191 };
192 if let Err(err) = run_native_installer(&component, &ctx) {
193 install_success = false;
194 let chain = err
195 .chain()
196 .map(|e| e.to_string())
197 .collect::<Vec<_>>()
198 .join(": ");
199 error_msg = format_user_friendly_error(&component.id, &chain);
200 let err_label = if component.category == Category::Performance {
202 "benchmarks"
203 } else {
204 "installation"
205 };
206 let _ = sender.send(InstallerEvent::Log(
207 format!("[ERROR] {} {} failed: {}", component.name, err_label, chain),
208 false,
209 ));
210 }
211 } else {
212 let script_path = format!("{}/{}", scripts_dir, component.script);
215 if fs::metadata(&script_path).is_err() {
216 overall_success = false;
217 let _ = sender.send(InstallerEvent::ComponentComplete {
218 component_id: component.id.clone(),
219 success: false,
220 message: format!("Script not found: {}", script_path),
221 });
222 continue;
223 }
224
225 if let Err(err) = run_script(
226 &component,
227 &script_path,
228 sudo_password.clone(),
229 batch_mode,
230 &install_method,
231 &sender,
232 Arc::clone(&input_rx),
233 ) {
234 install_success = false;
235 let chain = err
236 .chain()
237 .map(|e| e.to_string())
238 .collect::<Vec<_>>()
239 .join(": ");
240 error_msg = format_user_friendly_error(&component.id, &chain);
241 let err_label = if component.category == Category::Performance {
243 "benchmarks"
244 } else {
245 "installation"
246 };
247 let _ = sender.send(InstallerEvent::Log(
248 format!("[ERROR] {} {} failed: {}", component.name, err_label, chain),
249 false,
250 ));
251 }
252 }
253
254 let verify_label = if component.category == Category::Performance {
255 "Validating"
256 } else {
257 "Verifying"
258 };
259 let _ = sender.send(InstallerEvent::Progress {
260 component_id: component.id.clone(),
261 progress: 0.8,
262 message: format!("{} {}", verify_label, component.name),
263 });
264
265 let verification_outcome =
266 run_component_verification(&component, &python_candidates, &sender, &user_home);
267 let _ = sender.send(InstallerEvent::VerificationReport {
268 component_id: component.id.clone(),
269 lines: verification_outcome.report_lines.clone(),
270 });
271
272 if !verification_outcome.success {
273 overall_success = false;
274 }
275
276 let final_success = verification_outcome.success;
282 if !install_success && verification_outcome.success {
283 let _ = sender.send(InstallerEvent::Log(
284 format!(
285 "{} verification passed despite install warnings; marking as installed",
286 component.name
287 ),
288 false,
289 ));
290 }
291 if !final_success {
292 overall_success = false;
293 }
294
295 if component.id == "permanent-env"
297 || component.id == "enhanced-env"
298 || component.id == "basic-env"
299 {
300 let _ = sender.send(InstallerEvent::Log(
301 format!(
302 "[DEBUG] Component {}: install_success={}, verification_success={}, final={}",
303 component.id, install_success, verification_outcome.success, final_success
304 ),
305 false,
306 ));
307 if !install_success {
308 let _ = sender.send(InstallerEvent::Log(
309 format!("[DEBUG] Install error: {}", error_msg),
310 false,
311 ));
312 }
313 }
314
315 let completion_msg = if !install_success {
316 error_msg
317 } else if component.category == Category::Performance {
318 if verification_outcome.success {
319 format!("{} benchmarks completed", component.name)
320 } else {
321 format!("{} benchmarks completed with errors", component.name)
322 }
323 } else if verification_outcome.success {
324 format!("{} completed", component.name)
325 } else {
326 format!("{} completed with verification errors", component.name)
327 };
328
329 let _ = sender.send(InstallerEvent::ComponentComplete {
330 component_id: component.id.clone(),
331 success: final_success,
332 message: completion_msg,
333 });
334
335 let overall_progress = index as f32 / total;
336 let _ = sender.send(InstallerEvent::Progress {
337 component_id: "__overall__".into(),
338 progress: overall_progress,
339 message: format!(
340 "Finished {} ({:.0}%)",
341 component.name,
342 overall_progress * 100.0
343 ),
344 });
345 }
346
347 if overall_success && config.star_repos {
348 let _ = star_repo(&config, &sender);
349 }
350
351 if std::path::Path::new("/tmp/mlstack-reboot-required").exists() {
353 let _ = sender.send(InstallerEvent::Log(
354 "⚠ REBOOT REQUIRED: ROCm packages were reinstalled. A reboot is needed for kernel module changes to take effect.".to_string(),
355 false,
356 ));
357 let _ = std::fs::remove_file("/tmp/mlstack-reboot-required");
358 }
359
360 {
365 use crate::installers::common::fix_pip_cache_ownership;
366 let result = fix_pip_cache_ownership(&user_home);
367 match &result {
368 crate::installers::common::PipCacheFixResult::Fixed => {
369 let _ = sender.send(InstallerEvent::Log(
370 "[post-install] Fixed pip cache directory ownership (~/.cache/pip)".into(),
371 false,
372 ));
373 }
374 crate::installers::common::PipCacheFixResult::FixFailed(err) => {
375 let _ = sender.send(InstallerEvent::Log(
376 format!(
377 "[post-install] Warning: could not fix pip cache ownership: {}",
378 err
379 ),
380 false,
381 ));
382 }
383 _ => {
384 }
386 }
387 }
388
389 let _ = sender.send(InstallerEvent::Finished {
390 success: overall_success,
391 });
392}
393
394fn validate_sudo(password: String) -> Result<()> {
395 let mut child = Command::new("sudo")
396 .arg("-S")
397 .arg("-v")
398 .stdin(Stdio::piped())
399 .stdout(Stdio::null())
400 .stderr(Stdio::piped())
401 .spawn()
402 .context("Failed to spawn sudo")?;
403
404 let mut stdin = child.stdin.take().context("Failed to open stdin")?;
405 stdin.write_all(password.as_bytes())?;
406 stdin.write_all(b"\n")?;
407 stdin.flush()?;
408 drop(stdin);
409
410 let output = child.wait_with_output()?;
411 if output.status.success() {
412 return Ok(());
413 }
414
415 let stderr = String::from_utf8_lossy(&output.stderr);
416 let detail = stderr.lines().next().unwrap_or("sudo validation failed");
417 anyhow::bail!(detail.to_string())
418}
419
420#[derive(Debug, Clone, Copy, PartialEq, Eq)]
421enum EnvUpdate {
422 Created,
423 Updated,
424 Unchanged,
425}
426
427struct VerificationCheckResult {
428 status: VerificationResult,
429 output: Vec<String>,
430}
431
432#[derive(Debug, Clone)]
433struct VerificationOutcome {
434 success: bool,
435 report_lines: Vec<String>,
436}
437
438fn fix_libdrm_amdgpu_ids() {
448 let target = Path::new("/usr/share/libdrm/amdgpu.ids");
449 let link = Path::new("/opt/amdgpu/share/libdrm/amdgpu.ids");
450
451 if !target.exists() {
452 return;
453 }
454 if link.exists() {
455 return; }
457
458 if let Some(parent) = link.parent() {
460 if fs::create_dir_all(parent).is_ok() {
461 #[cfg(unix)]
462 {
463 let relative = compute_relative_symlink(link, target);
464 if std::os::unix::fs::symlink(&relative, link).is_ok() {
465 tracing::info!(
466 "Created libdrm symlink: {} -> {}",
467 link.display(),
468 relative.display()
469 );
470 return;
471 }
472 }
473 #[cfg(not(unix))]
474 {
475 let _ = (link, target);
476 }
477 }
478 }
479
480 let result = Command::new("sudo")
482 .args(["mkdir", "-p", "/opt/amdgpu/share/libdrm/"])
483 .stdin(std::process::Stdio::null())
484 .output();
485
486 if let Ok(out) = result {
487 if out.status.success() {
488 let result2 = Command::new("sudo")
489 .args([
490 "ln",
491 "-sf",
492 "/usr/share/libdrm/amdgpu.ids",
493 "/opt/amdgpu/share/libdrm/amdgpu.ids",
494 ])
495 .stdin(std::process::Stdio::null())
496 .output();
497 if let Ok(out2) = result2 {
498 if out2.status.success() {
499 tracing::info!(
500 "Created libdrm symlink via sudo: {} -> {}",
501 link.display(),
502 target.display()
503 );
504 return;
505 }
506 }
507 }
508 }
509
510 tracing::warn!(
512 "Cannot create /opt/amdgpu/share/libdrm/amdgpu.ids symlink (need sudo). \
513 Run manually: sudo mkdir -p /opt/amdgpu/share/libdrm && \
514 sudo ln -sf /usr/share/libdrm/amdgpu.ids /opt/amdgpu/share/libdrm/amdgpu.ids"
515 );
516}
517
518fn compute_relative_symlink(link: &Path, target: &Path) -> PathBuf {
520 let mut ups = 0u32;
522 for _ in link.parent().unwrap_or(Path::new(".")).ancestors().skip(1) {
523 ups += 1;
524 }
525 let mut result = PathBuf::new();
526 for _ in 0..ups.saturating_sub(1) {
527 result.push("..");
528 }
529 result.push(target.strip_prefix("/").unwrap_or(target));
530 result
531}
532
533fn ensure_mlstack_env(user_home: &str, install_method: &str) -> Result<EnvUpdate> {
534 let env_path = PathBuf::from(user_home).join(".mlstack_env");
535 let normalized_install_method = match install_method.trim().to_ascii_lowercase().as_str() {
536 "global" => "global",
537 "venv" => "venv",
538 "auto" => "auto",
539 _ => "auto",
540 };
541 #[allow(clippy::needless_borrow)] let env_exports = load_mlstack_env_exports(&user_home);
543 let persistent_python = env_exports
544 .get("MLSTACK_PYTHON_BIN")
545 .map(|v| v.trim().to_string())
546 .filter(|v| !v.is_empty());
547 let python_bin = persistent_python.unwrap_or_else(resolve_python_bin);
548 if env_path.exists() {
549 let contents = fs::read_to_string(&env_path).context("Failed to read .mlstack_env")?;
550 let rocm_version = detect_rocm_version();
551 let gpu_arch = detect_gpu_arch();
552 let gpu_list = detect_gpu_list();
553 let primary_gpu = first_gpu_index(&gpu_list);
554 let rocm_home = "/opt/rocm";
555 let rocm_lib = "/opt/rocm/lib";
556 let hsa_override_val =
558 hsa_override_from_gpu_arch(&gpu_arch).unwrap_or_else(|| "11.0.0".to_string());
559 let defaults = [
560 ("ROCM_VERSION", rocm_version.as_str()),
561 ("ROCM_CHANNEL", "latest"),
562 ("GPU_ARCH", gpu_arch.as_str()),
563 ("PYTORCH_ROCM_ARCH", gpu_arch.as_str()),
564 ("GPU_ARCHS", gpu_arch.as_str()),
565 ("ROCM_HOME", rocm_home),
566 ("ROCM_PATH", rocm_home),
567 ("HIP_PATH", rocm_home),
568 ("HSA_OVERRIDE_GFX_VERSION", hsa_override_val.as_str()),
569 ("HIP_VISIBLE_DEVICES", gpu_list.as_str()),
570 ("CUDA_VISIBLE_DEVICES", gpu_list.as_str()),
571 ("PYTORCH_ROCM_DEVICE", primary_gpu.as_str()),
572 ("MLSTACK_PYTHON_BIN", python_bin.as_str()),
573 ("UV_PYTHON", python_bin.as_str()),
574 ("MLSTACK_INSTALL_METHOD", normalized_install_method),
575 ("INSTALL_METHOD", normalized_install_method),
576 ("PYTHONPATH", rocm_lib),
577 ];
578 let (updated, changed) = sanitize_mlstack_env(
579 &contents,
580 &defaults,
581 gpu_list.as_str(),
582 rocm_home,
583 python_bin.as_str(),
584 normalized_install_method,
585 hsa_override_val.as_str(),
586 );
587 if changed {
588 fs::write(&env_path, updated).context("Failed to update .mlstack_env")?;
589 return Ok(EnvUpdate::Updated);
590 }
591 return Ok(EnvUpdate::Unchanged);
592 }
593
594 if let Some(parent) = env_path.parent() {
595 fs::create_dir_all(parent).context("Failed to create home directory")?;
596 }
597
598 let rocm_version = detect_rocm_version();
599 let gpu_arch = detect_gpu_arch();
600 let gpu_list = detect_gpu_list();
601 let primary_gpu = first_gpu_index(&gpu_list);
602 let rocm_home = "/opt/rocm";
603 let rocm_lib = "/opt/rocm/lib";
604 let hsa_override =
605 hsa_override_from_gpu_arch(&gpu_arch).unwrap_or_else(|| "11.0.0".to_string());
606
607 let content = format!(
608 "# ML Stack Environment File (generated by Rusty-Stack)\n\
609export ROCM_VERSION={}\n\
610export ROCM_CHANNEL=latest\n\
611export GPU_ARCH={}\n\
612export PYTORCH_ROCM_ARCH={}\n\
613export GPU_ARCHS={}\n\
614export ROCM_HOME={}\n\
615export ROCM_PATH={}\n\
616export HIP_PATH={}\n\
617export HSA_OVERRIDE_GFX_VERSION={}\n\
618export HIP_VISIBLE_DEVICES={}\n\
619export CUDA_VISIBLE_DEVICES={}\n\
620export PYTORCH_ROCM_DEVICE={}\n\
621export MLSTACK_PYTHON_BIN={}\n\
622export UV_PYTHON={}\n\
623export MLSTACK_INSTALL_METHOD={}\n\
624export INSTALL_METHOD={}\n\
625export PYTHONPATH={}:$PYTHONPATH\n\
626export TORCH_CUDA_ARCH_LIST=\"7.0;8.0;9.0\"\n\
627export PYTORCH_CUDA_ALLOC_CONF=\"max_split_size_mb:512\"\n\
628export PATH=\"/usr/local/bin:/usr/bin:/bin:/usr/local/games:/usr/games:{}/bin:{}/hip/bin:$PATH\"\n\
629export LD_LIBRARY_PATH=\"$HOME/.mlstack/libmpi-compat:$HOME/.mlstack/libmpi-compat-user-$(id -u):{}/lib:{}/hip/lib:{}/opencl/lib:$LD_LIBRARY_PATH\"\n",
630 rocm_version,
631 gpu_arch,
632 gpu_arch,
633 gpu_arch,
634 rocm_home,
635 rocm_home,
636 rocm_home,
637 hsa_override,
638 gpu_list,
639 gpu_list,
640 primary_gpu,
641 python_bin,
642 python_bin,
643 normalized_install_method,
644 normalized_install_method,
645 rocm_lib,
646 rocm_home,
647 rocm_home,
648 rocm_home,
649 rocm_home,
650 rocm_home
651 );
652
653 fs::write(&env_path, content).context("Failed to create .mlstack_env")?;
654
655 fix_libdrm_amdgpu_ids();
659
660 Ok(EnvUpdate::Created)
661}
662
663fn sanitize_mlstack_env(
664 contents: &str,
665 defaults: &[(&str, &str)],
666 gpu_list: &str,
667 rocm_home: &str,
668 python_bin: &str,
669 install_method: &str,
670 hsa_override: &str,
671) -> (String, bool) {
672 let mut changed = false;
673 let mut lines = Vec::new();
674 let mut seen = std::collections::HashSet::new();
675
676 for line in contents.lines() {
677 let trimmed = line.trim_start();
678 if let Some(key) = extract_env_key(trimmed) {
679 seen.insert(key.to_string());
680 }
681 let (line, changed_arch) = fix_env_assignment(line, "TORCH_CUDA_ARCH_LIST");
682 let (line, changed_alloc) = fix_env_assignment(&line, "PYTORCH_CUDA_ALLOC_CONF");
683 let (line, changed_hip) = normalize_env_value(&line, "HIP_PATH", "/opt/rocm");
684 let (line, changed_rocm) = normalize_env_value(&line, "ROCM_PATH", "/opt/rocm");
685 let (line, changed_python) = normalize_env_value(&line, "MLSTACK_PYTHON_BIN", python_bin);
686 let (line, changed_uv_python) = normalize_env_value(&line, "UV_PYTHON", python_bin);
687 let (line, changed_method) =
688 normalize_env_value(&line, "MLSTACK_INSTALL_METHOD", install_method);
689 let (line, changed_method_alias) =
690 normalize_env_value(&line, "INSTALL_METHOD", install_method);
691 let (line, changed_hsa) =
692 normalize_env_value(&line, "HSA_OVERRIDE_GFX_VERSION", hsa_override);
693 let (line, changed_hip_vis) =
694 normalize_visible_devices(&line, "HIP_VISIBLE_DEVICES", gpu_list);
695 let (line, changed_cuda_vis) =
696 normalize_visible_devices(&line, "CUDA_VISIBLE_DEVICES", gpu_list);
697 let (line, changed_pythonpath) = if line.starts_with("export PYTHONPATH=") {
698 let desired = format!("export PYTHONPATH={}/lib:$PYTHONPATH", rocm_home);
699 if line.trim() != desired {
700 (desired, true)
701 } else {
702 (line, false)
703 }
704 } else {
705 (line, false)
706 };
707
708 let (line, changed_path) = if line.starts_with("export PATH=") {
710 let desired = format!("export PATH=\"/usr/local/bin:/usr/bin:/bin:/usr/local/games:/usr/games:{}/bin:{}/hip/bin:$PATH\"", rocm_home, rocm_home);
711 if line.trim() != desired {
712 (desired, true)
713 } else {
714 (line, false)
715 }
716 } else {
717 (line, false)
718 };
719
720 let (line, changed_ld) = if line.starts_with("export LD_LIBRARY_PATH=") {
721 let desired = format!(
722 "export LD_LIBRARY_PATH=\"$HOME/.mlstack/libmpi-compat:$HOME/.mlstack/libmpi-compat-user-$(id -u):{}/lib:{}/hip/lib:{}/opencl/lib:$LD_LIBRARY_PATH\"",
723 rocm_home, rocm_home, rocm_home
724 );
725 if line.trim() != desired {
726 (desired, true)
727 } else {
728 (line, false)
729 }
730 } else {
731 (line, false)
732 };
733
734 changed |= changed_arch
735 || changed_alloc
736 || changed_hip
737 || changed_rocm
738 || changed_python
739 || changed_uv_python
740 || changed_method
741 || changed_method_alias
742 || changed_hsa
743 || changed_hip_vis
744 || changed_cuda_vis
745 || changed_pythonpath
746 || changed_path
747 || changed_ld;
748 lines.push(line);
749 }
750
751 for (key, value) in defaults {
752 if !seen.contains(*key) {
753 lines.push(format!("export {}={}", key, value));
754 changed = true;
755 }
756 }
757
758 let mut output = lines.join("\n");
759 output.push('\n');
760 (output, changed)
761}
762
763fn extract_env_key(line: &str) -> Option<&str> {
764 let export_idx = line.rfind("export ")?;
766 let rest = &line[export_idx + 7..]; rest.split('=').next().map(str::trim)
768}
769
770fn fix_env_assignment(line: &str, key: &str) -> (String, bool) {
771 let marker = format!("{}=", key);
772 let Some(idx) = line.find(&marker) else {
773 return (line.to_string(), false);
774 };
775
776 let before = &line[..idx + marker.len()];
777 let after = &line[idx + marker.len()..];
778
779 if after.starts_with('"') {
780 return (line.to_string(), false);
781 }
782
783 let clean_after = strip_shell_suffixes_from_value(after);
785 (format!("{}\"{}\"", before, clean_after), true)
786}
787
788fn normalize_env_value(line: &str, key: &str, desired: &str) -> (String, bool) {
789 let marker = format!("export {}=", key);
790 if let Some(export_idx) = line.rfind(&marker) {
792 let after_marker = &line[export_idx + marker.len()..];
793 let current = after_marker.trim().trim_matches('"');
794 let current = strip_shell_suffixes_from_value(current);
796 if current != desired {
797 return (format!("export {}={}", key, desired), true);
799 }
800 }
801 (line.to_string(), false)
802}
803
804fn normalize_visible_devices(line: &str, key: &str, desired: &str) -> (String, bool) {
805 let desired = desired.trim();
806 if desired.is_empty() {
807 return (line.to_string(), false);
808 }
809
810 let marker = format!("export {}=", key);
811 if let Some(export_idx) = line.rfind(&marker) {
813 let after_marker = &line[export_idx + marker.len()..];
814 let current = after_marker.trim().trim_matches('"');
815 let current = strip_shell_suffixes_from_value(current);
817
818 if current != desired {
824 return (format!("export {}={}", key, desired), true);
826 }
827 }
828 (line.to_string(), false)
829}
830
831fn detect_rocm_version() -> String {
832 if let Ok(version) = fs::read_to_string("/opt/rocm/.info/version") {
833 return version.trim().to_string();
834 }
835 "7.2.0".to_string()
836}
837
838fn run_verification(
839 component: &Component,
840 python_candidates: &[String],
841 sender: &Sender<InstallerEvent>,
842 user_home: &str,
843) -> VerificationOutcome {
844 let steps = verification_commands(&component.id, python_candidates);
845 let mut results: HashMap<String, VerificationCheckResult> = HashMap::new();
846 let mut executed = 0usize;
847 let mut failed = 0usize;
848
849 for step in &steps {
850 if let Some((benchmark_ok, benchmark_output)) = verify_benchmark_target(&step.target_id) {
851 executed += 1;
852 if benchmark_ok {
853 results.insert(
854 step.target_id.clone(),
855 VerificationCheckResult {
856 status: VerificationResult::Verified,
857 output: benchmark_output,
858 },
859 );
860 } else {
861 failed += 1;
862 results.insert(
863 step.target_id.clone(),
864 VerificationCheckResult {
865 status: VerificationResult::Failed,
866 output: benchmark_output,
867 },
868 );
869 let _ = sender.send(InstallerEvent::Log(
870 format!("Verification failed for {}", step.label),
871 false,
872 ));
873 }
874 continue;
875 }
876
877 if !is_component_installed_by_id(&step.target_id, python_candidates) {
878 results.insert(
879 step.target_id.clone(),
880 VerificationCheckResult {
881 status: VerificationResult::Missing,
882 output: vec!["Component not found in path or site-packages".to_string()],
883 },
884 );
885 let _ = sender.send(InstallerEvent::Log(
886 format!("Skipping {} (not installed)", step.label),
887 false,
888 ));
889 continue;
890 }
891 if !modules_available(&step.modules, python_candidates) {
892 results.insert(
893 step.target_id.clone(),
894 VerificationCheckResult {
895 status: VerificationResult::Missing,
896 output: vec![format!(
897 "Required Python modules missing: {}",
898 step.modules.join(", ")
899 )],
900 },
901 );
902 let _ = sender.send(InstallerEvent::Log(
903 format!("Skipping {} (module not available)", step.label),
904 false,
905 ));
906 continue;
907 }
908 executed += 1;
909 let _ = sender.send(InstallerEvent::Log(
910 format!("Verifying {}...", step.label),
911 false,
912 ));
913 match run_verification_command(step, sender) {
914 Ok((success, output)) => {
915 if success {
916 results.insert(
917 step.target_id.clone(),
918 VerificationCheckResult {
919 status: VerificationResult::Verified,
920 output,
921 },
922 );
923 } else {
924 failed += 1;
925 results.insert(
926 step.target_id.clone(),
927 VerificationCheckResult {
928 status: VerificationResult::Failed,
929 output,
930 },
931 );
932 let _ = sender.send(InstallerEvent::Log(
933 format!("Verification failed for {}", step.label),
934 false,
935 ));
936 }
937 }
938 Err(err) => {
939 failed += 1;
940 results.insert(
941 step.target_id.clone(),
942 VerificationCheckResult {
943 status: VerificationResult::Failed,
944 output: vec![err.to_string()],
945 },
946 );
947 let _ = sender.send(InstallerEvent::Log(
948 format!("Failed to execute verification for {}: {}", step.label, err),
949 false,
950 ));
951 }
952 }
953 }
954
955 let report_lines = build_verification_report(python_candidates, user_home, &results, &steps);
956 let success = executed > 0 && failed == 0;
957
958 if executed == 0 {
959 let _ = sender.send(InstallerEvent::Log(
960 "No installed components found to verify".into(),
961 false,
962 ));
963 }
964
965 if failed > 0 {
966 let _ = sender.send(InstallerEvent::Log(
967 format!("Verification completed with {} failed checks", failed),
968 false,
969 ));
970 }
971
972 VerificationOutcome {
973 success,
974 report_lines,
975 }
976}
977
978fn build_verification_report(
979 python_candidates: &[String],
980 user_home: &str,
981 results: &HashMap<String, VerificationCheckResult>,
982 steps: &[VerificationCommand],
983) -> Vec<String> {
984 let mut lines = Vec::new();
985 lines.push("Environment:".to_string());
986 for (key, value) in collect_env_info(user_home) {
987 lines.push(format!(" {}: {}", key, value));
988 }
989
990 lines.push("".to_string());
991 lines.push("Python Interpreters:".to_string());
992 for py in python_candidates {
993 lines.push(format!(" - {}", py));
994 }
995
996 lines.push("".to_string());
997 lines.push("Checks:".to_string());
998 for step in steps {
999 if let Some(res) = results.get(&step.target_id) {
1000 lines.push(format!(" - {}: {}", step.label, res.status.label()));
1001 for out_line in &res.output {
1002 lines.push(format!(" > {}", out_line));
1003 }
1004 } else {
1005 lines.push(format!(" - {}: Skipped", step.label));
1006 }
1007 }
1008
1009 lines
1010}
1011
1012fn run_component_verification(
1013 component: &Component,
1014 python_candidates: &[String],
1015 sender: &Sender<InstallerEvent>,
1016 user_home: &str,
1017) -> VerificationOutcome {
1018 let steps = component_verification_commands(&component.id, python_candidates);
1019 let mut results: HashMap<String, VerificationCheckResult> = HashMap::new();
1020 let mut executed = 0usize;
1021 let mut failed = 0usize;
1022
1023 for step in &steps {
1024 if let Some((benchmark_ok, benchmark_output)) = verify_benchmark_target(&step.target_id) {
1025 executed += 1;
1026 if benchmark_ok {
1027 results.insert(
1028 step.target_id.clone(),
1029 VerificationCheckResult {
1030 status: VerificationResult::Verified,
1031 output: benchmark_output,
1032 },
1033 );
1034 } else {
1035 failed += 1;
1036 results.insert(
1037 step.target_id.clone(),
1038 VerificationCheckResult {
1039 status: VerificationResult::Failed,
1040 output: benchmark_output,
1041 },
1042 );
1043 }
1044 continue;
1045 }
1046
1047 if !modules_available(&step.modules, python_candidates) {
1048 results.insert(
1049 step.target_id.clone(),
1050 VerificationCheckResult {
1051 status: VerificationResult::Missing,
1052 output: vec![format!(
1053 "Required Python modules missing: {}",
1054 step.modules.join(", ")
1055 )],
1056 },
1057 );
1058 continue;
1059 }
1060 executed += 1;
1061 match run_verification_command(step, sender) {
1062 Ok((success, output)) => {
1063 if success {
1064 results.insert(
1065 step.target_id.clone(),
1066 VerificationCheckResult {
1067 status: VerificationResult::Verified,
1068 output,
1069 },
1070 );
1071 } else {
1072 failed += 1;
1073 results.insert(
1074 step.target_id.clone(),
1075 VerificationCheckResult {
1076 status: VerificationResult::Failed,
1077 output,
1078 },
1079 );
1080 }
1081 }
1082 Err(err) => {
1083 failed += 1;
1084 results.insert(
1085 step.target_id.clone(),
1086 VerificationCheckResult {
1087 status: VerificationResult::Failed,
1088 output: vec![err.to_string()],
1089 },
1090 );
1091 }
1092 }
1093 }
1094
1095 let report_lines = build_component_report(component, user_home, &results, &steps);
1096 let success = executed > 0 && failed == 0;
1097
1098 VerificationOutcome {
1099 success,
1100 report_lines,
1101 }
1102}
1103
1104fn build_component_report(
1105 component: &Component,
1106 user_home: &str,
1107 results: &HashMap<String, VerificationCheckResult>,
1108 steps: &[VerificationCommand],
1109) -> Vec<String> {
1110 let mut lines = Vec::new();
1111 lines.push(format!("Component: {}", component.name));
1112
1113 for (key, value) in collect_env_info(user_home) {
1114 lines.push(format!(" {}: {}", key, value));
1115 }
1116
1117 if let Some(benchmark_lines) = extract_benchmark_results(&component.id) {
1119 lines.push("".to_string());
1120 lines.push("Benchmark Results:".to_string());
1121 lines.extend(benchmark_lines);
1122 }
1123
1124 lines.push("".to_string());
1125 lines.push("Verification Checks:".to_string());
1126
1127 for step in steps {
1128 if let Some(res) = results.get(&step.target_id) {
1129 lines.push(format!(" - {}: {}", step.label, res.status.label()));
1130 for out_line in &res.output {
1131 lines.push(format!(" > {}", out_line));
1132 }
1133 } else {
1134 lines.push(format!(" - {}: Skipped", step.label));
1135 }
1136 }
1137
1138 lines
1139}
1140
1141fn extract_benchmark_results(component_id: &str) -> Option<Vec<String>> {
1142 let log_dirs = benchmark_log_directories();
1143 if log_dirs.is_empty() {
1144 return None;
1145 }
1146
1147 let pattern = match component_id {
1149 "mlperf-inference" => "mlperf_inference",
1150 "rocm-benchmarks" => "rocm_benchmarks",
1151 "gpu-memory-bandwidth" => "gpu_memory_bandwidth",
1152 "rocm-smi-bench" => "rocm_smi_benchmarks",
1153 "pytorch-performance" => "pytorch_performance",
1154 "vllm-performance" => "vllm_benchmarks",
1155 "deepspeed-performance" => "deepspeed_benchmarks",
1156 "megatron-performance" => "megatron_benchmarks",
1157 "all-benchmarks" => "full_benchmarks",
1158 _ => return None,
1159 };
1160
1161 let log_file = find_latest_log_in_dirs(&log_dirs, pattern)?;
1162 let contents = std::fs::read_to_string(log_file).ok()?;
1163
1164 let mut result_lines: Vec<String> = Vec::new();
1166
1167 let Some(json) = extract_benchmark_json_value(&contents) else {
1168 result_lines.push(" No metrics found in log".to_string());
1169 return Some(result_lines);
1170 };
1171
1172 if let Some(results) = json.get("results").and_then(|r| r.as_object()) {
1174 for (bench_name, bench_data) in results {
1175 if bench_data.is_object() {
1176 let status = bench_data
1177 .get("success")
1178 .and_then(|v| v.as_bool())
1179 .map(|ok| if ok { "SUCCESS" } else { "FAILED" })
1180 .unwrap_or("UNKNOWN");
1181 let time_ms = bench_data
1182 .get("execution_time_ms")
1183 .and_then(|v| v.as_u64())
1184 .unwrap_or(0);
1185 result_lines.push(format!(
1186 " === {} ({}, {} ms) ===",
1187 bench_name.replace('_', " ").to_uppercase(),
1188 status,
1189 time_ms
1190 ));
1191 if let Some(metrics) = bench_data.get("metrics").and_then(|m| m.as_object()) {
1192 for (key, value) in metrics {
1193 let value_str = match value {
1194 serde_json::Value::Number(n) => n.to_string(),
1195 serde_json::Value::String(s) => s.clone(),
1196 serde_json::Value::Bool(b) => b.to_string(),
1197 serde_json::Value::Array(arr) => format!("{} values", arr.len()),
1198 _ => continue,
1199 };
1200 let formatted_key = key.replace('_', " ").to_uppercase();
1201 result_lines.push(format!(" {}: {}", formatted_key, value_str));
1202 }
1203 }
1204 if let Some(errors) = bench_data.get("errors").and_then(|v| v.as_array()) {
1205 for err in errors.iter().filter_map(|v| v.as_str()).take(3) {
1206 result_lines.push(format!(" ERROR: {}", err));
1207 }
1208 }
1209 result_lines.push(String::new());
1210 }
1211 }
1212 } else if let Some(metrics) = json.get("metrics").and_then(|m| m.as_object()) {
1213 for (key, value) in metrics {
1215 let value_str = match value {
1216 serde_json::Value::Number(n) => n.to_string(),
1217 serde_json::Value::String(s) => s.clone(),
1218 serde_json::Value::Bool(b) => b.to_string(),
1219 serde_json::Value::Array(arr) => format!("{} values", arr.len()),
1220 _ => continue,
1221 };
1222 let formatted_key = key.replace('_', " ").to_uppercase();
1223 result_lines.push(format!(" {}: {}", formatted_key, value_str));
1224 }
1225 }
1226
1227 if result_lines.is_empty() {
1228 result_lines.push(" No metrics found in log".to_string());
1229 }
1230
1231 Some(result_lines)
1232}
1233
1234fn benchmark_pattern_for_target(target_id: &str) -> Option<&'static str> {
1235 match target_id {
1236 "mlperf-inference" => Some("mlperf_inference"),
1237 "rocm-benchmarks" => Some("rocm_benchmarks"),
1238 "gpu-memory-bandwidth" => Some("gpu_memory_bandwidth"),
1239 "rocm-smi-bench" => Some("rocm_smi_benchmarks"),
1240 "pytorch-performance" => Some("pytorch_performance"),
1241 "vllm-performance" => Some("vllm_benchmarks"),
1242 "deepspeed-performance" => Some("deepspeed_benchmarks"),
1243 "megatron-performance" => Some("megatron_benchmarks"),
1244 "all-benchmarks" => Some("full_benchmarks"),
1245 _ => None,
1246 }
1247}
1248
1249fn benchmark_success_and_errors(value: &serde_json::Value) -> (bool, Vec<String>) {
1250 if let Some(success) = value.get("success").and_then(|v| v.as_bool()) {
1251 let mut details = Vec::new();
1252 if !success {
1253 if let Some(reason) = value.get("reason").and_then(|v| v.as_str()) {
1254 let reason = reason.trim();
1255 if !reason.is_empty() {
1256 details.push(reason.to_string());
1257 }
1258 }
1259 if let Some(errors) = value.get("errors").and_then(|v| v.as_array()) {
1260 for err in errors.iter().filter_map(|entry| entry.as_str()).take(4) {
1261 let trimmed = err.trim();
1262 if !trimmed.is_empty() {
1263 details.push(trimmed.to_string());
1264 }
1265 }
1266 }
1267 if details.is_empty() {
1268 details.push("Benchmark JSON reported success=false".to_string());
1269 }
1270 }
1271 return (success, details);
1272 }
1273
1274 if let Some(results) = value.get("results").and_then(|v| v.as_object()) {
1275 let mut failed = Vec::new();
1276 let mut seen_result = false;
1277 for (name, item) in results {
1278 if let Some(item_success) = item.get("success").and_then(|v| v.as_bool()) {
1279 seen_result = true;
1280 if !item_success {
1281 let mut reason = item
1282 .get("reason")
1283 .and_then(|v| v.as_str())
1284 .map(|s| s.trim().to_string())
1285 .filter(|s| !s.is_empty())
1286 .unwrap_or_else(|| format!("{name}: success=false"));
1287 if reason == format!("{name}: success=false") {
1288 if let Some(errors) = item.get("errors").and_then(|v| v.as_array()) {
1289 if let Some(first) = errors.iter().filter_map(|v| v.as_str()).next() {
1290 let first = first.trim();
1291 if !first.is_empty() {
1292 reason = format!("{name}: {first}");
1293 }
1294 }
1295 }
1296 }
1297 failed.push(reason);
1298 }
1299 }
1300 }
1301 if seen_result {
1302 return (failed.is_empty(), failed);
1303 }
1304 }
1305
1306 (
1307 false,
1308 vec!["Benchmark JSON does not contain a recognizable success field".to_string()],
1309 )
1310}
1311
1312fn verify_benchmark_target(target_id: &str) -> Option<(bool, Vec<String>)> {
1313 let pattern = benchmark_pattern_for_target(target_id)?;
1314 let log_dirs = benchmark_log_directories();
1315 if log_dirs.is_empty() {
1316 return Some((
1317 false,
1318 vec!["No benchmark log directories found".to_string()],
1319 ));
1320 }
1321
1322 let Some(log_file) = find_latest_log_in_dirs(&log_dirs, pattern) else {
1323 return Some((
1324 false,
1325 vec![format!("No benchmark logs found for pattern '{pattern}'")],
1326 ));
1327 };
1328
1329 let mut output = vec![format!("Latest benchmark log: {}", log_file.display())];
1330 let contents = match std::fs::read_to_string(&log_file) {
1331 Ok(contents) => contents,
1332 Err(err) => {
1333 output.push(format!("Could not read benchmark log: {err}"));
1334 return Some((false, output));
1335 }
1336 };
1337
1338 let Some(json) = extract_benchmark_json_value(&contents) else {
1339 output.push("Could not parse benchmark JSON payload from log".to_string());
1340 return Some((false, output));
1341 };
1342
1343 let (success, mut details) = benchmark_success_and_errors(&json);
1344 output.append(&mut details);
1345 Some((success, output))
1346}
1347
1348fn collect_env_info(user_home: &str) -> Vec<(String, String)> {
1349 let mut values = load_mlstack_env_exports(user_home);
1350
1351 let keys = [
1352 "ROCM_VERSION",
1353 "ROCM_CHANNEL",
1354 "GPU_ARCH",
1355 "GPU_ARCHS",
1356 "PYTORCH_ROCM_ARCH",
1357 "HIP_VISIBLE_DEVICES",
1358 "CUDA_VISIBLE_DEVICES",
1359 "PYTORCH_ROCM_DEVICE",
1360 "HSA_OVERRIDE_GFX_VERSION",
1361 "ROCM_HOME",
1362 "ROCM_PATH",
1363 "HIP_PATH",
1364 "MLSTACK_PYTHON_BIN",
1365 "PYTHONPATH",
1366 ];
1367
1368 for key in keys {
1369 if !values.contains_key(key) {
1370 if let Ok(value) = std::env::var(key) {
1371 values.insert(key.to_string(), value);
1372 }
1373 }
1374 }
1375
1376 keys.iter()
1377 .map(|key| {
1378 let default_val = if *key == "ROCM_HOME" || *key == "ROCM_PATH" || *key == "HIP_PATH" {
1379 "/opt/rocm".to_string()
1380 } else {
1381 "unknown".to_string()
1382 };
1383 (
1384 key.to_string(),
1385 values.get(*key).cloned().unwrap_or(default_val),
1386 )
1387 })
1388 .collect()
1389}
1390
1391fn load_mlstack_env_exports(user_home: &str) -> HashMap<String, String> {
1392 let mut values: HashMap<String, String> = HashMap::new();
1393 for home in candidate_mlstack_homes(user_home) {
1394 let env_path = PathBuf::from(home).join(".mlstack_env");
1395 if let Ok(contents) = fs::read_to_string(&env_path) {
1396 for line in contents.lines() {
1397 if let Some((key, value)) = parse_env_export(line) {
1398 values.entry(key).or_insert(value);
1399 }
1400 }
1401 }
1402 }
1403
1404 values
1405}
1406
1407fn parse_env_export(line: &str) -> Option<(String, String)> {
1408 let export_idx = line.rfind("export ")?;
1409 let rest = &line[export_idx + 7..];
1410 let mut parts = rest.splitn(2, '=');
1411 let key = parts.next()?.trim().to_string();
1412 let value = parts.next()?.trim().trim_matches('"').to_string();
1413 let value = strip_shell_suffixes_from_value(&value);
1415 Some((key, value))
1416}
1417
1418fn strip_shell_suffixes_from_value(value: &str) -> String {
1426 let mut v = value.to_string();
1427 if let Some(idx) = v.find("; fi") {
1429 v.truncate(idx);
1430 }
1431 if let Some(idx) = v.find("&& ") {
1433 v.truncate(idx);
1434 }
1435 v.trim().to_string()
1436}
1437
1438fn detect_onnx_rocm_provider(python: &str) -> Option<String> {
1439 let output = Command::new(python)
1440 .arg("-c")
1441 .arg(
1442 "import onnxruntime, pathlib; base=pathlib.Path(onnxruntime.__file__).parent; libs=list(base.rglob('libonnxruntime_providers_rocm.so')); print(libs[0] if libs else '')",
1443 )
1444 .output()
1445 .ok()?;
1446 if !output.status.success() {
1447 return None;
1448 }
1449 let value = String::from_utf8_lossy(&output.stdout).trim().to_string();
1450 if value.is_empty() {
1451 None
1452 } else {
1453 Some(value)
1454 }
1455}
1456
1457fn push_unique_env_path(paths: &mut Vec<String>, value: &str) {
1458 let candidate = value.trim();
1459 if candidate.is_empty() {
1460 return;
1461 }
1462 if paths.iter().any(|existing| existing == candidate) {
1463 return;
1464 }
1465 paths.push(candidate.to_string());
1466}
1467
1468fn user_home_from_passwd(user_name: &str) -> Option<String> {
1469 let target = user_name.trim();
1470 if target.is_empty() {
1471 return None;
1472 }
1473 let passwd = fs::read_to_string("/etc/passwd").ok()?;
1474 passwd.lines().find_map(|line| {
1475 if line.trim().is_empty() || line.starts_with('#') {
1476 return None;
1477 }
1478 let fields: Vec<&str> = line.split(':').collect();
1479 if fields.len() < 6 || fields[0] != target {
1480 return None;
1481 }
1482 let home = fields[5].trim();
1483 if home.is_empty() {
1484 None
1485 } else {
1486 Some(home.to_string())
1487 }
1488 })
1489}
1490
1491fn passwd_homes_with_mlstack() -> Vec<String> {
1492 let mut homes = Vec::new();
1493 let Ok(passwd) = fs::read_to_string("/etc/passwd") else {
1494 return homes;
1495 };
1496
1497 for line in passwd.lines() {
1498 if line.trim().is_empty() || line.starts_with('#') {
1499 continue;
1500 }
1501 let fields: Vec<&str> = line.split(':').collect();
1502 if fields.len() < 6 {
1503 continue;
1504 }
1505 let home = fields[5].trim();
1506 if home.is_empty() {
1507 continue;
1508 }
1509 if Path::new(home).join(".mlstack").is_dir()
1510 || Path::new(home).join(".mlstack_env").is_file()
1511 {
1512 push_unique_env_path(&mut homes, home);
1513 }
1514 }
1515
1516 homes
1517}
1518
1519fn candidate_mlstack_homes(preferred_home: &str) -> Vec<String> {
1520 let mut homes = Vec::new();
1521 push_unique_env_path(&mut homes, preferred_home);
1522
1523 if let Ok(value) = std::env::var("MLSTACK_USER_HOME") {
1524 push_unique_env_path(&mut homes, &value);
1525 }
1526 if let Ok(value) = std::env::var("HOME") {
1527 push_unique_env_path(&mut homes, &value);
1528 }
1529
1530 for key in ["SUDO_USER", "USER", "LOGNAME"] {
1531 if let Ok(user_name) = std::env::var(key) {
1532 if let Some(home) = user_home_from_passwd(&user_name) {
1533 push_unique_env_path(&mut homes, &home);
1534 }
1535 }
1536 }
1537
1538 for home in passwd_homes_with_mlstack() {
1539 push_unique_env_path(&mut homes, &home);
1540 }
1541
1542 homes
1543}
1544
1545fn resolve_mlstack_user_home() -> String {
1546 let fallback = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
1547 let candidates = candidate_mlstack_homes(&fallback);
1548
1549 for home in &candidates {
1550 if Path::new(home).join(".mlstack_env").is_file() {
1551 return home.clone();
1552 }
1553 }
1554 for home in &candidates {
1555 if Path::new(home).join(".mlstack").is_dir() {
1556 return home.clone();
1557 }
1558 }
1559
1560 candidates.into_iter().next().unwrap_or(fallback)
1561}
1562
1563fn mlstack_mpi_compat_dirs(user_home: &str) -> Vec<String> {
1564 let mut dirs = Vec::new();
1565 for home in candidate_mlstack_homes(user_home) {
1566 let mlstack_dir = PathBuf::from(home).join(".mlstack");
1567
1568 let primary = mlstack_dir.join("libmpi-compat");
1569 if primary.is_dir() {
1570 dirs.push(primary.to_string_lossy().to_string());
1571 }
1572
1573 if let Ok(entries) = fs::read_dir(&mlstack_dir) {
1574 for entry in entries.flatten() {
1575 let path = entry.path();
1576 if !path.is_dir() {
1577 continue;
1578 }
1579 let Some(name) = path.file_name().and_then(|value| value.to_str()) else {
1580 continue;
1581 };
1582 if name.starts_with("libmpi-compat-user-") {
1583 dirs.push(path.to_string_lossy().to_string());
1584 }
1585 }
1586 }
1587 }
1588
1589 dirs.sort();
1590 dirs.dedup();
1591 dirs
1592}
1593
1594fn build_verification_ld_library_path(user_home: &str) -> String {
1595 let mut paths: Vec<String> = Vec::new();
1596
1597 for compat_dir in mlstack_mpi_compat_dirs(user_home) {
1598 push_unique_env_path(&mut paths, &compat_dir);
1599 }
1600
1601 for rocm_path in ["/opt/rocm/lib", "/opt/rocm/hip/lib", "/opt/rocm/opencl/lib"] {
1602 if Path::new(rocm_path).exists() {
1603 push_unique_env_path(&mut paths, rocm_path);
1604 }
1605 }
1606
1607 if let Ok(existing) = std::env::var("LD_LIBRARY_PATH") {
1608 for part in existing.split(':') {
1609 push_unique_env_path(&mut paths, part);
1610 }
1611 }
1612
1613 paths.join(":")
1614}
1615
1616fn run_verification_command(
1617 step: &VerificationCommand,
1618 sender: &Sender<InstallerEvent>,
1619) -> Result<(bool, Vec<String>)> {
1620 let user_home = resolve_mlstack_user_home();
1621 let user_name = std::env::var("USER").unwrap_or_else(|_| "user".into());
1622 #[allow(clippy::needless_borrow)] let env_exports = load_mlstack_env_exports(&user_home);
1624
1625 let mut command = Command::new(&step.program);
1626 command.args(&step.args);
1627
1628 let extra_paths = python_search_paths();
1629 let mut pythonpath = std::env::var("PYTHONPATH").unwrap_or_default();
1630 if !extra_paths.is_empty() {
1631 if !pythonpath.is_empty() {
1632 pythonpath.push(':');
1633 }
1634 pythonpath.push_str(&extra_paths.join(":"));
1635 }
1636
1637 pythonpath = pythonpath
1639 .split(':')
1640 .filter(|part| !part.starts_with("/tmp/"))
1641 .collect::<Vec<_>>()
1642 .join(":");
1643
1644 if !pythonpath.is_empty() {
1645 command.env("PYTHONPATH", pythonpath);
1646 }
1647
1648 let verification_ld = build_verification_ld_library_path(&user_home);
1649 if !verification_ld.is_empty() {
1650 command.env("LD_LIBRARY_PATH", verification_ld);
1651 }
1652
1653 if step.target_id == "onnx" {
1654 if let Some(lib) = detect_onnx_rocm_provider(&step.program) {
1655 command.env("ORT_ROCM_EP_PROVIDER_PATH", lib);
1656 }
1657 }
1658
1659 let selected_python = if !step.modules.is_empty() {
1660 step.program.clone()
1661 } else {
1662 resolve_python_bin()
1663 };
1664 let gpu_arch = env_exports
1665 .get("GPU_ARCH")
1666 .map(|v| v.trim().to_string())
1667 .filter(|v| !v.is_empty())
1668 .unwrap_or_else(detect_gpu_arch);
1669 let py_rocm_arch = env_exports
1670 .get("PYTORCH_ROCM_ARCH")
1671 .map(|v| v.trim().to_string())
1672 .filter(|v| !v.is_empty())
1673 .unwrap_or_else(|| gpu_arch.clone());
1674 let mut gpu_archs = env_exports
1675 .get("GPU_ARCHS")
1676 .map(|v| v.trim().to_string())
1677 .filter(|v| !v.is_empty())
1678 .unwrap_or_else(|| gpu_arch.clone());
1679 if step.target_id == "aiter" {
1680 gpu_archs = py_rocm_arch
1681 .split(';')
1682 .next()
1683 .map(|v| v.trim().to_string())
1684 .filter(|v| !v.is_empty())
1685 .unwrap_or_else(|| gpu_arch.clone());
1686 }
1687 let mut hip_visible = env_exports
1688 .get("HIP_VISIBLE_DEVICES")
1689 .map(|v| v.trim().to_string())
1690 .filter(|v| !v.is_empty());
1691 let mut cuda_visible = env_exports
1692 .get("CUDA_VISIBLE_DEVICES")
1693 .map(|v| v.trim().to_string())
1694 .filter(|v| !v.is_empty())
1695 .or_else(|| hip_visible.clone());
1696 if step.target_id == "aiter" {
1697 let base_list = hip_visible
1698 .clone()
1699 .or_else(|| cuda_visible.clone())
1700 .unwrap_or_else(detect_gpu_list);
1701 let ordered = prioritize_gpu_list_for_arch(&base_list, &gpu_archs);
1702 hip_visible = Some(ordered.clone());
1703 cuda_visible = Some(ordered);
1704 }
1705 let hsa_override = env_exports
1706 .get("HSA_OVERRIDE_GFX_VERSION")
1707 .map(|v| v.trim().to_string())
1708 .filter(|v| !v.is_empty());
1709
1710 command
1711 .env("HOME", &user_home)
1712 .env("USER", &user_name)
1713 .env("LOGNAME", &user_name)
1714 .env("MLSTACK_USER_HOME", &user_home)
1715 .env("MLSTACK_PYTHON_BIN", &selected_python)
1716 .env("UV_PYTHON", &selected_python)
1717 .env("PYTHONWARNINGS", "ignore")
1718 .env("PYTHONUNBUFFERED", "1")
1719 .env("ROCM_HOME", "/opt/rocm")
1720 .env("HIP_PATH", "/opt/rocm")
1721 .env("GPU_ARCH", &gpu_arch)
1722 .env("PYTORCH_ROCM_ARCH", &py_rocm_arch)
1723 .env("GPU_ARCHS", &gpu_archs)
1724 .stdout(Stdio::piped())
1725 .stderr(Stdio::piped());
1726
1727 if let Some(value) = hip_visible {
1728 command.env("HIP_VISIBLE_DEVICES", value);
1729 }
1730 if let Some(value) = cuda_visible {
1731 command.env("CUDA_VISIBLE_DEVICES", value);
1732 }
1733 if let Some(value) = hsa_override {
1734 command.env("HSA_OVERRIDE_GFX_VERSION", value);
1735 }
1736 if step.target_id == "aiter" {
1737 let aiter_jit_dir = PathBuf::from(&user_home)
1738 .join(".mlstack")
1739 .join("aiter")
1740 .join("jit");
1741 let _ = fs::create_dir_all(&aiter_jit_dir);
1742 command.env("AITER_JIT_DIR", aiter_jit_dir);
1743 }
1744
1745 if step.target_id == "mpi4py" {
1746 command
1748 .env("OMPI_MCA_btl", "^smcuda")
1749 .env("OMPI_MCA_opal_cuda_support", "0")
1750 .env("OMPI_MCA_opal_accelerator_rocm_enable", "0");
1751 }
1752
1753 let output = command
1754 .output()
1755 .with_context(|| format!("Failed to run verification for {}", step.label))?;
1756
1757 let mut captured_output = Vec::new();
1758 const MAX_CAPTURED_OUTPUT_LINES: usize = 24;
1759
1760 let stdout = String::from_utf8_lossy(&output.stdout);
1761 for raw in stdout.replace('\r', "\n").lines() {
1762 let clean_line = sanitize_output_line(raw);
1763 if should_suppress_verification_line(&clean_line, &step.target_id) {
1764 continue;
1765 }
1766 if !clean_line.trim().is_empty() && captured_output.len() < MAX_CAPTURED_OUTPUT_LINES {
1767 let line = format!("{}: {}", step.label, clean_line);
1768 let _ = sender.send(InstallerEvent::Log(line.clone(), false));
1769 captured_output.push(clean_line);
1770 }
1771 }
1772
1773 let stderr = String::from_utf8_lossy(&output.stderr);
1774 for raw in stderr.replace('\r', "\n").lines() {
1775 let clean_line = sanitize_output_line(raw);
1776 if should_suppress_verification_line(&clean_line, &step.target_id) {
1777 continue;
1778 }
1779 if !clean_line.trim().is_empty() && captured_output.len() < MAX_CAPTURED_OUTPUT_LINES {
1780 let line = format!("{}: {}", step.label, clean_line);
1781 let _ = sender.send(InstallerEvent::Log(line.clone(), false));
1782 captured_output.push(clean_line);
1783 }
1784 }
1785
1786 if output.status.success() && captured_output.len() == MAX_CAPTURED_OUTPUT_LINES {
1787 captured_output.push("Output truncated".to_string());
1788 }
1789
1790 Ok((output.status.success(), captured_output))
1791}
1792
1793fn sort_components_by_dependencies(components: Vec<Component>) -> Vec<Component> {
1806 use crate::installers::components::topological_sort;
1807
1808 let mut installer_comps: Vec<Component> = Vec::new();
1810 let mut other_comps: Vec<Component> = Vec::new();
1811
1812 for comp in components {
1813 if comp.category == Category::Verification || comp.category == Category::Performance {
1814 other_comps.push(comp);
1815 } else {
1816 installer_comps.push(comp);
1817 }
1818 }
1819
1820 let ids: Vec<String> = installer_comps.iter().map(|c| c.id.clone()).collect();
1822
1823 match topological_sort(&ids) {
1824 Ok(sorted_ids) => {
1825 let mut reordered = Vec::with_capacity(installer_comps.len());
1827 for id in &sorted_ids {
1828 if let Some(pos) = installer_comps.iter().position(|c| c.id == *id) {
1829 reordered.push(installer_comps.remove(pos));
1830 }
1831 }
1832 reordered.extend(installer_comps);
1834
1835 reordered.extend(other_comps);
1837 reordered
1838 }
1839 Err(_) => {
1840 installer_comps.extend(other_comps);
1843 installer_comps
1844 }
1845 }
1846}
1847
1848struct NativeInstallerContext<'a> {
1855 sudo_password: Option<String>,
1856 batch_mode: bool,
1857 install_method: &'a str,
1858 sender: &'a Sender<InstallerEvent>,
1859 #[allow(dead_code)] input_rx: Arc<Mutex<Receiver<String>>>,
1861 user_home: &'a str,
1862 env_exports: &'a HashMap<String, String>,
1863}
1864
1865enum NativeCommand {
1875 Shell {
1877 program: String,
1878 args: Vec<String>,
1879 env: Vec<(String, String)>,
1880 working_dir: Option<PathBuf>,
1881 },
1882 Pip { program: String, args: Vec<String> },
1884 System { program: String, args: Vec<String> },
1886 Package { program: String, args: Vec<String> },
1888}
1889
1890impl NativeCommand {
1891 fn from_shell_cmd(program: &str, args: &[String], env: &[(String, String)]) -> Self {
1895 NativeCommand::Shell {
1896 program: program.to_string(),
1897 args: args.to_vec(),
1898 env: env.to_vec(),
1899 working_dir: None,
1900 }
1901 }
1902
1903 fn from_shell_cmd_with_dir(
1905 program: &str,
1906 args: &[String],
1907 env: &[(String, String)],
1908 working_dir: Option<PathBuf>,
1909 ) -> Self {
1910 NativeCommand::Shell {
1911 program: program.to_string(),
1912 args: args.to_vec(),
1913 env: env.to_vec(),
1914 working_dir,
1915 }
1916 }
1917}
1918
1919fn is_up_to_date_output(output: &str) -> bool {
1939 let lower = output.to_lowercase();
1940 lower.contains("up to date -- skipping")
1941 || lower.contains("up to date, skipping")
1942 || lower.contains("up to date -- reinstalling")
1943 || lower.contains("there is nothing to do")
1944 || lower.contains("already installed")
1945 || (lower.contains("warning:") && lower.contains("up to date"))
1946}
1947
1948fn execute_native_command(
1954 cmd: &NativeCommand,
1955 sudo_pw: Option<&str>,
1956 sender: &Sender<InstallerEvent>,
1957 component_name: &str,
1958) -> Result<()> {
1959 let (program, args, envs, working_dir) = match cmd {
1960 NativeCommand::Shell {
1961 program,
1962 args,
1963 env,
1964 working_dir,
1965 } => (
1966 program.clone(),
1967 args.clone(),
1968 env.clone(),
1969 working_dir.clone(),
1970 ),
1971 NativeCommand::Pip { program, args } => (program.clone(), args.clone(), vec![], None),
1972 NativeCommand::System { program, args } => (program.clone(), args.clone(), vec![], None),
1973 NativeCommand::Package { program, args } => (program.clone(), args.clone(), vec![], None),
1974 };
1975
1976 let already_sudo = program == "sudo";
1979 let mut command = if let Some(_pw) = sudo_pw {
1980 if already_sudo {
1981 let mut c = Command::new(&program);
1983 c.args(&args);
1984 c
1985 } else {
1986 let mut c = Command::new("sudo");
1987 c.arg("-S").arg("-p").arg("").arg("--").arg(&program);
1988 c.args(&args);
1989 let _ = sender.send(InstallerEvent::Log(
1991 format!("[native] Running with sudo: {} {}", program, args.join(" ")),
1992 false,
1993 ));
1994 c
1995 }
1996 } else {
1997 let mut c = Command::new(&program);
1998 c.args(&args);
1999 c
2000 };
2001
2002 for (key, value) in &envs {
2004 command.env(key, value);
2005 }
2006
2007 if let Some(ref dir) = working_dir {
2009 command.current_dir(dir);
2010 }
2011
2012 command
2014 .stdin(Stdio::piped())
2015 .stdout(Stdio::piped())
2016 .stderr(Stdio::piped());
2017
2018 let cmd_str = if args.is_empty() {
2020 program.clone()
2021 } else {
2022 format!("{} {}", program, args.join(" "))
2023 };
2024 let _ = sender.send(InstallerEvent::Log(
2025 format!("[native] $ {}", cmd_str),
2026 false,
2027 ));
2028
2029 let mut child = command.spawn().with_context(|| {
2030 let mut hint = String::new();
2031 if cmd_str.starts_with("sudo") {
2033 hint = " — ensure sudo is installed and you have the correct permissions".to_string();
2034 } else if cmd_str.starts_with("pip")
2035 || cmd_str.starts_with("uv")
2036 || cmd_str.starts_with("python")
2037 {
2038 hint = " — ensure Python/pip is installed and accessible".to_string();
2039 } else if cmd_str.starts_with("git") {
2040 hint = " — ensure git is installed".to_string();
2041 } else if cmd_str.starts_with("cmake") || cmd_str.starts_with("make") {
2042 hint = " — ensure build tools (cmake, make, gcc) are installed".to_string();
2043 } else if cmd_str.starts_with("rocminfo") || cmd_str.starts_with("rocm-smi") {
2044 hint = " — ensure ROCm is installed and /opt/rocm/bin is in PATH".to_string();
2045 }
2046 format!(
2047 "Failed to spawn command '{}' for {}.{}",
2048 cmd_str, component_name, hint
2049 )
2050 })?;
2051
2052 if let Some(pw) = sudo_pw {
2055 if let Some(mut stdin) = child.stdin.take() {
2056 let _ = stdin.write_all(pw.as_bytes());
2057 let _ = stdin.write_all(b"\n");
2058 let _ = stdin.flush();
2059 drop(stdin);
2060 }
2061 } else {
2062 drop(child.stdin.take());
2065 }
2066
2067 let sender_stdout = sender.clone();
2070 let stdout_stream = match child.stdout.take() {
2071 Some(s) => s,
2072 None => {
2073 let _ = sender.send(InstallerEvent::Log(
2074 format!(
2075 "[ERROR] {} — failed to capture stdout from spawned process",
2076 component_name
2077 ),
2078 false,
2079 ));
2080 bail!(
2081 "{} — failed to capture stdout from spawned process. \
2082 This usually indicates a system resource issue.",
2083 component_name
2084 );
2085 }
2086 };
2087 let stdout_captured: Arc<Mutex<String>> = Arc::new(Mutex::new(String::new()));
2088 let stdout_captured_clone = stdout_captured.clone();
2089 let stdout_handle = thread::spawn(move || {
2090 let reader = BufReader::new(stdout_stream);
2091 for line in reader.lines() {
2092 match line {
2093 Ok(l) => {
2094 let clean = strip_ansi_codes(&l);
2095 if !clean.trim().is_empty() {
2096 let _ = sender_stdout.send(InstallerEvent::Log(clean.clone(), false));
2097 if let Ok(mut buf) = stdout_captured_clone.lock() {
2098 buf.push_str(&clean);
2099 buf.push('\n');
2100 }
2101 }
2102 }
2103 Err(_) => break,
2104 }
2105 }
2106 });
2107
2108 let sender_stderr = sender.clone();
2109 let stderr_stream = match child.stderr.take() {
2110 Some(s) => s,
2111 None => {
2112 let _ = sender.send(InstallerEvent::Log(
2113 format!(
2114 "[ERROR] {} — failed to capture stderr from spawned process",
2115 component_name
2116 ),
2117 false,
2118 ));
2119 bail!(
2120 "{} — failed to capture stderr from spawned process. \
2121 This usually indicates a system resource issue.",
2122 component_name
2123 );
2124 }
2125 };
2126 let stderr_captured: Arc<Mutex<String>> = Arc::new(Mutex::new(String::new()));
2127 let stderr_captured_clone = stderr_captured.clone();
2128 let stderr_handle = thread::spawn(move || {
2129 let reader = BufReader::new(stderr_stream);
2130 for line in reader.lines() {
2131 match line {
2132 Ok(l) => {
2133 let clean = strip_ansi_codes(&l);
2134 if !clean.trim().is_empty() {
2135 let _ = sender_stderr.send(InstallerEvent::Log(clean.clone(), false));
2136 if let Ok(mut buf) = stderr_captured_clone.lock() {
2137 buf.push_str(&clean);
2138 buf.push('\n');
2139 }
2140 }
2141 }
2142 Err(_) => break,
2143 }
2144 }
2145 });
2146
2147 let status = child.wait().with_context(|| {
2149 format!(
2150 "{} — command process wait failed. \
2151 The process may have been killed or the system is under heavy load.",
2152 component_name
2153 )
2154 })?;
2155 let _ = stdout_handle.join();
2156 let _ = stderr_handle.join();
2157
2158 if !status.success() {
2159 let combined_output = {
2164 let stdout_buf = stdout_captured.lock().unwrap_or_else(|e| e.into_inner());
2165 let stderr_buf = stderr_captured.lock().unwrap_or_else(|e| e.into_inner());
2166 format!("{}{}", *stdout_buf, *stderr_buf)
2167 };
2168
2169 if is_up_to_date_output(&combined_output) {
2170 let _ = sender.send(InstallerEvent::Log(
2171 format!(
2172 "[native] {} — package already up to date, treating as success",
2173 component_name
2174 ),
2175 false,
2176 ));
2177 return Ok(());
2178 }
2179
2180 let code = status.code().unwrap_or(-1);
2181 let error_detail = match code {
2182 1 => "General error — check the command output above for details".to_string(),
2183 2 => "Misuse of shell builtins — check command arguments".to_string(),
2184 126 => "Command not executable — check file permissions".to_string(),
2185 127 => "Command not found — ensure the required tool is installed".to_string(),
2186 130 => "Process terminated by Ctrl+C (SIGINT)".to_string(),
2187 137 => "Process killed (SIGKILL) — possibly out of memory".to_string(),
2188 139 => "Segmentation fault — this is a bug in the invoked program".to_string(),
2189 n if n > 128 => format!(
2190 "Process terminated by signal {} — check system logs",
2191 n - 128
2192 ),
2193 _ => format!(
2194 "Command exited with code {} — check the output above for details",
2195 code
2196 ),
2197 };
2198 let _ = sender.send(InstallerEvent::Log(
2199 format!(
2200 "[ERROR] {} failed: {} (exit code {} at {})",
2201 component_name,
2202 error_detail,
2203 code,
2204 Local::now().format("%H:%M:%S")
2205 ),
2206 false,
2207 ));
2208 bail!(
2209 "{} failed: {} (exit code {})",
2210 component_name,
2211 error_detail,
2212 code
2213 );
2214 }
2215
2216 Ok(())
2217}
2218
2219#[allow(dead_code)] fn execute_command_sequence(
2225 commands: &[NativeCommand],
2226 sudo_pw: Option<&str>,
2227 sender: &Sender<InstallerEvent>,
2228 component_name: &str,
2229) -> Result<()> {
2230 for cmd in commands {
2231 execute_native_command(cmd, sudo_pw, sender, component_name)?;
2232 }
2233 Ok(())
2234}
2235
2236fn build_common_env(ctx: &NativeInstallerContext) -> Vec<(String, String)> {
2245 let mut env = Vec::new();
2246
2247 env.push(("HOME".into(), ctx.user_home.to_string()));
2249 if let Ok(user) = std::env::var("USER") {
2250 env.push(("USER".into(), user.clone()));
2251 env.push(("LOGNAME".into(), user));
2252 }
2253 env.push(("MLSTACK_USER_HOME".into(), ctx.user_home.to_string()));
2254 env.push((
2255 "MLSTACK_BATCH_MODE".into(),
2256 if ctx.batch_mode { "1" } else { "0" }.into(),
2257 ));
2258 env.push((
2259 "MLSTACK_INSTALL_METHOD".into(),
2260 ctx.install_method.to_string(),
2261 ));
2262 env.push(("INSTALL_METHOD".into(), ctx.install_method.to_string()));
2263 env.push(("RUSTY_STACK".into(), "true".into()));
2264
2265 for key in &[
2267 "ROCM_VERSION",
2268 "ROCM_CHANNEL",
2269 "GPU_ARCH",
2270 "GPU_ARCHS",
2271 "PYTORCH_ROCM_ARCH",
2272 "HIP_VISIBLE_DEVICES",
2273 "CUDA_VISIBLE_DEVICES",
2274 "PYTORCH_ROCM_DEVICE",
2275 "HSA_OVERRIDE_GFX_VERSION",
2276 "MLSTACK_PYTHON_BIN",
2277 "UV_PYTHON",
2278 ] {
2279 if let Some(value) = ctx.env_exports.get(*key) {
2280 let trimmed = value.trim();
2281 if !trimmed.is_empty() {
2282 env.push((key.to_string(), trimmed.to_string()));
2283 }
2284 }
2285 }
2286
2287 env.push(("ROCM_HOME".into(), "/opt/rocm".into()));
2289 env.push(("ROCM_PATH".into(), "/opt/rocm".into()));
2290 env.push(("HIP_PATH".into(), "/opt/rocm".into()));
2291 env.push(("HIP_ROOT_DIR".into(), "/opt/rocm".into()));
2292 env.push(("ROCM_ROOT".into(), "/opt/rocm".into()));
2293 env.push(("HIPCC_BIN_DIR".into(), "/opt/rocm/bin".into()));
2294
2295 env.push(("PIP_BREAK_SYSTEM_PACKAGES".into(), "1".into()));
2297 env.push(("PIP_ROOT_USER_ACTION".into(), "ignore".into()));
2298 env.push(("UV_PIP_BREAK_SYSTEM_PACKAGES".into(), "1".into()));
2299 env.push(("UV_SYSTEM_PYTHON".into(), "1".into()));
2300 env.push(("PYTHONUNBUFFERED".into(), "1".into()));
2301
2302 env
2303}
2304
2305fn is_existing_git_repo(target_dir: &str) -> bool {
2315 let path = Path::new(target_dir);
2316 match std::fs::symlink_metadata(target_dir) {
2319 Ok(meta) if meta.is_dir() => {
2320 path.join(".git")
2322 .symlink_metadata()
2323 .map(|m| m.is_dir() || m.is_file())
2324 .unwrap_or(false)
2325 }
2326 _ => false,
2327 }
2328}
2329
2330fn fix_directory_ownership(dir: &str, sudo_pw: Option<&str>, sender: &Sender<InstallerEvent>) {
2347 let stat_output = std::process::Command::new("stat")
2349 .args(["-c", "%U", dir])
2350 .output();
2351
2352 let needs_fix = match stat_output {
2353 Ok(out) if out.status.success() => {
2354 let owner = String::from_utf8_lossy(&out.stdout).trim().to_string();
2355 owner == "root"
2356 }
2357 _ => true, };
2359
2360 if !needs_fix {
2361 return;
2362 }
2363
2364 let real_user = std::env::var("SUDO_USER")
2365 .or_else(|_| std::env::var("USER"))
2366 .unwrap_or_else(|_| "scooter".to_string());
2367
2368 let _ = sender.send(InstallerEvent::Log(
2369 format!("[native] Fixing ownership of {} -> {}", dir, real_user),
2370 false,
2371 ));
2372
2373 if sudo_pw.is_some() {
2374 let chown_cmd = NativeCommand::from_shell_cmd(
2375 "chown",
2376 &[
2377 "-R".to_string(),
2378 format!("{}:{}", real_user, real_user),
2379 dir.to_string(),
2380 ],
2381 &[],
2382 );
2383 let _ = execute_native_command(&chown_cmd, sudo_pw, sender, "ownership-fix");
2384 }
2385 }
2387
2388fn git_clone_or_pull(
2389 repo_url: &str,
2390 target_dir: &str,
2391 extra_clone_args: &[&str],
2392 sudo_pw: Option<&str>,
2393 sender: &Sender<InstallerEvent>,
2394 component_name: &str,
2395) -> Result<()> {
2396 if is_existing_git_repo(target_dir) {
2397 let _ = sender.send(InstallerEvent::Log(
2398 format!(
2399 "[native] {} — directory '{}' already exists, pulling updates instead of cloning",
2400 component_name, target_dir
2401 ),
2402 false,
2403 ));
2404
2405 let pull_cmd = NativeCommand::from_shell_cmd(
2407 "git",
2408 &[
2409 "-C".to_string(),
2410 target_dir.to_string(),
2411 "pull".to_string(),
2412 "--ff-only".to_string(),
2413 ],
2414 &[],
2415 );
2416 let pull_result = execute_native_command(&pull_cmd, sudo_pw, sender, component_name);
2417
2418 if pull_result.is_ok() {
2419 return Ok(());
2420 }
2421
2422 let _ = sender.send(InstallerEvent::Log(
2424 format!(
2425 "[native] {} — git pull --ff-only failed, trying fetch + reset --hard origin/HEAD",
2426 component_name
2427 ),
2428 false,
2429 ));
2430
2431 let fetch_cmd = NativeCommand::from_shell_cmd(
2432 "git",
2433 &[
2434 "-C".to_string(),
2435 target_dir.to_string(),
2436 "fetch".to_string(),
2437 "--all".to_string(),
2438 ],
2439 &[],
2440 );
2441 execute_native_command(&fetch_cmd, sudo_pw, sender, component_name)?;
2442
2443 let reset_cmd = NativeCommand::from_shell_cmd(
2444 "git",
2445 &[
2446 "-C".to_string(),
2447 target_dir.to_string(),
2448 "reset".to_string(),
2449 "--hard".to_string(),
2450 "origin/HEAD".to_string(),
2451 ],
2452 &[],
2453 );
2454 execute_native_command(&reset_cmd, sudo_pw, sender, component_name)?;
2455
2456 Ok(())
2457 } else {
2458 let target_path = Path::new(target_dir);
2463 if target_path.is_dir() && !target_path.join(".git").exists() {
2464 let _ = sender.send(InstallerEvent::Log(
2465 format!(
2466 "[native] {} — directory '{}' exists but is not a git repo; removing incomplete clone",
2467 component_name, target_dir
2468 ),
2469 false,
2470 ));
2471 if std::fs::remove_dir_all(target_dir).is_err() {
2473 let rm_cmd = NativeCommand::from_shell_cmd(
2475 "rm",
2476 &["-rf".to_string(), target_dir.to_string()],
2477 &[],
2478 );
2479 let rm_result = execute_native_command(&rm_cmd, sudo_pw, sender, component_name);
2480 if let Err(e) = rm_result {
2481 let _ = sender.send(InstallerEvent::Log(
2482 format!(
2483 "[native] [WARN] {} — could not remove stale dir: {}. Clone may fail.",
2484 component_name, e
2485 ),
2486 false,
2487 ));
2488 }
2489 }
2490 }
2491
2492 let _ = sender.send(InstallerEvent::Log(
2493 format!(
2494 "[native] {} — cloning '{}' into '{}'",
2495 component_name, repo_url, target_dir
2496 ),
2497 false,
2498 ));
2499
2500 let mut args: Vec<String> = vec!["clone".to_string()];
2501 for extra in extra_clone_args {
2502 args.push(extra.to_string());
2503 }
2504 args.push(repo_url.to_string());
2505 args.push(target_dir.to_string());
2506
2507 let clone_cmd = NativeCommand::from_shell_cmd("git", &args, &[]);
2508 execute_native_command(&clone_cmd, sudo_pw, sender, component_name)
2509 }
2510}
2511
2512fn filter_cuda_requirements(req_path: &str) -> String {
2544 let content = match std::fs::read_to_string(req_path) {
2545 Ok(c) => c,
2546 Err(_) => return String::new(),
2547 };
2548
2549 let mut filtered = Vec::new();
2550
2551 for line in content.lines() {
2552 let trimmed = line.trim();
2553
2554 if trimmed.is_empty() || trimmed.starts_with('#') {
2556 continue;
2557 }
2558
2559 let pkg_name = trimmed
2561 .split(['=', '>', '<', '!', ';', '[', ' ', '\t'])
2562 .next()
2563 .unwrap_or("")
2564 .trim();
2565
2566 let pkg_lower = pkg_name.to_lowercase();
2568 if pkg_lower.starts_with("nvidia-")
2569 || pkg_lower.starts_with("cuda")
2570 || pkg_lower.starts_with("tensorrt")
2571 || pkg_lower == "triton"
2572 || pkg_lower == "xformers"
2573 || pkg_lower == "flash-attn"
2574 || pkg_lower == "torch"
2575 || pkg_lower == "torchvision"
2576 || pkg_lower == "torchaudio"
2577 {
2578 continue;
2579 }
2580
2581 if pkg_lower.starts_with("triton-")
2583 || pkg_lower.starts_with("triton==")
2584 || pkg_lower.starts_with("triton[")
2585 {
2586 continue;
2587 }
2588
2589 let line_lower = trimmed.to_lowercase();
2592 if line_lower.contains("http://") || line_lower.contains("https://") {
2593 if line_lower.contains("+cu1")
2595 || line_lower.contains("+cu2")
2596 || line_lower.contains("+cu3")
2597 || line_lower.contains("cu124")
2598 || line_lower.contains("cu128")
2599 || line_lower.contains("cu121")
2600 || line_lower.contains("cu118")
2601 {
2602 continue;
2603 }
2604 if line_lower.contains("ik_llama") {
2606 continue;
2607 }
2608 if line_lower.contains("exllamav3")
2610 && (line_lower.contains("+cu") || line_lower.contains("cu1"))
2611 {
2612 continue;
2613 }
2614 if line_lower.contains("flash_attn")
2616 && (line_lower.contains("+cu") || line_lower.contains("cu1"))
2617 {
2618 continue;
2619 }
2620 }
2621
2622 filtered.push(line.to_string());
2623 }
2624
2625 filtered.join("\n")
2626}
2627
2628fn is_force_reinstall() -> bool {
2633 std::env::var("MLSTACK_FORCE_REINSTALL")
2634 .or_else(|_| std::env::var("FORCE"))
2635 .map(|v| {
2636 matches!(
2637 v.trim().to_ascii_lowercase().as_str(),
2638 "1" | "true" | "yes" | "on"
2639 )
2640 })
2641 .unwrap_or(false)
2642}
2643
2644fn build_pip_uninstall(packages: &[&str]) -> NativeCommand {
2649 let use_uv = crate::installers::common::utils::command_exists("uv");
2650 if use_uv {
2651 let mut args = vec!["pip".to_string(), "uninstall".to_string()];
2652 for pkg in packages {
2653 args.push(pkg.to_string());
2654 }
2655 NativeCommand::Pip {
2656 program: "uv".to_string(),
2657 args,
2658 }
2659 } else {
2660 let mut args = vec!["uninstall".to_string(), "-y".to_string()];
2661 for pkg in packages {
2662 args.push(pkg.to_string());
2663 }
2664 NativeCommand::Pip {
2665 program: "pip3".to_string(),
2666 args,
2667 }
2668 }
2669}
2670
2671fn purge_component_packages(
2677 packages: &[&str],
2678 sender: &Sender<InstallerEvent>,
2679 component_name: &str,
2680) {
2681 let _ = sender.send(InstallerEvent::Log(
2682 format!(
2683 "[native] {} — purging packages: {}",
2684 component_name,
2685 packages.join(", ")
2686 ),
2687 false,
2688 ));
2689
2690 let uninstall_cmd = build_pip_uninstall(packages);
2692 let _ = execute_native_command(&uninstall_cmd, None, sender, component_name);
2693}
2694
2695fn purge_pytorch(sender: &Sender<InstallerEvent>, component_name: &str) {
2697 purge_component_packages(
2698 &[
2699 "torch",
2700 "torchvision",
2701 "torchaudio",
2702 "pytorch-triton-rocm",
2703 "pytorch-triton",
2704 ],
2705 sender,
2706 component_name,
2707 );
2708}
2709
2710fn purge_triton(sender: &Sender<InstallerEvent>, component_name: &str) {
2712 purge_component_packages(
2713 &["triton", "triton-rocm", "pytorch-triton-rocm"],
2714 sender,
2715 component_name,
2716 );
2717}
2718
2719fn purge_aiter(sender: &Sender<InstallerEvent>, component_name: &str) {
2721 purge_component_packages(&["aiter", "amd-aiter"], sender, component_name);
2722}
2723
2724fn purge_flash_attn(sender: &Sender<InstallerEvent>, component_name: &str) {
2726 purge_component_packages(
2727 &["flash-attn", "flash_attn", "flash_attention_amd"],
2728 sender,
2729 component_name,
2730 );
2731}
2732
2733fn purge_vllm(sender: &Sender<InstallerEvent>, component_name: &str) {
2735 purge_component_packages(&["vllm"], sender, component_name);
2736}
2737
2738fn run_native_installer(component: &Component, ctx: &NativeInstallerContext) -> Result<()> {
2744 let sender = ctx.sender;
2745 let _ = sender.send(InstallerEvent::Log(
2746 format!("[native] Installing {} via Rust module...", component.name),
2747 false,
2748 ));
2749
2750 let _ = sender.send(InstallerEvent::Progress {
2751 component_id: component.id.clone(),
2752 progress: 0.1,
2753 message: format!("Starting {} installation", component.name),
2754 });
2755
2756 let component_needs_sudo = component.needs_sudo && needs_sudo();
2758 let sudo_pw: Option<&str> = if component_needs_sudo {
2759 ctx.sudo_password.as_deref()
2760 } else {
2761 None
2762 };
2763
2764 if component_needs_sudo && ctx.sudo_password.is_none() {
2766 bail!("{} requires sudo but no password provided", component.name);
2767 }
2768
2769 let common_env = build_common_env(ctx);
2771
2772 match component.id.as_str() {
2776 "permanent-env" => {
2778 use crate::installers::components::permanent_env::{
2779 PermanentEnvConfig, PermanentEnvInstaller,
2780 };
2781 let detected_gpu_arch = detect_gpu_arch();
2782 let detected_gpu_list = detect_gpu_list();
2783 let detected_rocm_ver = detect_rocm_version();
2784 let hsa_override = hsa_override_from_gpu_arch(&detected_gpu_arch)
2785 .unwrap_or_else(|| "11.0.0".to_string());
2786 let resolved_python = resolve_python_bin();
2787 let inst = PermanentEnvInstaller::new(PermanentEnvConfig {
2788 python_bin: resolved_python,
2789 gpu_arch: detected_gpu_arch,
2790 hsa_override_gfx_version: hsa_override,
2791 discrete_gpu_list: detected_gpu_list,
2792 rocm_version: detected_rocm_ver,
2793 ..PermanentEnvConfig::default()
2794 });
2795 let cmd = inst.build_mkdir_triton_dirs_command();
2796 execute_native_command(
2797 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
2798 sudo_pw,
2799 sender,
2800 &component.name,
2801 )?;
2802
2803 let cmd = inst.build_rocminfo_gpu_detect_command();
2804 execute_native_command(
2805 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
2806 sudo_pw,
2807 sender,
2808 &component.name,
2809 )?;
2810
2811 let cmd = inst.build_rocm_version_detect_command();
2812 execute_native_command(
2813 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
2814 sudo_pw,
2815 sender,
2816 &component.name,
2817 )?;
2818 }
2819
2820 "rocm" => {
2822 use crate::installers::components::rocm::{PackageCommand, RocmConfig, RocmInstaller};
2823 let rocm_force = is_force_reinstall();
2824 let inst = RocmInstaller::new(RocmConfig {
2825 force_reinstall: rocm_force,
2826 ..RocmConfig::default()
2827 });
2828 let distro = crate::installers::common::DistroFacade::detect();
2829 let commands = if distro.uses_apt() {
2830 inst.apt_install_commands()
2831 } else if distro.uses_dnf() || distro.uses_yum() {
2832 let rhel_ver = distro.version().split('.').next().unwrap_or("9");
2833 inst.dnf_install_commands(rhel_ver)
2834 } else if distro.uses_zypper() {
2835 inst.zypper_install_commands()
2836 } else {
2837 let all_pkgs = inst.pacman_rocm_packages();
2840 let need_install = RocmInstaller::filter_already_installed_pacman(&all_pkgs);
2841
2842 if need_install.is_empty() {
2843 let force_reinstall = is_force_reinstall();
2845 if force_reinstall {
2846 let _ = sender.send(InstallerEvent::Log(
2847 format!(
2848 "[native] {} — all ROCm packages already installed, force-reinstalling",
2849 component.name
2850 ),
2851 false,
2852 ));
2853 } else {
2856 let _ = sender.send(InstallerEvent::Log(
2857 format!(
2858 "[native] {} — all ROCm packages already installed, skipping",
2859 component.name
2860 ),
2861 false,
2862 ));
2863 return Ok(());
2864 }
2865 }
2866 let force_reinstall = is_force_reinstall();
2868 let pkgs_to_install = if force_reinstall {
2869 all_pkgs.clone()
2870 } else {
2871 need_install
2872 };
2873
2874 let _ = sender.send(InstallerEvent::Log(
2875 format!(
2876 "[native] {} — {} of {} packages need installation: {}",
2877 component.name,
2878 pkgs_to_install.len(),
2879 all_pkgs.len(),
2880 pkgs_to_install.join(", ")
2881 ),
2882 false,
2883 ));
2884
2885 if is_force_reinstall() {
2888 let mut remove_args = vec!["-Rns".to_string(), "--noconfirm".to_string()];
2889 remove_args.extend(all_pkgs.iter().filter(|p| !p.is_empty()).cloned());
2890 let _ = sender.send(InstallerEvent::Log(
2891 format!(
2892 "[native] {} — force reinstall: removing existing packages",
2893 component.name
2894 ),
2895 false,
2896 ));
2897 let remove_cmd = NativeCommand::Package {
2898 program: "sudo".to_string(),
2899 args: {
2900 let mut a = vec!["pacman".to_string()];
2901 a.extend(remove_args);
2902 a
2903 },
2904 };
2905 let _ = execute_native_command(&remove_cmd, sudo_pw, sender, &component.name);
2906 }
2907
2908 let aur_helper = "yay";
2911 let mut args = if is_force_reinstall() {
2912 vec!["-S".to_string(), "--noconfirm".to_string()]
2913 } else {
2914 vec![
2915 "-S".to_string(),
2916 "--needed".to_string(),
2917 "--noconfirm".to_string(),
2918 ]
2919 };
2920 args.extend(pkgs_to_install);
2921
2922 vec![
2923 PackageCommand {
2924 program: "sudo".to_string(),
2925 args: vec![aur_helper.to_string()],
2926 },
2927 PackageCommand {
2928 program: aur_helper.to_string(),
2929 args,
2930 },
2931 ]
2932 };
2933 for pkg_cmd in &commands {
2934 let native_cmd = NativeCommand::Package {
2935 program: pkg_cmd.program.clone(),
2936 args: pkg_cmd.args.clone(),
2937 };
2938 execute_native_command(&native_cmd, sudo_pw, sender, &component.name)?;
2941 }
2942
2943 if is_force_reinstall() {
2945 let marker_path = std::path::Path::new("/tmp/mlstack-reboot-required");
2946 let _ = std::fs::write(
2947 marker_path,
2948 "ROCm packages were reinstalled. Reboot required for kernel module changes.\n",
2949 );
2950 let _ = sender.send(InstallerEvent::Log(
2951 "[native] ROCm — reboot recommended for kernel module changes".to_string(),
2952 false,
2953 ));
2954 }
2955 }
2956
2957 "pytorch" => {
2959 use crate::installers::components::pytorch::{PyTorchConfig, PyTorchInstaller};
2960 let force = is_force_reinstall();
2961 let config = PyTorchConfig {
2962 force_reinstall: force,
2963 ..PyTorchConfig::default()
2964 };
2965 let inst = PyTorchInstaller::new(config);
2966 let use_uv = crate::installers::common::utils::command_exists("uv");
2967 let rocm_mm = ctx
2968 .env_exports
2969 .get("ROCM_VERSION")
2970 .map(|v| {
2971 let parts: Vec<&str> = v.split('.').collect();
2972 if parts.len() >= 2 {
2973 format!("{}.{}", parts[0], parts[1])
2974 } else {
2975 v.clone()
2976 }
2977 })
2978 .unwrap_or_else(|| "7.2".to_string());
2979
2980 if force {
2982 purge_pytorch(sender, &component.name);
2983 }
2984
2985 let deps_cmd = inst.build_common_deps_command(use_uv);
2987 let _ = common_env.clone(); execute_native_command(
2989 &NativeCommand::Pip {
2990 program: deps_cmd.program.clone(),
2991 args: deps_cmd.args.clone(),
2992 },
2993 None, sender,
2995 &component.name,
2996 )?;
2997
2998 let install_cmd = inst.build_install_command(&rocm_mm, use_uv);
3000 execute_native_command(
3001 &NativeCommand::Pip {
3002 program: install_cmd.program.clone(),
3003 args: install_cmd.args.clone(),
3004 },
3005 None,
3006 sender,
3007 &component.name,
3008 )?;
3009 }
3010
3011 "triton" => {
3013 use crate::installers::components::triton::{TritonConfig, TritonInstaller};
3014 let triton_force = is_force_reinstall();
3015 let inst = TritonInstaller::new(TritonConfig::default());
3016 let target_dir = format!("{}/.mlstack/triton", ctx.user_home);
3017
3018 if triton_force {
3020 purge_triton(sender, &component.name);
3021 }
3022
3023 git_clone_or_pull(
3027 inst.git_clone_url(),
3028 &target_dir,
3029 &[],
3030 sudo_pw,
3031 sender,
3032 &component.name,
3033 )?;
3034
3035 fix_directory_ownership(&target_dir, sudo_pw, sender);
3038
3039 let safe_dir_cmd = NativeCommand::from_shell_cmd(
3041 "git",
3042 &[
3043 "config".to_string(),
3044 "--global".to_string(),
3045 "--add".to_string(),
3046 "safe.directory".to_string(),
3047 target_dir.clone(),
3048 ],
3049 &[],
3050 );
3051 let _ = execute_native_command(&safe_dir_cmd, None, sender, &component.name);
3052
3053 let branch = inst.select_branch().to_string();
3055 let fetch_branch_cmd = NativeCommand::from_shell_cmd(
3056 "git",
3057 &[
3058 "-C".to_string(),
3059 target_dir.clone(),
3060 "fetch".to_string(),
3061 "origin".to_string(),
3062 branch.clone(),
3063 ],
3064 &[],
3065 );
3066 let _ = execute_native_command(&fetch_branch_cmd, None, sender, &component.name);
3067
3068 let checkout_cmd = inst.build_git_checkout_command();
3070 let mut checkout_args = vec!["-C".to_string(), target_dir.clone()];
3071 checkout_args.extend(checkout_cmd.args.iter().cloned());
3072 let checkout_result = execute_native_command(
3073 &NativeCommand::from_shell_cmd(
3074 &checkout_cmd.program,
3075 &checkout_args,
3076 &checkout_cmd.env,
3077 ),
3078 None,
3079 sender,
3080 &component.name,
3081 );
3082 if let Err(err) = checkout_result {
3083 let _ = sender.send(InstallerEvent::Log(
3084 format!(
3085 "[native] [WARN] {} git checkout failed (non-fatal): {}",
3086 component.name, err
3087 ),
3088 false,
3089 ));
3090 }
3091
3092 let cmd = inst.build_prerequisites_command();
3094 execute_native_command(
3095 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
3096 None,
3097 sender,
3098 &component.name,
3099 )?;
3100
3101 let cmd = inst.build_pip_install_command(&target_dir);
3103 execute_native_command(
3104 &NativeCommand::from_shell_cmd_with_dir(
3105 &cmd.program,
3106 &cmd.args,
3107 &cmd.env,
3108 cmd.working_dir.clone(),
3109 ),
3110 None,
3111 sender,
3112 &component.name,
3113 )?;
3114 }
3115
3116 "mpi4py" => {
3118 use crate::installers::components::mpi4py::{Mpi4PyConfig, Mpi4PyInstaller};
3119 let force = is_force_reinstall();
3120 let inst = Mpi4PyInstaller::new(Mpi4PyConfig {
3121 force_reinstall: force,
3122 ..Mpi4PyConfig::default()
3123 });
3124 let distro = crate::installers::common::DistroFacade::detect();
3125
3126 if force {
3128 purge_component_packages(&["mpi4py"], sender, &component.name);
3129 }
3130
3131 let cmd = inst.build_system_mpi_install_command(&distro);
3133 execute_native_command(
3134 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
3135 sudo_pw,
3136 sender,
3137 &component.name,
3138 )?;
3139
3140 let mpi_path = std::path::PathBuf::from("/usr");
3142 let cmd = inst.build_pip_install_command(&mpi_path);
3143 execute_native_command(
3144 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
3145 None,
3146 sender,
3147 &component.name,
3148 )?;
3149 }
3150
3151 "deepspeed" => {
3153 use crate::installers::components::deepspeed::{DeepSpeedConfig, DeepSpeedInstaller};
3154 let force = is_force_reinstall();
3155 let config = DeepSpeedConfig {
3156 force_reinstall: force,
3157 ..DeepSpeedConfig::default()
3158 };
3159 let inst = DeepSpeedInstaller::new(config);
3160
3161 if force {
3163 purge_component_packages(&["deepspeed"], sender, &component.name);
3164 }
3165
3166 let deps_cmd = inst.build_deps_install_command();
3168 execute_native_command(
3169 &NativeCommand::Pip {
3170 program: deps_cmd.program.clone(),
3171 args: deps_cmd.args.clone(),
3172 },
3173 None,
3174 sender,
3175 &component.name,
3176 )?;
3177
3178 let install_cmd = inst.build_install_command();
3180 execute_native_command(
3181 &NativeCommand::Pip {
3182 program: install_cmd.program.clone(),
3183 args: install_cmd.args.clone(),
3184 },
3185 None,
3186 sender,
3187 &component.name,
3188 )?;
3189 }
3190
3191 "ml-stack-core" => {
3193 use crate::installers::components::ml_stack::{MlStackConfig, MlStackInstaller};
3194 let inst = MlStackInstaller::new(MlStackConfig::default());
3195 let distro = crate::installers::common::DistroFacade::detect();
3196
3197 let migraphx_cmd = inst.build_migraphx_install_command(&distro);
3199 execute_native_command(
3200 &NativeCommand::System {
3201 program: migraphx_cmd.program.clone(),
3202 args: migraphx_cmd.args.clone(),
3203 },
3204 sudo_pw,
3205 sender,
3206 &component.name,
3207 )?;
3208
3209 let rccl_cmd = inst.build_rccl_install_command(&distro);
3210 execute_native_command(
3211 &NativeCommand::System {
3212 program: rccl_cmd.program.clone(),
3213 args: rccl_cmd.args.clone(),
3214 },
3215 sudo_pw,
3216 sender,
3217 &component.name,
3218 )?;
3219
3220 let mpi_cmd = inst.build_mpi_install_command(&distro);
3221 execute_native_command(
3222 &NativeCommand::System {
3223 program: mpi_cmd.program.clone(),
3224 args: mpi_cmd.args.clone(),
3225 },
3226 sudo_pw,
3227 sender,
3228 &component.name,
3229 )?;
3230
3231 let megatron_clone = inst.build_megatron_clone_command();
3233 execute_native_command(
3234 &NativeCommand::System {
3235 program: megatron_clone.program.clone(),
3236 args: megatron_clone.args.clone(),
3237 },
3238 sudo_pw,
3239 sender,
3240 &component.name,
3241 )?;
3242
3243 let megatron_install = inst.build_megatron_install_command();
3244 execute_native_command(
3245 &NativeCommand::Pip {
3246 program: megatron_install.program.clone(),
3247 args: megatron_install.args.clone(),
3248 },
3249 None,
3250 sender,
3251 &component.name,
3252 )?;
3253 }
3254
3255 "flash-attn" => {
3257 use crate::installers::components::flash_attention_ck::{
3258 FlashAttentionConfig, FlashAttentionInstaller,
3259 };
3260 let force = is_force_reinstall();
3261 let inst = FlashAttentionInstaller::new(FlashAttentionConfig {
3262 force_reinstall: force,
3263 ..FlashAttentionConfig::default()
3264 });
3265
3266 if force {
3268 purge_flash_attn(sender, &component.name);
3269 }
3270
3271 let rocm_backend =
3273 std::env::var("ROCM_BACKEND").unwrap_or_else(|_| "triton".to_string());
3274
3275 if rocm_backend == "triton" {
3276 let target_dir = format!("{}/.mlstack/flash-attention", ctx.user_home);
3279
3280 git_clone_or_pull(
3281 "https://github.com/ROCm/flash-attention.git",
3282 &target_dir,
3283 &[],
3284 sudo_pw,
3285 sender,
3286 &component.name,
3287 )?;
3288
3289 fix_directory_ownership(&target_dir, sudo_pw, sender);
3291
3292 let safe_dir_cmd = NativeCommand::from_shell_cmd(
3294 "git",
3295 &[
3296 "config".to_string(),
3297 "--global".to_string(),
3298 "--add".to_string(),
3299 "safe.directory".to_string(),
3300 target_dir.clone(),
3301 ],
3302 &[],
3303 );
3304 let _ = execute_native_command(&safe_dir_cmd, None, sender, &component.name);
3305
3306 let gpu_arch = ctx
3308 .env_exports
3309 .get("GPU_ARCH")
3310 .cloned()
3311 .unwrap_or_else(|| "gfx1100".to_string());
3312 let hsa_version = ctx
3313 .env_exports
3314 .get("HSA_OVERRIDE_GFX_VERSION")
3315 .cloned()
3316 .unwrap_or_else(|| "11.0.0".to_string());
3317
3318 let mut pip_args = vec![
3320 "-m".to_string(),
3321 "pip".to_string(),
3322 "install".to_string(),
3323 "--no-build-isolation".to_string(),
3324 "--break-system-packages".to_string(),
3325 ".".to_string(),
3326 ];
3327 if force {
3329 pip_args.push("--force-reinstall".to_string());
3330 }
3331 let pip_env = vec![
3332 (
3333 "FLASH_ATTENTION_TRITON_AMD_ENABLE".to_string(),
3334 "TRUE".to_string(),
3335 ),
3336 ("GPU_ARCH".to_string(), gpu_arch.clone()),
3337 ("PYTORCH_ROCM_ARCH".to_string(), gpu_arch),
3338 ("ROCM_PATH".to_string(), "/opt/rocm".to_string()),
3339 ("HSA_OVERRIDE_GFX_VERSION".to_string(), hsa_version),
3340 ];
3341
3342 execute_native_command(
3343 &NativeCommand::Shell {
3344 program: "python3".to_string(),
3345 args: pip_args,
3346 env: pip_env,
3347 working_dir: Some(PathBuf::from(&target_dir)),
3348 },
3349 None, sender,
3351 &component.name,
3352 )?;
3353 } else {
3354 let clone_cmd = inst.build_git_clone_command();
3357 let target_dir_str = clone_cmd.args.last().cloned().unwrap_or_default();
3358 git_clone_or_pull(
3359 "https://github.com/ROCmSoftwarePlatform/flash-attention.git",
3360 &target_dir_str,
3361 &[],
3362 sudo_pw,
3363 sender,
3364 &component.name,
3365 )?;
3366
3367 fix_directory_ownership(&target_dir_str, sudo_pw, sender);
3369
3370 let cmd = inst.build_git_checkout_command();
3372 execute_native_command(
3373 &NativeCommand::from_shell_cmd_with_dir(
3374 &cmd.program,
3375 &cmd.args,
3376 &cmd.env,
3377 cmd.working_dir.clone(),
3378 ),
3379 None,
3380 sender,
3381 &component.name,
3382 )?;
3383
3384 if force {
3386 let build_dir = PathBuf::from(&target_dir_str).join("build");
3387 if build_dir.exists() {
3388 let _ = std::fs::remove_dir_all(&build_dir);
3389 let _ = sender.send(InstallerEvent::Log(
3390 format!(
3391 "[native] {} — force reinstall: cleaned build directory",
3392 component.name
3393 ),
3394 false,
3395 ));
3396 }
3397 }
3398 let rocm_env = crate::installers::common::RocmEnv::detect();
3399 let torch_prefix = format!(
3400 "{}/rocm_venv/lib/python3.12/site-packages/torch/share/cmake/Torch",
3401 ctx.user_home
3402 );
3403 let cmd = inst.build_cmake_command(&rocm_env, &torch_prefix);
3404 execute_native_command(
3405 &NativeCommand::from_shell_cmd_with_dir(
3406 &cmd.program,
3407 &cmd.args,
3408 &cmd.env,
3409 cmd.working_dir.clone(),
3410 ),
3411 None,
3412 sender,
3413 &component.name,
3414 )?;
3415
3416 let cmd = inst.build_make_command();
3418 execute_native_command(
3419 &NativeCommand::from_shell_cmd_with_dir(
3420 &cmd.program,
3421 &cmd.args,
3422 &cmd.env,
3423 cmd.working_dir.clone(),
3424 ),
3425 None,
3426 sender,
3427 &component.name,
3428 )?;
3429
3430 let cmd = inst.build_setup_install_command();
3432 execute_native_command(
3433 &NativeCommand::from_shell_cmd_with_dir(
3434 &cmd.program,
3435 &cmd.args,
3436 &cmd.env,
3437 cmd.working_dir.clone(),
3438 ),
3439 None,
3440 sender,
3441 &component.name,
3442 )?;
3443 }
3444 }
3445
3446 "repair-stack" => {
3450 use crate::installers::components::repair::RepairInstaller;
3451 let repair_ids = RepairInstaller::native_component_ids();
3452
3453 let _ = sender.send(InstallerEvent::Log(
3454 "[native] Starting ML Stack repair (8 steps via native installers)...".to_string(),
3455 false,
3456 ));
3457
3458 for step_id in &repair_ids {
3459 let sub_component = Component {
3462 id: (*step_id).to_string(),
3463 name: (*step_id).to_string(),
3464 description: String::new(),
3465 script: String::new(),
3466 category: Category::Core,
3467 required: true,
3468 selected: true,
3469 installed: false,
3470 progress: 0.0,
3471 estimate: String::new(),
3472 needs_sudo: component.needs_sudo,
3473 };
3474
3475 let _ = sender.send(InstallerEvent::Progress {
3476 component_id: component.id.clone(),
3477 progress: 0.5,
3478 message: format!("Repairing {}...", step_id),
3479 });
3480
3481 if let Err(err) = run_native_installer(&sub_component, ctx) {
3482 let _ = sender.send(InstallerEvent::Log(
3484 format!("[native] Repair step '{}' failed: {}", step_id, err),
3485 false,
3486 ));
3487 }
3488 }
3489 }
3490
3491 "megatron" => {
3493 use crate::installers::components::megatron::{MegatronConfig, MegatronInstaller};
3494 let force = is_force_reinstall();
3495 let inst = MegatronInstaller::new(MegatronConfig {
3496 force_reinstall: force,
3497 ..MegatronConfig::default()
3498 });
3499
3500 if force {
3502 purge_component_packages(&["megatron-core", "megatron"], sender, &component.name);
3503 }
3504
3505 for dep_pkg in &["ninja", "packaging", "wheel"] {
3507 let cmd = inst.build_dep_install_command(dep_pkg);
3508 execute_native_command(
3509 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
3510 None,
3511 sender,
3512 &component.name,
3513 )?;
3514 }
3515
3516 let clone_cmd = inst.build_git_clone_command();
3518 let clone_target_str = clone_cmd.args.last().cloned().unwrap_or_default();
3519 git_clone_or_pull(
3520 "https://github.com/NVIDIA/Megatron-LM.git",
3521 &clone_target_str,
3522 &[],
3523 sudo_pw,
3524 sender,
3525 &component.name,
3526 )?;
3527
3528 let clone_target_path = PathBuf::from(&clone_target_str);
3532 if let Ok(entries) = std::fs::read_dir(&clone_target_path) {
3533 for entry in entries.flatten() {
3534 let name = entry.file_name();
3535 let name_str = name.to_string_lossy();
3536 if name_str.ends_with(".egg-info") {
3537 let egg_path = entry.path();
3538 if std::fs::remove_dir_all(&egg_path).is_err() {
3540 let rm_cmd = NativeCommand::from_shell_cmd(
3542 "rm",
3543 &["-rf".to_string(), egg_path.to_string_lossy().to_string()],
3544 &[],
3545 );
3546 let _ =
3547 execute_native_command(&rm_cmd, sudo_pw, sender, &component.name);
3548 }
3549 }
3550 }
3551 }
3552
3553 let cmd = inst.build_pip_install_command();
3555 execute_native_command(
3556 &NativeCommand::from_shell_cmd_with_dir(
3557 &cmd.program,
3558 &cmd.args,
3559 &cmd.env,
3560 cmd.working_dir.clone(),
3561 ),
3562 None,
3563 sender,
3564 &component.name,
3565 )?;
3566
3567 let requirements_candidates = [
3569 clone_target_path
3570 .join("requirements")
3571 .join("requirements.txt"),
3572 clone_target_path.join("requirements.txt"),
3573 clone_target_path
3574 .join("requirements")
3575 .join("requirements_amd.txt"),
3576 ];
3577 for req_path in requirements_candidates.iter().filter(|p| p.exists()) {
3578 let filtered = filter_cuda_requirements(&req_path.to_string_lossy());
3579 if filtered.is_empty() {
3580 continue;
3581 }
3582 let tmp_dir = std::env::temp_dir().join("rusty-stack-megatron");
3583 let _ = std::fs::create_dir_all(&tmp_dir);
3584 let filtered_path = tmp_dir.join("requirements_filtered.txt");
3585 std::fs::write(&filtered_path, &filtered).with_context(|| {
3586 format!(
3587 "Failed to write filtered Megatron requirements to {:?}",
3588 filtered_path
3589 )
3590 })?;
3591 let filtered_path_str = filtered_path.to_string_lossy().to_string();
3592 let filtered_cmd = inst.build_requirements_install_command(&filtered_path_str);
3593 execute_native_command(
3594 &NativeCommand::from_shell_cmd(
3595 &filtered_cmd.program,
3596 &filtered_cmd.args,
3597 &filtered_cmd.env,
3598 ),
3599 None,
3600 sender,
3601 &component.name,
3602 )?;
3603 let _ = std::fs::remove_file(&filtered_path);
3604 break;
3605 }
3606 }
3607
3608 "vllm" => {
3610 use crate::installers::components::vllm_multi::{VllmConfig, VllmInstaller};
3611 let vllm_force = is_force_reinstall();
3612 let inst = VllmInstaller::new(VllmConfig::default());
3613
3614 if vllm_force {
3616 purge_vllm(sender, &component.name);
3617 }
3618
3619 let cmd = inst.build_vllm_install_command();
3621 execute_native_command(
3622 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
3623 None,
3624 sender,
3625 &component.name,
3626 )?;
3627
3628 let cmd = inst.build_deps_install_command();
3630 execute_native_command(
3631 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
3632 None,
3633 sender,
3634 &component.name,
3635 )?;
3636 }
3637
3638 "aiter" => {
3640 use crate::installers::components::aiter::{AiterConfig, AiterInstaller};
3641 let inst = AiterInstaller::new(AiterConfig::default());
3642 let target_dir = format!("{}/.mlstack/aiter", ctx.user_home);
3643
3644 if is_force_reinstall() {
3646 purge_aiter(sender, &component.name);
3647 }
3648
3649 git_clone_or_pull(
3651 inst.git_clone_url(),
3652 &target_dir,
3653 &["--recursive"],
3654 sudo_pw,
3655 sender,
3656 &component.name,
3657 )?;
3658
3659 let run_user = std::env::var("SUDO_USER")
3663 .or_else(|_| std::env::var("USER"))
3664 .unwrap_or_default();
3665 if !run_user.is_empty() {
3666 let chown_cmd = NativeCommand::from_shell_cmd(
3667 "chown",
3668 &[
3669 "-R".to_string(),
3670 format!("{run_user}:{run_user}"),
3671 target_dir.clone(),
3672 ],
3673 &[],
3674 );
3675 let _ = execute_native_command(&chown_cmd, sudo_pw, sender, &component.name);
3676 }
3677
3678 let cmd = inst.build_deps_install_command();
3680 execute_native_command(
3681 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
3682 sudo_pw,
3683 sender,
3684 &component.name,
3685 )?;
3686
3687 let cmd = inst.build_pip_install_command(&target_dir);
3689 execute_native_command(
3690 &NativeCommand::from_shell_cmd_with_dir(
3691 &cmd.program,
3692 &cmd.args,
3693 &cmd.env,
3694 cmd.working_dir.clone(),
3695 ),
3696 None,
3697 sender,
3698 &component.name,
3699 )?;
3700
3701 let run_user = std::env::var("SUDO_USER")
3705 .or_else(|_| std::env::var("USER"))
3706 .unwrap_or_default();
3707 if !run_user.is_empty() {
3708 let user_home = &ctx.user_home;
3709 for aiter_dir in [
3710 format!("{user_home}/.aiter"),
3711 format!("{user_home}/.mlstack/aiter"),
3712 ] {
3713 let dir_path = PathBuf::from(&aiter_dir);
3714 if dir_path.exists() {
3715 let needs_fix = std::fs::read_dir(&dir_path)
3717 .map(|mut entries| entries.next().is_some())
3718 .unwrap_or(false)
3719 && std::fs::write(dir_path.join(".permission_check"), b"").is_err();
3720 if needs_fix {
3721 let chown_cmd = NativeCommand::from_shell_cmd(
3722 "chown",
3723 &[
3724 "-R".to_string(),
3725 format!("{run_user}:{run_user}"),
3726 aiter_dir.clone(),
3727 ],
3728 &[],
3729 );
3730 let _ = execute_native_command(
3731 &chown_cmd,
3732 sudo_pw,
3733 sender,
3734 &component.name,
3735 );
3736 } else {
3737 let _ = std::fs::remove_file(dir_path.join(".permission_check"));
3739 }
3740 }
3741 }
3742 }
3743 }
3744 "vllm-studio" => {
3745 use crate::installers::components::vllm_studio::{
3746 VllmStudioConfig, VllmStudioInstaller,
3747 };
3748 let inst = VllmStudioInstaller::new(VllmStudioConfig::default());
3749
3750 let clone_cmd = inst.build_git_clone_command();
3752 let target_dir = clone_cmd.args.last().cloned().unwrap_or_default();
3753 git_clone_or_pull(
3754 inst.repo_url(),
3755 &target_dir,
3756 &[],
3757 None,
3758 sender,
3759 &component.name,
3760 )?;
3761
3762 let pkg_mgr = if crate::installers::common::utils::command_exists("npm") {
3764 "npm"
3765 } else {
3766 "yarn"
3767 };
3768 let cmd = inst.build_pkg_install_command(pkg_mgr);
3769 execute_native_command(
3770 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
3771 None,
3772 sender,
3773 &component.name,
3774 )?;
3775
3776 let frontend_dir = PathBuf::from(&target_dir).join("frontend");
3779 let cmd = inst.build_frontend_build_command(pkg_mgr);
3780 let build_result = execute_native_command(
3781 &NativeCommand::from_shell_cmd_with_dir(
3782 &cmd.program,
3783 &cmd.args,
3784 &cmd.env,
3785 Some(frontend_dir),
3786 ),
3787 None,
3788 sender,
3789 &component.name,
3790 );
3791 if let Err(err) = build_result {
3792 let _ = sender.send(InstallerEvent::Log(
3793 format!(
3794 "[native] [WARN] {} frontend build failed (non-fatal): {}",
3795 component.name, err
3796 ),
3797 false,
3798 ));
3799 }
3800 }
3801
3802 "comfyui" => {
3804 use crate::installers::components::comfyui::{ComfyuiConfig, ComfyuiInstaller};
3805 let inst = ComfyuiInstaller::new(ComfyuiConfig::default());
3806
3807 let cmd = inst.build_pytorch_check_command();
3809 execute_native_command(
3810 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
3811 None,
3812 sender,
3813 &component.name,
3814 )?;
3815
3816 let clone_cmd = inst.build_git_clone_command();
3818 let target_dir = clone_cmd.args.last().cloned().unwrap_or_default();
3819 git_clone_or_pull(
3820 inst.repo_url(),
3821 &target_dir,
3822 &[],
3823 None,
3824 sender,
3825 &component.name,
3826 )?;
3827
3828 let req_path = format!("{}/requirements.txt", target_dir);
3831 let cmd = inst.build_pip_install_command(&req_path);
3832 execute_native_command(
3833 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
3834 None,
3835 sender,
3836 &component.name,
3837 )?;
3838 }
3839
3840 "textgen" => {
3842 use crate::installers::components::textgen::{TextgenConfig, TextgenInstaller};
3843 let textgen_python = ctx
3844 .env_exports
3845 .get("MLSTACK_PYTHON_BIN")
3846 .map(|v| v.trim().to_string())
3847 .filter(|v| !v.is_empty())
3848 .unwrap_or_else(|| "python3".to_string());
3849 let inst = TextgenInstaller::new(TextgenConfig {
3853 python_bin: textgen_python.clone(),
3854 use_uv: false,
3855 break_system_packages: true,
3856 ..Default::default()
3857 });
3858
3859 let clone_cmd = inst.build_git_clone_command();
3861 let target_dir = clone_cmd.args.last().cloned().unwrap_or_default();
3862 git_clone_or_pull(
3863 inst.repo_url(),
3864 &target_dir,
3865 &[],
3866 None,
3867 sender,
3868 &component.name,
3869 )?;
3870
3871 let amd_req_full = format!("{}/requirements/full/requirements_amd.txt", target_dir);
3876 let amd_req_subdir = format!("{}/requirements/requirements_amd.txt", target_dir);
3877 let amd_req_path = if std::path::Path::new(&amd_req_full).exists() {
3878 amd_req_full
3879 } else if std::path::Path::new(&amd_req_subdir).exists() {
3880 amd_req_subdir
3881 } else {
3882 String::new()
3883 };
3884 if !amd_req_path.is_empty() {
3885 let _ = sender.send(InstallerEvent::Log(
3886 "[native] Installing AMD-specific requirements first (ROCm wheels)".to_string(),
3887 false,
3888 ));
3889 let cmd = inst.build_amd_requirements_command(&amd_req_path);
3890 execute_native_command(
3891 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
3892 None,
3893 sender,
3894 &component.name,
3895 )?;
3896 }
3897
3898 let req_path_root = format!("{}/requirements.txt", target_dir);
3905 let req_path_full = format!("{}/requirements/full/requirements.txt", target_dir);
3906 let req_path_subdir = format!("{}/requirements/requirements.txt", target_dir);
3907 let raw_req_path = if std::path::Path::new(&req_path_full).exists() {
3908 req_path_full
3909 } else if std::path::Path::new(&req_path_root).exists() {
3910 req_path_root
3911 } else if std::path::Path::new(&req_path_subdir).exists() {
3912 req_path_subdir
3913 } else {
3914 String::new()
3915 };
3916
3917 if !raw_req_path.is_empty() {
3918 let filtered = filter_cuda_requirements(&raw_req_path);
3920 if !filtered.is_empty() {
3921 let tmp_dir = std::env::temp_dir().join("rusty-stack-textgen");
3923 let _ = std::fs::create_dir_all(&tmp_dir);
3924 let filtered_path = tmp_dir.join("requirements_filtered.txt");
3925 std::fs::write(&filtered_path, &filtered).with_context(|| {
3926 format!(
3927 "Failed to write filtered requirements to {:?}",
3928 filtered_path
3929 )
3930 })?;
3931 let filtered_path_str = filtered_path.to_string_lossy().to_string();
3932 let _ = sender.send(InstallerEvent::Log(
3933 format!(
3934 "[native] Installing filtered requirements (CUDA deps excluded) from {}",
3935 raw_req_path
3936 ),
3937 false,
3938 ));
3939 let cmd = inst.build_pip_install_command(&filtered_path_str);
3940 execute_native_command(
3941 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
3942 None,
3943 sender,
3944 &component.name,
3945 )?;
3946 let _ = std::fs::remove_file(&filtered_path);
3947 } else {
3948 let _ = sender.send(InstallerEvent::Log(
3949 "[native] [WARN] All requirements were filtered out, nothing to install"
3950 .to_string(),
3951 false,
3952 ));
3953 }
3954 } else {
3955 let _ = sender.send(InstallerEvent::Log(
3956 "[native] [WARN] requirements.txt not found, skipping dependency installation"
3957 .to_string(),
3958 false,
3959 ));
3960 }
3961 }
3962
3963 "onnx" => {
3965 use crate::installers::components::onnxruntime::{
3966 OnnxInstallMethod, OnnxRuntimeConfig, OnnxRuntimeInstaller,
3967 };
3968
3969 let detected_rocm_version = detect_rocm_version();
3970 let detected_gpu_arch = detect_gpu_arch();
3971
3972 let config = OnnxRuntimeConfig {
3974 rocm_version: Some(detected_rocm_version.clone()),
3975 rocm_release: Some({
3976 let parts: Vec<&str> = detected_rocm_version.split('.').collect();
3979 if parts.len() >= 3 {
3980 detected_rocm_version.clone()
3981 } else if parts.len() >= 2 && parts[0] == "7" && parts[1] == "2" {
3982 "7.2.4".to_string()
3983 } else if parts.len() >= 2 {
3984 format!("{}.{}.3", parts[0], parts[1])
3985 } else {
3986 "7.2.4".to_string()
3987 }
3988 }),
3989 gpu_arch: Some(detected_gpu_arch),
3990 install_method: OnnxInstallMethod::MigraphxWheel,
3991 ..Default::default()
3992 };
3993 let inst = OnnxRuntimeInstaller::new(config);
3994
3995 let cmd = inst.build_uninstall_command();
3997 let _ = execute_native_command(
3998 &NativeCommand::from_shell_cmd_with_dir(
3999 &cmd.program,
4000 &cmd.args,
4001 &cmd.env,
4002 cmd.working_dir.clone(),
4003 ),
4004 None,
4005 sender,
4006 &component.name,
4007 );
4008
4009 let _ = sender.send(InstallerEvent::Log(
4011 format!(
4012 "[native] {} — installing onnxruntime-migraphx from AMD repo (ROCm {})",
4013 component.name, detected_rocm_version
4014 ),
4015 false,
4016 ));
4017 let cmd = inst.build_migraphx_install_command();
4018 execute_native_command(
4019 &NativeCommand::from_shell_cmd_with_dir(
4020 &cmd.program,
4021 &cmd.args,
4022 &cmd.env,
4023 cmd.working_dir.clone(),
4024 ),
4025 None,
4026 sender,
4027 &component.name,
4028 )?;
4029
4030 let model_dirs = [
4032 dirs::home_dir().map(|h| h.join(".mlstack/models")),
4033 dirs::home_dir().map(|h| h.join(".local/share/models")),
4034 ];
4035 for model_dir in model_dirs.iter().flatten() {
4036 if model_dir.exists() {
4037 if let Ok(entries) = std::fs::read_dir(model_dir) {
4038 for entry in entries.flatten() {
4039 let path = entry.path();
4040 if path.extension().map(|e| e == "onnx").unwrap_or(false) {
4041 let model_path = path.to_string_lossy().to_string();
4042 if model_path.ends_with(".onnx.optimized") {
4044 continue;
4045 }
4046 let _ = sender.send(InstallerEvent::Log(
4047 format!(
4048 "[native] {} — optimizing model: {}",
4049 component.name, model_path
4050 ),
4051 false,
4052 ));
4053 let cmd = inst.build_model_optimizer_command(&model_path);
4054 let _ = execute_native_command(
4055 &NativeCommand::from_shell_cmd_with_dir(
4056 &cmd.program,
4057 &cmd.args,
4058 &cmd.env,
4059 cmd.working_dir.clone(),
4060 ),
4061 None,
4062 sender,
4063 &component.name,
4064 );
4065 }
4066 }
4067 }
4068 }
4069 }
4070
4071 let _ = sender.send(InstallerEvent::Log(
4072 format!("[native] {} — ONNX Runtime (MIGraphX) installed successfully. \
4073 Source build available via OnnxInstallMethod::SourceBuild for ROCMExecutionProvider.",
4074 component.name),
4075 false,
4076 ));
4077 }
4078
4079 "bitsandbytes" => {
4081 use crate::installers::components::bitsandbytes_multi::{
4082 BitsAndBytesConfig, BitsAndBytesInstaller,
4083 };
4084 let inst = BitsAndBytesInstaller::new(BitsAndBytesConfig::default());
4085 let target_dir = format!("{}/.mlstack/bitsandbytes", ctx.user_home);
4086
4087 if is_force_reinstall() {
4089 purge_component_packages(&["bitsandbytes"], sender, &component.name);
4090 }
4091
4092 git_clone_or_pull(
4094 "https://github.com/ROCm/bitsandbytes.git",
4095 &target_dir,
4096 &["--recursive"],
4097 sudo_pw,
4098 sender,
4099 &component.name,
4100 )?;
4101
4102 let cmd = inst.build_pypi_install_command();
4104 execute_native_command(
4105 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4106 None,
4107 sender,
4108 &component.name,
4109 )?;
4110 }
4111
4112 "rocm-smi" => {
4114 use crate::installers::components::rocm_smi::{RocmSmiConfig, RocmSmiInstaller};
4115 let inst = RocmSmiInstaller::new(RocmSmiConfig::default());
4116 let distro = crate::installers::common::DistroFacade::detect();
4117 let target_dir = format!("{}/.mlstack/rocm_smi", ctx.user_home);
4118
4119 git_clone_or_pull(
4121 inst.git_clone_url(),
4122 &target_dir,
4123 &[],
4124 sudo_pw,
4125 sender,
4126 &component.name,
4127 )?;
4128
4129 if !(distro.family() == crate::platform::detection::DistroFamily::Arch
4131 && crate::installers::common::utils::command_exists("rocm-smi"))
4132 {
4133 let cmd = inst.build_package_install_command(&distro);
4134 execute_native_command(
4135 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4136 sudo_pw,
4137 sender,
4138 &component.name,
4139 )?;
4140 } else {
4141 let _ = sender.send(InstallerEvent::Log(
4142 format!(
4143 "[native] {} — rocm-smi already on PATH, skipping package install",
4144 component.name
4145 ),
4146 false,
4147 ));
4148 }
4149
4150 let cmd = inst.build_pip_install_command(&target_dir);
4152 execute_native_command(
4153 &NativeCommand::from_shell_cmd_with_dir(
4154 &cmd.program,
4155 &cmd.args,
4156 &cmd.env,
4157 cmd.working_dir.clone(),
4158 ),
4159 None,
4160 sender,
4161 &component.name,
4162 )?;
4163 }
4164
4165 "migraphx" => {
4167 use crate::installers::components::migraphx_multi::{
4168 MigraphxConfig, MigraphxInstaller,
4169 };
4170 let inst = MigraphxInstaller::new(MigraphxConfig::default());
4171 let distro = crate::installers::common::DistroFacade::detect();
4172 let packages = inst.required_packages(&distro);
4173
4174 let cmd = inst.build_package_install_command(&distro, &packages);
4176 execute_native_command(
4177 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4178 sudo_pw,
4179 sender,
4180 &component.name,
4181 )?;
4182
4183 if inst.should_skip_pip_install(&distro) {
4185 if let Some(msg) = inst.build_limitation_message(&distro) {
4186 let _ = sender.send(InstallerEvent::Log(
4187 format!("[native] [INFO] {}", msg),
4188 false,
4189 ));
4190 }
4191 } else {
4192 let cmd = inst.build_pip_install_command();
4193 execute_native_command(
4194 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4195 None,
4196 sender,
4197 &component.name,
4198 )?;
4199 }
4200 }
4201
4202 "pytorch-profiler" => {
4204 use crate::installers::components::pytorch_profiler::{
4205 PytorchProfilerConfig, PytorchProfilerInstaller,
4206 };
4207 let inst = PytorchProfilerInstaller::new(PytorchProfilerConfig::default());
4208
4209 let cmd = inst.build_install_command();
4211 execute_native_command(
4212 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4213 None,
4214 sender,
4215 &component.name,
4216 )?;
4217 }
4218
4219 "wandb" => {
4221 use crate::installers::components::wandb::{WandbConfig, WandbInstaller};
4222 let inst = WandbInstaller::new(WandbConfig::default());
4223
4224 let cmd = inst.build_install_command();
4226 execute_native_command(
4227 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4228 None,
4229 sender,
4230 &component.name,
4231 )?;
4232 }
4233
4234 "amdgpu-drivers" => {
4236 use crate::installers::components::amdgpu_drivers::{AmdgpuConfig, AmdgpuInstaller};
4237 let inst = AmdgpuInstaller::new(AmdgpuConfig::default());
4238 let distro = crate::installers::common::DistroFacade::detect();
4239 let packages = inst.required_packages(&distro);
4240
4241 let cmd = inst.build_package_install_command(&distro, &packages);
4243 execute_native_command(
4244 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4245 sudo_pw,
4246 sender,
4247 &component.name,
4248 )?;
4249
4250 let gpu_arch = ctx
4252 .env_exports
4253 .get("GPU_ARCH")
4254 .map(|v| v.trim().to_string())
4255 .filter(|v| !v.is_empty())
4256 .unwrap_or_else(detect_gpu_arch);
4257 for cmd in inst.build_env_setup_commands(&gpu_arch) {
4258 execute_native_command(
4259 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4260 sudo_pw,
4261 sender,
4262 &component.name,
4263 )?;
4264 }
4265
4266 let cmd = inst.build_verify_command();
4268 execute_native_command(
4269 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4270 None,
4271 sender,
4272 &component.name,
4273 )?;
4274 }
4275
4276 "migraphx-python" => {
4278 use crate::installers::components::migraphx_python::{
4279 MigraphxPythonConfig, MigraphxPythonInstaller,
4280 };
4281 let inst = MigraphxPythonInstaller::new(MigraphxPythonConfig::default());
4282
4283 let distro = crate::installers::common::DistroFacade::detect();
4285 if !inst.is_available_on_distro(&distro) {
4286 if let Some(msg) = inst.build_unavailable_message(&distro) {
4287 let _ = sender.send(InstallerEvent::Log(
4288 format!("[native] [SKIP] {}", msg),
4289 false,
4290 ));
4291 }
4292 } else {
4293 let cmd = inst.build_install_command();
4295 execute_native_command(
4296 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4297 None,
4298 sender,
4299 &component.name,
4300 )?;
4301 }
4302 }
4303
4304 "enhanced-env" => {
4306 let _ = sender.send(InstallerEvent::Log(
4310 "[native] Enhanced environment setup handled during pre-install phase".into(),
4311 false,
4312 ));
4313 }
4314
4315 "mlperf-inference"
4317 | "rocm-benchmarks"
4318 | "gpu-memory-bandwidth"
4319 | "rocm-smi-bench"
4320 | "pytorch-performance"
4321 | "vllm-performance"
4322 | "deepspeed-performance"
4323 | "megatron-performance"
4324 | "all-benchmarks" => {
4325 run_native_benchmark(&component.id, sender)?;
4326 }
4327
4328 "fastvideo" => {
4330 use crate::installers::components::fastvideo::{FastVideoConfig, FastVideoInstaller};
4331 let gpu_arch = ctx
4332 .env_exports
4333 .get("GPU_ARCH")
4334 .map(|v| v.trim().to_string())
4335 .filter(|v| !v.is_empty())
4336 .unwrap_or_else(detect_gpu_arch);
4337 let python_bin = ctx
4338 .env_exports
4339 .get("MLSTACK_PYTHON_BIN")
4340 .map(|v| v.trim().to_string())
4341 .filter(|v| !v.is_empty())
4342 .unwrap_or_else(|| "python3".to_string());
4343 let inst = FastVideoInstaller::new(FastVideoConfig {
4344 gpu_arch: gpu_arch.clone(),
4345 python_bin: python_bin.clone(),
4346 });
4347 let build_dir = PathBuf::from("/tmp/FastVideo_ROCm_build");
4348 let kernel_dir = build_dir.join("fastvideo-kernel");
4349
4350 let mkdir_cmd = inst.mkdir_build_dir();
4352 execute_native_command(
4353 &NativeCommand::from_shell_cmd(&mkdir_cmd.program, &mkdir_cmd.args, &mkdir_cmd.env),
4354 sudo_pw,
4355 sender,
4356 &component.name,
4357 )?;
4358
4359 git_clone_or_pull(
4361 "https://github.com/scooter-lacroix/FastVideo.git",
4362 &build_dir.to_string_lossy(),
4363 &[],
4364 sudo_pw,
4365 sender,
4366 &component.name,
4367 )?;
4368
4369 let checkout_cmd = inst.git_checkout();
4371 execute_native_command(
4372 &NativeCommand::from_shell_cmd_with_dir(
4373 &checkout_cmd.program,
4374 &checkout_cmd.args,
4375 &checkout_cmd.env,
4376 Some(build_dir.clone()),
4377 ),
4378 sudo_pw,
4379 sender,
4380 &component.name,
4381 )?;
4382
4383 let submod_cmd = inst.git_submodule_init();
4385 execute_native_command(
4386 &NativeCommand::from_shell_cmd_with_dir(
4387 &submod_cmd.program,
4388 &submod_cmd.args,
4389 &submod_cmd.env,
4390 Some(kernel_dir.clone()),
4391 ),
4392 sudo_pw,
4393 sender,
4394 &component.name,
4395 )?;
4396
4397 let patch_cmd = inst.patch_cmake();
4399 execute_native_command(
4400 &NativeCommand::from_shell_cmd_with_dir(
4401 &patch_cmd.program,
4402 &patch_cmd.args,
4403 &patch_cmd.env,
4404 Some(kernel_dir.clone()),
4405 ),
4406 sudo_pw,
4407 sender,
4408 &component.name,
4409 )?;
4410
4411 let deps_cmd = inst.install_build_deps();
4413 execute_native_command(
4414 &NativeCommand::from_shell_cmd(&deps_cmd.program, &deps_cmd.args, &deps_cmd.env),
4415 sudo_pw,
4416 sender,
4417 &component.name,
4418 )?;
4419
4420 let pip_cmd = inst.pip_install_kernel();
4422 execute_native_command(
4423 &NativeCommand::from_shell_cmd_with_dir(
4424 &pip_cmd.program,
4425 &pip_cmd.args,
4426 &pip_cmd.env,
4427 Some(kernel_dir),
4428 ),
4429 sudo_pw,
4430 sender,
4431 &component.name,
4432 )?;
4433
4434 let cleanup_cmd = inst.cleanup();
4436 execute_native_command(
4437 &NativeCommand::from_shell_cmd(
4438 &cleanup_cmd.program,
4439 &cleanup_cmd.args,
4440 &cleanup_cmd.env,
4441 ),
4442 sudo_pw,
4443 sender,
4444 &component.name,
4445 )?;
4446 }
4447 "llama-cpp" => {
4449 use crate::installers::components::llama_cpp::{LlamaCppConfig, LlamaCppInstaller};
4450 let gpu_arch = ctx
4451 .env_exports
4452 .get("GPU_ARCH")
4453 .map(|v| v.trim().to_string())
4454 .filter(|v| !v.is_empty())
4455 .unwrap_or_else(detect_gpu_arch);
4456 let channel = ctx
4457 .env_exports
4458 .get("ROCM_CHANNEL")
4459 .map(|v| v.trim().to_string())
4460 .filter(|v| !v.is_empty())
4461 .unwrap_or_else(|| "latest".to_string());
4462 let inst = LlamaCppInstaller::new(LlamaCppConfig {
4463 gpu_arch: gpu_arch.clone(),
4464 channel: channel.clone(),
4465 ..Default::default()
4466 });
4467 let home = ctx.user_home;
4468 let manifest = inst.check_latest_release();
4469 let strategy = inst.resolve_install_strategy(&home, manifest.as_ref());
4470 match &strategy {
4471 crate::installers::components::llama_cpp::InstallStrategy::Prebuilt(plan) => {
4472 let _ = sender.send(InstallerEvent::Log(
4473 format!(
4474 "llama-cpp install strategy selected: prebuilt ({})",
4475 plan.arch
4476 ),
4477 false,
4478 ));
4479 if let Err(err) = inst.download_prebuilt_binary(plan) {
4480 let _ = sender.send(InstallerEvent::Log(
4481 format!(
4482 "llama-cpp prebuilt download failed; falling back to source: {}",
4483 err
4484 ),
4485 false,
4486 ));
4487 inst.run_source_install(&home).map_err(|source_err| {
4488 anyhow::anyhow!(
4489 "llama-cpp prebuilt install failed and source fallback failed: {}",
4490 source_err
4491 )
4492 })?;
4493 }
4494 }
4495 crate::installers::components::llama_cpp::InstallStrategy::Source(plan) => {
4496 let _ = sender.send(InstallerEvent::Log(
4497 format!(
4498 "llama-cpp install strategy selected: source ({})",
4499 plan.reason
4500 ),
4501 false,
4502 ));
4503 inst.run_source_install(&home).map_err(|err| {
4504 anyhow::anyhow!("llama-cpp source install failed: {}", err)
4505 })?;
4506 }
4507 }
4508
4509 let fork_dir = format!(
4511 "{}/Documents/Product/Stan-s-ML-Stack/Fork/llama.cpp-turboquant-hip",
4512 home
4513 );
4514 let verification =
4515 crate::installers::components::llama_cpp::verify_installed_binary(&home, &fork_dir);
4516 let _ = sender.send(InstallerEvent::Log(
4517 format!(
4518 "llama-cpp post-install verification: {}",
4519 verification.summary
4520 ),
4521 false,
4522 ));
4523 if verification.stronger_check_passed {
4524 let _ = sender.send(InstallerEvent::Log(
4525 format!(
4526 "llama-cpp llama-bench: prefill={:.0} t/s, decode={:.0} t/s",
4527 verification.bench_prefill_tps.unwrap_or(0.0),
4528 verification.bench_decode_tps.unwrap_or(0.0),
4529 ),
4530 false,
4531 ));
4532 }
4533 if verification.rdna3_is_valid.unwrap_or(false) {
4534 let _ = sender.send(InstallerEvent::Log(
4535 format!(
4536 "llama-cpp RDNA3 proof: device=RDNA3, WMMA={:.0} t/s",
4537 verification.rdna3_wmma_tps.unwrap_or(0.0),
4538 ),
4539 false,
4540 ));
4541 }
4542 }
4543
4544 _ => {
4546 bail!(
4547 "No native installer dispatch for component '{}'. \
4548 This should not happen — is_native_component() should return false for this ID.",
4549 component.id
4550 );
4551 }
4552 }
4553
4554 let _ = sender.send(InstallerEvent::Progress {
4555 component_id: component.id.clone(),
4556 progress: 0.9,
4557 message: format!("{} installation complete", component.name),
4558 });
4559
4560 Ok(())
4561}
4562
4563fn run_script(
4564 component: &Component,
4565 script_path: &str,
4566 sudo_password: Option<String>,
4567 batch_mode: bool,
4568 install_method: &str,
4569 sender: &Sender<InstallerEvent>,
4570 input_rx: Arc<Mutex<Receiver<String>>>,
4571) -> Result<()> {
4572 let user_home = resolve_mlstack_user_home();
4573 let user_name = std::env::var("USER").unwrap_or_else(|_| "user".into());
4574 let preserve_env = "HOME,USER,LOGNAME,PATH,MLSTACK_USER_HOME,MLSTACK_SKIP_TORCH_INSTALL,MLSTACK_PYTHON_BIN,MLSTACK_INSTALL_METHOD,INSTALL_METHOD,PIP_BREAK_SYSTEM_PACKAGES,PIP_ROOT_USER_ACTION,UV_PIP_BREAK_SYSTEM_PACKAGES,UV_SYSTEM_PYTHON,PYTHONPATH,LD_LIBRARY_PATH,ROCM_HOME,ROCM_PATH,ROCM_VERSION,ROCM_CHANNEL,GPU_ARCH,GPU_ARCHS,PYTORCH_ROCM_ARCH,HSA_OVERRIDE_GFX_VERSION,HIP_VISIBLE_DEVICES,CUDA_VISIBLE_DEVICES,AITER_JIT_DIR,HIP_PATH,HIP_ROOT_DIR,ROCM_ROOT,HIPCC_BIN_DIR,VLLM_TARGET_DEVICE,VLLM_USE_ROCM,USE_ROCM,VLLM_VERSION,UV_PYTHON,DS_ACCELERATOR,FORCE,PYTORCH_REINSTALL,MLSTACK_FORCE_REINSTALL,MLSTACK_SUDO_PASSWORD,MLSTACK_TARGET_UID,MLSTACK_TARGET_GID";
4575
4576 let venv_bin = PathBuf::from(&user_home).join("rocm_venv").join("bin");
4577 let local_bin = PathBuf::from(&user_home).join(".local").join("bin");
4578 let mut path = std::env::var("PATH").unwrap_or_default();
4579
4580 if local_bin.exists() {
4582 let local_str = local_bin.to_string_lossy();
4583 if !path.split(':').any(|entry| entry == local_str) {
4584 if path.is_empty() {
4585 path = local_str.to_string();
4586 } else {
4587 path = format!("{}:{}", path, local_str);
4588 }
4589 }
4590 }
4591
4592 if venv_bin.exists() {
4593 let venv_str = venv_bin.to_string_lossy();
4594 if !path.split(':').any(|entry| entry == venv_str) {
4595 if path.is_empty() {
4596 path = venv_str.to_string();
4597 } else {
4598 path = format!("{}:{}", venv_str, path);
4599 }
4600 }
4601 }
4602
4603 #[allow(clippy::needless_borrow)] let env_exports = load_mlstack_env_exports(&user_home);
4605 let persistent_python = env_exports
4606 .get("MLSTACK_PYTHON_BIN")
4607 .map(|v| v.trim().to_string())
4608 .filter(|v| !v.is_empty());
4609 let python_bin = persistent_python.unwrap_or_else(resolve_python_bin);
4610 let force_reinstall = is_force_reinstall();
4611 let force_env_value = if force_reinstall { "true" } else { "false" };
4612 let sudo_password_env = sudo_password.clone();
4613 let user_home_meta = fs::metadata(&user_home).ok();
4614 #[cfg(unix)]
4615 let (mlstack_target_uid, mlstack_target_gid) = {
4616 use std::os::unix::fs::MetadataExt;
4617 (
4618 user_home_meta.as_ref().map(|meta| meta.uid().to_string()),
4619 user_home_meta.as_ref().map(|meta| meta.gid().to_string()),
4620 )
4621 };
4622 #[cfg(not(unix))]
4623 let (mlstack_target_uid, mlstack_target_gid): (Option<String>, Option<String>) = (None, None);
4624
4625 let component_needs_sudo = component.needs_sudo && needs_sudo();
4627
4628 let mut command = if component_needs_sudo {
4629 let mut cmd = Command::new("sudo");
4630 cmd.arg("-S")
4631 .arg("-p")
4632 .arg("")
4633 .arg(format!("--preserve-env={}", preserve_env))
4634 .arg("bash")
4635 .arg(script_path);
4636 cmd
4637 } else {
4638 let mut cmd = Command::new("bash");
4639 cmd.arg(script_path);
4640 cmd
4641 };
4642
4643 let mut pythonpath = std::env::var("PYTHONPATH").unwrap_or_default();
4644 let rocm_lib = "/opt/rocm/lib";
4645 if !pythonpath.split(':').any(|entry| entry == rocm_lib) {
4646 if pythonpath.is_empty() {
4647 pythonpath = rocm_lib.to_string();
4648 } else {
4649 pythonpath = format!("{}:{}", rocm_lib, pythonpath);
4650 }
4651 }
4652
4653 pythonpath = pythonpath
4655 .split(':')
4656 .filter(|part| !part.starts_with("/tmp/"))
4657 .collect::<Vec<_>>()
4658 .join(":");
4659
4660 let rocm_version = env_exports
4661 .get("ROCM_VERSION")
4662 .map(|v| v.trim().to_string())
4663 .filter(|v| !v.is_empty())
4664 .unwrap_or_else(detect_rocm_version);
4665 let gpu_arch = env_exports
4666 .get("GPU_ARCH")
4667 .map(|v| v.trim().to_string())
4668 .filter(|v| !v.is_empty())
4669 .or_else(|| {
4670 env_exports
4671 .get("PYTORCH_ROCM_ARCH")
4672 .map(|v| v.split(';').next().unwrap_or("").trim().to_string())
4673 .filter(|v| !v.is_empty())
4674 })
4675 .unwrap_or_else(detect_gpu_arch);
4676 let py_rocm_arch = env_exports
4677 .get("PYTORCH_ROCM_ARCH")
4678 .map(|v| v.trim().to_string())
4679 .filter(|v| !v.is_empty())
4680 .unwrap_or_else(|| gpu_arch.clone());
4681 let mut gpu_archs = env_exports
4682 .get("GPU_ARCHS")
4683 .map(|v| v.trim().to_string())
4684 .filter(|v| !v.is_empty())
4685 .unwrap_or_else(|| gpu_arch.clone());
4686 let mut gpu_list = env_exports
4687 .get("HIP_VISIBLE_DEVICES")
4688 .map(|v| v.trim().to_string())
4689 .filter(|v| !v.is_empty())
4690 .or_else(|| {
4691 env_exports
4692 .get("CUDA_VISIBLE_DEVICES")
4693 .map(|v| v.trim().to_string())
4694 .filter(|v| !v.is_empty())
4695 })
4696 .unwrap_or_else(detect_gpu_list);
4697 let hsa_override = env_exports
4698 .get("HSA_OVERRIDE_GFX_VERSION")
4699 .map(|v| v.trim().to_string())
4700 .filter(|v| !v.is_empty());
4701 let aiter_jit_dir = PathBuf::from(&user_home)
4702 .join(".mlstack")
4703 .join("aiter")
4704 .join("jit");
4705
4706 if component.id == "aiter" {
4707 gpu_archs = py_rocm_arch
4708 .split(';')
4709 .next()
4710 .map(|v| v.trim().to_string())
4711 .filter(|v| !v.is_empty())
4712 .unwrap_or_else(|| gpu_arch.clone());
4713 gpu_list = prioritize_gpu_list_for_arch(&gpu_list, &gpu_archs);
4714 let _ = fs::create_dir_all(&aiter_jit_dir);
4715 }
4716
4717 command
4718 .stdin(Stdio::piped())
4719 .stdout(Stdio::piped())
4720 .stderr(Stdio::piped())
4721 .env("HOME", &user_home)
4722 .env("USER", &user_name)
4723 .env("LOGNAME", &user_name)
4724 .env("PATH", &path)
4725 .env("PYTHONPATH", &pythonpath)
4726 .env("MLSTACK_USER_HOME", &user_home)
4727 .env("MLSTACK_PYTHON_BIN", &python_bin)
4728 .env("UV_PYTHON", &python_bin)
4729 .env("ROCM_HOME", "/opt/rocm")
4730 .env("ROCM_PATH", "/opt/rocm")
4731 .env("ROCM_VERSION", &rocm_version)
4732 .env("ROCM_CHANNEL", "latest")
4733 .env("GPU_ARCH", &gpu_arch)
4734 .env("GPU_ARCHS", &gpu_archs)
4735 .env("PYTORCH_ROCM_ARCH", &py_rocm_arch)
4736 .env("HIP_PATH", "/opt/rocm")
4737 .env("HIP_ROOT_DIR", "/opt/rocm")
4738 .env("ROCM_ROOT", "/opt/rocm")
4739 .env("HIPCC_BIN_DIR", "/opt/rocm/bin")
4740 .env("HIP_VISIBLE_DEVICES", &gpu_list)
4741 .env("CUDA_VISIBLE_DEVICES", &gpu_list)
4742 .env("MLSTACK_BATCH_MODE", if batch_mode { "1" } else { "0" })
4743 .env("MLSTACK_INSTALL_METHOD", install_method)
4744 .env("INSTALL_METHOD", install_method)
4745 .env("PIP_BREAK_SYSTEM_PACKAGES", "1")
4746 .env("PIP_ROOT_USER_ACTION", "ignore")
4747 .env("UV_PIP_BREAK_SYSTEM_PACKAGES", "1")
4748 .env("UV_SYSTEM_PYTHON", "1")
4749 .env("FORCE", force_env_value)
4750 .env("PYTORCH_REINSTALL", force_env_value)
4751 .env("MLSTACK_FORCE_REINSTALL", force_env_value)
4752 .env(
4753 "MLSTACK_SKIP_TORCH_INSTALL",
4754 if component.id == "pytorch" { "0" } else { "1" },
4755 )
4756 .env("RUSTY_STACK", "true");
4757
4758 if let Some(uid) = mlstack_target_uid.as_deref() {
4759 command.env("MLSTACK_TARGET_UID", uid);
4760 }
4761 if let Some(gid) = mlstack_target_gid.as_deref() {
4762 command.env("MLSTACK_TARGET_GID", gid);
4763 }
4764
4765 if let Some(password) = sudo_password_env {
4766 command.env("MLSTACK_SUDO_PASSWORD", password);
4767 }
4768
4769 if force_reinstall {
4770 if let Ok(script_contents) = fs::read_to_string(script_path) {
4771 let supports_force_flag = script_contents.contains("--force)")
4772 || script_contents.contains("--force ")
4773 || script_contents.contains("\"--force\"")
4774 || script_contents.contains("'--force'");
4775 if supports_force_flag {
4776 command.arg("--force");
4777 }
4778 }
4779 }
4780
4781 if let Some(value) = hsa_override {
4782 command.env("HSA_OVERRIDE_GFX_VERSION", value);
4783 }
4784 if component.id == "aiter" {
4785 command.env("AITER_JIT_DIR", aiter_jit_dir);
4786 }
4787
4788 let mut child = command.spawn().with_context(|| {
4789 format!(
4790 "Failed to spawn installer script '{}' for {}. \
4791 Ensure bash is installed and the script file is accessible.",
4792 script_path, component.name
4793 )
4794 })?;
4795
4796 let mut child_stdin = child.stdin.take().with_context(|| {
4797 format!(
4798 "Failed to open stdin for script '{}' — \
4799 this usually indicates a system resource issue.",
4800 script_path
4801 )
4802 })?;
4803 if let Some(password) = sudo_password {
4804 let _ = child_stdin.write_all(password.as_bytes());
4805 let _ = child_stdin.write_all(b"\n");
4806 let _ = child_stdin.flush();
4807 }
4808
4809 if let Some(default_input) = default_component_input(&component.id) {
4811 let _ = child_stdin.write_all(default_input.as_bytes());
4812 let _ = child_stdin.flush();
4813 }
4814
4815 let stop_input = Arc::new(AtomicBool::new(false));
4816
4817 let stop_input_clone = Arc::clone(&stop_input);
4819 let input_thread = thread::spawn(move || {
4820 while !stop_input_clone.load(Ordering::Relaxed) {
4821 let rx = match input_rx.lock() {
4822 Ok(guard) => guard,
4823 Err(_) => break, };
4825 match rx.recv_timeout(Duration::from_millis(100)) {
4826 Ok(input) => {
4827 let _ = child_stdin.write_all(input.as_bytes());
4828 let _ = child_stdin.flush();
4829 }
4830 Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {
4831 continue;
4832 }
4833 Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => break,
4834 }
4835 }
4836 });
4837
4838 let sender_stdout = sender.clone();
4839 let component_id = component.id.clone();
4840 let component_name = component_name(component);
4841 let stdout_stream = match child.stdout.take() {
4842 Some(s) => s,
4843 None => {
4844 stop_input.store(true, Ordering::Relaxed);
4845 let _ = input_thread.join();
4846 let _ = sender.send(InstallerEvent::Log(
4847 format!(
4848 "[ERROR] {} — failed to capture stdout from script process",
4849 component.name
4850 ),
4851 false,
4852 ));
4853 bail!(
4854 "{} — failed to capture stdout from script process. \
4855 This usually indicates a system resource issue.",
4856 component.name
4857 );
4858 }
4859 };
4860 let stdout_handle = thread::spawn(move || {
4861 let mut reader = BufReader::new(stdout_stream);
4862 let mut buffer = Vec::new();
4863 let mut chunk = [0u8; 1024];
4864
4865 while let Ok(n) = reader.read(&mut chunk) {
4866 if n == 0 {
4867 break;
4868 }
4869 for &b in chunk.iter().take(n) {
4870 if b == b'\n' || b == b'\r' {
4871 let is_transient = b == b'\r';
4872 if !buffer.is_empty() {
4873 let line = String::from_utf8_lossy(&buffer).to_string();
4874 let clean_line = strip_ansi_codes(&line);
4875 if !clean_line.trim().is_empty() {
4876 let _ = sender_stdout
4877 .send(InstallerEvent::Log(clean_line.clone(), is_transient));
4878 if let Some(progress) = parse_progress(&clean_line) {
4879 let _ = sender_stdout.send(InstallerEvent::Progress {
4880 component_id: component_id.clone(),
4881 progress,
4882 message: format!(
4883 "{}: {}%",
4884 component_name,
4885 (progress * 100.0) as i32
4886 ),
4887 });
4888 }
4889 }
4890 buffer.clear();
4891 }
4892 } else {
4893 buffer.push(b);
4894 if buffer.len() > 8192 {
4895 buffer.clear();
4896 }
4897 }
4898 }
4899 }
4900 });
4901
4902 let sender_stderr = sender.clone();
4903 let stderr_stream = match child.stderr.take() {
4904 Some(s) => s,
4905 None => {
4906 stop_input.store(true, Ordering::Relaxed);
4907 let _ = input_thread.join();
4908 let _ = stdout_handle.join();
4909 let _ = sender.send(InstallerEvent::Log(
4910 format!(
4911 "[ERROR] {} — failed to capture stderr from script process",
4912 component.name
4913 ),
4914 false,
4915 ));
4916 bail!(
4917 "{} — failed to capture stderr from script process. \
4918 This usually indicates a system resource issue.",
4919 component.name
4920 );
4921 }
4922 };
4923 let stderr_handle = thread::spawn(move || {
4924 let mut reader = BufReader::new(stderr_stream);
4925 let mut buffer = Vec::new();
4926 let mut chunk = [0u8; 1024];
4927
4928 while let Ok(n) = reader.read(&mut chunk) {
4929 if n == 0 {
4930 break;
4931 }
4932 for &b in chunk.iter().take(n) {
4933 if b == b'\n' || b == b'\r' {
4934 let is_transient = b == b'\r';
4935 if !buffer.is_empty() {
4936 let line = String::from_utf8_lossy(&buffer).to_string();
4937 let clean_line = strip_ansi_codes(&line);
4938 if !clean_line.trim().is_empty() {
4939 let _ =
4940 sender_stderr.send(InstallerEvent::Log(clean_line, is_transient));
4941 }
4942 buffer.clear();
4943 }
4944 } else {
4945 buffer.push(b);
4946 if buffer.len() > 8192 {
4947 buffer.clear();
4948 }
4949 }
4950 }
4951 }
4952 });
4953
4954 let status = child.wait().with_context(|| {
4955 format!(
4956 "{} — script process wait failed. \
4957 The process may have been killed or the system is under heavy load.",
4958 component.name
4959 )
4960 })?;
4961 stop_input.store(true, Ordering::Relaxed);
4962 let _ = input_thread.join();
4963 let _ = stdout_handle.join();
4964 let _ = stderr_handle.join();
4965
4966 if !status.success() {
4967 let code = status.code().unwrap_or(-1);
4968 let error_detail = match code {
4969 1 => "General error — check the script output above for details".to_string(),
4970 2 => "Misuse of shell builtins — check script arguments".to_string(),
4971 126 => "Script not executable — check file permissions (chmod +x)".to_string(),
4972 127 => "Command not found in script — ensure required tools are installed".to_string(),
4973 130 => "Script terminated by Ctrl+C (SIGINT)".to_string(),
4974 137 => "Script killed (SIGKILL) — possibly out of memory".to_string(),
4975 139 => "Segmentation fault in script subprocess".to_string(),
4976 n if n > 128 => format!(
4977 "Script process terminated by signal {} — check system logs",
4978 n - 128
4979 ),
4980 _ => format!(
4981 "Script exited with code {} — check the output above for details",
4982 code
4983 ),
4984 };
4985 let _ = sender.send(InstallerEvent::Log(
4986 format!(
4987 "[ERROR] {} failed: {} (exit code {} at {})",
4988 component.name,
4989 error_detail,
4990 code,
4991 Local::now().format("%H:%M:%S")
4992 ),
4993 false,
4994 ));
4995 bail!(
4996 "{} failed: {} (exit code {})",
4997 component.name,
4998 error_detail,
4999 code
5000 );
5001 }
5002
5003 Ok(())
5004}
5005
5006fn star_repo(config: &InstallerConfig, sender: &Sender<InstallerEvent>) -> Result<()> {
5007 if !config.star_repos {
5008 return Ok(());
5009 }
5010
5011 let repo_url = "https://github.com/scooter-lacroix/Stan-s-ML-Stack";
5012
5013 let gh_check = Command::new("gh")
5015 .arg("auth")
5016 .arg("status")
5017 .stdout(Stdio::null())
5018 .stderr(Stdio::null())
5019 .status();
5020
5021 if let Ok(status) = gh_check {
5022 if status.success() {
5023 let _ = sender.send(InstallerEvent::Log(
5024 "Starring ML Stack repository on GitHub...".into(),
5025 false,
5026 ));
5027 let star_status = Command::new("gh")
5028 .arg("repo")
5029 .arg("star")
5030 .arg(repo_url)
5031 .stdout(Stdio::null())
5032 .stderr(Stdio::null())
5033 .status();
5034
5035 if let Ok(s) = star_status {
5036 if s.success() {
5037 let _ = sender.send(InstallerEvent::Log(
5038 "Thank you for starring the repo! ✦".into(),
5039 false,
5040 ));
5041 }
5042 }
5043 }
5044 }
5045
5046 Ok(())
5047}
5048
5049fn component_name(comp: &Component) -> String {
5050 comp.name.clone()
5051}
5052
5053fn format_user_friendly_error(component_id: &str, raw_error: &str) -> String {
5062 let error_lower = raw_error.to_lowercase();
5063
5064 match component_id {
5067 "permanent-env" => {
5068 if error_lower.contains("rocminfo") {
5069 return format!(
5070 "{}: Failed to detect GPU — rocminfo command failed. \
5071 Ensure ROCm is properly installed and /opt/rocm/bin is in your PATH. \
5072 Original error: {}",
5073 component_id, raw_error
5074 );
5075 }
5076 return format!(
5077 "{}: Environment setup failed. \
5078 Ensure ROCm is installed and GPU is accessible. \
5079 Original error: {}",
5080 component_id, raw_error
5081 );
5082 }
5083 "rocm" if error_lower.contains("no such file") || error_lower.contains("not found") => {
5084 return format!(
5085 "{}: ROCm installation failed — required packages or repositories not found. \
5086 Ensure your system's package manager is configured correctly and ROCm repositories are added. \
5087 Original error: {}",
5088 component_id, raw_error
5089 );
5090 }
5091 "pytorch" => {
5092 if error_lower.contains("no such file") || error_lower.contains("not found") {
5093 return format!(
5094 "{}: PyTorch installation failed — required packages not found. \
5095 Ensure ROCm is installed first and Python 3.10+ is available. \
5096 Original error: {}",
5097 component_id, raw_error
5098 );
5099 }
5100 return format!(
5101 "{}: PyTorch installation failed. \
5102 Ensure ROCm is installed first, Python 3.10-3.13 is available, \
5103 and pip/uv can reach the PyTorch package index. \
5104 Original error: {}",
5105 component_id, raw_error
5106 );
5107 }
5108 "triton" | "flash-attn" | "vllm" | "aiter" => {
5109 return format!(
5110 "{}: Installation failed. \
5111 Ensure PyTorch and ROCm are properly installed first \
5112 (these are required dependencies). \
5113 Original error: {}",
5114 component_id, raw_error
5115 );
5116 }
5117 "megatron" => {
5118 return format!(
5119 "{}: Megatron-LM installation failed. \
5120 Ensure PyTorch and MPI4Py are installed first (required dependencies). \
5121 Original error: {}",
5122 component_id, raw_error
5123 );
5124 }
5125 "deepspeed" => {
5126 return format!(
5127 "{}: DeepSpeed installation failed. \
5128 Ensure PyTorch is installed first (required dependency). \
5129 Original error: {}",
5130 component_id, raw_error
5131 );
5132 }
5133 _ => {}
5134 }
5135
5136 if error_lower.contains("permission denied") || error_lower.contains("sudo") {
5139 return format!(
5140 "{}: Permission denied. Try running with elevated privileges or check file/directory permissions. \
5141 Original error: {}",
5142 component_id, raw_error
5143 );
5144 }
5145
5146 if error_lower.contains("network")
5147 || error_lower.contains("connection")
5148 || error_lower.contains("timeout")
5149 || error_lower.contains("resolve")
5150 {
5151 return format!(
5152 "{}: Network error — failed to download packages or clone repositories. \
5153 Check your internet connection and try again. \
5154 Original error: {}",
5155 component_id, raw_error
5156 );
5157 }
5158
5159 if error_lower.contains("out of memory")
5160 || error_lower.contains("cannot allocate")
5161 || error_lower.contains("killed")
5162 {
5163 return format!(
5164 "{}: Out of memory — the build process was killed due to insufficient memory. \
5165 Try closing other applications or increasing swap space. \
5166 Original error: {}",
5167 component_id, raw_error
5168 );
5169 }
5170
5171 if error_lower.contains("exit code") {
5172 return raw_error.to_string();
5174 }
5175
5176 if error_lower.contains("no such file") || error_lower.contains("not found") {
5177 return format!(
5178 "{}: Required file or command not found. \
5179 Ensure all dependencies are installed before this component. \
5180 Original error: {}",
5181 component_id, raw_error
5182 );
5183 }
5184
5185 format!(
5188 "{}: Installation failed. Check the log output above for details. \
5189 Original error: {}",
5190 component_id, raw_error
5191 )
5192}
5193
5194fn detect_gpu_arch() -> String {
5195 let rocminfo_path = crate::installers::common::utils::resolve_rocminfo_path();
5196 if let Ok(output) = Command::new(&rocminfo_path).output() {
5197 let stdout = String::from_utf8_lossy(&output.stdout);
5198 let mut best_arch = String::new();
5199 let mut best_value = 0u32;
5200 for token in stdout.split_whitespace() {
5201 let cleaned = token.trim_matches(|c: char| !c.is_alphanumeric());
5202 if !cleaned.starts_with("gfx") {
5203 continue;
5204 }
5205 let raw = cleaned.trim_start_matches("gfx");
5206 if let Ok(value) = raw.parse::<u32>() {
5207 if value >= best_value {
5208 best_value = value;
5209 best_arch = cleaned.to_string();
5210 }
5211 }
5212 }
5213 if !best_arch.is_empty() {
5214 return best_arch;
5215 }
5216 }
5217
5218 "gfx000".to_string()
5219}
5220
5221const IGPU_PATTERNS: &[&str] = &[
5224 "Cezanne",
5226 "Rembrandt",
5227 "Phoenix",
5228 "Raphael",
5229 "Barcelo",
5230 "Pink Sardine",
5231 "Yellow Carp",
5232 "Green Sardine",
5233 "Integrated",
5235 "iGPU",
5236 "APU",
5237 "AMD Radeon Graphics",
5239 "Radeon Graphics",
5240 " 5600G",
5242 " 5700G",
5243 " 5600GE",
5244 " 5700GE",
5245 " 5700G",
5246 " 5600GT",
5247 " 5700GT",
5248 " 8600G",
5249 " 8700G",
5250 " 8500G",
5251 " 8300G",
5252 "7800X3D",
5254 "7950X3D",
5255 "7900X3D",
5256 "7700X3D",
5257 "7600X3D",
5258 "Ryzen 7 7735",
5260 "Ryzen 7 7840",
5261 "Ryzen 7 8840",
5262 "Ryzen 5 7535",
5263 "Ryzen 5 7640",
5264 "Ryzen 5 8640",
5265 "Ryzen 9 7945",
5266 "Ryzen 9 8945",
5267 "Hawk Point",
5269 "Dragon Range",
5270 "Fire Range",
5271 "Stella",
5273 "Krackan",
5275];
5276
5277fn is_igpu_name(marketing_name: &str) -> bool {
5279 let name_upper = marketing_name.to_uppercase();
5280
5281 for pattern in IGPU_PATTERNS {
5283 if marketing_name.contains(pattern) {
5284 return true;
5285 }
5286 }
5287
5288 if name_upper.contains("RYZEN") && !name_upper.contains("RX") {
5300 return true;
5301 }
5302
5303 if name_upper.contains("RYZEN") {
5306 if let Some(pos) = name_upper.find(" RYZEN ") {
5308 let after_ryzen = &name_upper[pos + 7..]; let words: Vec<&str> = after_ryzen.split_whitespace().collect();
5311 if !words.is_empty() {
5312 let first_word = words[0];
5313 if first_word.ends_with('G') && first_word.len() >= 5 {
5314 if !first_word.starts_with("RX") && !first_word.starts_with("RADEON") {
5316 if !first_word.contains("XT") && !first_word.contains("XTX") {
5318 return true;
5319 }
5320 }
5321 }
5322 }
5323 }
5324 }
5325
5326 if name_upper.contains("RADEON GRAPHICS") && !name_upper.contains("RX") {
5328 let has_model = marketing_name.chars().any(|c| c.is_ascii_digit());
5330 if !has_model || marketing_name.contains("AMD Radeon Graphics") {
5331 return true;
5332 }
5333 }
5334
5335 false
5336}
5337
5338fn detect_gpu_list() -> String {
5341 let rocminfo_path = crate::installers::common::utils::resolve_rocminfo_path();
5342 if let Ok(output) = Command::new(&rocminfo_path).output() {
5343 let stdout = String::from_utf8_lossy(&output.stdout);
5344
5345 let discrete_indices = parse_rocminfo_for_discrete_gpus(&stdout);
5347
5348 if !discrete_indices.is_empty() {
5349 return discrete_indices.join(",");
5350 }
5351 }
5352
5353 if let Ok(output) = Command::new("rocm-smi")
5355 .args(["--showproductname", "--showbus", "--json"])
5356 .output()
5357 {
5358 let stdout = String::from_utf8_lossy(&output.stdout);
5359 let discrete_indices = parse_rocm_smi_for_discrete_gpus(&stdout);
5360
5361 if !discrete_indices.is_empty() {
5362 return discrete_indices.join(",");
5363 }
5364 }
5365
5366 if let Ok(output) = Command::new("lspci").output() {
5368 let stdout = String::from_utf8_lossy(&output.stdout);
5369 let discrete_indices = parse_lspci_for_discrete_gpus(&stdout);
5370
5371 if !discrete_indices.is_empty() {
5372 return discrete_indices.join(",");
5373 }
5374 }
5375
5376 if let Some(count) = detect_gpu_count_sysfs() {
5378 if count > 0 {
5379 return (0..count)
5380 .map(|idx| idx.to_string())
5381 .collect::<Vec<_>>()
5382 .join(",");
5383 }
5384 }
5385
5386 "0".to_string()
5387}
5388
5389fn detect_gpu_arch_by_id() -> HashMap<String, String> {
5391 let mut mapping: HashMap<String, String> = HashMap::new();
5392
5393 let rocminfo_path = crate::installers::common::utils::resolve_rocminfo_path();
5394 let Ok(output) = Command::new(&rocminfo_path).output() else {
5395 return mapping;
5396 };
5397 let stdout = String::from_utf8_lossy(&output.stdout);
5398 let mut current_gpu_id: Option<String> = None;
5399
5400 for line in stdout.lines() {
5401 let line = line.trim();
5402 if line.contains("GPU ID:") {
5403 if let Some(id_part) = line.split(':').nth(1) {
5404 let id: String = id_part.chars().filter(|c| c.is_ascii_digit()).collect();
5405 if !id.is_empty() {
5406 current_gpu_id = Some(id);
5407 }
5408 }
5409 continue;
5410 }
5411
5412 if !line.contains("Name:") || !line.contains("gfx") {
5413 continue;
5414 }
5415
5416 if let Some(id) = current_gpu_id.clone() {
5417 for token in line.split_whitespace() {
5418 let cleaned = token.trim_matches(|c: char| !c.is_alphanumeric());
5419 if cleaned.starts_with("gfx") {
5420 mapping.entry(id).or_insert_with(|| cleaned.to_string());
5421 break;
5422 }
5423 }
5424 }
5425 }
5426
5427 mapping
5428}
5429
5430fn prioritize_gpu_list_for_arch(gpu_list: &str, target_arch: &str) -> String {
5433 let target_arch = target_arch.trim();
5434 if target_arch.is_empty() {
5435 return gpu_list.to_string();
5436 }
5437
5438 let arch_by_id = detect_gpu_arch_by_id();
5439 if arch_by_id.is_empty() {
5440 return gpu_list.to_string();
5441 }
5442
5443 let mut ids: Vec<String> = gpu_list
5444 .split(',')
5445 .map(|v| v.trim().to_string())
5446 .filter(|v| !v.is_empty())
5447 .collect();
5448 if ids.len() < 2 {
5449 return gpu_list.to_string();
5450 }
5451
5452 let preferred_pos = ids.iter().position(|id| {
5453 arch_by_id
5454 .get(id)
5455 .map(|v| v == target_arch)
5456 .unwrap_or(false)
5457 });
5458 if let Some(pos) = preferred_pos {
5459 if pos > 0 {
5460 let preferred = ids.remove(pos);
5461 ids.insert(0, preferred);
5462 return ids.join(",");
5463 }
5464 }
5465
5466 gpu_list.to_string()
5467}
5468
5469fn parse_rocminfo_for_discrete_gpus(rocminfo_output: &str) -> Vec<String> {
5472 let mut discrete_indices: Vec<String> = Vec::new();
5473 let mut current_gpu_id: Option<usize> = None;
5474 let mut current_marketing_name = String::new();
5475 let mut current_device_type = String::new();
5476 let mut is_igpu = false;
5477
5478 for line in rocminfo_output.lines() {
5479 let line = line.trim();
5480
5481 if line.contains("GPU ID:") {
5483 if let Some(id) = current_gpu_id {
5485 if !is_igpu && !current_marketing_name.is_empty() {
5486 discrete_indices.push(id.to_string());
5487 }
5488 }
5489
5490 if let Some(id_part) = line.split(':').nth(1) {
5493 let id_str: String = id_part.chars().filter(|c| c.is_ascii_digit()).collect();
5495 current_gpu_id = id_str.parse().ok();
5496 current_marketing_name.clear();
5497 current_device_type.clear();
5498 is_igpu = false;
5499 }
5500 }
5501 else if line.contains("Marketing Name:") {
5503 if let Some(name) = line.split(':').nth(1) {
5504 let name = name.trim_end_matches('*').trim();
5506 current_marketing_name = name.to_string();
5507 is_igpu = is_igpu_name(¤t_marketing_name);
5509 }
5510 }
5511 else if line.contains("Device Type:") {
5513 if let Some(dev_type) = line.split(':').nth(1) {
5514 let dev_type = dev_type.trim_end_matches('*').trim();
5516 current_device_type = dev_type.to_string();
5517 if current_device_type.to_uppercase().contains("CPU") {
5519 is_igpu = true;
5520 }
5521 }
5522 }
5523 }
5524
5525 if let Some(id) = current_gpu_id {
5527 if !is_igpu && !current_marketing_name.is_empty() {
5528 discrete_indices.push(id.to_string());
5529 }
5530 }
5531
5532 discrete_indices
5533}
5534
5535fn parse_rocm_smi_for_discrete_gpus(rocm_smi_output: &str) -> Vec<String> {
5538 fn card_index(key: &str) -> Option<usize> {
5539 let digits_rev: String = key
5540 .chars()
5541 .rev()
5542 .take_while(|c| c.is_ascii_digit())
5543 .collect();
5544 if digits_rev.is_empty() {
5545 return None;
5546 }
5547 let digits: String = digits_rev.chars().rev().collect();
5548 digits.parse().ok()
5549 }
5550
5551 fn value_as_string(value: &serde_json::Value) -> Option<String> {
5552 if let Some(s) = value.as_str() {
5553 return Some(s.to_string());
5554 }
5555 if let Some(obj) = value.as_object() {
5556 if let Some(v) = obj.get("value").and_then(|v| v.as_str()) {
5557 return Some(v.to_string());
5558 }
5559 }
5560 None
5561 }
5562
5563 fn field_is_placeholder(raw: &str) -> bool {
5564 matches!(
5565 raw.trim().to_ascii_lowercase().as_str(),
5566 "" | "n/a" | "na" | "none" | "unknown" | "not available" | "-"
5567 )
5568 }
5569
5570 fn meaningful_field(value: Option<String>) -> Option<String> {
5571 let value = value?;
5572 let trimmed = value.trim();
5573 if trimmed.is_empty() || field_is_placeholder(trimmed) {
5574 None
5575 } else {
5576 Some(trimmed.to_string())
5577 }
5578 }
5579
5580 fn normalize_bus_id(raw: &str) -> Option<String> {
5581 let mut normalized = raw.trim().to_ascii_lowercase();
5582 if normalized.is_empty() {
5583 return None;
5584 }
5585 if let Some(idx) = normalized.find(' ') {
5586 normalized.truncate(idx);
5587 }
5588 normalized = normalized.trim_start_matches("0000:").to_string();
5589 if normalized.len() == "00:00.0".len() {
5590 Some(normalized)
5591 } else {
5592 None
5593 }
5594 }
5595
5596 fn descriptor_from_lspci_bus(bus: &str) -> Option<String> {
5597 let bus = normalize_bus_id(bus)?;
5598 let output = Command::new("lspci").args(["-s", &bus]).output().ok()?;
5599 if !output.status.success() {
5600 return None;
5601 }
5602 let line = String::from_utf8_lossy(&output.stdout).trim().to_string();
5603 if line.is_empty() {
5604 None
5605 } else {
5606 Some(line)
5607 }
5608 }
5609
5610 fn bus_looks_integrated(raw: &str) -> bool {
5611 let Some(bus) = normalize_bus_id(raw) else {
5612 return false;
5613 };
5614
5615 if let Ok(output) = Command::new("lspci").args(["-s", &bus]).output() {
5616 let line = String::from_utf8_lossy(&output.stdout).trim().to_string();
5617 if !line.is_empty() {
5618 if is_igpu_name(&line) {
5619 return true;
5620 }
5621 let lower = line.to_ascii_lowercase();
5622 if lower.contains("integrated")
5623 || lower.contains("ryzen")
5624 || lower.contains("apu")
5625 || lower.contains("raphael")
5626 || lower.contains("phoenix")
5627 || lower.contains("rembrandt")
5628 || lower.contains("renoir")
5629 || lower.contains("raven")
5630 || lower.contains("picasso")
5631 || lower.contains("cezanne")
5632 || lower.contains("mendocino")
5633 || lower.contains("hawk point")
5634 || lower.contains("strix")
5635 {
5636 return true;
5637 }
5638 }
5639 }
5640
5641 let sysfs_path = format!("/sys/bus/pci/devices/0000:{}/mem_info_vram_total", bus);
5642 if let Ok(value) = fs::read_to_string(sysfs_path) {
5643 if let Ok(vram_bytes) = value.trim().parse::<u64>() {
5644 return vram_bytes < 4 * 1024 * 1024 * 1024;
5645 }
5646 }
5647 false
5648 }
5649
5650 fn parse_rocm_smi_json(raw: &str) -> Option<serde_json::Value> {
5651 if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(raw) {
5652 return Some(parsed);
5653 }
5654 let start = raw.find('{')?;
5655 let end = raw.rfind('}')?;
5656 if end <= start {
5657 return None;
5658 }
5659 serde_json::from_str::<serde_json::Value>(&raw[start..=end]).ok()
5660 }
5661
5662 let mut cards: Vec<(usize, String)> = Vec::new();
5663 let Some(parsed) = parse_rocm_smi_json(rocm_smi_output) else {
5664 return Vec::new();
5665 };
5666 let Some(obj) = parsed.as_object() else {
5667 return Vec::new();
5668 };
5669
5670 for (card_key, payload) in obj {
5671 let Some(index) = card_index(card_key) else {
5672 continue;
5673 };
5674 let Some(payload_obj) = payload.as_object() else {
5675 continue;
5676 };
5677
5678 let series = meaningful_field(
5679 payload_obj
5680 .get("Card Series")
5681 .and_then(value_as_string)
5682 .or_else(|| payload_obj.get("Card series").and_then(value_as_string)),
5683 );
5684 let model = meaningful_field(
5685 payload_obj
5686 .get("Card Model")
5687 .and_then(value_as_string)
5688 .or_else(|| payload_obj.get("Card model").and_then(value_as_string)),
5689 );
5690 let sku = meaningful_field(
5691 payload_obj
5692 .get("Card SKU")
5693 .and_then(value_as_string)
5694 .or_else(|| {
5695 payload_obj
5696 .get("Card Product Name")
5697 .and_then(value_as_string)
5698 }),
5699 );
5700 let bus = payload_obj
5701 .get("PCI Bus")
5702 .and_then(value_as_string)
5703 .or_else(|| payload_obj.get("PCI Bus Address").and_then(value_as_string))
5704 .or_else(|| payload_obj.get("PCIe Bus").and_then(value_as_string))
5705 .or_else(|| payload_obj.get("Bus").and_then(value_as_string))
5706 .unwrap_or_default();
5707 let mut descriptor_parts = Vec::new();
5708 if let Some(value) = series {
5709 descriptor_parts.push(value);
5710 }
5711 if let Some(value) = model {
5712 descriptor_parts.push(value);
5713 }
5714 if let Some(value) = sku {
5715 descriptor_parts.push(value);
5716 }
5717 if let Some(value) = descriptor_from_lspci_bus(&bus) {
5718 if !value.trim().is_empty() {
5719 descriptor_parts.push(value);
5720 }
5721 }
5722 let resolved_descriptor = descriptor_parts.join(" ");
5723
5724 if !resolved_descriptor.trim().is_empty()
5725 && !is_igpu_name(&resolved_descriptor)
5726 && !bus_looks_integrated(&bus)
5727 {
5728 cards.push((index, index.to_string()));
5729 }
5730 }
5731
5732 cards.sort_by_key(|(idx, _)| *idx);
5733 cards.into_iter().map(|(_, idx)| idx).collect()
5734}
5735
5736fn parse_lspci_for_discrete_gpus(lspci_output: &str) -> Vec<String> {
5739 fn lspci_bus_looks_integrated(bus_id: &str) -> bool {
5740 let bus_id = bus_id.trim().trim_start_matches("0000:");
5741 if bus_id.len() != "00:00.0".len() {
5742 return false;
5743 }
5744 let sysfs_path = format!("/sys/bus/pci/devices/0000:{}/mem_info_vram_total", bus_id);
5745 if let Ok(value) = fs::read_to_string(sysfs_path) {
5746 if let Ok(vram_bytes) = value.trim().parse::<u64>() {
5747 return vram_bytes < 4 * 1024 * 1024 * 1024;
5748 }
5749 }
5750 false
5751 }
5752
5753 let mut discrete_indices: Vec<String> = Vec::new();
5754 let mut gpu_index = 0usize;
5755
5756 for line in lspci_output.lines() {
5757 let line_lower = line.to_lowercase();
5758 let bus_id = line.split_whitespace().next().unwrap_or_default();
5759
5760 if (line_lower.contains("amd")
5762 || line_lower.contains("radeon")
5763 || line_lower.contains("advanced micro devices"))
5764 && (line_lower.contains("vga")
5765 || line_lower.contains("3d")
5766 || line_lower.contains("display"))
5767 {
5768 if !is_igpu_name(line) && !lspci_bus_looks_integrated(bus_id) {
5770 discrete_indices.push(gpu_index.to_string());
5771 }
5772 gpu_index += 1;
5773 }
5774 }
5775
5776 discrete_indices
5777}
5778
5779fn detect_gpu_count_sysfs() -> Option<usize> {
5780 let entries = fs::read_dir("/sys/class/drm").ok()?;
5781 let mut count = 0usize;
5782 for entry in entries.flatten() {
5783 let name = entry.file_name();
5784 let name = name.to_string_lossy();
5785 if !name.starts_with("card") || name.contains("render") {
5786 continue;
5787 }
5788 let vendor_path = entry.path().join("device/vendor");
5789 let vendor = match fs::read_to_string(&vendor_path) {
5790 Ok(value) => value,
5791 Err(_) => continue,
5792 };
5793 if vendor.trim().eq_ignore_ascii_case("0x1002") {
5794 count += 1;
5795 }
5796 }
5797 if count > 0 {
5798 Some(count)
5799 } else {
5800 None
5801 }
5802}
5803
5804fn first_gpu_index(list: &str) -> String {
5805 list.split(',').next().unwrap_or("0").to_string()
5806}
5807
5808fn hsa_override_from_gpu_arch(gfx: &str) -> Option<String> {
5815 let digits: String = gfx
5816 .trim_start_matches("gfx")
5817 .chars()
5818 .take_while(|c| c.is_ascii_digit())
5819 .collect();
5820 if digits.len() < 3 {
5821 return None;
5822 }
5823 let major: u32 = digits[..2].parse().ok()?;
5824 if major >= 10 {
5826 Some(format!("{}.0.0", major))
5827 } else {
5828 None
5829 }
5830}
5831
5832#[derive(Debug, Clone, Copy, PartialEq, Eq)]
5833#[allow(dead_code)] enum VerificationResult {
5835 Verified,
5836 Failed,
5837 Missing,
5838 Warning,
5840}
5841
5842impl VerificationResult {
5843 fn label(&self) -> &'static str {
5844 match self {
5845 VerificationResult::Verified => "Verified",
5846 VerificationResult::Failed => "Failed",
5847 VerificationResult::Missing => "Missing",
5848 VerificationResult::Warning => "Warning",
5849 }
5850 }
5851}
5852
5853fn resolve_python_bin() -> String {
5854 if let Ok(value) = std::env::var("MLSTACK_PYTHON_BIN") {
5855 let value = value.trim();
5856 if !value.is_empty() {
5857 return value.to_string();
5858 }
5859 }
5860
5861 if let Ok(value) = std::env::var("UV_PYTHON") {
5862 let value = value.trim();
5863 if !value.is_empty() {
5864 return value.to_string();
5865 }
5866 }
5867
5868 let home = std::env::var("HOME").unwrap_or_default();
5869
5870 let uv_python_dir = PathBuf::from(&home).join(".local/share/uv/python");
5875 if let Ok(entries) = std::fs::read_dir(&uv_python_dir) {
5876 let mut uv_pythons: Vec<PathBuf> = entries
5877 .filter_map(|e| e.ok())
5878 .filter_map(|e| {
5879 let bin = e.path().join("bin/python3");
5880 if bin.exists() {
5881 Some(bin)
5882 } else {
5883 None
5884 }
5885 })
5886 .collect();
5887 uv_pythons.sort_by(|a, b| {
5889 let va = python_version_tuple(a);
5890 let vb = python_version_tuple(b);
5891 python_ml_priority(&va).cmp(&python_ml_priority(&vb))
5892 });
5893 for python in &uv_pythons {
5895 if python_has_rocm_torch(&python.to_string_lossy()) {
5896 return python.to_string_lossy().to_string();
5897 }
5898 }
5899 if let Some(first) = uv_pythons.first() {
5901 return first.to_string_lossy().to_string();
5902 }
5903 }
5904
5905 let local_bin_python = format!("{home}/.local/bin/python3");
5907 if Path::new(&local_bin_python).exists() {
5908 if python_has_rocm_torch(&local_bin_python) {
5909 return local_bin_python;
5910 }
5911 return local_bin_python;
5912 }
5913
5914 let candidates = [
5916 Path::new(&home)
5917 .join("rocm_venv")
5918 .join("bin")
5919 .join("python"),
5920 PathBuf::from("/usr/local/bin/python3"),
5921 PathBuf::from("/usr/bin/python3"),
5922 PathBuf::from("python3"),
5923 PathBuf::from("python"),
5924 ];
5925
5926 for candidate in &candidates {
5927 let candidate_str = candidate.to_string_lossy();
5928 if candidate_str.is_empty() {
5929 continue;
5930 }
5931 if candidate.is_absolute() && !candidate.exists() {
5932 continue;
5933 }
5934 if python_has_rocm_torch(&candidate_str) {
5935 return candidate_str.to_string();
5936 }
5937 }
5938
5939 for candidate in &candidates {
5940 let candidate_str = candidate.to_string_lossy();
5941 if candidate.is_absolute() && !candidate.exists() {
5942 continue;
5943 }
5944 return candidate_str.to_string();
5945 }
5946
5947 "python3".to_string()
5948}
5949
5950fn python_version_tuple(path: &Path) -> (u32, u32) {
5952 let name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
5953 let parent = path
5955 .parent()
5956 .and_then(|p| p.parent())
5957 .and_then(|p| p.file_name())
5958 .and_then(|n| n.to_str())
5959 .unwrap_or("");
5960 if let Some(v) = parse_cpython_version(parent) {
5962 return v;
5963 }
5964 if let Some(rest) = name.strip_prefix("python") {
5966 let parts: Vec<&str> = rest.split('.').collect();
5967 if parts.len() >= 2 {
5968 if let (Ok(major), Ok(minor)) = (parts[0].parse::<u32>(), parts[1].parse::<u32>()) {
5969 return (major, minor);
5970 }
5971 }
5972 }
5973 (0, 0)
5974}
5975
5976fn parse_cpython_version(dir_name: &str) -> Option<(u32, u32)> {
5978 let rest = dir_name.strip_prefix("cpython-")?;
5979 let parts: Vec<&str> = rest.split('.').collect();
5980 if parts.len() >= 2 {
5981 let major = parts[0].parse::<u32>().ok()?;
5982 let minor = parts[1].parse::<u32>().ok()?;
5983 return Some((major, minor));
5984 }
5985 None
5986}
5987
5988fn python_ml_priority(version: &(u32, u32)) -> (i32, u32) {
5990 match version {
5991 (3, 12) => (0, version.1), (3, 13) => (1, version.1),
5993 (3, 11) => (2, version.1),
5994 (3, 10) => (3, version.1),
5995 _ => (4, version.1), }
5997}
5998
5999fn python_has_rocm_torch(python: &str) -> bool {
6000 Command::new(python)
6001 .arg("-c")
6002 .arg(
6003 "import torch, sys; ok=hasattr(torch,'__version__') and hasattr(torch,'version') and getattr(torch.version,'hip', None); sys.exit(0 if ok else 1)",
6004 )
6005 .stdout(Stdio::null())
6006 .stderr(Stdio::null())
6007 .status()
6008 .map(|status| status.success())
6009 .unwrap_or(false)
6010}
6011
6012fn default_component_input(component_id: &str) -> Option<&'static str> {
6013 match component_id {
6014 "ml-stack-core" => Some("9\n0\n"),
6015 _ => None,
6016 }
6017}
6018
6019fn parse_progress(line: &str) -> Option<f32> {
6020 let lower = line.to_lowercase();
6021 if let Some(percent_idx) = lower.find('%') {
6022 let prefix = &lower[..percent_idx];
6023 let digits: String = prefix
6024 .chars()
6025 .rev()
6026 .take_while(|c| c.is_ascii_digit())
6027 .collect();
6028 if !digits.is_empty() {
6029 let value: String = digits.chars().rev().collect();
6030 if let Ok(percent) = value.parse::<f32>() {
6031 return Some((percent / 100.0).min(1.0));
6032 }
6033 }
6034 }
6035 if lower.contains("progress") && lower.contains("100") {
6036 return Some(1.0);
6037 }
6038 None
6039}
6040
6041fn needs_sudo() -> bool {
6042 #[cfg(unix)]
6043 {
6044 unsafe { libc::geteuid() != 0 }
6045 }
6046 #[cfg(not(unix))]
6047 {
6048 false
6049 }
6050}
6051
6052fn strip_ansi_codes(s: &str) -> String {
6053 let mut clean = String::with_capacity(s.len());
6054 let mut in_escape = false;
6055 for c in s.chars() {
6056 if c == '\x1b' {
6057 in_escape = true;
6058 } else if in_escape {
6059 if ('@'..='~').contains(&c) || c == 'm' {
6060 in_escape = false;
6061 }
6062 } else {
6063 clean.push(c);
6064 }
6065 }
6066 clean
6067}
6068
6069fn sanitize_output_line(line: &str) -> String {
6070 let stripped = strip_ansi_codes(line);
6071 let mut clean = String::with_capacity(stripped.len());
6072 for c in stripped.chars() {
6073 if c == '\t' {
6074 clean.push(' ');
6075 } else if !c.is_control() {
6076 clean.push(c);
6077 }
6078 }
6079 clean
6080}
6081
6082fn should_suppress_verification_line(line: &str, target_id: &str) -> bool {
6083 let trimmed = line.trim();
6084 if trimmed.is_empty() {
6085 return true;
6086 }
6087
6088 let lower = trimmed.to_lowercase();
6089 if is_exception_or_traceback_line(&lower) {
6090 return false;
6091 }
6092
6093 if lower.contains("amdgpu.ids")
6094 || lower.contains("tool-lib")
6095 || lower.contains("cannot load tool")
6096 || lower.contains("libdrm/amdgpu.ids: no such file or directory")
6097 {
6098 return true;
6099 }
6100
6101 if (target_id == "rocm" || target_id == "rocm-smi")
6102 && (lower.starts_with("======")
6103 || lower.starts_with("====")
6104 || lower.starts_with("---")
6105 || lower.contains("the info is not accurate")
6106 || lower.contains("================================")
6107 || lower.starts_with('{')
6108 || lower.starts_with('[')
6109 || lower.starts_with(']')
6110 || lower.starts_with('}')
6111 || lower.starts_with("\"card")
6112 || lower.starts_with("device ")
6113 || lower.starts_with("*****"))
6114 {
6115 return true;
6116 }
6117
6118 false
6119}
6120
6121fn is_exception_or_traceback_line(lowercase_line: &str) -> bool {
6122 lowercase_line.contains("traceback")
6123 || lowercase_line.contains("exception")
6124 || lowercase_line.contains("segmentation fault")
6125 || lowercase_line.contains("fatal:")
6126 || lowercase_line.contains("error:")
6127 || lowercase_line.starts_with("error ")
6128}
6129
6130fn run_native_benchmark(component_id: &str, sender: &Sender<InstallerEvent>) -> Result<()> {
6135 use crate::benchmark_runners;
6136
6137 let bench_name = match component_id {
6138 "mlperf-inference" => "all",
6139 "rocm-benchmarks" => "gpu-capability",
6140 "gpu-memory-bandwidth" => "memory-bandwidth",
6141 "rocm-smi-bench" => "gpu-capability",
6142 "pytorch-performance" => "pytorch",
6143 "vllm-performance" => "vllm",
6144 "deepspeed-performance" => "deepspeed",
6145 "megatron-performance" => "megatron",
6146 "all-benchmarks" => "all",
6147 _ => bail!("Unknown benchmark component: {}", component_id),
6148 };
6149
6150 let _ = sender.send(InstallerEvent::Log(
6151 format!("[native] Running {} benchmark...", component_id),
6152 false,
6153 ));
6154
6155 let _ = sender.send(InstallerEvent::Progress {
6156 component_id: component_id.to_string(),
6157 progress: 0.3,
6158 message: format!("Running {} benchmark", component_id),
6159 });
6160
6161 match benchmark_runners::run_benchmark(bench_name) {
6162 Ok(output) => {
6163 let _ = sender.send(InstallerEvent::Log(
6164 format!(
6165 "[native] {} benchmark completed: {} ({} ms)",
6166 component_id,
6167 if output.success { "SUCCESS" } else { "FAILED" },
6168 output.execution_time_ms
6169 ),
6170 false,
6171 ));
6172 if !output.errors.is_empty() {
6173 for err in &output.errors {
6174 let _ = sender.send(InstallerEvent::Log(
6175 format!("[native] {} error: {}", component_id, err),
6176 false,
6177 ));
6178 }
6179 }
6180 if !output.success {
6181 let error_details = output.errors.join("; ");
6182 bail!(
6183 "{} benchmark failed: {}",
6184 component_id,
6185 if error_details.is_empty() {
6186 "unknown error".to_string()
6187 } else {
6188 error_details
6189 }
6190 );
6191 }
6192 Ok(())
6193 }
6194 Err(err) => {
6195 let _ = sender.send(InstallerEvent::Log(
6196 format!("[native] {} benchmark error: {}", component_id, err),
6197 false,
6198 ));
6199 bail!("{} benchmark error: {}", component_id, err)
6200 }
6201 }
6202}
6203
6204#[cfg(test)]
6205mod tests {
6206 use super::*;
6207
6208 #[test]
6211 fn test_is_existing_git_repo_detects_git_dir() {
6212 let tmp = tempfile::tempdir().unwrap();
6213 let repo_dir = tmp.path().join("my-repo");
6214 fs::create_dir_all(repo_dir.join(".git")).unwrap();
6215 assert!(is_existing_git_repo(repo_dir.to_str().unwrap()));
6216 }
6217
6218 #[test]
6219 fn test_is_existing_git_repo_false_when_no_git() {
6220 let tmp = tempfile::tempdir().unwrap();
6221 let plain_dir = tmp.path().join("plain-dir");
6222 fs::create_dir_all(&plain_dir).unwrap();
6223 assert!(!is_existing_git_repo(plain_dir.to_str().unwrap()));
6224 }
6225
6226 #[test]
6227 fn test_is_existing_git_repo_false_when_no_dir() {
6228 assert!(!is_existing_git_repo(
6229 "/nonexistent/path/that/does/not/exist"
6230 ));
6231 }
6232
6233 #[test]
6234 fn test_git_clone_or_pull_clones_when_missing() {
6235 let tmp = tempfile::tempdir().unwrap();
6236 let target = tmp.path().join("new-repo");
6237 let (tx, rx) = std::sync::mpsc::channel();
6238
6239 let result = git_clone_or_pull(
6241 "https://github.com/nonexistent/repo.git",
6242 target.to_str().unwrap(),
6243 &[],
6244 None,
6245 &tx,
6246 "test-component",
6247 );
6248 assert!(result.is_err());
6250 let mut found_clone_log = false;
6252 while let Ok(event) = rx.try_recv() {
6253 if let InstallerEvent::Log(msg, _) = event {
6254 found_clone_log = found_clone_log || msg.contains("cloning");
6255 }
6256 }
6257 assert!(found_clone_log, "Should log cloning intent");
6258 }
6259
6260 #[test]
6261 fn test_git_clone_or_pull_pulls_when_existing() {
6262 let tmp = tempfile::tempdir().unwrap();
6263 let repo_dir = tmp.path().join("existing-repo");
6264 fs::create_dir_all(repo_dir.join(".git")).unwrap();
6265 let (tx, rx) = std::sync::mpsc::channel();
6266
6267 let _ = git_clone_or_pull(
6268 "https://github.com/ROCm/aiter.git",
6269 repo_dir.to_str().unwrap(),
6270 &["--recursive"],
6271 None,
6272 &tx,
6273 "test-component",
6274 );
6275 let mut found_pull_log = false;
6277 while let Ok(event) = rx.try_recv() {
6278 if let InstallerEvent::Log(msg, _) = event {
6279 found_pull_log = found_pull_log || msg.contains("pulling updates");
6280 }
6281 }
6282 assert!(found_pull_log, "Should log pulling intent, not cloning");
6283 }
6284
6285 #[test]
6286 fn test_hsa_override_from_gpu_arch() {
6287 assert_eq!(
6288 hsa_override_from_gpu_arch("gfx1100"),
6289 Some("11.0.0".to_string())
6290 );
6291 assert_eq!(
6292 hsa_override_from_gpu_arch("gfx1101"),
6293 Some("11.0.0".to_string())
6294 );
6295 assert_eq!(
6296 hsa_override_from_gpu_arch("gfx1200"),
6297 Some("12.0.0".to_string())
6298 );
6299 assert_eq!(
6300 hsa_override_from_gpu_arch("gfx1030"),
6301 Some("10.0.0".to_string())
6302 );
6303 assert!(hsa_override_from_gpu_arch("gfx99").is_none());
6305 assert!(hsa_override_from_gpu_arch("").is_none());
6306 }
6307
6308 #[test]
6309 fn test_is_igpu_name_raphael() {
6310 assert!(is_igpu_name("AMD Radeon Graphics (Raphael)"));
6312 assert!(is_igpu_name("Raphael"));
6313 }
6314
6315 #[test]
6316 fn test_is_igpu_name_apu_models() {
6317 assert!(is_igpu_name("Ryzen 5 5600G"));
6319 assert!(is_igpu_name("Ryzen 7 5700G"));
6320 assert!(is_igpu_name("Ryzen 7 8700G"));
6321 assert!(is_igpu_name("Ryzen 5 5600GE"));
6322 }
6323
6324 #[test]
6325 fn test_is_igpu_name_codenames() {
6326 assert!(is_igpu_name(
6328 "AMD Ryzen 9 7945HS with Radeon Graphics (Phoenix)"
6329 ));
6330 assert!(is_igpu_name("Rembrandt"));
6331 assert!(is_igpu_name("Cezanne"));
6332 }
6333
6334 #[test]
6335 fn test_is_igpu_name_discrete_gpus() {
6336 assert!(!is_igpu_name("Radeon RX 7900 XTX"));
6338 assert!(!is_igpu_name("Radeon RX 7800 XT"));
6339 assert!(!is_igpu_name("Radeon RX 7700 XT"));
6340 assert!(!is_igpu_name("AMD Radeon RX 6800 XT"));
6341 }
6342
6343 #[test]
6344 fn test_is_igpu_name_x3d_processors() {
6345 assert!(is_igpu_name("AMD Ryzen 7 7800X3D 8-Core Processor"));
6347 assert!(is_igpu_name("AMD Ryzen 9 7950X3D 16-Core Processor"));
6348 assert!(is_igpu_name("AMD Ryzen 9 7900X3D 12-Core Processor"));
6349 assert!(is_igpu_name("AMD Ryzen 5 7600X3D"));
6351 }
6352
6353 #[test]
6354 fn test_is_igpu_name_ryzen_heuristic() {
6355 assert!(is_igpu_name("AMD Ryzen 7 7800X3D 8-Core Processor"));
6357 assert!(is_igpu_name("AMD Ryzen 5 8600G"));
6358 assert!(is_igpu_name("AMD Ryzen 9 7945HS"));
6359 assert!(!is_igpu_name("AMD Radeon RX 7900 XTX"));
6361 assert!(!is_igpu_name("AMD Radeon RX 7800 XT Radeon RX 7900 XTX"));
6362 }
6363
6364 #[test]
6365 fn test_parse_rocminfo_mixed_gpus() {
6366 let rocminfo_output = r#"
6369 ****!**** The info is not accurate, please use it carefully !****
6370
6371===================== ROCm Information =========================
6372ROCm Version: 7.2.0
6373
6374===================== System Info =======================
6375Kernel Version: 6.12.63+deb13-rt-amd64
6376...
6377
6378===================== ASICs Info =========================
6379**** 2: Agent ****
6380** GPU ID: 0 **
6381** Marketing Name: AMD Radeon RX 7900 XTX **
6382** Device Type: GPU **
6383...
6384**** 2: Agent ****
6385** GPU ID: 1 **
6386** Marketing Name: AMD Radeon RX 7800 XT **
6387** Device Type: GPU **
6388...
6389**** 2: Agent ****
6390** GPU ID: 2 **
6391** Marketing Name: AMD Radeon Graphics (Raphael) **
6392** Device Type: GPU **
6393...
6394"#;
6395
6396 let result = parse_rocminfo_for_discrete_gpus(rocminfo_output);
6397 assert_eq!(result, vec!["0", "1"]);
6398 }
6399
6400 #[test]
6401 fn test_parse_rocminfo_x3d_igpu() {
6402 let rocminfo_output = r#"
6404Agent 1
6405*******
6406 Name: AMD Ryzen 7 7800X3D 8-Core Processor
6407 Marketing Name: AMD Ryzen 7 7800X3D 8-Core Processor
6408 Device Type: CPU
6409Agent 2
6410*******
6411** GPU ID: 0 **
6412 Name: gfx1100
6413 Marketing Name: AMD Radeon RX 7900 XTX
6414 Device Type: GPU
6415Agent 3
6416*******
6417** GPU ID: 1 **
6418 Name: gfx1100
6419 Marketing Name: AMD Radeon RX 7800 XT
6420 Device Type: GPU
6421Agent 4
6422*******
6423** GPU ID: 2 **
6424 Name: gfx1100
6425 Marketing Name: AMD Ryzen 7 7800X3D 8-Core Processor
6426 Device Type: GPU
6427"#;
6428 let result = parse_rocminfo_for_discrete_gpus(rocminfo_output);
6429 assert_eq!(result, vec!["0", "1"], "7800X3D iGPU should be excluded");
6430 }
6431
6432 #[test]
6433 fn test_parse_rocminfo_only_igpu() {
6434 let rocminfo_output = r#"
6436===================== ASICs Info =========================
6437**** 2: Agent ****
6438** GPU ID: 0 **
6439** Marketing Name: AMD Radeon Graphics (Raphael) **
6440** Device Type: GPU **
6441...
6442"#;
6443
6444 let result = parse_rocminfo_for_discrete_gpus(rocminfo_output);
6445 assert!(result.is_empty());
6446 }
6447
6448 #[test]
6449 fn test_parse_rocminfo_only_discrete() {
6450 let rocminfo_output = r#"
6452===================== ASICs Info =========================
6453**** 2: Agent ****
6454** GPU ID: 0 **
6455** Marketing Name: AMD Radeon RX 7900 XTX **
6456** Device Type: GPU **
6457...
6458**** 2: Agent ****
6459** GPU ID: 1 **
6460** Marketing Name: AMD Radeon RX 7800 XT **
6461** Device Type: GPU **
6462...
6463"#;
6464
6465 let result = parse_rocminfo_for_discrete_gpus(rocminfo_output);
6466 assert_eq!(result, vec!["0", "1"]);
6467 }
6468
6469 #[test]
6470 fn test_parse_rocm_smi_mixed_gpus() {
6471 let rocm_smi_output = r#"
6472{
6473 "card0": {
6474 "Card Series": "AMD Radeon RX 7900 XTX",
6475 "Card Model": "Navi 31"
6476 },
6477 "card1": {
6478 "Card Series": "AMD Radeon RX 7800 XT",
6479 "Card Model": "Navi 32"
6480 },
6481 "card2": {
6482 "Card Series": "AMD Radeon Graphics",
6483 "Card Model": "Raphael"
6484 }
6485}
6486"#;
6487 let result = parse_rocm_smi_for_discrete_gpus(rocm_smi_output);
6488 assert_eq!(result, vec!["0", "1"]);
6489 }
6490
6491 #[test]
6492 fn test_parse_rocm_smi_only_igpu() {
6493 let rocm_smi_output = r#"
6494{
6495 "card0": {
6496 "Card Series": "AMD Radeon Graphics",
6497 "Card Model": "Raphael"
6498 }
6499}
6500"#;
6501 let result = parse_rocm_smi_for_discrete_gpus(rocm_smi_output);
6502 assert!(result.is_empty());
6503 }
6504
6505 #[test]
6506 fn test_parse_rocm_smi_with_prefix_and_na_fields() {
6507 let rocm_smi_output = r#"
6508****!**** The info is not accurate, please use it carefully !****
6509{
6510 "card0": {
6511 "Card Series": "N/A",
6512 "Card Model": "N/A",
6513 "Card SKU": "AMD Radeon RX 7900 XTX"
6514 },
6515 "card1": {
6516 "Card Series": "N/A",
6517 "Card Model": "Navi 32",
6518 "Card SKU": "AMD Radeon RX 7800 XT"
6519 },
6520 "card2": {
6521 "Card Series": "N/A",
6522 "Card Model": "N/A",
6523 "Card SKU": "Raphael"
6524 }
6525}
6526"#;
6527 let result = parse_rocm_smi_for_discrete_gpus(rocm_smi_output);
6528 assert_eq!(result, vec!["0", "1"]);
6529 }
6530
6531 #[test]
6536 fn test_format_user_friendly_error_permission_denied() {
6537 let msg = format_user_friendly_error("rocm", "Permission denied: /opt/rocm");
6538 assert!(
6539 msg.contains("Permission denied"),
6540 "Should mention permission issue"
6541 );
6542 assert!(
6543 msg.contains("elevated privileges"),
6544 "Should suggest running with sudo"
6545 );
6546 }
6547
6548 #[test]
6549 fn test_format_user_friendly_error_not_found() {
6550 let msg = format_user_friendly_error("pytorch", "No such file or directory: python3");
6551 assert!(msg.contains("not found"), "Should mention file not found");
6552 assert!(
6553 msg.contains("ROCm") || msg.contains("Python"),
6554 "Should mention ROCm or Python dependency, got: {}",
6555 msg
6556 );
6557 }
6558
6559 #[test]
6560 fn test_format_user_friendly_error_network() {
6561 let msg = format_user_friendly_error(
6563 "onnxruntime",
6564 "network error: Could not resolve host: github.com",
6565 );
6566 assert!(
6567 msg.contains("Network error"),
6568 "Should mention network issue, got: {}",
6569 msg
6570 );
6571 assert!(
6572 msg.contains("internet connection"),
6573 "Should suggest checking internet, got: {}",
6574 msg
6575 );
6576 }
6577
6578 #[test]
6579 fn test_format_user_friendly_error_out_of_memory() {
6580 let msg = format_user_friendly_error("onnxruntime", "out of memory during build");
6582 assert!(
6583 msg.contains("Out of memory"),
6584 "Should mention memory issue, got: {}",
6585 msg
6586 );
6587 assert!(
6588 msg.contains("swap space"),
6589 "Should suggest increasing swap, got: {}",
6590 msg
6591 );
6592 }
6593
6594 #[test]
6595 fn test_format_user_friendly_error_permanent_env_rocminfo() {
6596 let msg = format_user_friendly_error("permanent-env", "rocminfo command not found");
6597 assert!(
6598 msg.contains("rocminfo"),
6599 "Should mention rocminfo specifically"
6600 );
6601 assert!(
6602 msg.contains("ROCm is properly installed"),
6603 "Should suggest ROCm installation"
6604 );
6605 }
6606
6607 #[test]
6608 fn test_format_user_friendly_error_pytorch_specific() {
6609 let msg = format_user_friendly_error("pytorch", "pip install failed: build error");
6611 assert!(msg.contains("PyTorch"), "Should mention PyTorch by name");
6612 assert!(
6613 msg.contains("ROCm is installed first"),
6614 "Should mention ROCm dependency"
6615 );
6616 assert!(
6617 msg.contains("Python 3.10-3.13"),
6618 "Should mention Python version, got: {}",
6619 msg
6620 );
6621 }
6622
6623 #[test]
6624 fn test_format_user_friendly_error_triton_dependency() {
6625 let msg = format_user_friendly_error("triton", "git clone failed");
6626 assert!(
6627 msg.contains("PyTorch and ROCm"),
6628 "Should mention required dependencies"
6629 );
6630 }
6631
6632 #[test]
6633 fn test_format_user_friendly_error_megatron_dependency() {
6634 let msg = format_user_friendly_error("megatron", "pip install failed");
6635 assert!(
6636 msg.contains("PyTorch and MPI4Py"),
6637 "Should mention Megatron's specific dependencies"
6638 );
6639 }
6640
6641 #[test]
6642 fn test_format_user_friendly_error_deepspeed_dependency() {
6643 let msg = format_user_friendly_error("deepspeed", "build failed");
6644 assert!(
6645 msg.contains("PyTorch"),
6646 "Should mention DeepSpeed's dependency on PyTorch"
6647 );
6648 }
6649
6650 #[test]
6651 fn test_format_user_friendly_error_exit_code_passthrough() {
6652 let msg = format_user_friendly_error("rocm", "ROCm failed: General error (exit code 1)");
6653 assert!(
6655 msg.contains("exit code 1"),
6656 "Should preserve exit code context"
6657 );
6658 }
6659
6660 #[test]
6661 fn test_format_user_friendly_error_unknown_component() {
6662 let msg = format_user_friendly_error("some-unknown-component", "something went wrong");
6663 assert!(
6664 msg.contains("Installation failed"),
6665 "Should have generic failure message"
6666 );
6667 assert!(
6668 msg.contains("something went wrong"),
6669 "Should include original error"
6670 );
6671 }
6672
6673 #[test]
6674 fn test_format_user_friendly_error_rocm_not_found() {
6675 let msg = format_user_friendly_error("rocm", "package not found: amdgpu-dkms");
6676 assert!(msg.contains("ROCm"), "Should mention ROCm specifically");
6677 assert!(msg.contains("repositories"), "Should mention repositories");
6678 }
6679
6680 #[test]
6681 fn test_execute_native_command_nonexistent_program() {
6682 let (tx, rx) = std::sync::mpsc::channel();
6685 let cmd = NativeCommand::Shell {
6686 program: "/nonexistent/program/that/does/not/exist".to_string(),
6687 args: vec![],
6688 env: vec![],
6689 working_dir: None,
6690 };
6691 let result = execute_native_command(&cmd, None, &tx, "test-component");
6692 assert!(result.is_err(), "Should fail for non-existent program");
6693 let err_msg = result.unwrap_err().to_string();
6694 assert!(!err_msg.is_empty(), "Error message should not be empty");
6695
6696 drop(tx);
6698 let events: Vec<_> = rx.try_iter().collect();
6699 assert!(
6700 !events.is_empty(),
6701 "Should have sent at least one event before failing"
6702 );
6703 }
6704
6705 #[test]
6706 fn test_execute_native_command_captures_exit_code() {
6707 let (tx, rx) = std::sync::mpsc::channel();
6709 let cmd = NativeCommand::Shell {
6710 program: "bash".to_string(),
6711 args: vec!["-c".to_string(), "echo test_output; exit 42".to_string()],
6712 env: vec![],
6713 working_dir: None,
6714 };
6715 let result = execute_native_command(&cmd, None, &tx, "test-component");
6716 assert!(
6717 result.is_err(),
6718 "Should fail for non-zero exit code, got: {:?}",
6719 result
6720 );
6721
6722 let err_msg = result.unwrap_err().to_string();
6723 assert!(
6724 err_msg.contains("42"),
6725 "Error should include exit code 42, got: {}",
6726 err_msg
6727 );
6728
6729 drop(tx);
6731 let events: Vec<_> = rx.try_iter().collect();
6732 let log_messages: Vec<&str> = events
6733 .iter()
6734 .filter_map(|e| match e {
6735 InstallerEvent::Log(msg, _) => Some(msg.as_str()),
6736 _ => None,
6737 })
6738 .collect();
6739 assert!(
6740 log_messages.iter().any(|m| m.contains("test_output")),
6741 "Should have captured stdout output before failure"
6742 );
6743 assert!(
6744 log_messages
6745 .iter()
6746 .any(|m| m.contains("[ERROR]") || m.contains("exit code")),
6747 "Should have logged error with exit code"
6748 );
6749 }
6750
6751 #[test]
6752 fn test_execute_native_command_success_sends_events() {
6753 let (tx, rx) = std::sync::mpsc::channel();
6755 let cmd = NativeCommand::Shell {
6756 program: "bash".to_string(),
6757 args: vec![
6758 "-c".to_string(),
6759 "echo 'hello world'; echo 'second line'".to_string(),
6760 ],
6761 env: vec![],
6762 working_dir: None,
6763 };
6764 let result = execute_native_command(&cmd, None, &tx, "test-component");
6765 assert!(result.is_ok(), "Should succeed for zero exit code");
6766
6767 drop(tx);
6768 let events: Vec<_> = rx.try_iter().collect();
6769 let log_messages: Vec<&str> = events
6770 .iter()
6771 .filter_map(|e| match e {
6772 InstallerEvent::Log(msg, _) => Some(msg.as_str()),
6773 _ => None,
6774 })
6775 .collect();
6776 assert!(
6777 log_messages.iter().any(|m| m.contains("hello world")),
6778 "Should have captured first line of output"
6779 );
6780 assert!(
6781 log_messages.iter().any(|m| m.contains("second line")),
6782 "Should have captured second line of output"
6783 );
6784 }
6785
6786 #[test]
6791 fn test_no_sudo_wrap_when_password_is_none() {
6792 let (tx, rx) = std::sync::mpsc::channel();
6796 let cmd = NativeCommand::Shell {
6797 program: "echo".to_string(),
6798 args: vec!["hello".to_string()],
6799 env: vec![],
6800 working_dir: None,
6801 };
6802 let result = execute_native_command(&cmd, None, &tx, "test-component");
6803 assert!(result.is_ok(), "Should succeed without sudo: {:?}", result);
6804
6805 drop(tx);
6806 let events: Vec<_> = rx.try_iter().collect();
6807 let log_messages: Vec<String> = events
6808 .iter()
6809 .filter_map(|e| match e {
6810 InstallerEvent::Log(msg, _) => Some(msg.clone()),
6811 _ => None,
6812 })
6813 .collect();
6814
6815 let cmd_log = log_messages.iter().find(|m| m.contains("[native] $"));
6817 assert!(cmd_log.is_some(), "Should have logged the command");
6818 let cmd_log = cmd_log.unwrap();
6819 assert!(
6820 !cmd_log.contains("sudo"),
6821 "Command should NOT be wrapped with sudo when sudo_pw is None, got: {}",
6822 cmd_log
6823 );
6824 assert!(
6825 cmd_log.contains("echo hello"),
6826 "Command should be 'echo hello', got: {}",
6827 cmd_log
6828 );
6829 }
6830
6831 #[test]
6832 fn test_sudo_wrap_when_password_is_some() {
6833 let (tx, rx) = std::sync::mpsc::channel();
6837 let cmd = NativeCommand::Shell {
6838 program: "echo".to_string(),
6839 args: vec!["hello".to_string()],
6840 env: vec![],
6841 working_dir: None,
6842 };
6843 let _result = execute_native_command(&cmd, Some("dummy_pw"), &tx, "test-component");
6845 drop(tx);
6848 let events: Vec<_> = rx.try_iter().collect();
6849 let log_messages: Vec<String> = events
6850 .iter()
6851 .filter_map(|e| match e {
6852 InstallerEvent::Log(msg, _) => Some(msg.clone()),
6853 _ => None,
6854 })
6855 .collect();
6856
6857 let has_sudo_log = log_messages.iter().any(|m| m.contains("Running with sudo"));
6859 assert!(
6860 has_sudo_log,
6861 "Should have logged 'Running with sudo' when sudo_pw is Some"
6862 );
6863 }
6864
6865 #[test]
6866 fn test_git_clone_without_sudo_works() {
6867 let (tx, rx) = std::sync::mpsc::channel();
6872 let cmd = NativeCommand::Shell {
6873 program: "git".to_string(),
6874 args: vec!["--version".to_string()],
6875 env: vec![],
6876 working_dir: None,
6877 };
6878 let result = execute_native_command(&cmd, None, &tx, "textgen");
6879 assert!(
6880 result.is_ok(),
6881 "git --version should succeed without sudo wrapping: {:?}",
6882 result
6883 );
6884
6885 drop(tx);
6886 let events: Vec<_> = rx.try_iter().collect();
6887 let log_messages: Vec<String> = events
6888 .iter()
6889 .filter_map(|e| match e {
6890 InstallerEvent::Log(msg, _) => Some(msg.clone()),
6891 _ => None,
6892 })
6893 .collect();
6894
6895 assert!(
6897 log_messages.iter().any(|m| m.contains("git version")),
6898 "Should have captured 'git version' output, got: {:?}",
6899 log_messages
6900 );
6901
6902 let cmd_log = log_messages.iter().find(|m| m.contains("[native] $"));
6904 assert!(cmd_log.is_some(), "Should have logged the command");
6905 assert!(
6906 !cmd_log.unwrap().contains("sudo"),
6907 "git command should NOT be wrapped with sudo"
6908 );
6909 }
6910
6911 #[test]
6912 fn test_textgen_git_clone_args_not_empty() {
6913 use crate::installers::components::textgen::{TextgenConfig, TextgenInstaller};
6916 let inst = TextgenInstaller::new(TextgenConfig {
6917 install_dir: "/tmp/test-textgen".to_string(),
6918 python_bin: "python3".to_string(),
6919 use_uv: false,
6920 break_system_packages: true,
6921 dry_run: false,
6922 });
6923 let cmd = inst.build_git_clone_command();
6924
6925 for (i, arg) in cmd.args.iter().enumerate() {
6927 assert!(
6928 !arg.is_empty(),
6929 "Argument {} should not be empty in git clone command",
6930 i
6931 );
6932 }
6933
6934 assert_eq!(cmd.program, "git");
6936 assert_eq!(
6937 cmd.args.len(),
6938 3,
6939 "git clone should have 3 args: clone, URL, dir"
6940 );
6941 assert_eq!(cmd.args[0], "clone");
6942 assert_eq!(
6943 cmd.args[1],
6944 "https://github.com/oobabooga/text-generation-webui.git"
6945 );
6946 assert_eq!(cmd.args[2], "/tmp/test-textgen");
6947
6948 assert!(
6950 !cmd.args.contains(&"--".to_string()),
6951 "git clone args should not contain '--' separator (from sudo wrapping)"
6952 );
6953 }
6954
6955 #[test]
6960 fn test_verification_result_labels() {
6961 assert_eq!(VerificationResult::Verified.label(), "Verified");
6962 assert_eq!(VerificationResult::Failed.label(), "Failed");
6963 assert_eq!(VerificationResult::Missing.label(), "Missing");
6964 assert_eq!(VerificationResult::Warning.label(), "Warning");
6965 }
6966
6967 #[test]
6968 fn test_verification_result_warning_is_distinct() {
6969 assert_ne!(VerificationResult::Warning, VerificationResult::Failed);
6971 assert_ne!(VerificationResult::Warning, VerificationResult::Verified);
6972 assert_ne!(VerificationResult::Warning, VerificationResult::Missing);
6973 }
6974}