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