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