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 run_installation_with_version(components, config, sudo_password, sender, input_rx, None);
59}
60
61pub fn run_installation_with_version(
68 components: Vec<Component>,
69 config: InstallerConfig,
70 sudo_password: Option<String>,
71 sender: Sender<InstallerEvent>,
72 input_rx: Receiver<String>,
73 target_version: Option<String>,
74) {
75 let scripts_dir = &config.scripts_dir;
76 let batch_mode = config.batch_mode;
77 let install_method = config.install_method.clone();
78 let total = components.len() as f32;
79 let mut index = 0usize;
80 let mut python_candidates = python_interpreters();
81 let user_home = resolve_mlstack_user_home();
82 let input_rx = Arc::new(Mutex::new(input_rx));
83
84 match ensure_mlstack_env(&user_home, &install_method, false) {
85 Ok(EnvUpdate::Created) => {
86 let _ = sender.send(InstallerEvent::Log(
87 format!("Created .mlstack_env in {}", user_home),
88 false,
89 ));
90 }
91 Ok(EnvUpdate::Updated) => {
92 let _ = sender.send(InstallerEvent::Log(
93 format!("Updated .mlstack_env in {}", user_home),
94 false,
95 ));
96 }
97 Ok(EnvUpdate::Unchanged) => {}
98 Err(err) => {
99 let _ = sender.send(InstallerEvent::Log(
100 format!("Failed to ensure .mlstack_env: {}", err),
101 false,
102 ));
103 }
104 }
105
106 let any_component_needs_sudo = components.iter().any(|c| c.needs_sudo);
108
109 if needs_sudo() && any_component_needs_sudo {
110 if let Some(password) = sudo_password.clone() {
111 if let Err(err) = validate_sudo(password) {
112 let _ = sender.send(InstallerEvent::Log(
113 format!("Sudo validation failed: {}", err),
114 false,
115 ));
116 let _ = sender.send(InstallerEvent::Finished { success: false });
117 return;
118 }
119 } else if can_sudo_non_interactive() {
120 let _ = sender.send(InstallerEvent::Log(
121 "Using cached/passwordless sudo credentials".into(),
122 false,
123 ));
124 } else {
125 let _ = sender.send(InstallerEvent::Log(
126 "Sudo password missing; cannot continue".into(),
127 false,
128 ));
129 let _ = sender.send(InstallerEvent::Finished { success: false });
130 return;
131 }
132 }
133
134 let no_named_env = std::env::var("MLSTACK_ENV_NAME")
141 .map(|v| v.trim().is_empty())
142 .unwrap_or(true);
143 let no_explicit_python = std::env::var("MLSTACK_PYTHON_BIN")
144 .map(|v| v.trim().is_empty())
145 .unwrap_or(true);
146 let is_global_install = no_named_env && no_explicit_python;
150 if is_global_install {
151 let bootstrap = resolve_python_bin();
152 if let Ok(global_py) = crate::platform::environment::ensure_global_venv(&bootstrap) {
153 std::env::set_var(
154 "MLSTACK_PYTHON_BIN",
155 global_py.to_string_lossy().to_string(),
156 );
157 } else if std::env::var("MLSTACK_PYTHON_BIN").is_err() {
158 let _ = sender.send(InstallerEvent::Log(
161 "[native] could not create managed global venv (~/.mlstack/global); \
162 falling back to a discovered interpreter"
163 .into(),
164 false,
165 ));
166 }
167 }
168
169 let mut overall_success = true;
170 let env_exports = load_mlstack_env_exports(&user_home);
171 let persistent_python = env_exports
172 .get("MLSTACK_PYTHON_BIN")
173 .map(|v| v.trim().to_string())
174 .filter(|v| !v.is_empty());
175 let python_bin = std::env::var("MLSTACK_ENV_NAME")
179 .ok()
180 .map(|n| n.trim().to_string())
181 .filter(|n| !n.is_empty())
182 .map(|n| {
183 crate::platform::environment::mlstack_env_dir(&n)
184 .join("bin")
185 .join("python")
186 .to_string_lossy()
187 .to_string()
188 })
189 .or(persistent_python)
190 .unwrap_or_else(resolve_python_bin);
191 python_candidates.retain(|candidate| candidate != &python_bin);
192 python_candidates.insert(0, python_bin.clone());
193 let _ = sender.send(InstallerEvent::Log(
194 format!("Selected python interpreter: {}", python_bin),
195 false,
196 ));
197
198 let components = sort_components_by_dependencies(components);
202
203 for component in components {
204 index += 1;
205 let _ = sender.send(InstallerEvent::ComponentStart {
206 component_id: component.id.clone(),
207 name: component.name.clone(),
208 });
209
210 if component.category == Category::Maintenance && component.id.starts_with("verify-") {
211 let outcome = run_verification(&component, &python_candidates, &sender, &user_home);
212 let _ = sender.send(InstallerEvent::VerificationReport {
213 component_id: component.id.clone(),
214 lines: outcome.report_lines.clone(),
215 });
216 if !outcome.success {
217 overall_success = false;
218 }
219 let _ = sender.send(InstallerEvent::ComponentComplete {
220 component_id: component.id.clone(),
221 success: outcome.success,
222 message: if outcome.success {
223 format!("{} verification complete", component.name)
224 } else {
225 format!("{} verification completed with issues", component.name)
226 },
227 });
228
229 let overall_progress = index as f32 / total;
230 let _ = sender.send(InstallerEvent::Progress {
231 component_id: "__overall__".into(),
232 progress: overall_progress,
233 message: format!(
234 "Finished {} ({:.0}%)",
235 component.name,
236 overall_progress * 100.0
237 ),
238 });
239 continue;
240 }
241
242 let is_native = crate::installers::components::is_native_component(&component.id);
245
246 let mut install_success = true;
247 let mut error_msg = String::new();
248
249 if is_native {
250 let ctx = NativeInstallerContext {
255 sudo_password: sudo_password.clone(),
256 batch_mode,
257 install_method: &install_method,
258 sender: &sender,
259 user_home: &user_home,
260 env_exports: &env_exports,
261 vllm_version_lag: config.vllm_version_lag,
262 vllm_version_min_age_days: config.vllm_version_min_age_days,
263 target_version: target_version.clone(),
264 };
265 match registry_gate(&component, &sender) {
267 Ok(false) => { }
268 Ok(true) => {
269 if let Err(err) = run_native_installer(&component, &ctx) {
270 install_success = false;
271 let chain = err
272 .chain()
273 .map(|e| e.to_string())
274 .collect::<Vec<_>>()
275 .join(": ");
276 error_msg = format_user_friendly_error(&component.id, &chain);
277 let err_label = if component.category == Category::Performance {
279 "benchmarks"
280 } else {
281 "installation"
282 };
283 let _ = sender.send(InstallerEvent::Log(
284 format!("[ERROR] {} {} failed: {}", component.name, err_label, chain),
285 false,
286 ));
287 }
288 }
289 Err(blocked) => {
290 install_success = false;
291 error_msg = format_user_friendly_error(&component.id, &blocked.to_string());
292 let _ = sender.send(InstallerEvent::Log(
293 format!("[ERROR] {} blocked: {}", component.name, blocked),
294 false,
295 ));
296 }
297 }
298 } else {
299 let script_path = format!("{}/{}", scripts_dir, component.script);
302 if fs::metadata(&script_path).is_err() {
303 overall_success = false;
304 let _ = sender.send(InstallerEvent::ComponentComplete {
305 component_id: component.id.clone(),
306 success: false,
307 message: format!("Script not found: {}", script_path),
308 });
309 continue;
310 }
311
312 if let Err(err) = run_script(
313 &component,
314 &script_path,
315 sudo_password.clone(),
316 batch_mode,
317 &install_method,
318 &sender,
319 Arc::clone(&input_rx),
320 ) {
321 install_success = false;
322 let chain = err
323 .chain()
324 .map(|e| e.to_string())
325 .collect::<Vec<_>>()
326 .join(": ");
327 error_msg = format_user_friendly_error(&component.id, &chain);
328 let err_label = if component.category == Category::Performance {
330 "benchmarks"
331 } else {
332 "installation"
333 };
334 let _ = sender.send(InstallerEvent::Log(
335 format!("[ERROR] {} {} failed: {}", component.name, err_label, chain),
336 false,
337 ));
338 }
339 }
340
341 let verify_label = if component.category == Category::Performance {
342 "Validating"
343 } else {
344 "Verifying"
345 };
346 let _ = sender.send(InstallerEvent::Progress {
347 component_id: component.id.clone(),
348 progress: 0.8,
349 message: format!("{} {}", verify_label, component.name),
350 });
351
352 let verification_outcome = if component.category == Category::Performance {
359 VerificationOutcome {
360 success: install_success,
361 report_lines: vec![format!(
362 "Benchmark {}",
363 if install_success {
364 "completed successfully"
365 } else {
366 "reported failures"
367 }
368 )],
369 }
370 } else if is_native_maintenance_action(&component) {
371 native_maintenance_action_outcome(&component, install_success)
372 } else {
373 run_component_verification(&component, &python_candidates, &sender, &user_home)
374 };
375 let _ = sender.send(InstallerEvent::VerificationReport {
376 component_id: component.id.clone(),
377 lines: verification_outcome.report_lines.clone(),
378 });
379
380 if !verification_outcome.success {
381 overall_success = false;
382 }
383
384 let final_success = install_success && verification_outcome.success;
393 if !install_success && verification_outcome.success {
394 let _ = sender.send(InstallerEvent::Log(
395 format!(
396 "[{}] install step failed — NOT marking installed despite passing \
397 verification (verification likely ran against pre-existing/stale artifacts)",
398 component.name
399 ),
400 true,
401 ));
402 }
403 if !final_success {
404 overall_success = false;
405 } else {
406 registry_record(&component);
409 }
410
411 if component.id == "permanent-env"
413 || component.id == "enhanced-env"
414 || component.id == "basic-env"
415 {
416 let _ = sender.send(InstallerEvent::Log(
417 format!(
418 "[DEBUG] Component {}: install_success={}, verification_success={}, final={}",
419 component.id, install_success, verification_outcome.success, final_success
420 ),
421 false,
422 ));
423 if !install_success {
424 let _ = sender.send(InstallerEvent::Log(
425 format!("[DEBUG] Install error: {}", error_msg),
426 false,
427 ));
428 }
429 }
430
431 let completion_msg = if !install_success {
432 error_msg
433 } else if component.category == Category::Performance {
434 if verification_outcome.success {
435 format!("{} benchmarks completed", component.name)
436 } else {
437 format!("{} benchmarks completed with errors", component.name)
438 }
439 } else if verification_outcome.success {
440 format!("{} completed", component.name)
441 } else {
442 format!("{} completed with verification errors", component.name)
443 };
444
445 let _ = sender.send(InstallerEvent::ComponentComplete {
446 component_id: component.id.clone(),
447 success: final_success,
448 message: completion_msg,
449 });
450
451 let overall_progress = index as f32 / total;
452 let _ = sender.send(InstallerEvent::Progress {
453 component_id: "__overall__".into(),
454 progress: overall_progress,
455 message: format!(
456 "Finished {} ({:.0}%)",
457 component.name,
458 overall_progress * 100.0
459 ),
460 });
461 }
462
463 match ensure_mlstack_env(&user_home, &install_method, false) {
471 Ok(EnvUpdate::Created) | Ok(EnvUpdate::Updated) => {
472 let _ = sender.send(InstallerEvent::Log(
473 "[post-install] Refreshed ~/.mlstack_env from installed state \
474 (ROCM_VERSION / GPU_ARCH / etc.)"
475 .into(),
476 false,
477 ));
478 }
479 Ok(EnvUpdate::Unchanged) => {}
480 Err(e) => {
481 let _ = sender.send(InstallerEvent::Log(
482 format!("[post-install] Warning: could not refresh ~/.mlstack_env: {e}"),
483 false,
484 ));
485 }
486 }
487
488 if is_global_install {
493 if let Err(e) = crate::platform::environment::offer_mlstack_env_source() {
494 let _ = sender.send(InstallerEvent::Log(
495 format!(
496 "[post-install] Warning: could not auto-source ~/.mlstack_env into shell rc: {e}"
497 ),
498 false,
499 ));
500 }
501 }
502
503 if overall_success && config.star_repos {
504 let _ = star_repo(&config, &sender);
505 }
506
507 if std::path::Path::new("/tmp/mlstack-reboot-required").exists() {
509 let _ = sender.send(InstallerEvent::Log(
510 "⚠ REBOOT REQUIRED: ROCm packages were reinstalled. A reboot is needed for kernel module changes to take effect.".to_string(),
511 false,
512 ));
513 let _ = std::fs::remove_file("/tmp/mlstack-reboot-required");
514 }
515
516 {
521 use crate::installers::common::fix_pip_cache_ownership;
522 let result = fix_pip_cache_ownership(&user_home);
523 match &result {
524 crate::installers::common::PipCacheFixResult::Fixed => {
525 let _ = sender.send(InstallerEvent::Log(
526 "[post-install] Fixed pip cache directory ownership (~/.cache/pip)".into(),
527 false,
528 ));
529 }
530 crate::installers::common::PipCacheFixResult::FixFailed(err) => {
531 let _ = sender.send(InstallerEvent::Log(
532 format!(
533 "[post-install] Warning: could not fix pip cache ownership: {}",
534 err
535 ),
536 false,
537 ));
538 }
539 _ => {
540 }
542 }
543 }
544
545 let _ = sender.send(InstallerEvent::Finished {
546 success: overall_success,
547 });
548}
549
550fn validate_sudo(password: String) -> Result<()> {
551 let mut child = Command::new("sudo")
552 .arg("-S")
553 .arg("-v")
554 .stdin(Stdio::piped())
555 .stdout(Stdio::null())
556 .stderr(Stdio::piped())
557 .spawn()
558 .context("Failed to spawn sudo")?;
559
560 let mut stdin = child.stdin.take().context("Failed to open stdin")?;
561 stdin.write_all(password.as_bytes())?;
562 stdin.write_all(b"\n")?;
563 stdin.flush()?;
564 drop(stdin);
565
566 let output = child.wait_with_output()?;
567 if output.status.success() {
568 return Ok(());
569 }
570
571 let stderr = String::from_utf8_lossy(&output.stderr);
572 let detail = stderr.lines().next().unwrap_or("sudo validation failed");
573 anyhow::bail!(detail.to_string())
574}
575
576fn can_sudo_non_interactive() -> bool {
577 Command::new("sudo")
578 .arg("-n")
579 .arg("true")
580 .stdin(Stdio::null())
581 .stdout(Stdio::null())
582 .stderr(Stdio::null())
583 .status()
584 .map(|s| s.success())
585 .unwrap_or(false)
586}
587
588#[derive(Debug, Clone, Copy, PartialEq, Eq)]
589enum EnvUpdate {
590 Created,
591 Updated,
592 Unchanged,
593}
594
595struct VerificationCheckResult {
596 status: VerificationResult,
597 output: Vec<String>,
598}
599
600#[derive(Debug, Clone)]
601struct VerificationOutcome {
602 success: bool,
603 report_lines: Vec<String>,
604}
605
606fn fix_libdrm_amdgpu_ids() {
616 let target = Path::new("/usr/share/libdrm/amdgpu.ids");
617 let link = Path::new("/opt/amdgpu/share/libdrm/amdgpu.ids");
618
619 if !target.exists() {
620 return;
621 }
622 if link.exists() {
623 return; }
625
626 if let Some(parent) = link.parent() {
628 if fs::create_dir_all(parent).is_ok() {
629 #[cfg(unix)]
630 {
631 let relative = compute_relative_symlink(link, target);
632 if std::os::unix::fs::symlink(&relative, link).is_ok() {
633 tracing::info!(
634 "Created libdrm symlink: {} -> {}",
635 link.display(),
636 relative.display()
637 );
638 return;
639 }
640 }
641 #[cfg(not(unix))]
642 {
643 let _ = (link, target);
644 }
645 }
646 }
647
648 let result = Command::new("sudo")
654 .args(["-n", "mkdir", "-p", "/opt/amdgpu/share/libdrm/"])
655 .stdin(std::process::Stdio::null())
656 .output();
657
658 if let Ok(out) = result {
659 if out.status.success() {
660 let result2 = Command::new("sudo")
661 .args([
662 "-n",
663 "ln",
664 "-sf",
665 "/usr/share/libdrm/amdgpu.ids",
666 "/opt/amdgpu/share/libdrm/amdgpu.ids",
667 ])
668 .stdin(std::process::Stdio::null())
669 .output();
670 if let Ok(out2) = result2 {
671 if out2.status.success() {
672 tracing::info!(
673 "Created libdrm symlink via sudo: {} -> {}",
674 link.display(),
675 target.display()
676 );
677 return;
678 }
679 }
680 }
681 }
682
683 tracing::warn!(
685 "Cannot create /opt/amdgpu/share/libdrm/amdgpu.ids symlink (need sudo). \
686 Run manually: sudo mkdir -p /opt/amdgpu/share/libdrm && \
687 sudo ln -sf /usr/share/libdrm/amdgpu.ids /opt/amdgpu/share/libdrm/amdgpu.ids"
688 );
689}
690
691fn compute_relative_symlink(link: &Path, target: &Path) -> PathBuf {
693 let mut ups = 0u32;
695 for _ in link.parent().unwrap_or(Path::new(".")).ancestors().skip(1) {
696 ups += 1;
697 }
698 let mut result = PathBuf::new();
699 for _ in 0..ups.saturating_sub(1) {
700 result.push("..");
701 }
702 result.push(target.strip_prefix("/").unwrap_or(target));
703 result
704}
705
706fn ensure_mlstack_env(user_home: &str, install_method: &str, force: bool) -> Result<EnvUpdate> {
707 let env_path = PathBuf::from(user_home).join(".mlstack_env");
708 let normalized_install_method = match install_method.trim().to_ascii_lowercase().as_str() {
709 "global" => "global",
710 "venv" => "venv",
711 "auto" => "auto",
712 _ => "auto",
713 };
714 let env_exports = load_mlstack_env_exports(user_home);
715 let persistent_python = env_exports
716 .get("MLSTACK_PYTHON_BIN")
717 .map(|v| v.trim().to_string())
718 .filter(|v| !v.is_empty());
719 let python_bin = std::env::var("MLSTACK_ENV_NAME")
723 .ok()
724 .map(|n| n.trim().to_string())
725 .filter(|n| !n.is_empty())
726 .map(|n| {
727 crate::platform::environment::mlstack_env_dir(&n)
728 .join("bin")
729 .join("python")
730 .to_string_lossy()
731 .to_string()
732 })
733 .or(persistent_python)
734 .unwrap_or_else(|| {
735 let global_py = crate::platform::environment::mlstack_global_python();
737 if global_py.exists() {
738 global_py.to_string_lossy().to_string()
739 } else {
740 resolve_python_bin()
741 }
742 });
743
744 let dgpus = crate::gpu::detect_discrete_amd_gpus();
746 if dgpus.is_empty() {
747 bail!("cannot detect AMD dGPU arch via sysfs — ensure amdgpu driver loaded");
748 }
749 let primary = &dgpus[0];
750 let gpu_arch = primary.gfx_arch.as_ref().ok_or_else(|| {
751 anyhow::anyhow!(
752 "could not resolve gfx arch for PCI device {} — set GPU_ARCH manually",
753 primary.pci_device_id
754 )
755 })?;
756 let discrete_indices = if dgpus.iter().all(|g| g.rocm_index.is_some()) {
768 dgpus
769 .iter()
770 .map(|g| g.rocm_index.expect("checked Some above").to_string())
771 .collect::<Vec<_>>()
772 .join(",")
773 } else {
774 dgpus
775 .iter()
776 .enumerate()
777 .map(|(i, _)| i.to_string())
778 .collect::<Vec<_>>()
779 .join(",")
780 };
781 let primary_gpu = "0";
785 let gpu_vram_gb_csv = dgpus
790 .iter()
791 .map(|g| match g.vram_bytes {
792 Some(b) => format!("{:.0}", b as f64 / (1024.0 * 1024.0 * 1024.0)),
793 None => "0".to_string(),
794 })
795 .collect::<Vec<_>>()
796 .join(",");
797 let (gpu_names_csv, gpu_cus_csv) = gpu_identity_from_rocminfo(&dgpus);
798 let gpu_pairs = [
799 ("MLSTACK_GPU_NAMES", gpu_names_csv.as_str()),
800 ("MLSTACK_GPU_VRAM_GB", gpu_vram_gb_csv.as_str()),
801 ("MLSTACK_GPU_CUS", gpu_cus_csv.as_str()),
802 ];
803
804 let hsa_override =
806 crate::bootstrap::env_setup::GpuArchInfo::from_arch(gpu_arch).hsa_override_gfx_version;
807
808 if env_path.exists() && !force {
814 let contents = fs::read_to_string(&env_path).context("Failed to read .mlstack_env")?;
815 let rocm_version = detect_rocm_version();
816 let rocm_home = "/opt/rocm";
817 let rocm_lib = "/opt/rocm/lib";
818 let defaults = [
819 ("ROCM_VERSION", rocm_version.as_str()),
820 ("ROCM_CHANNEL", "latest"),
821 ("GPU_ARCH", gpu_arch),
822 ("PYTORCH_ROCM_ARCH", gpu_arch),
823 ("GPU_ARCHS", gpu_arch),
824 ("ROCM_HOME", rocm_home),
825 ("ROCM_PATH", rocm_home),
826 ("HIP_PATH", rocm_home),
827 ("HSA_OVERRIDE_GFX_VERSION", hsa_override.as_str()),
828 ("HIP_VISIBLE_DEVICES", discrete_indices.as_str()),
829 ("CUDA_VISIBLE_DEVICES", discrete_indices.as_str()),
830 ("ROCR_VISIBLE_DEVICES", discrete_indices.as_str()),
835 ("PYTORCH_ROCM_DEVICE", primary_gpu),
836 ("MLSTACK_PYTHON_BIN", python_bin.as_str()),
837 ("UV_PYTHON", python_bin.as_str()),
838 ("MLSTACK_INSTALL_METHOD", normalized_install_method),
839 ("INSTALL_METHOD", normalized_install_method),
840 ("PYTHONPATH", rocm_lib),
841 ("PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:512"),
847 ("PYTORCH_HIP_ALLOC_CONF", "max_split_size_mb:512"),
851 ("AMDGPU_ASIC_ID_TABLE_PATH", "/usr/share/libdrm/amdgpu.ids"),
853 ("AMDGPU_ASIC_ID_TABLE_PATHS", "/usr/share/libdrm"),
854 ];
855 let (sanitized, _changed) = sanitize_mlstack_env(
856 &contents,
857 &defaults,
858 discrete_indices.as_str(),
859 rocm_home,
860 python_bin.as_str(),
861 normalized_install_method,
862 hsa_override.as_str(),
863 );
864 let updated = ensure_gpu_metadata_exports(&sanitized, &gpu_pairs);
867 if updated != contents {
868 fs::write(&env_path, &updated).context("Failed to update .mlstack_env")?;
869 return Ok(EnvUpdate::Updated);
870 }
871 return Ok(EnvUpdate::Unchanged);
872 }
873
874 if let Some(parent) = env_path.parent() {
875 fs::create_dir_all(parent).context("Failed to create home directory")?;
876 }
877
878 let rocm_version = detect_rocm_version();
879 let rocm_home = "/opt/rocm";
880
881 let mut content = crate::bootstrap::env_setup::generate_env_file_content(
885 &discrete_indices,
886 rocm_home,
887 &rocm_version,
888 "latest",
889 gpu_arch,
890 &hsa_override,
891 &python_bin,
892 );
893 content.push_str(&format!(
894 "\n# Rusty-Stack install metadata\nexport MLSTACK_INSTALL_METHOD={}\nexport INSTALL_METHOD={}\n",
895 normalized_install_method, normalized_install_method,
896 ));
897 content = ensure_gpu_metadata_exports(&content, &gpu_pairs);
899
900 fs::write(&env_path, &content).context("Failed to create .mlstack_env")?;
901
902 let fish_dir = PathBuf::from(user_home).join(".config/fish/conf.d");
905 if fs::create_dir_all(&fish_dir).is_ok() {
906 let fish_content = crate::bootstrap::env_setup::generate_fish_env_file_content(
907 &discrete_indices,
908 rocm_home,
909 &rocm_version,
910 "latest",
911 gpu_arch,
912 &hsa_override,
913 &python_bin,
914 );
915 let fish_path = fish_dir.join("mlstack_env.fish");
916 fs::write(&fish_path, fish_content)
917 .context("Failed to write fish conf.d/mlstack_env.fish")?;
918 }
919
920 fix_libdrm_amdgpu_ids();
924
925 Ok(EnvUpdate::Created)
926}
927
928fn gpu_identity_from_rocminfo(dgpus: &[crate::gpu::AmdGpu]) -> (String, String) {
935 let count = dgpus.len();
936 let text =
937 match std::process::Command::new(crate::installers::common::utils::resolve_rocminfo_path())
938 .output()
939 {
940 Ok(o) if o.status.success() => String::from_utf8_lossy(&o.stdout).to_string(),
941 _ => return gpu_identity_fallback(dgpus),
942 };
943 let mut entries: Vec<(String, u32)> = Vec::new();
944 let mut name = String::new();
945 let mut cus: u32 = 0;
946 let mut is_gpu = false;
947 let mut started = false;
948 for line in text.lines() {
949 let t = line.trim();
950 if t.starts_with("Agent ") && t.contains(char::is_numeric) {
951 if started && is_gpu && !name.is_empty() && !crate::gpu::is_integrated_gpu_name(&name) {
952 entries.push((name.clone(), cus));
953 }
954 name.clear();
955 cus = 0;
956 is_gpu = false;
957 started = false;
958 } else if let Some(v) = t.strip_prefix("Marketing Name:") {
959 name = v.trim().to_string();
960 started = true;
961 } else if let Some(v) = t.strip_prefix("Device Type:") {
962 is_gpu = v.trim().contains("GPU");
963 started = true;
964 } else if let Some(v) = t.strip_prefix("Compute Unit:") {
965 cus = v.trim().parse().unwrap_or(0);
966 started = true;
967 }
968 }
969 if started && is_gpu && !name.is_empty() && !crate::gpu::is_integrated_gpu_name(&name) {
970 entries.push((name, cus));
971 }
972 entries.truncate(count);
973 if entries.len() == count {
974 let names = entries
975 .iter()
976 .map(|(n, _)| n.clone())
977 .collect::<Vec<_>>()
978 .join(",");
979 let cus = entries
980 .iter()
981 .map(|(_, c)| c.to_string())
982 .collect::<Vec<_>>()
983 .join(",");
984 (names, cus)
985 } else {
986 gpu_identity_fallback(dgpus)
987 }
988}
989
990fn gpu_identity_fallback(dgpus: &[crate::gpu::AmdGpu]) -> (String, String) {
993 let names = dgpus
994 .iter()
995 .map(|g| {
996 g.marketing_name
997 .clone()
998 .unwrap_or_else(|| format!("AMD GPU {}", g.pci_device_id))
999 })
1000 .collect::<Vec<_>>()
1001 .join(",");
1002 let cus = dgpus
1003 .iter()
1004 .map(|_| "0".to_string())
1005 .collect::<Vec<_>>()
1006 .join(",");
1007 (names, cus)
1008}
1009
1010fn ensure_gpu_metadata_exports(contents: &str, pairs: &[(&str, &str)]) -> String {
1018 let mut lines: Vec<String> = contents.lines().map(|l| l.to_string()).collect();
1019 for (key, value) in pairs {
1020 let quoted = format!("export {key}=\"{value}\"");
1021 let prefix = format!("export {key}=");
1022 match lines
1023 .iter()
1024 .position(|l| l.trim_start().starts_with(&prefix))
1025 {
1026 Some(idx) => lines[idx] = quoted,
1027 None => lines.push(quoted),
1028 }
1029 }
1030 let mut out = lines.join("\n");
1031 if !out.ends_with('\n') {
1032 out.push('\n');
1033 }
1034 out
1035}
1036
1037fn sanitize_mlstack_env(
1038 contents: &str,
1039 defaults: &[(&str, &str)],
1040 gpu_list: &str,
1041 rocm_home: &str,
1042 python_bin: &str,
1043 install_method: &str,
1044 hsa_override: &str,
1045) -> (String, bool) {
1046 let mut changed = false;
1047 let mut lines = Vec::new();
1048 let mut seen = std::collections::HashSet::new();
1049 let default_for = |key: &str| -> &str {
1054 defaults
1055 .iter()
1056 .find(|(k, _)| *k == key)
1057 .map(|(_, v)| *v)
1058 .unwrap_or("")
1059 };
1060
1061 for line in contents.lines() {
1062 let trimmed = line.trim_start();
1063 if let Some(key) = extract_env_key(trimmed) {
1064 seen.insert(key.to_string());
1065 }
1066 if trimmed.starts_with("export TORCH_CUDA_ARCH_LIST=")
1076 || trimmed.starts_with("export PYTORCH_ALLOC_CONF=")
1077 {
1078 changed = true;
1079 continue;
1080 }
1081 let (line, changed_hip) = normalize_env_value(line, "HIP_PATH", "/opt/rocm");
1082 let (line, changed_rocm) = normalize_env_value(&line, "ROCM_PATH", "/opt/rocm");
1083 let (line, changed_python) = normalize_env_value(&line, "MLSTACK_PYTHON_BIN", python_bin);
1084 let (line, changed_uv_python) = normalize_env_value(&line, "UV_PYTHON", python_bin);
1085 let (line, changed_method) =
1086 normalize_env_value(&line, "MLSTACK_INSTALL_METHOD", install_method);
1087 let (line, changed_method_alias) =
1088 normalize_env_value(&line, "INSTALL_METHOD", install_method);
1089 let (line, changed_hsa) =
1090 normalize_env_value(&line, "HSA_OVERRIDE_GFX_VERSION", hsa_override);
1091 let (line, changed_hip_vis) =
1092 normalize_visible_devices(&line, "HIP_VISIBLE_DEVICES", gpu_list);
1093 let (line, changed_cuda_vis) =
1094 normalize_visible_devices(&line, "CUDA_VISIBLE_DEVICES", gpu_list);
1095 let (line, changed_rocr_vis) =
1096 normalize_visible_devices(&line, "ROCR_VISIBLE_DEVICES", gpu_list);
1097 let (line, changed_rocm_ver) =
1101 normalize_env_value(&line, "ROCM_VERSION", default_for("ROCM_VERSION"));
1102 let (line, changed_gpu_arch) =
1103 normalize_env_value(&line, "GPU_ARCH", default_for("GPU_ARCH"));
1104 let (line, changed_pytorch_arch) =
1105 normalize_env_value(&line, "PYTORCH_ROCM_ARCH", default_for("PYTORCH_ROCM_ARCH"));
1106 let (line, changed_gpu_archs) =
1107 normalize_env_value(&line, "GPU_ARCHS", default_for("GPU_ARCHS"));
1108 let (line, changed_rocm_home) =
1109 normalize_env_value(&line, "ROCM_HOME", default_for("ROCM_HOME"));
1110 let (line, changed_pytorch_dev) = normalize_env_value(
1111 &line,
1112 "PYTORCH_ROCM_DEVICE",
1113 default_for("PYTORCH_ROCM_DEVICE"),
1114 );
1115 let (line, changed_pythonpath) = if line.starts_with("export PYTHONPATH=") {
1116 let desired = format!("export PYTHONPATH={}/lib:$PYTHONPATH", rocm_home);
1117 if line.trim() != desired {
1118 (desired, true)
1119 } else {
1120 (line, false)
1121 }
1122 } else {
1123 (line, false)
1124 };
1125
1126 let (line, changed_path) = if line.starts_with("export PATH=") {
1128 let desired = format!("export PATH=\"/usr/local/bin:/usr/bin:/bin:/usr/local/games:/usr/games:{}/bin:{}/hip/bin:$PATH\"", rocm_home, rocm_home);
1129 if line.trim() != desired {
1130 (desired, true)
1131 } else {
1132 (line, false)
1133 }
1134 } else {
1135 (line, false)
1136 };
1137
1138 let (line, changed_ld) = if line.starts_with("export LD_LIBRARY_PATH=") {
1139 let desired = format!(
1140 "export LD_LIBRARY_PATH=\"$HOME/.mlstack/libmpi-compat:$HOME/.mlstack/libmpi-compat-user-$(id -u):{}/lib:{}/hip/lib:{}/opencl/lib:$LD_LIBRARY_PATH\"",
1141 rocm_home, rocm_home, rocm_home
1142 );
1143 if line.trim() != desired {
1144 (desired, true)
1145 } else {
1146 (line, false)
1147 }
1148 } else {
1149 (line, false)
1150 };
1151
1152 changed |= changed_hip
1153 || changed_rocm
1154 || changed_python
1155 || changed_uv_python
1156 || changed_method
1157 || changed_method_alias
1158 || changed_hsa
1159 || changed_hip_vis
1160 || changed_cuda_vis
1161 || changed_rocr_vis
1162 || changed_rocm_ver
1163 || changed_gpu_arch
1164 || changed_pytorch_arch
1165 || changed_gpu_archs
1166 || changed_rocm_home
1167 || changed_pytorch_dev
1168 || changed_pythonpath
1169 || changed_path
1170 || changed_ld;
1171 lines.push(line);
1172 }
1173
1174 for (key, value) in defaults {
1175 if !seen.contains(*key) {
1176 lines.push(format!("export {}={}", key, value));
1177 changed = true;
1178 }
1179 }
1180
1181 let mut output = lines.join("\n");
1182 output.push('\n');
1183 (output, changed)
1184}
1185
1186fn extract_env_key(line: &str) -> Option<&str> {
1187 let export_idx = line.rfind("export ")?;
1189 let rest = &line[export_idx + 7..]; rest.split('=').next().map(str::trim)
1191}
1192
1193fn normalize_env_value(line: &str, key: &str, desired: &str) -> (String, bool) {
1194 let marker = format!("export {}=", key);
1195 if let Some(export_idx) = line.rfind(&marker) {
1197 let after_marker = &line[export_idx + marker.len()..];
1198 let current = after_marker.trim().trim_matches('"');
1199 let current = strip_shell_suffixes_from_value(current);
1201 if current != desired {
1202 return (format!("export {}={}", key, desired), true);
1204 }
1205 }
1206 (line.to_string(), false)
1207}
1208
1209fn normalize_visible_devices(line: &str, key: &str, desired: &str) -> (String, bool) {
1210 let desired = desired.trim();
1211 if desired.is_empty() {
1212 return (line.to_string(), false);
1213 }
1214
1215 let marker = format!("export {}=", key);
1216 if let Some(export_idx) = line.rfind(&marker) {
1218 let after_marker = &line[export_idx + marker.len()..];
1219 let current = after_marker.trim().trim_matches('"');
1220 let current = strip_shell_suffixes_from_value(current);
1222
1223 if current != desired {
1229 return (format!("export {}={}", key, desired), true);
1231 }
1232 }
1233 (line.to_string(), false)
1234}
1235
1236fn detect_rocm_version() -> String {
1237 if let Ok(version) = fs::read_to_string("/opt/rocm/.info/version") {
1238 return version.trim().to_string();
1239 }
1240 "7.2.0".to_string()
1241}
1242
1243fn run_verification(
1244 component: &Component,
1245 python_candidates: &[String],
1246 sender: &Sender<InstallerEvent>,
1247 user_home: &str,
1248) -> VerificationOutcome {
1249 let steps = verification_commands(&component.id, python_candidates);
1250 let mut results: HashMap<String, VerificationCheckResult> = HashMap::new();
1251 let mut executed = 0usize;
1252 let mut failed = 0usize;
1253
1254 for step in &steps {
1255 if let Some((benchmark_ok, benchmark_output)) = verify_benchmark_target(&step.target_id) {
1256 executed += 1;
1257 if benchmark_ok {
1258 results.insert(
1259 step.target_id.clone(),
1260 VerificationCheckResult {
1261 status: VerificationResult::Verified,
1262 output: benchmark_output,
1263 },
1264 );
1265 } else {
1266 failed += 1;
1267 results.insert(
1268 step.target_id.clone(),
1269 VerificationCheckResult {
1270 status: VerificationResult::Failed,
1271 output: benchmark_output,
1272 },
1273 );
1274 let _ = sender.send(InstallerEvent::Log(
1275 format!("Verification failed for {}", step.label),
1276 false,
1277 ));
1278 }
1279 continue;
1280 }
1281
1282 if !is_component_installed_by_id(&step.target_id, python_candidates) {
1283 results.insert(
1284 step.target_id.clone(),
1285 VerificationCheckResult {
1286 status: VerificationResult::Missing,
1287 output: vec!["Component not found in path or site-packages".to_string()],
1288 },
1289 );
1290 let _ = sender.send(InstallerEvent::Log(
1291 format!("Skipping {} (not installed)", step.label),
1292 false,
1293 ));
1294 continue;
1295 }
1296 if !modules_available(&step.modules, python_candidates) {
1297 results.insert(
1298 step.target_id.clone(),
1299 VerificationCheckResult {
1300 status: VerificationResult::Missing,
1301 output: vec![format!(
1302 "Required Python modules missing: {}",
1303 step.modules.join(", ")
1304 )],
1305 },
1306 );
1307 let _ = sender.send(InstallerEvent::Log(
1308 format!("Skipping {} (module not available)", step.label),
1309 false,
1310 ));
1311 continue;
1312 }
1313 executed += 1;
1314 let _ = sender.send(InstallerEvent::Log(
1315 format!("Verifying {}...", step.label),
1316 false,
1317 ));
1318 match run_verification_command(step, sender) {
1319 Ok((success, output)) => {
1320 if success {
1321 results.insert(
1322 step.target_id.clone(),
1323 VerificationCheckResult {
1324 status: VerificationResult::Verified,
1325 output,
1326 },
1327 );
1328 } else {
1329 failed += 1;
1330 results.insert(
1331 step.target_id.clone(),
1332 VerificationCheckResult {
1333 status: VerificationResult::Failed,
1334 output,
1335 },
1336 );
1337 let _ = sender.send(InstallerEvent::Log(
1338 format!("Verification failed for {}", step.label),
1339 false,
1340 ));
1341 }
1342 }
1343 Err(err) => {
1344 failed += 1;
1345 results.insert(
1346 step.target_id.clone(),
1347 VerificationCheckResult {
1348 status: VerificationResult::Failed,
1349 output: vec![err.to_string()],
1350 },
1351 );
1352 let _ = sender.send(InstallerEvent::Log(
1353 format!("Failed to execute verification for {}: {}", step.label, err),
1354 false,
1355 ));
1356 }
1357 }
1358 }
1359
1360 let report_lines = build_verification_report(python_candidates, user_home, &results, &steps);
1361 let success = executed > 0 && failed == 0;
1362
1363 if executed == 0 {
1364 let _ = sender.send(InstallerEvent::Log(
1365 "No installed components found to verify".into(),
1366 false,
1367 ));
1368 }
1369
1370 if failed > 0 {
1371 let _ = sender.send(InstallerEvent::Log(
1372 format!("Verification completed with {} failed checks", failed),
1373 false,
1374 ));
1375 }
1376
1377 VerificationOutcome {
1378 success,
1379 report_lines,
1380 }
1381}
1382
1383fn build_verification_report(
1384 python_candidates: &[String],
1385 user_home: &str,
1386 results: &HashMap<String, VerificationCheckResult>,
1387 steps: &[VerificationCommand],
1388) -> Vec<String> {
1389 let mut lines = Vec::new();
1390 lines.push("Environment:".to_string());
1391 for (key, value) in collect_env_info(user_home) {
1392 lines.push(format!(" {}: {}", key, value));
1393 }
1394
1395 lines.push("".to_string());
1396 lines.push("Python Interpreters:".to_string());
1397 for py in python_candidates {
1398 lines.push(format!(" - {}", py));
1399 }
1400
1401 lines.push("".to_string());
1402 lines.push("Checks:".to_string());
1403 for step in steps {
1404 if let Some(res) = results.get(&step.target_id) {
1405 lines.push(format!(" - {}: {}", step.label, res.status.label()));
1406 for out_line in &res.output {
1407 lines.push(format!(" > {}", out_line));
1408 }
1409 } else {
1410 lines.push(format!(" - {}: Skipped", step.label));
1411 }
1412 }
1413
1414 lines
1415}
1416
1417fn run_component_verification(
1418 component: &Component,
1419 python_candidates: &[String],
1420 sender: &Sender<InstallerEvent>,
1421 user_home: &str,
1422) -> VerificationOutcome {
1423 let steps = component_verification_commands(&component.id, python_candidates);
1424 let mut results: HashMap<String, VerificationCheckResult> = HashMap::new();
1425 let mut executed = 0usize;
1426 let mut failed = 0usize;
1427
1428 for step in &steps {
1429 if let Some((benchmark_ok, benchmark_output)) = verify_benchmark_target(&step.target_id) {
1430 executed += 1;
1431 if benchmark_ok {
1432 results.insert(
1433 step.target_id.clone(),
1434 VerificationCheckResult {
1435 status: VerificationResult::Verified,
1436 output: benchmark_output,
1437 },
1438 );
1439 } else {
1440 failed += 1;
1441 results.insert(
1442 step.target_id.clone(),
1443 VerificationCheckResult {
1444 status: VerificationResult::Failed,
1445 output: benchmark_output,
1446 },
1447 );
1448 }
1449 continue;
1450 }
1451
1452 if !modules_available(&step.modules, python_candidates) {
1453 results.insert(
1454 step.target_id.clone(),
1455 VerificationCheckResult {
1456 status: VerificationResult::Missing,
1457 output: vec![format!(
1458 "Required Python modules missing: {}",
1459 step.modules.join(", ")
1460 )],
1461 },
1462 );
1463 continue;
1464 }
1465 executed += 1;
1466 match run_verification_command(step, sender) {
1467 Ok((success, output)) => {
1468 if success {
1469 results.insert(
1470 step.target_id.clone(),
1471 VerificationCheckResult {
1472 status: VerificationResult::Verified,
1473 output,
1474 },
1475 );
1476 } else {
1477 failed += 1;
1478 results.insert(
1479 step.target_id.clone(),
1480 VerificationCheckResult {
1481 status: VerificationResult::Failed,
1482 output,
1483 },
1484 );
1485 }
1486 }
1487 Err(err) => {
1488 failed += 1;
1489 results.insert(
1490 step.target_id.clone(),
1491 VerificationCheckResult {
1492 status: VerificationResult::Failed,
1493 output: vec![err.to_string()],
1494 },
1495 );
1496 }
1497 }
1498 }
1499
1500 let report_lines = build_component_report(component, user_home, &results, &steps);
1501 let success = executed > 0 && failed == 0;
1502
1503 VerificationOutcome {
1504 success,
1505 report_lines,
1506 }
1507}
1508
1509fn is_native_maintenance_action(component: &Component) -> bool {
1510 component.category == Category::Maintenance && !component.id.starts_with("verify-")
1511}
1512
1513fn native_maintenance_action_outcome(
1514 component: &Component,
1515 install_success: bool,
1516) -> VerificationOutcome {
1517 VerificationOutcome {
1518 success: install_success,
1519 report_lines: vec![format!(
1520 "{} native action {}",
1521 component.name,
1522 if install_success {
1523 "reported success"
1524 } else {
1525 "reported failure"
1526 }
1527 )],
1528 }
1529}
1530
1531fn build_component_report(
1532 component: &Component,
1533 user_home: &str,
1534 results: &HashMap<String, VerificationCheckResult>,
1535 steps: &[VerificationCommand],
1536) -> Vec<String> {
1537 let mut lines = Vec::new();
1538 lines.push(format!("Component: {}", component.name));
1539
1540 for (key, value) in collect_env_info(user_home) {
1541 lines.push(format!(" {}: {}", key, value));
1542 }
1543
1544 if let Some(benchmark_lines) = extract_benchmark_results(&component.id) {
1546 lines.push("".to_string());
1547 lines.push("Benchmark Results:".to_string());
1548 lines.extend(benchmark_lines);
1549 }
1550
1551 lines.push("".to_string());
1552 lines.push("Verification Checks:".to_string());
1553
1554 for step in steps {
1555 if let Some(res) = results.get(&step.target_id) {
1556 lines.push(format!(" - {}: {}", step.label, res.status.label()));
1557 for out_line in &res.output {
1558 lines.push(format!(" > {}", out_line));
1559 }
1560 } else {
1561 lines.push(format!(" - {}: Skipped", step.label));
1562 }
1563 }
1564
1565 lines
1566}
1567
1568fn extract_benchmark_results(component_id: &str) -> Option<Vec<String>> {
1569 let log_dirs = benchmark_log_directories();
1570 if log_dirs.is_empty() {
1571 return None;
1572 }
1573
1574 let pattern = match component_id {
1576 "mlperf-inference" => "all_benchmarks",
1577 "rocm-benchmarks" => "gpu-capability_benchmarks",
1578 "gpu-memory-bandwidth" => "memory-bandwidth_benchmarks",
1579 "rocm-smi-bench" => "gpu-capability_benchmarks",
1580 "pytorch-performance" => "pytorch_performance",
1581 "vllm-performance" => "vllm_benchmarks",
1582 "deepspeed-performance" => "deepspeed_benchmarks",
1583 "megatron-performance" => "megatron_benchmarks",
1584 "onnx-performance" => "onnx_benchmarks",
1585 "all-benchmarks" => "all_benchmarks",
1586 _ => return None,
1587 };
1588
1589 let log_file = find_latest_log_in_dirs(&log_dirs, pattern)?;
1590 let contents = std::fs::read_to_string(log_file).ok()?;
1591
1592 let Some(json) = extract_benchmark_json_value(&contents) else {
1593 return Some(vec![" No metrics found in log".to_string()]);
1594 };
1595
1596 let mut result_lines = benchmark_result_lines_from_json(&json);
1597 if result_lines.is_empty() {
1598 result_lines.push(" No metrics found in log".to_string());
1599 }
1600
1601 Some(result_lines)
1602}
1603
1604fn benchmark_result_lines_from_json(json: &serde_json::Value) -> Vec<String> {
1605 let mut result_lines = Vec::new();
1606
1607 if let Some(results) = json.get("results").and_then(|r| r.as_object()) {
1610 let nested_suite = results.values().any(|value| {
1611 value.as_object().is_some_and(|obj| {
1612 obj.contains_key("success")
1613 || obj.contains_key("metrics")
1614 || obj.contains_key("execution_time_ms")
1615 })
1616 });
1617 if nested_suite {
1618 for (bench_name, bench_data) in results {
1619 if bench_data.is_object() {
1620 let status = bench_data
1621 .get("success")
1622 .and_then(|v| v.as_bool())
1623 .map(|ok| if ok { "SUCCESS" } else { "FAILED" })
1624 .unwrap_or("UNKNOWN");
1625 let time_ms = bench_data
1626 .get("execution_time_ms")
1627 .and_then(|v| v.as_u64())
1628 .unwrap_or(0);
1629 result_lines.push(format!(
1630 " === {} ({}, {} ms) ===",
1631 bench_name.replace('_', " ").to_uppercase(),
1632 status,
1633 time_ms
1634 ));
1635 if let Some(metrics) = bench_data.get("metrics").and_then(|m| m.as_object()) {
1636 append_benchmark_metric_lines(&mut result_lines, metrics, " ");
1637 }
1638 append_benchmark_error_lines(&mut result_lines, bench_data, " ");
1639 result_lines.push(String::new());
1640 }
1641 }
1642 } else {
1643 append_benchmark_metric_lines(&mut result_lines, results, " ");
1644 append_benchmark_error_lines(&mut result_lines, json, " ");
1645 }
1646 } else if let Some(metrics) = json.get("metrics").and_then(|m| m.as_object()) {
1647 append_benchmark_metric_lines(&mut result_lines, metrics, " ");
1649 append_benchmark_error_lines(&mut result_lines, json, " ");
1650 }
1651
1652 result_lines
1653}
1654
1655fn append_benchmark_metric_lines(
1656 lines: &mut Vec<String>,
1657 metrics: &serde_json::Map<String, serde_json::Value>,
1658 indent: &str,
1659) {
1660 for (key, value) in metrics {
1661 let value_str = match value {
1662 serde_json::Value::Number(n) => n.to_string(),
1663 serde_json::Value::String(s) => s.clone(),
1664 serde_json::Value::Bool(b) => b.to_string(),
1665 serde_json::Value::Array(arr) => format!("{} values", arr.len()),
1666 _ => continue,
1667 };
1668 let formatted_key = key.replace('_', " ").to_uppercase();
1669 lines.push(format!("{}{}: {}", indent, formatted_key, value_str));
1670 }
1671}
1672
1673fn append_benchmark_error_lines(lines: &mut Vec<String>, value: &serde_json::Value, indent: &str) {
1674 if let Some(errors) = value.get("errors").and_then(|v| v.as_array()) {
1675 for err in errors.iter().filter_map(|v| v.as_str()).take(3) {
1676 lines.push(format!("{}ERROR: {}", indent, err));
1677 }
1678 }
1679}
1680
1681fn benchmark_pattern_for_target(target_id: &str) -> Option<&'static str> {
1682 match target_id {
1683 "mlperf-inference" => Some("all_benchmarks"),
1684 "rocm-benchmarks" => Some("gpu-capability_benchmarks"),
1685 "gpu-memory-bandwidth" => Some("memory-bandwidth_benchmarks"),
1686 "rocm-smi-bench" => Some("gpu-capability_benchmarks"),
1687 "pytorch-performance" => Some("pytorch_performance"),
1688 "vllm-performance" => Some("vllm_benchmarks"),
1689 "deepspeed-performance" => Some("deepspeed_benchmarks"),
1690 "megatron-performance" => Some("megatron_benchmarks"),
1691 "onnx-performance" => Some("onnx_benchmarks"),
1692 "all-benchmarks" => Some("all_benchmarks"),
1693 _ => None,
1694 }
1695}
1696
1697fn benchmark_success_and_errors(value: &serde_json::Value) -> (bool, Vec<String>) {
1698 if let Some(success) = value.get("success").and_then(|v| v.as_bool()) {
1699 let mut details = Vec::new();
1700 if !success {
1701 if let Some(reason) = value.get("reason").and_then(|v| v.as_str()) {
1702 let reason = reason.trim();
1703 if !reason.is_empty() {
1704 details.push(reason.to_string());
1705 }
1706 }
1707 if let Some(errors) = value.get("errors").and_then(|v| v.as_array()) {
1708 for err in errors.iter().filter_map(|entry| entry.as_str()).take(4) {
1709 let trimmed = err.trim();
1710 if !trimmed.is_empty() {
1711 details.push(trimmed.to_string());
1712 }
1713 }
1714 }
1715 if details.is_empty() {
1716 details.push("Benchmark JSON reported success=false".to_string());
1717 }
1718 }
1719 return (success, details);
1720 }
1721
1722 if let Some(results) = value.get("results").and_then(|v| v.as_object()) {
1723 let mut failed = Vec::new();
1724 let mut seen_result = false;
1725 for (name, item) in results {
1726 if let Some(item_success) = item.get("success").and_then(|v| v.as_bool()) {
1727 seen_result = true;
1728 if !item_success {
1729 let mut reason = item
1730 .get("reason")
1731 .and_then(|v| v.as_str())
1732 .map(|s| s.trim().to_string())
1733 .filter(|s| !s.is_empty())
1734 .unwrap_or_else(|| format!("{name}: success=false"));
1735 if reason == format!("{name}: success=false") {
1736 if let Some(errors) = item.get("errors").and_then(|v| v.as_array()) {
1737 if let Some(first) = errors.iter().filter_map(|v| v.as_str()).next() {
1738 let first = first.trim();
1739 if !first.is_empty() {
1740 reason = format!("{name}: {first}");
1741 }
1742 }
1743 }
1744 }
1745 failed.push(reason);
1746 }
1747 }
1748 }
1749 if seen_result {
1750 return (failed.is_empty(), failed);
1751 }
1752 }
1753
1754 (
1755 false,
1756 vec!["Benchmark JSON does not contain a recognizable success field".to_string()],
1757 )
1758}
1759
1760fn verify_benchmark_target(target_id: &str) -> Option<(bool, Vec<String>)> {
1761 let pattern = benchmark_pattern_for_target(target_id)?;
1762 let log_dirs = benchmark_log_directories();
1763 if log_dirs.is_empty() {
1764 return Some((
1765 false,
1766 vec!["No benchmark log directories found".to_string()],
1767 ));
1768 }
1769
1770 let Some(log_file) = find_latest_log_in_dirs(&log_dirs, pattern) else {
1771 return Some((
1772 false,
1773 vec![format!("No benchmark logs found for pattern '{pattern}'")],
1774 ));
1775 };
1776
1777 let mut output = vec![format!("Latest benchmark log: {}", log_file.display())];
1778 let contents = match std::fs::read_to_string(&log_file) {
1779 Ok(contents) => contents,
1780 Err(err) => {
1781 output.push(format!("Could not read benchmark log: {err}"));
1782 return Some((false, output));
1783 }
1784 };
1785
1786 let Some(json) = extract_benchmark_json_value(&contents) else {
1787 output.push("Could not parse benchmark JSON payload from log".to_string());
1788 return Some((false, output));
1789 };
1790
1791 let (success, mut details) = benchmark_success_and_errors(&json);
1792 output.append(&mut details);
1793 Some((success, output))
1794}
1795
1796fn collect_env_info(user_home: &str) -> Vec<(String, String)> {
1797 let mut values = load_mlstack_env_exports(user_home);
1798
1799 let keys = [
1800 "ROCM_VERSION",
1801 "ROCM_CHANNEL",
1802 "GPU_ARCH",
1803 "GPU_ARCHS",
1804 "PYTORCH_ROCM_ARCH",
1805 "HIP_VISIBLE_DEVICES",
1806 "CUDA_VISIBLE_DEVICES",
1807 "ROCR_VISIBLE_DEVICES",
1808 "PYTORCH_ROCM_DEVICE",
1809 "HSA_OVERRIDE_GFX_VERSION",
1810 "ROCM_HOME",
1811 "ROCM_PATH",
1812 "HIP_PATH",
1813 "MLSTACK_PYTHON_BIN",
1814 "PYTHONPATH",
1815 ];
1816
1817 for key in keys {
1818 if !values.contains_key(key) {
1819 if let Ok(value) = std::env::var(key) {
1820 values.insert(key.to_string(), value);
1821 }
1822 }
1823 }
1824
1825 keys.iter()
1826 .map(|key| {
1827 let default_val = if *key == "ROCM_HOME" || *key == "ROCM_PATH" || *key == "HIP_PATH" {
1828 "/opt/rocm".to_string()
1829 } else {
1830 "unknown".to_string()
1831 };
1832 (
1833 key.to_string(),
1834 values.get(*key).cloned().unwrap_or(default_val),
1835 )
1836 })
1837 .collect()
1838}
1839
1840fn load_mlstack_env_exports(user_home: &str) -> HashMap<String, String> {
1841 let mut values: HashMap<String, String> = HashMap::new();
1842 for home in candidate_mlstack_homes(user_home) {
1843 let env_path = PathBuf::from(home).join(".mlstack_env");
1844 if let Ok(contents) = fs::read_to_string(&env_path) {
1845 for line in contents.lines() {
1846 if let Some((key, value)) = parse_env_export(line) {
1847 values.entry(key).or_insert(value);
1848 }
1849 }
1850 }
1851 }
1852
1853 values
1854}
1855
1856fn parse_env_export(line: &str) -> Option<(String, String)> {
1857 let export_idx = line.rfind("export ")?;
1858 let rest = &line[export_idx + 7..];
1859 let mut parts = rest.splitn(2, '=');
1860 let key = parts.next()?.trim().to_string();
1861 let value = parts.next()?.trim().trim_matches('"').to_string();
1862 let value = strip_shell_suffixes_from_value(&value);
1864 Some((key, value))
1865}
1866
1867fn strip_shell_suffixes_from_value(value: &str) -> String {
1875 let mut v = value.to_string();
1876 if let Some(idx) = v.find("; fi") {
1878 v.truncate(idx);
1879 }
1880 if let Some(idx) = v.find("&& ") {
1882 v.truncate(idx);
1883 }
1884 v.trim().to_string()
1885}
1886
1887fn detect_onnx_rocm_provider(python: &str) -> Option<String> {
1888 let output = Command::new(python)
1889 .arg("-c")
1890 .arg(
1891 "import onnxruntime, pathlib; base=pathlib.Path(onnxruntime.__file__).parent; libs=list(base.rglob('libonnxruntime_providers_rocm.so')); print(libs[0] if libs else '')",
1892 )
1893 .output()
1894 .ok()?;
1895 if !output.status.success() {
1896 return None;
1897 }
1898 let value = String::from_utf8_lossy(&output.stdout).trim().to_string();
1899 if value.is_empty() {
1900 None
1901 } else {
1902 Some(value)
1903 }
1904}
1905
1906fn push_unique_env_path(paths: &mut Vec<String>, value: &str) {
1907 let candidate = value.trim();
1908 if candidate.is_empty() {
1909 return;
1910 }
1911 if paths.iter().any(|existing| existing == candidate) {
1912 return;
1913 }
1914 paths.push(candidate.to_string());
1915}
1916
1917fn user_home_from_passwd(user_name: &str) -> Option<String> {
1918 let target = user_name.trim();
1919 if target.is_empty() {
1920 return None;
1921 }
1922 let passwd = fs::read_to_string("/etc/passwd").ok()?;
1923 passwd.lines().find_map(|line| {
1924 if line.trim().is_empty() || line.starts_with('#') {
1925 return None;
1926 }
1927 let fields: Vec<&str> = line.split(':').collect();
1928 if fields.len() < 6 || fields[0] != target {
1929 return None;
1930 }
1931 let home = fields[5].trim();
1932 if home.is_empty() {
1933 None
1934 } else {
1935 Some(home.to_string())
1936 }
1937 })
1938}
1939
1940fn passwd_homes_with_mlstack() -> Vec<String> {
1941 let mut homes = Vec::new();
1942 let Ok(passwd) = fs::read_to_string("/etc/passwd") else {
1943 return homes;
1944 };
1945
1946 for line in passwd.lines() {
1947 if line.trim().is_empty() || line.starts_with('#') {
1948 continue;
1949 }
1950 let fields: Vec<&str> = line.split(':').collect();
1951 if fields.len() < 6 {
1952 continue;
1953 }
1954 let home = fields[5].trim();
1955 if home.is_empty() {
1956 continue;
1957 }
1958 if Path::new(home).join(".mlstack").is_dir()
1959 || Path::new(home).join(".mlstack_env").is_file()
1960 {
1961 push_unique_env_path(&mut homes, home);
1962 }
1963 }
1964
1965 homes
1966}
1967
1968fn candidate_mlstack_homes(preferred_home: &str) -> Vec<String> {
1969 let mut homes = Vec::new();
1970 push_unique_env_path(&mut homes, preferred_home);
1971
1972 if let Ok(value) = std::env::var("MLSTACK_USER_HOME") {
1973 push_unique_env_path(&mut homes, &value);
1974 }
1975 if let Ok(value) = std::env::var("HOME") {
1976 push_unique_env_path(&mut homes, &value);
1977 }
1978
1979 for key in ["SUDO_USER", "USER", "LOGNAME"] {
1980 if let Ok(user_name) = std::env::var(key) {
1981 if let Some(home) = user_home_from_passwd(&user_name) {
1982 push_unique_env_path(&mut homes, &home);
1983 }
1984 }
1985 }
1986
1987 for home in passwd_homes_with_mlstack() {
1988 push_unique_env_path(&mut homes, &home);
1989 }
1990
1991 homes
1992}
1993
1994fn resolve_mlstack_user_home() -> String {
1995 let fallback = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
1996 let candidates = candidate_mlstack_homes(&fallback);
1997
1998 for home in &candidates {
1999 if Path::new(home).join(".mlstack_env").is_file() {
2000 return home.clone();
2001 }
2002 }
2003 for home in &candidates {
2004 if Path::new(home).join(".mlstack").is_dir() {
2005 return home.clone();
2006 }
2007 }
2008
2009 candidates.into_iter().next().unwrap_or(fallback)
2010}
2011
2012fn mlstack_mpi_compat_dirs(user_home: &str) -> Vec<String> {
2013 let mut dirs = Vec::new();
2014 for home in candidate_mlstack_homes(user_home) {
2015 let mlstack_dir = PathBuf::from(home).join(".mlstack");
2016
2017 let primary = mlstack_dir.join("libmpi-compat");
2018 if primary.is_dir() {
2019 dirs.push(primary.to_string_lossy().to_string());
2020 }
2021
2022 if let Ok(entries) = fs::read_dir(&mlstack_dir) {
2023 for entry in entries.flatten() {
2024 let path = entry.path();
2025 if !path.is_dir() {
2026 continue;
2027 }
2028 let Some(name) = path.file_name().and_then(|value| value.to_str()) else {
2029 continue;
2030 };
2031 if name.starts_with("libmpi-compat-user-") {
2032 dirs.push(path.to_string_lossy().to_string());
2033 }
2034 }
2035 }
2036 }
2037
2038 dirs.sort();
2039 dirs.dedup();
2040 dirs
2041}
2042
2043fn build_verification_ld_library_path(user_home: &str) -> String {
2044 let mut paths: Vec<String> = Vec::new();
2045
2046 for compat_dir in mlstack_mpi_compat_dirs(user_home) {
2047 push_unique_env_path(&mut paths, &compat_dir);
2048 }
2049
2050 for rocm_path in ["/opt/rocm/lib", "/opt/rocm/hip/lib", "/opt/rocm/opencl/lib"] {
2051 if Path::new(rocm_path).exists() {
2052 push_unique_env_path(&mut paths, rocm_path);
2053 }
2054 }
2055
2056 if let Ok(existing) = std::env::var("LD_LIBRARY_PATH") {
2057 for part in existing.split(':') {
2058 push_unique_env_path(&mut paths, part);
2059 }
2060 }
2061
2062 paths.join(":")
2063}
2064
2065fn run_verification_command(
2066 step: &VerificationCommand,
2067 sender: &Sender<InstallerEvent>,
2068) -> Result<(bool, Vec<String>)> {
2069 let user_home = resolve_mlstack_user_home();
2070 let user_name = std::env::var("USER").unwrap_or_else(|_| "user".into());
2071 let env_exports = load_mlstack_env_exports(&user_home);
2072
2073 let mut command = Command::new(&step.program);
2074 command.args(&step.args);
2075
2076 let extra_paths = python_search_paths();
2077 let mut pythonpath = std::env::var("PYTHONPATH").unwrap_or_default();
2078 if !extra_paths.is_empty() {
2079 if !pythonpath.is_empty() {
2080 pythonpath.push(':');
2081 }
2082 pythonpath.push_str(&extra_paths.join(":"));
2083 }
2084
2085 pythonpath = pythonpath
2087 .split(':')
2088 .filter(|part| !part.starts_with("/tmp/"))
2089 .collect::<Vec<_>>()
2090 .join(":");
2091
2092 if !pythonpath.is_empty() {
2093 command.env("PYTHONPATH", pythonpath);
2094 }
2095
2096 let verification_ld = build_verification_ld_library_path(&user_home);
2097 if !verification_ld.is_empty() {
2098 command.env("LD_LIBRARY_PATH", verification_ld);
2099 }
2100
2101 if step.target_id == "onnx" {
2102 if let Some(lib) = detect_onnx_rocm_provider(&step.program) {
2103 command.env("ORT_ROCM_EP_PROVIDER_PATH", lib);
2104 }
2105 }
2106
2107 let selected_python = if !step.modules.is_empty() {
2108 step.program.clone()
2109 } else {
2110 resolve_python_bin()
2111 };
2112 let gpu_arch = env_exports
2113 .get("GPU_ARCH")
2114 .map(|v| v.trim().to_string())
2115 .filter(|v| !v.is_empty())
2116 .unwrap_or_else(|| {
2117 detect_gpu_arch().unwrap_or_else(|| {
2118 eprintln!("Warning: no GPU arch detected, falling back to gfx1100");
2119 "gfx1100".to_string()
2120 })
2121 });
2122 let py_rocm_arch = env_exports
2123 .get("PYTORCH_ROCM_ARCH")
2124 .map(|v| v.trim().to_string())
2125 .filter(|v| !v.is_empty())
2126 .unwrap_or_else(|| gpu_arch.clone());
2127 let mut gpu_archs = env_exports
2128 .get("GPU_ARCHS")
2129 .map(|v| v.trim().to_string())
2130 .filter(|v| !v.is_empty())
2131 .unwrap_or_else(|| gpu_arch.clone());
2132 if step.target_id == "aiter" {
2133 gpu_archs = py_rocm_arch
2134 .split(';')
2135 .next()
2136 .map(|v| v.trim().to_string())
2137 .filter(|v| !v.is_empty())
2138 .unwrap_or_else(|| gpu_arch.clone());
2139 }
2140 let mut hip_visible = env_exports
2141 .get("HIP_VISIBLE_DEVICES")
2142 .map(|v| v.trim().to_string())
2143 .filter(|v| !v.is_empty());
2144 let mut cuda_visible = env_exports
2145 .get("CUDA_VISIBLE_DEVICES")
2146 .map(|v| v.trim().to_string())
2147 .filter(|v| !v.is_empty())
2148 .or_else(|| hip_visible.clone());
2149 let mut rocr_visible = env_exports
2150 .get("ROCR_VISIBLE_DEVICES")
2151 .map(|v| v.trim().to_string())
2152 .filter(|v| !v.is_empty())
2153 .or_else(|| hip_visible.clone());
2154 if step.target_id == "aiter" {
2155 let base_list = hip_visible
2156 .clone()
2157 .or_else(|| cuda_visible.clone())
2158 .unwrap_or_else(detect_gpu_list);
2159 let ordered = prioritize_gpu_list_for_arch(&base_list, &gpu_archs);
2160 hip_visible = Some(ordered.clone());
2161 cuda_visible = Some(ordered.clone());
2162 rocr_visible = Some(ordered);
2163 }
2164 let hsa_override = env_exports
2165 .get("HSA_OVERRIDE_GFX_VERSION")
2166 .map(|v| v.trim().to_string())
2167 .filter(|v| !v.is_empty());
2168
2169 command
2170 .env("HOME", &user_home)
2171 .env("USER", &user_name)
2172 .env("LOGNAME", &user_name)
2173 .env("MLSTACK_USER_HOME", &user_home)
2174 .env("MLSTACK_PYTHON_BIN", &selected_python)
2175 .env("UV_PYTHON", &selected_python)
2176 .env("PYTHONWARNINGS", "ignore")
2177 .env("PYTHONUNBUFFERED", "1")
2178 .env("ROCM_HOME", "/opt/rocm")
2179 .env("HIP_PATH", "/opt/rocm")
2180 .env("GPU_ARCH", &gpu_arch)
2181 .env("PYTORCH_ROCM_ARCH", &py_rocm_arch)
2182 .env("GPU_ARCHS", &gpu_archs)
2183 .stdout(Stdio::piped())
2184 .stderr(Stdio::piped());
2185
2186 if let Some(value) = hip_visible {
2187 command.env("HIP_VISIBLE_DEVICES", value);
2188 }
2189 if let Some(value) = cuda_visible {
2190 command.env("CUDA_VISIBLE_DEVICES", value);
2191 }
2192 if let Some(value) = rocr_visible {
2193 command.env("ROCR_VISIBLE_DEVICES", value);
2194 }
2195 if let Some(value) = hsa_override {
2196 command.env("HSA_OVERRIDE_GFX_VERSION", value);
2197 }
2198 if step.target_id == "aiter" {
2199 let aiter_jit_dir = PathBuf::from(&user_home)
2200 .join(".mlstack")
2201 .join("aiter")
2202 .join("jit");
2203 let _ = fs::create_dir_all(&aiter_jit_dir);
2204 command.env("AITER_JIT_DIR", aiter_jit_dir);
2205 }
2206
2207 if step.target_id == "mpi4py" {
2208 command
2210 .env("OMPI_MCA_btl", "^smcuda")
2211 .env("OMPI_MCA_opal_cuda_support", "0")
2212 .env("OMPI_MCA_opal_accelerator_rocm_enable", "0");
2213 }
2214
2215 let output = command
2216 .output()
2217 .with_context(|| format!("Failed to run verification for {}", step.label))?;
2218
2219 let mut captured_output = Vec::new();
2220 const MAX_CAPTURED_OUTPUT_LINES: usize = 24;
2221
2222 let stdout = String::from_utf8_lossy(&output.stdout);
2223 for raw in stdout.replace('\r', "\n").lines() {
2224 let clean_line = sanitize_output_line(raw);
2225 if should_suppress_verification_line(&clean_line, &step.target_id) {
2226 continue;
2227 }
2228 if !clean_line.trim().is_empty() && captured_output.len() < MAX_CAPTURED_OUTPUT_LINES {
2229 let line = format!("{}: {}", step.label, clean_line);
2230 let _ = sender.send(InstallerEvent::Log(line.clone(), false));
2231 captured_output.push(clean_line);
2232 }
2233 }
2234
2235 let stderr = String::from_utf8_lossy(&output.stderr);
2236 for raw in stderr.replace('\r', "\n").lines() {
2237 let clean_line = sanitize_output_line(raw);
2238 if should_suppress_verification_line(&clean_line, &step.target_id) {
2239 continue;
2240 }
2241 if !clean_line.trim().is_empty() && captured_output.len() < MAX_CAPTURED_OUTPUT_LINES {
2242 let line = format!("{}: {}", step.label, clean_line);
2243 let _ = sender.send(InstallerEvent::Log(line.clone(), false));
2244 captured_output.push(clean_line);
2245 }
2246 }
2247
2248 if output.status.success() && captured_output.len() == MAX_CAPTURED_OUTPUT_LINES {
2249 captured_output.push("Output truncated".to_string());
2250 }
2251
2252 Ok((output.status.success(), captured_output))
2253}
2254
2255fn sort_components_by_dependencies(components: Vec<Component>) -> Vec<Component> {
2268 use crate::installers::components::topological_sort;
2269
2270 let mut installer_comps: Vec<Component> = Vec::new();
2272 let mut other_comps: Vec<Component> = Vec::new();
2273
2274 for comp in components {
2275 if comp.category == Category::Maintenance || comp.category == Category::Performance {
2276 other_comps.push(comp);
2277 } else {
2278 installer_comps.push(comp);
2279 }
2280 }
2281
2282 let ids: Vec<String> = installer_comps.iter().map(|c| c.id.clone()).collect();
2284
2285 match topological_sort(&ids) {
2286 Ok(sorted_ids) => {
2287 let mut reordered = Vec::with_capacity(installer_comps.len());
2289 for id in &sorted_ids {
2290 if let Some(pos) = installer_comps.iter().position(|c| c.id == *id) {
2291 reordered.push(installer_comps.remove(pos));
2292 }
2293 }
2294 reordered.extend(installer_comps);
2296
2297 reordered.extend(other_comps);
2299 reordered
2300 }
2301 Err(_) => {
2302 installer_comps.extend(other_comps);
2305 installer_comps
2306 }
2307 }
2308}
2309
2310struct NativeInstallerContext<'a> {
2317 sudo_password: Option<String>,
2318 batch_mode: bool,
2319 install_method: &'a str,
2320 sender: &'a Sender<InstallerEvent>,
2321 user_home: &'a str,
2322 env_exports: &'a HashMap<String, String>,
2323 vllm_version_lag: u32,
2325 vllm_version_min_age_days: u32,
2326 target_version: Option<String>,
2328}
2329
2330enum NativeCommand {
2340 Shell {
2342 program: String,
2343 args: Vec<String>,
2344 env: Vec<(String, String)>,
2345 working_dir: Option<PathBuf>,
2346 },
2347 Pip { program: String, args: Vec<String> },
2349 System { program: String, args: Vec<String> },
2351 Package { program: String, args: Vec<String> },
2353}
2354
2355impl NativeCommand {
2356 fn from_shell_cmd(program: &str, args: &[String], env: &[(String, String)]) -> Self {
2360 NativeCommand::Shell {
2361 program: program.to_string(),
2362 args: args.to_vec(),
2363 env: env.to_vec(),
2364 working_dir: None,
2365 }
2366 }
2367
2368 fn from_shell_cmd_with_dir(
2370 program: &str,
2371 args: &[String],
2372 env: &[(String, String)],
2373 working_dir: Option<PathBuf>,
2374 ) -> Self {
2375 NativeCommand::Shell {
2376 program: program.to_string(),
2377 args: args.to_vec(),
2378 env: env.to_vec(),
2379 working_dir,
2380 }
2381 }
2382}
2383
2384fn is_up_to_date_output(output: &str) -> bool {
2404 let lower = output.to_lowercase();
2405 lower.contains("up to date -- skipping")
2406 || lower.contains("up to date, skipping")
2407 || lower.contains("up to date -- reinstalling")
2408 || lower.contains("there is nothing to do")
2409 || lower.contains("already installed")
2410 || (lower.contains("warning:") && lower.contains("up to date"))
2411}
2412
2413fn safe_rocm_path(inherited: &str) -> String {
2425 const ROCM_BINS: &[&str] = &["/opt/rocm/bin", "/opt/rocm/hip/bin"];
2426 if !Path::new("/opt/rocm/bin").exists() {
2427 return inherited.to_string();
2428 }
2429 let mut parts: Vec<&str> = ROCM_BINS.to_vec();
2430 for entry in inherited.split(':') {
2431 if !entry.is_empty() && !ROCM_BINS.contains(&entry) {
2432 parts.push(entry);
2433 }
2434 }
2435 parts.join(":")
2436}
2437
2438fn execute_native_command(
2444 cmd: &NativeCommand,
2445 sudo_pw: Option<&str>,
2446 sender: &Sender<InstallerEvent>,
2447 component_name: &str,
2448) -> Result<()> {
2449 let (program, args, envs, working_dir) = match cmd {
2450 NativeCommand::Shell {
2451 program,
2452 args,
2453 env,
2454 working_dir,
2455 } => (
2456 program.clone(),
2457 args.clone(),
2458 env.clone(),
2459 working_dir.clone(),
2460 ),
2461 NativeCommand::Pip { program, args } => (program.clone(), args.clone(), vec![], None),
2462 NativeCommand::System { program, args } => (program.clone(), args.clone(), vec![], None),
2463 NativeCommand::Package { program, args } => (program.clone(), args.clone(), vec![], None),
2464 };
2465
2466 let invokes_pip = program == "uv"
2475 || program.ends_with("pip")
2476 || program.ends_with("pip3")
2477 || args.iter().any(|a| a == "pip");
2478 let is_install = args.iter().any(|a| a == "install");
2479 if invokes_pip && is_install {
2480 let mut skip_next = false;
2481 for a in &args {
2482 if skip_next {
2483 skip_next = false;
2487 continue;
2488 }
2489 if a == "-r" || a == "--requirement" || a == "-c" || a == "--constraint" {
2490 skip_next = true;
2491 continue;
2492 }
2493 if a.starts_with('-') || a == "install" || a == "." || a == "-e" || a.is_empty() {
2494 continue;
2495 }
2496 if crate::installers::common::nvidia_blocklist::is_nvidia_cuda_runtime_package(a) {
2497 let _ = sender.send(InstallerEvent::Log(
2498 format!("[native] BLOCKED nvidia/cuda package in {component_name}: {a}"),
2499 true,
2500 ));
2501 anyhow::bail!(
2502 "NVIDIA/CUDA contamination blocked: component '{component_name}' attempted \
2503 to pip install '{a}' — violates the No-CUDA hard-prime tenet"
2504 );
2505 }
2506 if a.starts_with("http")
2507 && crate::installers::common::nvidia_blocklist::is_cuda_wheel_url(a)
2508 {
2509 let _ = sender.send(InstallerEvent::Log(
2510 format!("[native] BLOCKED cuda wheel URL in {component_name}: {a}"),
2511 true,
2512 ));
2513 anyhow::bail!("NVIDIA/CUDA wheel URL blocked in {component_name}: {a}");
2514 }
2515 }
2516 }
2517
2518 let already_sudo = program == "sudo";
2521 let mut command = if let Some(_pw) = sudo_pw {
2522 if already_sudo {
2523 let mut c = Command::new(&program);
2525 c.args(&args);
2526 c
2527 } else {
2528 let mut c = Command::new("sudo");
2529 c.arg("-S").arg("-p").arg("").arg("--").arg(&program);
2530 c.args(&args);
2531 let _ = sender.send(InstallerEvent::Log(
2533 format!("[native] Running with sudo: {} {}", program, args.join(" ")),
2534 false,
2535 ));
2536 c
2537 }
2538 } else {
2539 let mut c = Command::new(&program);
2540 c.args(&args);
2541 c
2542 };
2543
2544 for (key, value) in &envs {
2546 command.env(key, value);
2547 }
2548
2549 command.env(
2558 "PATH",
2559 safe_rocm_path(&std::env::var("PATH").unwrap_or_default()),
2560 );
2561
2562 if let Some(ref dir) = working_dir {
2564 command.current_dir(dir);
2565 }
2566
2567 command
2569 .stdin(Stdio::piped())
2570 .stdout(Stdio::piped())
2571 .stderr(Stdio::piped());
2572
2573 let cmd_str = if args.is_empty() {
2575 program.clone()
2576 } else {
2577 format!("{} {}", program, args.join(" "))
2578 };
2579 let _ = sender.send(InstallerEvent::Log(
2580 format!("[native] $ {}", cmd_str),
2581 false,
2582 ));
2583
2584 let mut child = command.spawn().with_context(|| {
2585 let mut hint = String::new();
2586 if cmd_str.starts_with("sudo") {
2588 hint = " — ensure sudo is installed and you have the correct permissions".to_string();
2589 } else if cmd_str.starts_with("pip")
2590 || cmd_str.starts_with("uv")
2591 || cmd_str.starts_with("python")
2592 {
2593 hint = " — ensure Python/pip is installed and accessible".to_string();
2594 } else if cmd_str.starts_with("git") {
2595 hint = " — ensure git is installed".to_string();
2596 } else if cmd_str.starts_with("cmake") || cmd_str.starts_with("make") {
2597 hint = " — ensure build tools (cmake, make, gcc) are installed".to_string();
2598 } else if cmd_str.starts_with("rocminfo") || cmd_str.starts_with("rocm-smi") {
2599 hint = " — ensure ROCm is installed and /opt/rocm/bin is in PATH".to_string();
2600 }
2601 format!(
2602 "Failed to spawn command '{}' for {}.{}",
2603 cmd_str, component_name, hint
2604 )
2605 })?;
2606
2607 if let Some(pw) = sudo_pw {
2610 if let Some(mut stdin) = child.stdin.take() {
2611 let _ = stdin.write_all(pw.as_bytes());
2612 let _ = stdin.write_all(b"\n");
2613 let _ = stdin.flush();
2614 drop(stdin);
2615 }
2616 } else {
2617 drop(child.stdin.take());
2620 }
2621
2622 let sender_stdout = sender.clone();
2625 let stdout_stream = match child.stdout.take() {
2626 Some(s) => s,
2627 None => {
2628 let _ = sender.send(InstallerEvent::Log(
2629 format!(
2630 "[ERROR] {} — failed to capture stdout from spawned process",
2631 component_name
2632 ),
2633 false,
2634 ));
2635 bail!(
2636 "{} — failed to capture stdout from spawned process. \
2637 This usually indicates a system resource issue.",
2638 component_name
2639 );
2640 }
2641 };
2642 let stdout_captured: Arc<Mutex<String>> = Arc::new(Mutex::new(String::new()));
2643 let stdout_captured_clone = stdout_captured.clone();
2644 let stdout_handle = thread::spawn(move || {
2645 let reader = BufReader::new(stdout_stream);
2646 for line in reader.lines() {
2647 match line {
2648 Ok(l) => {
2649 let clean = strip_ansi_codes(&l);
2650 if !clean.trim().is_empty() {
2651 let _ = sender_stdout.send(InstallerEvent::Log(clean.clone(), false));
2652 if let Ok(mut buf) = stdout_captured_clone.lock() {
2653 buf.push_str(&clean);
2654 buf.push('\n');
2655 }
2656 }
2657 }
2658 Err(_) => break,
2659 }
2660 }
2661 });
2662
2663 let sender_stderr = sender.clone();
2664 let stderr_stream = match child.stderr.take() {
2665 Some(s) => s,
2666 None => {
2667 let _ = sender.send(InstallerEvent::Log(
2668 format!(
2669 "[ERROR] {} — failed to capture stderr from spawned process",
2670 component_name
2671 ),
2672 false,
2673 ));
2674 bail!(
2675 "{} — failed to capture stderr from spawned process. \
2676 This usually indicates a system resource issue.",
2677 component_name
2678 );
2679 }
2680 };
2681 let stderr_captured: Arc<Mutex<String>> = Arc::new(Mutex::new(String::new()));
2682 let stderr_captured_clone = stderr_captured.clone();
2683 let stderr_handle = thread::spawn(move || {
2684 let reader = BufReader::new(stderr_stream);
2685 for line in reader.lines() {
2686 match line {
2687 Ok(l) => {
2688 let clean = strip_ansi_codes(&l);
2689 if !clean.trim().is_empty() {
2690 let _ = sender_stderr.send(InstallerEvent::Log(clean.clone(), false));
2691 if let Ok(mut buf) = stderr_captured_clone.lock() {
2692 buf.push_str(&clean);
2693 buf.push('\n');
2694 }
2695 }
2696 }
2697 Err(_) => break,
2698 }
2699 }
2700 });
2701
2702 let status = child.wait().with_context(|| {
2704 format!(
2705 "{} — command process wait failed. \
2706 The process may have been killed or the system is under heavy load.",
2707 component_name
2708 )
2709 })?;
2710 let _ = stdout_handle.join();
2711 let _ = stderr_handle.join();
2712
2713 if !status.success() {
2714 let combined_output = {
2719 let stdout_buf = stdout_captured.lock().unwrap_or_else(|e| e.into_inner());
2720 let stderr_buf = stderr_captured.lock().unwrap_or_else(|e| e.into_inner());
2721 format!("{}{}", *stdout_buf, *stderr_buf)
2722 };
2723
2724 if is_up_to_date_output(&combined_output) {
2725 let _ = sender.send(InstallerEvent::Log(
2726 format!(
2727 "[native] {} — package already up to date, treating as success",
2728 component_name
2729 ),
2730 false,
2731 ));
2732 return Ok(());
2733 }
2734
2735 let code = status.code().unwrap_or(-1);
2736 let error_detail = match code {
2737 1 => "General error — check the command output above for details".to_string(),
2738 2 => "Misuse of shell builtins — check command arguments".to_string(),
2739 126 => "Command not executable — check file permissions".to_string(),
2740 127 => "Command not found — ensure the required tool is installed".to_string(),
2741 130 => "Process terminated by Ctrl+C (SIGINT)".to_string(),
2742 137 => "Process killed (SIGKILL) — possibly out of memory".to_string(),
2743 139 => "Segmentation fault — this is a bug in the invoked program".to_string(),
2744 n if n > 128 => format!(
2745 "Process terminated by signal {} — check system logs",
2746 n - 128
2747 ),
2748 _ => format!(
2749 "Command exited with code {} — check the output above for details",
2750 code
2751 ),
2752 };
2753 let _ = sender.send(InstallerEvent::Log(
2754 format!(
2755 "[ERROR] {} failed: {} (exit code {} at {})",
2756 component_name,
2757 error_detail,
2758 code,
2759 Local::now().format("%H:%M:%S")
2760 ),
2761 false,
2762 ));
2763 bail!(
2764 "{} failed: {} (exit code {})",
2765 component_name,
2766 error_detail,
2767 code
2768 );
2769 }
2770
2771 Ok(())
2772}
2773
2774fn build_common_env(ctx: &NativeInstallerContext) -> Vec<(String, String)> {
2783 let mut env = Vec::new();
2784
2785 env.push(("HOME".into(), ctx.user_home.to_string()));
2787 if let Ok(user) = std::env::var("USER") {
2788 env.push(("USER".into(), user.clone()));
2789 env.push(("LOGNAME".into(), user));
2790 }
2791 env.push(("MLSTACK_USER_HOME".into(), ctx.user_home.to_string()));
2792 env.push((
2793 "MLSTACK_BATCH_MODE".into(),
2794 if ctx.batch_mode { "1" } else { "0" }.into(),
2795 ));
2796 env.push((
2797 "MLSTACK_INSTALL_METHOD".into(),
2798 ctx.install_method.to_string(),
2799 ));
2800 env.push(("INSTALL_METHOD".into(), ctx.install_method.to_string()));
2801 env.push(("RUSTY_STACK".into(), "true".into()));
2802
2803 for key in &[
2805 "ROCM_VERSION",
2806 "ROCM_CHANNEL",
2807 "GPU_ARCH",
2808 "GPU_ARCHS",
2809 "PYTORCH_ROCM_ARCH",
2810 "HIP_VISIBLE_DEVICES",
2811 "CUDA_VISIBLE_DEVICES",
2812 "ROCR_VISIBLE_DEVICES",
2813 "PYTORCH_ROCM_DEVICE",
2814 "HSA_OVERRIDE_GFX_VERSION",
2815 "MLSTACK_PYTHON_BIN",
2816 "UV_PYTHON",
2817 ] {
2818 if let Some(value) = ctx.env_exports.get(*key) {
2819 let trimmed = value.trim();
2820 if !trimmed.is_empty() {
2821 env.push((key.to_string(), trimmed.to_string()));
2822 }
2823 }
2824 }
2825
2826 env.push(("ROCM_HOME".into(), "/opt/rocm".into()));
2828 env.push(("ROCM_PATH".into(), "/opt/rocm".into()));
2829 env.push(("HIP_PATH".into(), "/opt/rocm".into()));
2830 env.push(("HIP_ROOT_DIR".into(), "/opt/rocm".into()));
2831 env.push(("ROCM_ROOT".into(), "/opt/rocm".into()));
2832 env.push(("HIPCC_BIN_DIR".into(), "/opt/rocm/bin".into()));
2833
2834 env.push(("PIP_BREAK_SYSTEM_PACKAGES".into(), "1".into()));
2836 env.push(("PIP_ROOT_USER_ACTION".into(), "ignore".into()));
2837 env.push(("UV_PIP_BREAK_SYSTEM_PACKAGES".into(), "1".into()));
2838 env.push(("UV_SYSTEM_PYTHON".into(), "1".into()));
2839 env.push(("PYTHONUNBUFFERED".into(), "1".into()));
2840
2841 env
2842}
2843
2844fn is_existing_git_repo(target_dir: &str) -> bool {
2854 let path = Path::new(target_dir);
2855 match std::fs::symlink_metadata(target_dir) {
2858 Ok(meta) if meta.is_dir() => {
2859 path.join(".git")
2861 .symlink_metadata()
2862 .map(|m| m.is_dir() || m.is_file())
2863 .unwrap_or(false)
2864 }
2865 _ => false,
2866 }
2867}
2868
2869fn fix_directory_ownership(dir: &str, sudo_pw: Option<&str>, sender: &Sender<InstallerEvent>) {
2886 let stat_output = std::process::Command::new("stat")
2888 .args(["-c", "%U", dir])
2889 .output();
2890
2891 let needs_fix = match stat_output {
2892 Ok(out) if out.status.success() => {
2893 let owner = String::from_utf8_lossy(&out.stdout).trim().to_string();
2894 owner == "root"
2895 }
2896 _ => true, };
2898
2899 if !needs_fix {
2900 return;
2901 }
2902
2903 let real_user = std::env::var("SUDO_USER")
2904 .or_else(|_| std::env::var("USER"))
2905 .unwrap_or_else(|_| "scooter".to_string());
2906
2907 let _ = sender.send(InstallerEvent::Log(
2908 format!("[native] Fixing ownership of {} -> {}", dir, real_user),
2909 false,
2910 ));
2911
2912 if sudo_pw.is_some() {
2913 let chown_cmd = NativeCommand::from_shell_cmd(
2914 "chown",
2915 &[
2916 "-R".to_string(),
2917 format!("{}:{}", real_user, real_user),
2918 dir.to_string(),
2919 ],
2920 &[],
2921 );
2922 let _ = execute_native_command(&chown_cmd, sudo_pw, sender, "ownership-fix");
2923 }
2924 }
2926
2927fn git_clone_or_pull(
2928 repo_url: &str,
2929 target_dir: &str,
2930 extra_clone_args: &[&str],
2931 sudo_pw: Option<&str>,
2932 sender: &Sender<InstallerEvent>,
2933 component_name: &str,
2934) -> Result<()> {
2935 if is_existing_git_repo(target_dir) {
2936 let _ = sender.send(InstallerEvent::Log(
2937 format!(
2938 "[native] {} — directory '{}' already exists, pulling updates instead of cloning",
2939 component_name, target_dir
2940 ),
2941 false,
2942 ));
2943
2944 let pull_cmd = NativeCommand::from_shell_cmd(
2946 "git",
2947 &[
2948 "-C".to_string(),
2949 target_dir.to_string(),
2950 "pull".to_string(),
2951 "--ff-only".to_string(),
2952 ],
2953 &[],
2954 );
2955 let pull_result = execute_native_command(&pull_cmd, sudo_pw, sender, component_name);
2956
2957 if pull_result.is_ok() {
2958 return Ok(());
2959 }
2960
2961 let _ = sender.send(InstallerEvent::Log(
2963 format!(
2964 "[native] {} — git pull --ff-only failed, trying fetch + reset --hard origin/HEAD",
2965 component_name
2966 ),
2967 false,
2968 ));
2969
2970 let fetch_cmd = NativeCommand::from_shell_cmd(
2971 "git",
2972 &[
2973 "-C".to_string(),
2974 target_dir.to_string(),
2975 "fetch".to_string(),
2976 "--all".to_string(),
2977 ],
2978 &[],
2979 );
2980 let fetch_ok = execute_native_command(&fetch_cmd, sudo_pw, sender, component_name).is_ok();
2981
2982 let reset_ok = if fetch_ok {
2983 let reset_cmd = NativeCommand::from_shell_cmd(
2984 "git",
2985 &[
2986 "-C".to_string(),
2987 target_dir.to_string(),
2988 "reset".to_string(),
2989 "--hard".to_string(),
2990 "origin/HEAD".to_string(),
2991 ],
2992 &[],
2993 );
2994 execute_native_command(&reset_cmd, sudo_pw, sender, component_name).is_ok()
2995 } else {
2996 false
2997 };
2998
2999 if fetch_ok && reset_ok {
3000 return Ok(());
3001 }
3002
3003 let _ = sender.send(InstallerEvent::Log(
3007 format!(
3008 "[native] {} — existing .git is corrupt/unusable (pull+fetch+reset failed); removing and re-cloning",
3009 component_name
3010 ),
3011 false,
3012 ));
3013 if std::fs::remove_dir_all(target_dir).is_err() {
3014 let rm_cmd = NativeCommand::from_shell_cmd(
3015 "rm",
3016 &["-rf".to_string(), target_dir.to_string()],
3017 &[],
3018 );
3019 let _ = execute_native_command(&rm_cmd, sudo_pw, sender, component_name);
3020 }
3021 if Path::new(target_dir).join(".git").exists() {
3028 bail!(
3029 "{} — could not remove corrupt repo at '{}' (.git still present); \
3030 aborting re-clone to avoid infinite recursion",
3031 component_name,
3032 target_dir
3033 );
3034 }
3035 git_clone_or_pull(
3036 repo_url,
3037 target_dir,
3038 extra_clone_args,
3039 sudo_pw,
3040 sender,
3041 component_name,
3042 )
3043 } else {
3044 let target_path = Path::new(target_dir);
3049 if target_path.is_dir() && !target_path.join(".git").exists() {
3050 let _ = sender.send(InstallerEvent::Log(
3051 format!(
3052 "[native] {} — directory '{}' exists but is not a git repo; removing incomplete clone",
3053 component_name, target_dir
3054 ),
3055 false,
3056 ));
3057 if std::fs::remove_dir_all(target_dir).is_err() {
3059 let rm_cmd = NativeCommand::from_shell_cmd(
3061 "rm",
3062 &["-rf".to_string(), target_dir.to_string()],
3063 &[],
3064 );
3065 let rm_result = execute_native_command(&rm_cmd, sudo_pw, sender, component_name);
3066 if let Err(e) = rm_result {
3067 let _ = sender.send(InstallerEvent::Log(
3068 format!(
3069 "[native] [WARN] {} — could not remove stale dir: {}. Clone may fail.",
3070 component_name, e
3071 ),
3072 false,
3073 ));
3074 }
3075 }
3076 }
3077
3078 let _ = sender.send(InstallerEvent::Log(
3079 format!(
3080 "[native] {} — cloning '{}' into '{}'",
3081 component_name, repo_url, target_dir
3082 ),
3083 false,
3084 ));
3085
3086 let mut args: Vec<String> = vec!["clone".to_string()];
3087 for extra in extra_clone_args {
3088 args.push(extra.to_string());
3089 }
3090 args.push(repo_url.to_string());
3091 args.push(target_dir.to_string());
3092
3093 let clone_cmd = NativeCommand::from_shell_cmd("git", &args, &[]);
3094 execute_native_command(&clone_cmd, sudo_pw, sender, component_name)
3095 }
3096}
3097
3098fn filter_cuda_requirements(req_path: &str) -> std::io::Result<String> {
3129 crate::installers::common::nvidia_blocklist::filter_requirements_file(req_path)
3130}
3131
3132fn is_force_reinstall() -> bool {
3137 std::env::var("MLSTACK_FORCE_REINSTALL")
3138 .or_else(|_| std::env::var("FORCE"))
3139 .map(|v| {
3140 matches!(
3141 v.trim().to_ascii_lowercase().as_str(),
3142 "1" | "true" | "yes" | "on"
3143 )
3144 })
3145 .unwrap_or(false)
3146}
3147
3148fn triton_installed_satisfies_target(python_bin: &str, target: Option<&str>) -> bool {
3163 let target = match target {
3164 Some(t) => t.trim(),
3165 None => return true, };
3167 if target.is_empty() || target == "latest" || target == "installed" {
3169 return true;
3170 }
3171 if is_force_reinstall() {
3173 return false;
3174 }
3175 let installed = std::process::Command::new(python_bin)
3177 .arg("-c")
3178 .arg(
3179 "try:\n import triton, importlib.metadata as im\n \
3180 print(getattr(triton, '__version__', None) or im.version('triton') \
3181 or im.version('triton-rocm') or im.version('pytorch-triton'))\n\
3182 except Exception:\n pass",
3183 )
3184 .output()
3185 .ok()
3186 .and_then(|o| {
3187 if o.status.success() {
3188 String::from_utf8(o.stdout).ok()
3189 } else {
3190 None
3191 }
3192 })
3193 .map(|s| s.trim().to_string())
3194 .filter(|s| !s.is_empty());
3195 match installed {
3196 Some(ver) => version_satisfies(&ver, target),
3197 None => false, }
3199}
3200
3201fn version_satisfies(installed: &str, target: &str) -> bool {
3207 fn parts(v: &str) -> Vec<u64> {
3208 v.split('-')
3209 .next()
3210 .unwrap_or("")
3211 .split('.')
3212 .filter_map(|s| s.parse::<u64>().ok())
3213 .collect()
3214 }
3215 let a = parts(installed.trim());
3216 let b = parts(target.trim());
3217 if a.is_empty() || b.is_empty() {
3218 return false;
3220 }
3221 for (ai, bi) in a.iter().zip(b.iter()) {
3222 match ai.cmp(bi) {
3223 std::cmp::Ordering::Greater => return true,
3224 std::cmp::Ordering::Less => return false,
3225 std::cmp::Ordering::Equal => continue,
3226 }
3227 }
3228 a.len() >= b.len()
3231}
3232
3233fn build_pip_uninstall(packages: &[&str]) -> NativeCommand {
3238 let use_uv = crate::installers::common::utils::command_exists("uv");
3239 if use_uv {
3240 let mut args = vec!["pip".to_string(), "uninstall".to_string()];
3241 for pkg in packages {
3242 args.push(pkg.to_string());
3243 }
3244 NativeCommand::Pip {
3245 program: "uv".to_string(),
3246 args,
3247 }
3248 } else {
3249 let mut args = vec!["uninstall".to_string(), "-y".to_string()];
3250 for pkg in packages {
3251 args.push(pkg.to_string());
3252 }
3253 NativeCommand::Pip {
3254 program: "pip3".to_string(),
3255 args,
3256 }
3257 }
3258}
3259
3260fn purge_component_packages(
3266 packages: &[&str],
3267 sender: &Sender<InstallerEvent>,
3268 component_name: &str,
3269) {
3270 let _ = sender.send(InstallerEvent::Log(
3271 format!(
3272 "[native] {} — purging packages: {}",
3273 component_name,
3274 packages.join(", ")
3275 ),
3276 false,
3277 ));
3278
3279 let uninstall_cmd = build_pip_uninstall(packages);
3281 let _ = execute_native_command(&uninstall_cmd, None, sender, component_name);
3282}
3283
3284fn purge_pytorch(sender: &Sender<InstallerEvent>, component_name: &str) {
3286 purge_component_packages(
3287 &[
3288 "torch",
3289 "torchvision",
3290 "torchaudio",
3291 "pytorch-triton-rocm",
3292 "pytorch-triton",
3293 ],
3294 sender,
3295 component_name,
3296 );
3297}
3298
3299fn purge_triton(sender: &Sender<InstallerEvent>, component_name: &str) {
3301 purge_component_packages(
3302 &["triton", "triton-rocm", "pytorch-triton-rocm"],
3303 sender,
3304 component_name,
3305 );
3306}
3307
3308fn purge_aiter(sender: &Sender<InstallerEvent>, component_name: &str) {
3310 purge_component_packages(&["aiter", "amd-aiter"], sender, component_name);
3311}
3312
3313fn purge_flash_attn(sender: &Sender<InstallerEvent>, component_name: &str) {
3315 purge_component_packages(
3316 &["flash-attn", "flash_attn", "flash_attention_amd"],
3317 sender,
3318 component_name,
3319 );
3320}
3321
3322fn purge_vllm(sender: &Sender<InstallerEvent>, component_name: &str) {
3324 purge_component_packages(&["vllm"], sender, component_name);
3325}
3326
3327fn unseal_core_requested() -> bool {
3330 std::env::var("MLSTACK_UNSEAL_CORE")
3331 .map(|v| {
3332 matches!(
3333 v.trim().to_ascii_lowercase().as_str(),
3334 "1" | "true" | "yes" | "on"
3335 )
3336 })
3337 .unwrap_or(false)
3338}
3339
3340fn registry_gate(component: &Component, sender: &Sender<InstallerEvent>) -> Result<bool> {
3349 let registry = crate::core::registry::InstalledComponentRegistry::load();
3350 if !registry.is_sealed(&component.id) {
3351 return Ok(true);
3352 }
3353 if is_force_reinstall() && !unseal_core_requested() {
3354 let _ = sender.send(InstallerEvent::Log(
3361 format!(
3362 "[native] {} is a sealed CORE component — force ignored, reusing the installed copy. \
3363 Set MLSTACK_UNSEAL_CORE=1 to intentionally override it.",
3364 component.name
3365 ),
3366 false,
3367 ));
3368 return Ok(false);
3369 }
3370 if !is_force_reinstall() {
3371 let _ = sender.send(InstallerEvent::Log(
3372 format!(
3373 "[native] {} already installed (sealed core) — reusing, NOT reinstalling",
3374 component.name
3375 ),
3376 false,
3377 ));
3378 return Ok(false);
3379 }
3380 Ok(true)
3381}
3382
3383fn registry_record(component: &Component) {
3389 use crate::core::registry::{InstalledComponent, InstalledComponentRegistry};
3390 let mut registry = InstalledComponentRegistry::load();
3391 let record = InstalledComponent {
3392 id: component.id.clone(),
3393 version: component_version(&component.id),
3394 source_index: component_source_index(&component.id),
3395 location: install_target_location(),
3396 sealed: crate::core::registry::should_seal_component(&component.id),
3397 installed_at: Some(chrono::Utc::now().to_rfc3339()),
3398 pip_packages: component_pip_packages(&component.id),
3399 };
3400 registry.mark_installed(record);
3401 if let Err(e) = registry.save() {
3402 eprintln!("[WARN] Failed to save installed component registry: {e}");
3406 }
3407}
3408
3409fn component_pip_packages(id: &str) -> Vec<String> {
3413 match id {
3414 "pytorch" => vec!["torch", "torchvision", "torchaudio", "pytorch-triton-rocm"],
3415 "triton" => vec!["triton", "pytorch-triton-rocm"],
3416 "vllm" | "vllm-multi" => vec!["vllm"],
3417 "deepspeed" => vec!["deepspeed", "einops"],
3418 "megatron" => vec!["megatron-core"],
3419 "bitsandbytes" => vec!["bitsandbytes"],
3420 "flash-attn" | "flash-attn-triton" | "flash-attn-ck" | "flash_attention" => {
3421 vec!["flash-attn", "flash_attn"]
3422 }
3423 "aiter" => vec!["aiter"],
3424 "onnx" | "onnxruntime" => vec!["onnxruntime", "onnxruntime-rocm"],
3425 "migraphx" => vec!["migraphx"],
3426 "mpi4py" => vec!["mpi4py"],
3427 "wandb" => vec!["wandb"],
3428 "fastvideo" => vec!["fastvideo"],
3429 _ => Vec::new(),
3430 }
3431 .into_iter()
3432 .map(String::from)
3433 .collect()
3434}
3435
3436fn install_target_location() -> Option<String> {
3440 if let Ok(name) = std::env::var("MLSTACK_ENV_NAME") {
3441 let name = name.trim();
3442 if !name.is_empty() {
3443 return Some(
3444 crate::platform::environment::mlstack_env_dir(name)
3445 .to_string_lossy()
3446 .into_owned(),
3447 );
3448 }
3449 }
3450 let global = crate::platform::environment::mlstack_global_dir();
3451 if global.exists() {
3452 return Some(global.to_string_lossy().into_owned());
3453 }
3454 std::env::var("MLSTACK_PYTHON_BIN")
3455 .ok()
3456 .map(|s| s.trim().to_string())
3457 .filter(|s| !s.is_empty())
3458}
3459
3460fn component_source_index(id: &str) -> Option<String> {
3466 match id {
3467 "pytorch" => Some("https://download.pytorch.org/whl/rocm".to_string()),
3468 "triton" => Some("triton-rocm (bundled with PyTorch)".to_string()),
3469 "rocm" => Some("/opt/rocm".to_string()),
3470 _ => None,
3471 }
3472}
3473
3474fn component_version(id: &str) -> Option<String> {
3478 match id {
3479 "rocm" => {
3480 let v = detect_rocm_version();
3481 if v.is_empty() {
3482 None
3483 } else {
3484 Some(v)
3485 }
3486 }
3487 _ => None,
3488 }
3489}
3490
3491fn ensure_user_in_groups(
3508 groups: &[&str],
3509 sudo_pw: Option<&str>,
3510 sender: &Sender<InstallerEvent>,
3511 component_name: &str,
3512) {
3513 use std::process::Command;
3514
3515 let user = std::env::var("SUDO_USER")
3518 .ok()
3519 .map(|u| u.trim().to_string())
3520 .filter(|u| !u.is_empty())
3521 .or_else(|| {
3522 std::env::var("USER")
3523 .ok()
3524 .map(|u| u.trim().to_string())
3525 .filter(|u| !u.is_empty())
3526 });
3527 let user = match user {
3528 Some(u) => u,
3529 None => {
3530 let _ = sender.send(InstallerEvent::Log(
3531 format!(
3532 "[native] {component_name} — could not determine target user; skipping \
3533 group check. Add yourself manually: sudo usermod -aG {} $USER",
3534 groups.join(",")
3535 ),
3536 false,
3537 ));
3538 return;
3539 }
3540 };
3541
3542 let current_groups: Vec<String> = Command::new("id")
3545 .args(["-nG", &user])
3546 .output()
3547 .ok()
3548 .filter(|o| o.status.success())
3549 .map(|o| {
3550 String::from_utf8_lossy(&o.stdout)
3551 .split_whitespace()
3552 .map(|s| s.to_string())
3553 .collect()
3554 })
3555 .unwrap_or_default();
3556
3557 let mut to_add: Vec<&str> = Vec::new();
3559 for g in groups {
3560 let exists = Command::new("getent")
3561 .args(["group", g])
3562 .output()
3563 .map(|o| o.status.success())
3564 .unwrap_or(false);
3565 if !exists {
3566 let _ = sender.send(InstallerEvent::Log(
3567 format!(
3568 "[native] {component_name} — group '{g}' not present on this system; skipping"
3569 ),
3570 false,
3571 ));
3572 continue;
3573 }
3574 if !current_groups.iter().any(|c| c == g) {
3575 to_add.push(*g);
3576 }
3577 }
3578
3579 if to_add.is_empty() {
3580 let _ = sender.send(InstallerEvent::Log(
3581 format!(
3582 "[native] {component_name} — user '{user}' already in all required group(s); nothing to add"
3583 ),
3584 false,
3585 ));
3586 return;
3587 }
3588
3589 let csv = to_add.join(",");
3590
3591 let askpass = sudo_pw.filter(|p| !p.is_empty()).and_then(|p| {
3593 match crate::installers::common::askpass::Askpass::new(p) {
3594 Ok(g) => Some(g),
3595 Err(e) => {
3596 let _ = sender.send(InstallerEvent::Log(
3597 format!(
3598 "[native] {component_name} — askpass setup failed ({e}); cannot add \
3599 groups non-interactively"
3600 ),
3601 true,
3602 ));
3603 None
3604 }
3605 }
3606 });
3607
3608 let added = if let Some(ap) = askpass {
3609 let args = vec![
3612 "-A".to_string(),
3613 "usermod".to_string(),
3614 "-aG".to_string(),
3615 csv.clone(),
3616 user.clone(),
3617 ];
3618 let envs = vec![("SUDO_ASKPASS".to_string(), ap.path().to_string())];
3619 let cmd = NativeCommand::from_shell_cmd("sudo", &args, &envs);
3620 execute_native_command(&cmd, None, sender, component_name).is_ok()
3621 } else {
3622 let _ = sender.send(InstallerEvent::Log(
3623 format!(
3624 "[native] {component_name} — no sudo password available; please run: \
3625 sudo usermod -aG {csv} {user}"
3626 ),
3627 true,
3628 ));
3629 return;
3630 };
3631
3632 if added {
3633 let _ = sender.send(InstallerEvent::Log(
3634 format!(
3635 "[native] {component_name} — added user '{user}' to group(s): {csv}. \
3636 A re-login (or reboot) is required for the new groups to take effect."
3637 ),
3638 false,
3639 ));
3640 }
3641}
3642
3643fn run_rccl_remediation(ctx: &NativeInstallerContext, component_name: &str) -> Result<()> {
3645 let sender = ctx.sender;
3646 let rccl_remediation = crate::installers::components::rccl::RcclInstaller::new(
3647 ctx.user_home,
3648 resolve_python_bin(),
3649 ctx.env_exports
3650 .get("ROCM_PATH")
3651 .cloned()
3652 .unwrap_or_else(|| "/opt/rocm".to_string()),
3653 ctx.env_exports
3654 .get("ROCM_VERSION")
3655 .cloned()
3656 .unwrap_or_else(|| "7.2.0".to_string()),
3657 ctx.env_exports
3658 .get("ROCM_CHANNEL")
3659 .cloned()
3660 .unwrap_or_else(|| "latest".to_string()),
3661 ctx.env_exports
3662 .get("ROCR_VISIBLE_DEVICES")
3663 .cloned()
3664 .or_else(|| ctx.env_exports.get("HIP_VISIBLE_DEVICES").cloned()),
3665 );
3666
3667 let _ = sender.send(InstallerEvent::Log(
3668 "[native] Probing RCCL multi-GPU collectives...".into(),
3669 false,
3670 ));
3671 if let Some(build) = rccl_remediation.prepare_remediation()? {
3672 let _ = sender.send(InstallerEvent::Log(
3673 "[native] Known RCCL collective failure detected; building sealed overlay...".into(),
3674 false,
3675 ));
3676 execute_native_command(
3677 &NativeCommand::from_shell_cmd(&build.program, &build.args, &build.env),
3678 None,
3679 sender,
3680 component_name,
3681 )?;
3682 rccl_remediation.activate_and_verify()?;
3683 let _ = sender.send(InstallerEvent::Log(
3684 "[native] RCCL sealed overlay activated and validated on both GPUs".into(),
3685 false,
3686 ));
3687 } else {
3688 let _ = sender.send(InstallerEvent::Log(
3689 "[native] RCCL probe healthy or existing sealed overlay revalidated; no repair needed"
3690 .into(),
3691 false,
3692 ));
3693 }
3694 Ok(())
3695}
3696
3697fn run_native_installer(component: &Component, ctx: &NativeInstallerContext) -> Result<()> {
3705 let sender = ctx.sender;
3706 let _ = sender.send(InstallerEvent::Log(
3707 format!("[native] Installing {} via Rust module...", component.name),
3708 false,
3709 ));
3710
3711 let _ = sender.send(InstallerEvent::Progress {
3712 component_id: component.id.clone(),
3713 progress: 0.1,
3714 message: format!("Starting {} installation", component.name),
3715 });
3716
3717 let component_needs_sudo = component.needs_sudo && needs_sudo();
3719 let cached_sudo =
3720 component_needs_sudo && ctx.sudo_password.is_none() && can_sudo_non_interactive();
3721 let sudo_pw: Option<&str> = if component_needs_sudo {
3722 ctx.sudo_password
3723 .as_deref()
3724 .or(if cached_sudo { Some("") } else { None })
3725 } else {
3726 None
3727 };
3728
3729 if component_needs_sudo && sudo_pw.is_none() {
3731 bail!("{} requires sudo but no password provided", component.name);
3732 }
3733
3734 let common_env = build_common_env(ctx);
3736
3737 match component.id.as_str() {
3741 "permanent-env" => {
3743 use crate::installers::components::permanent_env::PermanentEnvInstaller;
3744 match ensure_mlstack_env(ctx.user_home, ctx.install_method, is_force_reinstall()) {
3755 Ok(EnvUpdate::Created) | Ok(EnvUpdate::Updated) => {
3756 let _ = sender.send(InstallerEvent::Log(
3757 "[native] Permanent ROCm Env — regenerated ~/.mlstack_env from current state".into(),
3758 false,
3759 ));
3760 }
3761 Ok(EnvUpdate::Unchanged) => {
3762 let _ = sender.send(InstallerEvent::Log(
3763 "[native] Permanent ROCm Env — ~/.mlstack_env already up to date".into(),
3764 false,
3765 ));
3766 }
3767 Err(e) => {
3768 let _ = sender.send(InstallerEvent::Log(
3769 format!(
3770 "[native] Permanent ROCm Env — could not regenerate ~/.mlstack_env: {e}"
3771 ),
3772 true,
3773 ));
3774 }
3775 }
3776 crate::platform::environment::write_activate_snippets()?;
3778 crate::platform::environment::offer_global_prepend(true)?;
3779 let inst = PermanentEnvInstaller::default();
3780 let cmd = inst.build_mkdir_triton_dirs_command();
3781 execute_native_command(
3782 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
3783 sudo_pw,
3784 sender,
3785 &component.name,
3786 )?;
3787 }
3788
3789 "rocm" => {
3791 use crate::installers::components::rocm::{RocmConfig, RocmInstaller};
3792 let rocm_force = is_force_reinstall();
3793 let inst = RocmInstaller::new(RocmConfig {
3794 force_reinstall: rocm_force,
3795 ..RocmConfig::default()
3796 });
3797 let distro = crate::installers::common::DistroFacade::detect();
3798 let is_arch = distro.is_arch_family();
3799 if is_arch {
3800 let all_pkgs = inst.pacman_rocm_packages();
3806 let force_reinstall = is_force_reinstall();
3807 let pkgs_to_install = if force_reinstall {
3808 all_pkgs.clone()
3809 } else {
3810 let need_install = RocmInstaller::filter_already_installed_pacman(&all_pkgs);
3811 if need_install.is_empty() {
3812 let _ = sender.send(InstallerEvent::Log(
3813 format!(
3814 "[native] {} — all ROCm packages already installed, skipping",
3815 component.name
3816 ),
3817 false,
3818 ));
3819 return Ok(());
3820 }
3821 need_install
3822 };
3823
3824 let _ = sender.send(InstallerEvent::Log(
3825 format!(
3826 "[native] {} — installing {} of {} ROCm packages via yay \
3827 (user-space, non-interactive): {}",
3828 component.name,
3829 pkgs_to_install.len(),
3830 all_pkgs.len(),
3831 pkgs_to_install.join(", ")
3832 ),
3833 false,
3834 ));
3835
3836 let mut args = vec!["-S".to_string()];
3846 if !force_reinstall {
3847 args.push("--needed".to_string());
3848 }
3849 args.push("--noconfirm".to_string());
3850
3851 let askpass = sudo_pw.filter(|p| !p.is_empty()).and_then(|p| {
3853 match crate::installers::common::askpass::Askpass::new(p) {
3854 Ok(g) => Some(g),
3855 Err(e) => {
3856 let _ = sender.send(InstallerEvent::Log(
3857 format!(
3858 "[native] {} — askpass setup failed ({e}); yay's \
3859 privileged steps may not authenticate",
3860 component.name
3861 ),
3862 true,
3863 ));
3864 None
3865 }
3866 }
3867 });
3868 let mut envs: Vec<(String, String)> = Vec::new();
3869 if let Some(ref ap) = askpass {
3870 envs.push(("SUDO_ASKPASS".to_string(), ap.path().to_string()));
3871 args.push("--sudoflags=-A".to_string());
3872 } else {
3873 let _ = sender.send(InstallerEvent::Log(
3874 format!(
3875 "[native] {} — no sudo password supplied; yay will only \
3876 authenticate if sudo is cached / passwordless",
3877 component.name
3878 ),
3879 false,
3880 ));
3881 }
3882 args.extend(pkgs_to_install);
3883
3884 let yay_cmd = NativeCommand::Shell {
3885 program: "yay".to_string(),
3886 args,
3887 env: envs,
3888 working_dir: None,
3889 };
3890 execute_native_command(&yay_cmd, None, sender, &component.name)?;
3893 } else {
3894 let commands = if distro.uses_apt() {
3895 inst.apt_install_commands()
3896 } else if distro.uses_dnf() || distro.uses_yum() {
3897 let rhel_ver = distro.version().split('.').next().unwrap_or("9");
3898 inst.dnf_install_commands(rhel_ver)
3899 } else if distro.uses_zypper() {
3900 inst.zypper_install_commands()
3901 } else {
3902 bail!(
3903 "Unsupported distribution for ROCm installation: {} ({})",
3904 distro.name(),
3905 distro.id()
3906 );
3907 };
3908 for pkg_cmd in &commands {
3909 let native_cmd = NativeCommand::Package {
3910 program: pkg_cmd.program.clone(),
3911 args: pkg_cmd.args.clone(),
3912 };
3913 execute_native_command(&native_cmd, sudo_pw, sender, &component.name)?;
3916 }
3917 }
3918
3919 if is_force_reinstall() {
3921 let marker_path = std::path::Path::new("/tmp/mlstack-reboot-required");
3922 let _ = std::fs::write(
3923 marker_path,
3924 "ROCm packages were reinstalled. Reboot required for kernel module changes.\n",
3925 );
3926 let _ = sender.send(InstallerEvent::Log(
3927 "[native] ROCm — reboot recommended for kernel module changes".to_string(),
3928 false,
3929 ));
3930 }
3931
3932 ensure_user_in_groups(&["render", "video"], sudo_pw, sender, &component.name);
3937 }
3938
3939 "pytorch" => {
3941 use crate::installers::components::pytorch::{PyTorchConfig, PyTorchInstaller};
3942 let force = is_force_reinstall();
3943 let python_bin = resolve_python_bin();
3944 let config = PyTorchConfig {
3945 python_bin: python_bin.clone(),
3946 force_reinstall: force,
3947 ..PyTorchConfig::default()
3948 };
3949 let inst = PyTorchInstaller::new(config);
3950 let use_uv = crate::installers::common::utils::command_exists("uv");
3951 let rocm_mm = ctx
3952 .env_exports
3953 .get("ROCM_VERSION")
3954 .map(|v| {
3955 let parts: Vec<&str> = v.split('.').collect();
3956 if parts.len() >= 2 {
3957 format!("{}.{}", parts[0], parts[1])
3958 } else {
3959 v.clone()
3960 }
3961 })
3962 .unwrap_or_else(|| "7.2".to_string());
3963
3964 if force {
3966 purge_pytorch(sender, &component.name);
3967 }
3968
3969 let install_cmd =
3978 inst.build_install_command(&rocm_mm, use_uv, ctx.target_version.as_deref());
3979 execute_native_command(
3980 &NativeCommand::Pip {
3981 program: install_cmd.program.clone(),
3982 args: install_cmd.args.clone(),
3983 },
3984 None,
3985 sender,
3986 &component.name,
3987 )?;
3988
3989 let deps_cmd = inst.build_common_deps_command(use_uv);
3994 let _ = common_env.clone(); execute_native_command(
3996 &NativeCommand::Pip {
3997 program: deps_cmd.program.clone(),
3998 args: deps_cmd.args.clone(),
3999 },
4000 None, sender,
4002 &component.name,
4003 )?;
4004
4005 let managed_python = resolve_python_bin();
4010 let managed_install =
4011 crate::installers::common::managed_python_install_command(managed_python.clone());
4012 execute_native_command(
4013 &NativeCommand::Pip {
4014 program: managed_install.program,
4015 args: managed_install.args,
4016 },
4017 None,
4018 sender,
4019 &component.name,
4020 )?;
4021
4022 let managed_verify =
4023 crate::installers::common::managed_python_verify_command(managed_python);
4024 execute_native_command(
4025 &NativeCommand::from_shell_cmd(&managed_verify.program, &managed_verify.args, &[]),
4026 None,
4027 sender,
4028 &component.name,
4029 )?;
4030 }
4031
4032 "triton" => {
4034 use crate::installers::components::triton::{TritonConfig, TritonInstaller};
4035
4036 let python_bin = resolve_python_bin();
4053 let triton_present = std::process::Command::new(&python_bin)
4054 .arg("-c")
4055 .arg("import triton")
4056 .status()
4057 .map(|s| s.success())
4058 .unwrap_or(false);
4059 let triton_satisfies_target = triton_present
4060 && triton_installed_satisfies_target(&python_bin, ctx.target_version.as_deref());
4061 if triton_satisfies_target {
4062 let _ = sender.send(InstallerEvent::Log(
4063 "[native] Triton — already importable (triton-rocm bundled with PyTorch) \
4064 and satisfies the requested target; skipping the redundant from-source \
4065 ROCm/triton build"
4066 .into(),
4067 false,
4068 ));
4069 return Ok(());
4070 }
4071
4072 let triton_force = is_force_reinstall();
4073 let inst = TritonInstaller::new(TritonConfig {
4074 python_bin: resolve_python_bin(),
4075 ..Default::default()
4076 });
4077 let target_dir = format!("{}/.mlstack/triton", ctx.user_home);
4078
4079 if triton_force {
4081 purge_triton(sender, &component.name);
4082 }
4083
4084 git_clone_or_pull(
4088 inst.git_clone_url(),
4089 &target_dir,
4090 &[],
4091 sudo_pw,
4092 sender,
4093 &component.name,
4094 )?;
4095
4096 fix_directory_ownership(&target_dir, sudo_pw, sender);
4099
4100 let safe_dir_cmd = NativeCommand::from_shell_cmd(
4102 "git",
4103 &[
4104 "config".to_string(),
4105 "--global".to_string(),
4106 "--add".to_string(),
4107 "safe.directory".to_string(),
4108 target_dir.clone(),
4109 ],
4110 &[],
4111 );
4112 let _ = execute_native_command(&safe_dir_cmd, None, sender, &component.name);
4113
4114 let branch = inst.select_branch().to_string();
4116 let fetch_branch_cmd = NativeCommand::from_shell_cmd(
4117 "git",
4118 &[
4119 "-C".to_string(),
4120 target_dir.clone(),
4121 "fetch".to_string(),
4122 "origin".to_string(),
4123 branch.clone(),
4124 ],
4125 &[],
4126 );
4127 let _ = execute_native_command(&fetch_branch_cmd, None, sender, &component.name);
4128
4129 let checkout_cmd = inst.build_git_checkout_command();
4131 let mut checkout_args = vec!["-C".to_string(), target_dir.clone()];
4132 checkout_args.extend(checkout_cmd.args.iter().cloned());
4133 let checkout_result = execute_native_command(
4134 &NativeCommand::from_shell_cmd(
4135 &checkout_cmd.program,
4136 &checkout_args,
4137 &checkout_cmd.env,
4138 ),
4139 None,
4140 sender,
4141 &component.name,
4142 );
4143 if let Err(err) = checkout_result {
4144 let _ = sender.send(InstallerEvent::Log(
4145 format!(
4146 "[native] [WARN] {} git checkout failed (non-fatal): {}",
4147 component.name, err
4148 ),
4149 false,
4150 ));
4151 }
4152
4153 let cmd = inst.build_prerequisites_command();
4155 execute_native_command(
4156 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4157 None,
4158 sender,
4159 &component.name,
4160 )?;
4161
4162 let cmd = inst.build_pip_install_command(&target_dir);
4164 execute_native_command(
4165 &NativeCommand::from_shell_cmd_with_dir(
4166 &cmd.program,
4167 &cmd.args,
4168 &cmd.env,
4169 cmd.working_dir.clone(),
4170 ),
4171 None,
4172 sender,
4173 &component.name,
4174 )?;
4175 }
4176
4177 "mpi4py" => {
4179 use crate::installers::components::mpi4py::{Mpi4PyConfig, Mpi4PyInstaller};
4180 let force = is_force_reinstall();
4181 let inst = Mpi4PyInstaller::new(Mpi4PyConfig {
4182 python_bin: resolve_python_bin(),
4183 force_reinstall: force,
4184 ..Mpi4PyConfig::default()
4185 });
4186 let distro = crate::installers::common::DistroFacade::detect();
4187
4188 if force {
4190 purge_component_packages(&["mpi4py"], sender, &component.name);
4191 }
4192
4193 let cmd = inst.build_system_mpi_install_command(&distro);
4195 execute_native_command(
4196 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4197 sudo_pw,
4198 sender,
4199 &component.name,
4200 )?;
4201
4202 let mpi_path = std::path::PathBuf::from("/usr");
4204 let cmd = inst.build_pip_install_command(&mpi_path);
4205 execute_native_command(
4206 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4207 None,
4208 sender,
4209 &component.name,
4210 )?;
4211 }
4212
4213 "deepspeed" => {
4215 use crate::installers::components::deepspeed::{DeepSpeedConfig, DeepSpeedInstaller};
4216 let force = is_force_reinstall();
4217 let config = DeepSpeedConfig {
4218 python_bin: resolve_python_bin(),
4219 force_reinstall: force,
4220 ..DeepSpeedConfig::default()
4221 };
4222 let inst = DeepSpeedInstaller::new(config);
4223
4224 if force {
4226 purge_component_packages(&["deepspeed"], sender, &component.name);
4227 }
4228
4229 let deps_cmd = inst.build_deps_install_command();
4231 execute_native_command(
4232 &NativeCommand::Pip {
4233 program: deps_cmd.program.clone(),
4234 args: deps_cmd.args.clone(),
4235 },
4236 None,
4237 sender,
4238 &component.name,
4239 )?;
4240
4241 let install_cmd = inst.build_install_command();
4243 execute_native_command(
4244 &NativeCommand::Pip {
4245 program: install_cmd.program.clone(),
4246 args: install_cmd.args.clone(),
4247 },
4248 None,
4249 sender,
4250 &component.name,
4251 )?;
4252 }
4253
4254 "ml-stack-core" => {
4256 use crate::installers::components::ml_stack::{MlStackConfig, MlStackInstaller};
4257 let inst = MlStackInstaller::new(MlStackConfig {
4258 python_bin: resolve_python_bin(),
4259 ..Default::default()
4260 });
4261 let distro = crate::installers::common::DistroFacade::detect();
4262
4263 let migraphx_cmd = inst.build_migraphx_install_command(&distro);
4265 execute_native_command(
4266 &NativeCommand::System {
4267 program: migraphx_cmd.program.clone(),
4268 args: migraphx_cmd.args.clone(),
4269 },
4270 sudo_pw,
4271 sender,
4272 &component.name,
4273 )?;
4274
4275 let rccl_cmd = inst.build_rccl_install_command(&distro);
4276 execute_native_command(
4277 &NativeCommand::System {
4278 program: rccl_cmd.program.clone(),
4279 args: rccl_cmd.args.clone(),
4280 },
4281 sudo_pw,
4282 sender,
4283 &component.name,
4284 )?;
4285
4286 run_rccl_remediation(ctx, &component.name)?;
4289
4290 let mpi_cmd = inst.build_mpi_install_command(&distro);
4291 execute_native_command(
4292 &NativeCommand::System {
4293 program: mpi_cmd.program.clone(),
4294 args: mpi_cmd.args.clone(),
4295 },
4296 sudo_pw,
4297 sender,
4298 &component.name,
4299 )?;
4300
4301 let megatron_dir = format!("{}/.mlstack/src/Megatron-LM", ctx.user_home);
4306 let megatron_clone = inst.build_megatron_clone_command(&megatron_dir);
4307 execute_native_command(
4308 &NativeCommand::System {
4309 program: megatron_clone.program.clone(),
4310 args: megatron_clone.args.clone(),
4311 },
4312 sudo_pw,
4313 sender,
4314 &component.name,
4315 )?;
4316
4317 let megatron_install = inst.build_megatron_install_command();
4318 execute_native_command(
4319 &NativeCommand::from_shell_cmd_with_dir(
4320 &megatron_install.program,
4321 &megatron_install.args,
4322 &[],
4323 Some(PathBuf::from(&megatron_dir)),
4324 ),
4325 None,
4326 sender,
4327 &component.name,
4328 )?;
4329 }
4330
4331 "flash-attn" | "flash-attn-triton" | "flash-attn-ck" => {
4341 let force = is_force_reinstall();
4342 let is_triton = matches!(component.id.as_str(), "flash-attn" | "flash-attn-triton");
4345 let backend_name = if is_triton { "triton" } else { "ck" };
4346
4347 if force {
4349 purge_flash_attn(sender, &component.name);
4350 }
4351
4352 let target_dir = format!("{}/.mlstack/flash-attention", ctx.user_home);
4353
4354 git_clone_or_pull(
4355 "https://github.com/ROCm/flash-attention.git",
4356 &target_dir,
4357 &[],
4358 sudo_pw,
4359 sender,
4360 &component.name,
4361 )?;
4362
4363 fix_directory_ownership(&target_dir, sudo_pw, sender);
4365
4366 let safe_dir_cmd = NativeCommand::from_shell_cmd(
4368 "git",
4369 &[
4370 "config".to_string(),
4371 "--global".to_string(),
4372 "--add".to_string(),
4373 "safe.directory".to_string(),
4374 target_dir.clone(),
4375 ],
4376 &[],
4377 );
4378 let _ = execute_native_command(&safe_dir_cmd, None, sender, &component.name);
4379
4380 let gpu_arch = ctx
4382 .env_exports
4383 .get("GPU_ARCH")
4384 .cloned()
4385 .unwrap_or_else(|| "gfx1100".to_string());
4386 let hsa_version = ctx
4387 .env_exports
4388 .get("HSA_OVERRIDE_GFX_VERSION")
4389 .cloned()
4390 .unwrap_or_else(|| "11.0.0".to_string());
4391
4392 let fa_deps_args = vec![
4398 "-m".to_string(),
4399 "pip".to_string(),
4400 "install".to_string(),
4401 "--break-system-packages".to_string(),
4402 "einops".to_string(),
4403 "ninja".to_string(),
4404 "packaging".to_string(),
4405 "psutil".to_string(),
4406 ];
4407 execute_native_command(
4408 &NativeCommand::Shell {
4409 program: resolve_python_bin(),
4410 args: fa_deps_args,
4411 env: vec![],
4412 working_dir: None,
4413 },
4414 None, sender,
4416 &component.name,
4417 )?;
4418
4419 let mut pip_args = vec![
4426 "-m".to_string(),
4427 "pip".to_string(),
4428 "install".to_string(),
4429 "-v".to_string(),
4430 "--no-build-isolation".to_string(),
4431 "--no-deps".to_string(),
4432 "--break-system-packages".to_string(),
4433 ".".to_string(),
4434 ];
4435 if force {
4436 pip_args.push("--force-reinstall".to_string());
4437 }
4438
4439 let mut pip_env = vec![
4440 ("GPU_ARCH".to_string(), gpu_arch.clone()),
4441 ("PYTORCH_ROCM_ARCH".to_string(), gpu_arch),
4442 ("ROCM_PATH".to_string(), "/opt/rocm".to_string()),
4443 ("HSA_OVERRIDE_GFX_VERSION".to_string(), hsa_version),
4444 ];
4445 if is_triton {
4446 pip_env.push((
4448 "FLASH_ATTENTION_TRITON_AMD_ENABLE".to_string(),
4449 "TRUE".to_string(),
4450 ));
4451 } else {
4452 let jobs = std::thread::available_parallelism()
4456 .map(|n| n.get())
4457 .unwrap_or(4)
4458 .min(8);
4459 pip_env.push(("MAX_JOBS".to_string(), jobs.to_string()));
4460 }
4461
4462 execute_native_command(
4463 &NativeCommand::Shell {
4464 program: resolve_python_bin(),
4467 args: pip_args,
4468 env: pip_env,
4469 working_dir: Some(PathBuf::from(&target_dir)),
4470 },
4471 None, sender,
4473 &component.name,
4474 )?;
4475
4476 let marker = PathBuf::from(&target_dir).join(".backend");
4482 std::fs::write(&marker, backend_name).map_err(|err| {
4483 anyhow::anyhow!(
4484 "{}: failed to write Flash Attention backend marker {}: {} \
4485 — package installed but untrackable; failing install",
4486 component.name,
4487 marker.display(),
4488 err
4489 )
4490 })?;
4491 }
4492
4493 "repair-stack" => {
4497 use crate::installers::components::repair::RepairInstaller;
4498 let repair_ids = RepairInstaller::native_component_ids();
4499
4500 let _ = sender.send(InstallerEvent::Log(
4501 "[native] Starting ML Stack repair (8 steps via native installers)...".to_string(),
4502 false,
4503 ));
4504
4505 for step_id in &repair_ids {
4506 let sub_component = Component {
4509 id: (*step_id).to_string(),
4510 name: (*step_id).to_string(),
4511 description: String::new(),
4512 script: String::new(),
4513 category: Category::Core,
4514 required: true,
4515 selected: true,
4516 installed: false,
4517 progress: 0.0,
4518 estimate: String::new(),
4519 needs_sudo: component.needs_sudo,
4520 experimental: false,
4521 note: None,
4522 };
4523
4524 let _ = sender.send(InstallerEvent::Progress {
4525 component_id: component.id.clone(),
4526 progress: 0.5,
4527 message: format!("Repairing {}...", step_id),
4528 });
4529
4530 if let Err(err) = run_native_installer(&sub_component, ctx) {
4531 let _ = sender.send(InstallerEvent::Log(
4533 format!("[native] Repair step '{}' failed: {}", step_id, err),
4534 false,
4535 ));
4536 }
4537 }
4538 }
4539
4540 "rccl-repair" => {
4541 run_rccl_remediation(ctx, &component.name)?;
4542 }
4543
4544 "megatron" => {
4546 use crate::installers::components::megatron::{MegatronConfig, MegatronInstaller};
4547 let force = is_force_reinstall();
4548 let inst = MegatronInstaller::new(MegatronConfig {
4549 python_bin: resolve_python_bin(),
4550 force_reinstall: force,
4551 ..MegatronConfig::default()
4552 });
4553
4554 if force {
4556 purge_component_packages(&["megatron-core", "megatron"], sender, &component.name);
4557 }
4558
4559 for dep_pkg in &["ninja", "packaging", "wheel"] {
4561 let cmd = inst.build_dep_install_command(dep_pkg);
4562 execute_native_command(
4563 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4564 None,
4565 sender,
4566 &component.name,
4567 )?;
4568 }
4569
4570 let clone_cmd = inst.build_git_clone_command();
4572 let clone_target_str = clone_cmd.args.last().cloned().unwrap_or_default();
4573 git_clone_or_pull(
4574 "https://github.com/NVIDIA/Megatron-LM.git",
4575 &clone_target_str,
4576 &[],
4577 sudo_pw,
4578 sender,
4579 &component.name,
4580 )?;
4581
4582 let run_user = std::env::var("SUDO_USER")
4588 .or_else(|_| std::env::var("USER"))
4589 .unwrap_or_default();
4590 if !run_user.is_empty() {
4591 let chown_cmd = NativeCommand::from_shell_cmd(
4592 "chown",
4593 &[
4594 "-R".to_string(),
4595 format!("{run_user}:{run_user}"),
4596 clone_target_str.clone(),
4597 ],
4598 &[],
4599 );
4600 let _ = execute_native_command(&chown_cmd, sudo_pw, sender, &component.name);
4601 }
4602
4603 let clone_target_path = PathBuf::from(&clone_target_str);
4607 if let Ok(entries) = std::fs::read_dir(&clone_target_path) {
4608 for entry in entries.flatten() {
4609 let name = entry.file_name();
4610 let name_str = name.to_string_lossy();
4611 if name_str.ends_with(".egg-info") {
4612 let egg_path = entry.path();
4613 if std::fs::remove_dir_all(&egg_path).is_err() {
4615 let rm_cmd = NativeCommand::from_shell_cmd(
4617 "rm",
4618 &["-rf".to_string(), egg_path.to_string_lossy().to_string()],
4619 &[],
4620 );
4621 let _ =
4622 execute_native_command(&rm_cmd, sudo_pw, sender, &component.name);
4623 }
4624 }
4625 }
4626 }
4627
4628 let cmd = inst.build_pip_install_command();
4630 execute_native_command(
4631 &NativeCommand::from_shell_cmd_with_dir(
4632 &cmd.program,
4633 &cmd.args,
4634 &cmd.env,
4635 cmd.working_dir.clone(),
4636 ),
4637 None,
4638 sender,
4639 &component.name,
4640 )?;
4641
4642 let requirements_candidates = [
4644 clone_target_path
4645 .join("requirements")
4646 .join("requirements.txt"),
4647 clone_target_path.join("requirements.txt"),
4648 clone_target_path
4649 .join("requirements")
4650 .join("requirements_amd.txt"),
4651 ];
4652 for req_path in requirements_candidates.iter().filter(|p| p.exists()) {
4653 let filtered = match filter_cuda_requirements(&req_path.to_string_lossy()) {
4654 Ok(f) => f,
4655 Err(e) => {
4656 let _ = sender.send(InstallerEvent::Log(
4657 format!(
4658 "[native] WARN: could not read Megatron requirements {}: {e} — skipping requirements install",
4659 req_path.display()
4660 ),
4661 true,
4662 ));
4663 anyhow::bail!(
4664 "failed to read Megatron requirements {}: {e}",
4665 req_path.display()
4666 );
4667 }
4668 };
4669 if filtered.is_empty() {
4670 continue;
4671 }
4672 let tmp_dir = std::env::temp_dir().join("rusty-stack-megatron");
4673 let _ = std::fs::create_dir_all(&tmp_dir);
4674 let filtered_path = tmp_dir.join("requirements_filtered.txt");
4675 std::fs::write(&filtered_path, &filtered).with_context(|| {
4676 format!(
4677 "Failed to write filtered Megatron requirements to {:?}",
4678 filtered_path
4679 )
4680 })?;
4681 let filtered_path_str = filtered_path.to_string_lossy().to_string();
4682 let filtered_cmd = inst.build_requirements_install_command(&filtered_path_str);
4683 execute_native_command(
4684 &NativeCommand::from_shell_cmd(
4685 &filtered_cmd.program,
4686 &filtered_cmd.args,
4687 &filtered_cmd.env,
4688 ),
4689 None,
4690 sender,
4691 &component.name,
4692 )?;
4693 let _ = std::fs::remove_file(&filtered_path);
4694 break;
4695 }
4696 }
4697
4698 "vllm" => {
4700 use crate::installers::components::vllm_multi::{VllmConfig, VllmInstaller};
4701 let vllm_force = is_force_reinstall();
4702 let inst = VllmInstaller::new(VllmConfig {
4703 python_bin: resolve_python_bin(),
4704 ..Default::default()
4705 });
4706
4707 if vllm_force {
4709 purge_vllm(sender, &component.name);
4710 }
4711
4712 let cmd = inst.build_tools_install_command();
4715 execute_native_command(
4716 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4717 None,
4718 sender,
4719 &component.name,
4720 )?;
4721
4722 let vllm_version = inst
4727 .resolve_version_with_policy(ctx.vllm_version_lag, ctx.vllm_version_min_age_days);
4728 let _ = sender.send(InstallerEvent::Log(
4729 format!(
4730 "[native] vLLM: selected {vllm_version} (gate: lag={}, min_age_days={}) — \
4731 building from source (ROCm, env torch)",
4732 ctx.vllm_version_lag, ctx.vllm_version_min_age_days
4733 ),
4734 false,
4735 ));
4736 let cmd = inst.build_source_install_command(&vllm_version);
4741 execute_native_command(
4742 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4743 None,
4744 sender,
4745 &component.name,
4746 )?;
4747
4748 let cmd = inst.build_versioned_deps_command();
4752 execute_native_command(
4753 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4754 None,
4755 sender,
4756 &component.name,
4757 )?;
4758
4759 let cmd = inst.build_deps_install_command();
4761 execute_native_command(
4762 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4763 None,
4764 sender,
4765 &component.name,
4766 )?;
4767
4768 let cmd = inst.build_integrity_check_command();
4772 execute_native_command(
4773 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4774 None,
4775 sender,
4776 &component.name,
4777 )?;
4778 }
4779
4780 "aiter" => {
4782 use crate::installers::components::aiter::{AiterConfig, AiterInstaller};
4783 let detected_gpu_arch = ctx
4786 .env_exports
4787 .get("GPU_ARCH")
4788 .map(|v| v.trim().to_string())
4789 .filter(|v| !v.is_empty())
4790 .unwrap_or_else(|| detect_gpu_arch().unwrap_or_else(|| "gfx1100".to_string()));
4791 let inst = AiterInstaller::new(AiterConfig {
4792 python_bin: resolve_python_bin(),
4793 gpu_arch: detected_gpu_arch,
4794 ..Default::default()
4795 });
4796 let target_dir = format!("{}/.mlstack/aiter", ctx.user_home);
4797
4798 if is_force_reinstall() {
4800 purge_aiter(sender, &component.name);
4801 }
4802
4803 git_clone_or_pull(
4805 inst.git_clone_url(),
4806 &target_dir,
4807 &["--recursive"],
4808 sudo_pw,
4809 sender,
4810 &component.name,
4811 )?;
4812
4813 let run_user = std::env::var("SUDO_USER")
4817 .or_else(|_| std::env::var("USER"))
4818 .unwrap_or_default();
4819 if !run_user.is_empty() {
4820 let chown_cmd = NativeCommand::from_shell_cmd(
4821 "chown",
4822 &[
4823 "-R".to_string(),
4824 format!("{run_user}:{run_user}"),
4825 target_dir.clone(),
4826 ],
4827 &[],
4828 );
4829 let _ = execute_native_command(&chown_cmd, sudo_pw, sender, &component.name);
4830 }
4831
4832 let cmd = inst.build_deps_install_command();
4834 execute_native_command(
4835 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4836 sudo_pw,
4837 sender,
4838 &component.name,
4839 )?;
4840
4841 let cmd = inst.build_pip_install_command(&target_dir);
4843 execute_native_command(
4844 &NativeCommand::from_shell_cmd_with_dir(
4845 &cmd.program,
4846 &cmd.args,
4847 &cmd.env,
4848 cmd.working_dir.clone(),
4849 ),
4850 None,
4851 sender,
4852 &component.name,
4853 )?;
4854
4855 let run_user = std::env::var("SUDO_USER")
4859 .or_else(|_| std::env::var("USER"))
4860 .unwrap_or_default();
4861 if !run_user.is_empty() {
4862 let user_home = &ctx.user_home;
4863 for aiter_dir in [
4864 format!("{user_home}/.aiter"),
4865 format!("{user_home}/.mlstack/aiter"),
4866 ] {
4867 let dir_path = PathBuf::from(&aiter_dir);
4868 if dir_path.exists() {
4869 let needs_fix = std::fs::read_dir(&dir_path)
4871 .map(|mut entries| entries.next().is_some())
4872 .unwrap_or(false)
4873 && std::fs::write(dir_path.join(".permission_check"), b"").is_err();
4874 if needs_fix {
4875 let chown_cmd = NativeCommand::from_shell_cmd(
4876 "chown",
4877 &[
4878 "-R".to_string(),
4879 format!("{run_user}:{run_user}"),
4880 aiter_dir.clone(),
4881 ],
4882 &[],
4883 );
4884 let _ = execute_native_command(
4885 &chown_cmd,
4886 sudo_pw,
4887 sender,
4888 &component.name,
4889 );
4890 } else {
4891 let _ = std::fs::remove_file(dir_path.join(".permission_check"));
4893 }
4894 }
4895 }
4896 }
4897 }
4898 "vllm-studio" => {
4899 use crate::installers::components::vllm_studio::{
4900 VllmStudioConfig, VllmStudioInstaller,
4901 };
4902 let inst = VllmStudioInstaller::new(VllmStudioConfig {
4903 python_bin: resolve_python_bin(),
4904 ..Default::default()
4905 });
4906
4907 let clone_cmd = inst.build_git_clone_command();
4909 let target_dir = clone_cmd.args.last().cloned().unwrap_or_default();
4910 git_clone_or_pull(
4911 inst.repo_url(),
4912 &target_dir,
4913 &[],
4914 None,
4915 sender,
4916 &component.name,
4917 )?;
4918
4919 let pkg_mgr = if crate::installers::common::utils::command_exists("npm") {
4921 "npm"
4922 } else {
4923 "yarn"
4924 };
4925 let cmd = inst.build_pkg_install_command(pkg_mgr);
4926 execute_native_command(
4927 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4928 None,
4929 sender,
4930 &component.name,
4931 )?;
4932
4933 let frontend_dir = PathBuf::from(&target_dir).join("frontend");
4936 let cmd = inst.build_frontend_build_command(pkg_mgr);
4937 let build_result = execute_native_command(
4938 &NativeCommand::from_shell_cmd_with_dir(
4939 &cmd.program,
4940 &cmd.args,
4941 &cmd.env,
4942 Some(frontend_dir),
4943 ),
4944 None,
4945 sender,
4946 &component.name,
4947 );
4948 if let Err(err) = build_result {
4949 let _ = sender.send(InstallerEvent::Log(
4950 format!(
4951 "[native] [WARN] {} frontend build failed (non-fatal): {}",
4952 component.name, err
4953 ),
4954 false,
4955 ));
4956 }
4957 }
4958
4959 "comfyui" => {
4961 use crate::installers::components::comfyui::{ComfyuiConfig, ComfyuiInstaller};
4962 let inst = ComfyuiInstaller::new(ComfyuiConfig {
4963 python_bin: resolve_python_bin(),
4964 ..Default::default()
4965 });
4966
4967 let cmd = inst.build_pytorch_check_command();
4969 execute_native_command(
4970 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4971 None,
4972 sender,
4973 &component.name,
4974 )?;
4975
4976 let clone_cmd = inst.build_git_clone_command();
4978 let target_dir = clone_cmd.args.last().cloned().unwrap_or_default();
4979 git_clone_or_pull(
4980 inst.repo_url(),
4981 &target_dir,
4982 &[],
4983 None,
4984 sender,
4985 &component.name,
4986 )?;
4987
4988 let req_path = format!("{}/requirements.txt", target_dir);
4994 if std::path::Path::new(&req_path).exists() {
4995 let filtered = match filter_cuda_requirements(&req_path) {
4996 Ok(f) => f,
4997 Err(e) => {
4998 let _ = sender.send(InstallerEvent::Log(
4999 format!(
5000 "[native] WARN: could not read ComfyUI requirements {req_path}: {e} — skipping requirements install"
5001 ),
5002 true,
5003 ));
5004 anyhow::bail!("failed to read ComfyUI requirements {req_path}: {e}");
5005 }
5006 };
5007 if !filtered.is_empty() {
5008 let tmp_dir = std::env::temp_dir().join("rusty-stack-comfyui");
5009 let _ = std::fs::create_dir_all(&tmp_dir);
5010 let filtered_path = tmp_dir.join("requirements_filtered.txt");
5011 std::fs::write(&filtered_path, &filtered).with_context(|| {
5012 format!(
5013 "Failed to write filtered requirements to {:?}",
5014 filtered_path
5015 )
5016 })?;
5017 let filtered_path_str = filtered_path.to_string_lossy().to_string();
5018 let _ = sender.send(InstallerEvent::Log(
5019 format!(
5020 "[native] Installing filtered requirements (CUDA deps excluded) from {}",
5021 req_path
5022 ),
5023 false,
5024 ));
5025 let cmd = inst.build_pip_install_command(&filtered_path_str);
5026 execute_native_command(
5027 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5028 None,
5029 sender,
5030 &component.name,
5031 )?;
5032 let _ = std::fs::remove_file(&filtered_path);
5033 } else {
5034 let _ = sender.send(InstallerEvent::Log(
5035 "[native] [WARN] All ComfyUI requirements were filtered out, nothing to install".to_string(),
5036 false,
5037 ));
5038 }
5039 } else {
5040 let _ = sender.send(InstallerEvent::Log(
5041 "[native] [WARN] ComfyUI requirements.txt not found, skipping dependency installation".to_string(),
5042 false,
5043 ));
5044 }
5045 }
5046
5047 "textgen" => {
5049 use crate::installers::components::textgen::{TextgenConfig, TextgenInstaller};
5050 let textgen_python = resolve_python_bin();
5053 let inst = TextgenInstaller::new(TextgenConfig {
5058 python_bin: textgen_python.clone(),
5059 use_uv: false,
5060 break_system_packages: true,
5061 ..Default::default()
5062 });
5063
5064 let clone_cmd = inst.build_git_clone_command();
5066 let target_dir = clone_cmd.args.last().cloned().unwrap_or_default();
5067 git_clone_or_pull(
5068 inst.repo_url(),
5069 &target_dir,
5070 &[],
5071 None,
5072 sender,
5073 &component.name,
5074 )?;
5075
5076 let amd_req_full = format!("{}/requirements/full/requirements_amd.txt", target_dir);
5081 let amd_req_subdir = format!("{}/requirements/requirements_amd.txt", target_dir);
5082 let amd_req_path = if std::path::Path::new(&amd_req_full).exists() {
5083 amd_req_full
5084 } else if std::path::Path::new(&amd_req_subdir).exists() {
5085 amd_req_subdir
5086 } else {
5087 String::new()
5088 };
5089 if !amd_req_path.is_empty() {
5090 let _ = sender.send(InstallerEvent::Log(
5091 "[native] Installing AMD-specific requirements first (ROCm wheels)".to_string(),
5092 false,
5093 ));
5094 let cmd = inst.build_amd_requirements_command(&amd_req_path);
5095 execute_native_command(
5096 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5097 None,
5098 sender,
5099 &component.name,
5100 )?;
5101 }
5102
5103 let req_path_root = format!("{}/requirements.txt", target_dir);
5110 let req_path_full = format!("{}/requirements/full/requirements.txt", target_dir);
5111 let req_path_subdir = format!("{}/requirements/requirements.txt", target_dir);
5112 let raw_req_path = if std::path::Path::new(&req_path_full).exists() {
5113 req_path_full
5114 } else if std::path::Path::new(&req_path_root).exists() {
5115 req_path_root
5116 } else if std::path::Path::new(&req_path_subdir).exists() {
5117 req_path_subdir
5118 } else {
5119 String::new()
5120 };
5121
5122 if !raw_req_path.is_empty() {
5123 let filtered = match filter_cuda_requirements(&raw_req_path) {
5125 Ok(f) => f,
5126 Err(e) => {
5127 let _ = sender.send(InstallerEvent::Log(
5128 format!(
5129 "[native] WARN: could not read textgen requirements {raw_req_path}: {e} — skipping requirements install"
5130 ),
5131 true,
5132 ));
5133 anyhow::bail!("failed to read textgen requirements {raw_req_path}: {e}");
5134 }
5135 };
5136 if !filtered.is_empty() {
5137 let tmp_dir = std::env::temp_dir().join("rusty-stack-textgen");
5139 let _ = std::fs::create_dir_all(&tmp_dir);
5140 let filtered_path = tmp_dir.join("requirements_filtered.txt");
5141 std::fs::write(&filtered_path, &filtered).with_context(|| {
5142 format!(
5143 "Failed to write filtered requirements to {:?}",
5144 filtered_path
5145 )
5146 })?;
5147 let filtered_path_str = filtered_path.to_string_lossy().to_string();
5148 let _ = sender.send(InstallerEvent::Log(
5149 format!(
5150 "[native] Installing filtered requirements (CUDA deps excluded) from {}",
5151 raw_req_path
5152 ),
5153 false,
5154 ));
5155 let cmd = inst.build_pip_install_command(&filtered_path_str);
5156 execute_native_command(
5157 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5158 None,
5159 sender,
5160 &component.name,
5161 )?;
5162 let _ = std::fs::remove_file(&filtered_path);
5163 } else {
5164 let _ = sender.send(InstallerEvent::Log(
5165 "[native] [WARN] All requirements were filtered out, nothing to install"
5166 .to_string(),
5167 false,
5168 ));
5169 }
5170 } else {
5171 let _ = sender.send(InstallerEvent::Log(
5172 "[native] [WARN] requirements.txt not found, skipping dependency installation"
5173 .to_string(),
5174 false,
5175 ));
5176 }
5177 }
5178
5179 "onnx" => {
5181 use crate::installers::components::onnxruntime::{
5182 install_method_from_env, prebuilt_version_from_env, OnnxInstallMethod,
5183 OnnxRuntimeConfig, OnnxRuntimeInstaller,
5184 };
5185
5186 let detected_rocm_version = detect_rocm_version();
5187 let detected_gpu_arch = detect_gpu_arch().unwrap_or_else(|| {
5188 eprintln!("Warning: no GPU arch detected for ONNX, falling back to gfx1100");
5189 "gfx1100".to_string()
5190 });
5191
5192 let install_method = install_method_from_env();
5199 let prebuilt_version = prebuilt_version_from_env();
5200
5201 let config = OnnxRuntimeConfig {
5203 rocm_version: Some(detected_rocm_version.clone()),
5204 rocm_release: Some({
5205 let parts: Vec<&str> = detected_rocm_version.split('.').collect();
5208 if parts.len() >= 3 {
5209 detected_rocm_version.clone()
5210 } else if parts.len() >= 2 && parts[0] == "7" && parts[1] == "2" {
5211 "7.2.4".to_string()
5212 } else if parts.len() >= 2 {
5213 format!("{}.{}.3", parts[0], parts[1])
5214 } else {
5215 "7.2.4".to_string()
5216 }
5217 }),
5218 gpu_arch: Some(detected_gpu_arch),
5219 python_bin: resolve_python_bin(),
5220 install_method,
5221 runtime_version: ctx.target_version.clone(),
5222 prebuilt_version,
5223 ..Default::default()
5224 };
5225 let inst = OnnxRuntimeInstaller::new(config);
5226
5227 let cmd = inst.build_uninstall_command();
5229 let _ = execute_native_command(
5230 &NativeCommand::from_shell_cmd_with_dir(
5231 &cmd.program,
5232 &cmd.args,
5233 &cmd.env,
5234 cmd.working_dir.clone(),
5235 ),
5236 None,
5237 sender,
5238 &component.name,
5239 );
5240
5241 match install_method {
5243 OnnxInstallMethod::MigraphxWheel => {
5244 let _ = sender.send(InstallerEvent::Log(
5248 format!(
5249 "[native] {} — installing onnxruntime-migraphx=={} from PyPI (ROCm {})",
5250 component.name,
5251 inst.prebuilt_version(),
5252 detected_rocm_version
5253 ),
5254 false,
5255 ));
5256 let cmd = inst.build_migraphx_install_command();
5257 execute_native_command(
5258 &NativeCommand::from_shell_cmd_with_dir(
5259 &cmd.program,
5260 &cmd.args,
5261 &cmd.env,
5262 cmd.working_dir.clone(),
5263 ),
5264 None,
5265 sender,
5266 &component.name,
5267 )?;
5268 }
5269 OnnxInstallMethod::PrebuiltWheel => {
5270 let _ = sender.send(InstallerEvent::Log(
5273 format!(
5274 "[native] {} — installing prebuilt onnxruntime-rocm from PyPI \
5275 (legacy; may be ABI-incompatible with ROCm {})",
5276 component.name, detected_rocm_version
5277 ),
5278 false,
5279 ));
5280 let cmd = inst.build_prebuilt_install_command();
5281 execute_native_command(
5282 &NativeCommand::from_shell_cmd_with_dir(
5283 &cmd.program,
5284 &cmd.args,
5285 &cmd.env,
5286 cmd.working_dir.clone(),
5287 ),
5288 None,
5289 sender,
5290 &component.name,
5291 )?;
5292 }
5293 OnnxInstallMethod::SourceBuild => {
5294 let _ = sender.send(InstallerEvent::Log(
5296 format!(
5297 "[native] {} — building from source (target {}, ROCm {}). Heavy build.",
5298 component.name,
5299 inst.runtime_version(),
5300 detected_rocm_version
5301 ),
5302 false,
5303 ));
5304 let rocm_env = crate::installers::common::RocmEnv::from_known(
5305 Some(std::path::PathBuf::from("/opt/rocm")),
5306 detected_rocm_version.clone(),
5307 );
5308 for cmd in [
5309 inst.build_git_clone_command(),
5310 inst.build_git_checkout_command(),
5311 inst.build_git_submodule_command(),
5312 inst.build_build_command(&rocm_env),
5313 inst.build_wheel_install_command(),
5314 ] {
5315 execute_native_command(
5316 &NativeCommand::from_shell_cmd_with_dir(
5317 &cmd.program,
5318 &cmd.args,
5319 &cmd.env,
5320 cmd.working_dir.clone(),
5321 ),
5322 None,
5323 sender,
5324 &component.name,
5325 )?;
5326 }
5327 }
5328 }
5329
5330 let cmd = inst.build_provider_validation_command();
5332 execute_native_command(
5333 &NativeCommand::from_shell_cmd_with_dir(
5334 &cmd.program,
5335 &cmd.args,
5336 &cmd.env,
5337 cmd.working_dir.clone(),
5338 ),
5339 None,
5340 sender,
5341 &component.name,
5342 )?;
5343
5344 let model_dirs = [
5346 dirs::home_dir().map(|h| h.join(".mlstack/models")),
5347 dirs::home_dir().map(|h| h.join(".local/share/models")),
5348 ];
5349 for model_dir in model_dirs.iter().flatten() {
5350 if model_dir.exists() {
5351 if let Ok(entries) = std::fs::read_dir(model_dir) {
5352 for entry in entries.flatten() {
5353 let path = entry.path();
5354 if path.extension().map(|e| e == "onnx").unwrap_or(false) {
5355 let model_path = path.to_string_lossy().to_string();
5356 if model_path.ends_with(".onnx.optimized") {
5358 continue;
5359 }
5360 let _ = sender.send(InstallerEvent::Log(
5361 format!(
5362 "[native] {} — optimizing model: {}",
5363 component.name, model_path
5364 ),
5365 false,
5366 ));
5367 let cmd = inst.build_model_optimizer_command(&model_path);
5368 let _ = execute_native_command(
5369 &NativeCommand::from_shell_cmd_with_dir(
5370 &cmd.program,
5371 &cmd.args,
5372 &cmd.env,
5373 cmd.working_dir.clone(),
5374 ),
5375 None,
5376 sender,
5377 &component.name,
5378 );
5379 }
5380 }
5381 }
5382 }
5383 }
5384
5385 persist_onnx_install_status(
5386 inst.prebuilt_version(),
5387 &["MIGraphXExecutionProvider", "CPUExecutionProvider"],
5388 );
5389
5390 let _ = sender.send(InstallerEvent::Log(
5391 format!("[native] {} — ONNX Runtime (MIGraphX) installed successfully. \
5392 Source build available via OnnxInstallMethod::SourceBuild for ROCMExecutionProvider.",
5393 component.name),
5394 false,
5395 ));
5396 }
5397
5398 "bitsandbytes" => {
5400 use crate::installers::components::bitsandbytes_multi::{
5401 BitsAndBytesConfig, BitsAndBytesInstaller,
5402 };
5403 let inst = BitsAndBytesInstaller::new(BitsAndBytesConfig {
5404 python_bin: resolve_python_bin(),
5405 ..Default::default()
5406 });
5407 let target_dir = format!("{}/.mlstack/bitsandbytes", ctx.user_home);
5408
5409 if is_force_reinstall() {
5411 purge_component_packages(&["bitsandbytes"], sender, &component.name);
5412 }
5413
5414 git_clone_or_pull(
5416 "https://github.com/ROCm/bitsandbytes.git",
5417 &target_dir,
5418 &["--recursive"],
5419 sudo_pw,
5420 sender,
5421 &component.name,
5422 )?;
5423
5424 let cmd = inst.build_pypi_install_command();
5426 execute_native_command(
5427 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5428 None,
5429 sender,
5430 &component.name,
5431 )?;
5432 }
5433
5434 "rocm-smi" => {
5436 use crate::installers::components::rocm_smi::RocmSmiInstaller;
5437 let inst = RocmSmiInstaller::new();
5438 let distro = crate::installers::common::DistroFacade::detect();
5439 let target_dir = format!("{}/.mlstack/rocm_smi", ctx.user_home);
5440
5441 git_clone_or_pull(
5443 inst.git_clone_url(),
5444 &target_dir,
5445 &[],
5446 sudo_pw,
5447 sender,
5448 &component.name,
5449 )?;
5450
5451 if !(distro.family() == crate::platform::detection::DistroFamily::Arch
5453 && crate::installers::common::utils::command_exists("rocm-smi"))
5454 {
5455 let cmd = inst.build_package_install_command(&distro);
5456 execute_native_command(
5457 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5458 sudo_pw,
5459 sender,
5460 &component.name,
5461 )?;
5462 } else {
5463 let _ = sender.send(InstallerEvent::Log(
5464 format!(
5465 "[native] {} — rocm-smi already on PATH, skipping package install",
5466 component.name
5467 ),
5468 false,
5469 ));
5470 }
5471 }
5472
5473 "migraphx" => {
5475 use crate::installers::components::migraphx_multi::{
5476 MigraphxConfig, MigraphxInstaller,
5477 };
5478 let inst = MigraphxInstaller::new(MigraphxConfig {
5479 python_bin: resolve_python_bin(),
5480 ..Default::default()
5481 });
5482 let distro = crate::installers::common::DistroFacade::detect();
5483 let packages = inst.required_packages(&distro);
5484
5485 let cmd = inst.build_package_install_command(&distro, &packages);
5487 execute_native_command(
5488 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5489 sudo_pw,
5490 sender,
5491 &component.name,
5492 )?;
5493
5494 if inst.should_skip_pip_install(&distro) {
5496 if let Some(msg) = inst.build_limitation_message(&distro) {
5497 let _ = sender.send(InstallerEvent::Log(
5498 format!("[native] [INFO] {}", msg),
5499 false,
5500 ));
5501 }
5502 } else {
5503 let cmd = inst.build_pip_install_command();
5504 execute_native_command(
5505 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5506 None,
5507 sender,
5508 &component.name,
5509 )?;
5510 }
5511
5512 match crate::installers::components::migraphx_multi::migraphx_driver_status() {
5517 Ok(()) => {}
5518 Err(missing) => {
5519 let detail = if missing.is_empty() {
5520 "migraphx-driver would not run".to_string()
5521 } else {
5522 format!("missing shared libraries: {}", missing.join(", "))
5523 };
5524 let _ = sender.send(InstallerEvent::Log(
5525 format!(
5526 "[native] [WARN] {} installed but NOT functional — {detail}. \
5527 This is usually a library version desync from a partial upgrade. \
5528 Run `sudo pacman -Syu` (Arch/CachyOS) or the equivalent system \
5529 update to sync dependencies, then re-verify.",
5530 component.name
5531 ),
5532 true,
5533 ));
5534 }
5535 }
5536 }
5537
5538 "pytorch-profiler" => {
5540 use crate::installers::components::pytorch_profiler::{
5541 PytorchProfilerConfig, PytorchProfilerInstaller,
5542 };
5543 let inst = PytorchProfilerInstaller::new(PytorchProfilerConfig {
5544 python_bin: resolve_python_bin(),
5545 ..Default::default()
5546 });
5547
5548 let cmd = inst.build_install_command();
5550 execute_native_command(
5551 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5552 None,
5553 sender,
5554 &component.name,
5555 )?;
5556 }
5557
5558 "wandb" => {
5560 use crate::installers::components::wandb::{WandbConfig, WandbInstaller};
5561 let inst = WandbInstaller::new(WandbConfig {
5562 python_bin: resolve_python_bin(),
5563 ..Default::default()
5564 });
5565
5566 let cmd = inst.build_install_command();
5568 execute_native_command(
5569 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5570 None,
5571 sender,
5572 &component.name,
5573 )?;
5574 }
5575
5576 "amdgpu-drivers" => {
5578 use crate::installers::components::amdgpu_drivers::{AmdgpuConfig, AmdgpuInstaller};
5579 let inst = AmdgpuInstaller::new(AmdgpuConfig::default());
5580 let distro = crate::installers::common::DistroFacade::detect();
5581 let packages = inst.required_packages(&distro);
5582
5583 let cmd = inst.build_package_install_command(&distro, &packages);
5585 execute_native_command(
5586 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5587 sudo_pw,
5588 sender,
5589 &component.name,
5590 )?;
5591
5592 let cmd = inst.build_verify_command();
5600 execute_native_command(
5601 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5602 None,
5603 sender,
5604 &component.name,
5605 )?;
5606
5607 ensure_user_in_groups(&["render", "video"], sudo_pw, sender, &component.name);
5610 }
5611
5612 "migraphx-python" => {
5614 use crate::installers::components::migraphx_python::{
5615 build_python_source_requested, source_branch, MigraphxPythonConfig,
5616 MigraphxPythonInstaller,
5617 };
5618 let inst = MigraphxPythonInstaller::new(MigraphxPythonConfig {
5619 python_bin: resolve_python_bin(),
5620 ..Default::default()
5621 });
5622
5623 let distro = crate::installers::common::DistroFacade::detect();
5625 if !inst.is_available_on_distro(&distro) {
5626 if build_python_source_requested() {
5627 let _ = sender.send(InstallerEvent::Log(
5630 "[native] MLSTACK_MIGRAPHX_BUILD_PYTHON=1 — building Python \
5631 bindings from source (AMDMIGraphX). Heavy build (~20-40 min)."
5632 .into(),
5633 false,
5634 ));
5635 let workdir = std::env::temp_dir()
5636 .join("rusty-migraphx-python")
5637 .to_string_lossy()
5638 .to_string();
5639 let _ = std::fs::remove_dir_all(&workdir);
5640 let gpu_arch =
5641 std::env::var("GPU_ARCH").unwrap_or_else(|_| "gfx1100".to_string());
5642 let cmds = inst.build_source_commands(
5643 &workdir,
5644 "/opt/rocm",
5645 &gpu_arch,
5646 &source_branch(),
5647 );
5648 for cmd in cmds {
5649 execute_native_command(
5650 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5651 None,
5652 sender,
5653 &component.name,
5654 )?;
5655 }
5656 } else if let Some(msg) = inst.build_unavailable_message(&distro) {
5657 let _ = sender.send(InstallerEvent::Log(
5658 format!("[native] [SKIP] {}", msg),
5659 false,
5660 ));
5661 }
5662 } else {
5663 let cmd = inst.build_install_command();
5665 execute_native_command(
5666 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5667 None,
5668 sender,
5669 &component.name,
5670 )?;
5671 }
5672 }
5673
5674 "enhanced-env" => {
5676 let _ = sender.send(InstallerEvent::Log(
5680 "[native] Enhanced environment setup handled during pre-install phase".into(),
5681 false,
5682 ));
5683 }
5684
5685 "mlperf-inference"
5687 | "rocm-benchmarks"
5688 | "gpu-memory-bandwidth"
5689 | "rocm-smi-bench"
5690 | "pytorch-performance"
5691 | "vllm-performance"
5692 | "deepspeed-performance"
5693 | "megatron-performance"
5694 | "onnx-performance"
5695 | "rusty-llama-performance"
5696 | "flash-attention-ck-performance"
5697 | "all-benchmarks" => {
5698 run_native_benchmark(&component.id, sender)?;
5699 }
5700
5701 "fastvideo" => {
5703 use crate::installers::components::fastvideo::{
5704 detect_discrete_gpu_archs_from_sysfs, fastvideo_build_warning, post_checkout_plan,
5705 resolve_gpu_archs, resolve_install_user, select_visible_gpu_archs, FastVideoConfig,
5706 FastVideoInstallStep, FastVideoInstaller,
5707 };
5708 let detected_archs =
5709 detect_discrete_gpu_archs_from_sysfs(PathBuf::from("/sys/class/drm").as_path());
5710 let selected_archs = select_visible_gpu_archs(
5711 &detected_archs,
5712 ctx.env_exports
5713 .get("ROCR_VISIBLE_DEVICES")
5714 .map(String::as_str),
5715 ctx.env_exports
5716 .get("HIP_VISIBLE_DEVICES")
5717 .map(String::as_str),
5718 ctx.env_exports
5719 .get("CUDA_VISIBLE_DEVICES")
5720 .map(String::as_str),
5721 );
5722 let gpu_archs = resolve_gpu_archs(
5723 &selected_archs,
5724 ctx.env_exports.get("GPU_ARCHS").map(String::as_str),
5725 ctx.env_exports.get("GPU_ARCH").map(String::as_str),
5726 None,
5727 );
5728 let python_bin = ctx
5729 .env_exports
5730 .get("MLSTACK_PYTHON_BIN")
5731 .map(|v| v.trim().to_string())
5732 .filter(|v| !v.is_empty())
5733 .unwrap_or_else(resolve_python_bin);
5734 let inst = FastVideoInstaller::new(FastVideoConfig {
5735 gpu_archs: gpu_archs.clone(),
5736 python_bin: python_bin.clone(),
5737 });
5738 let build_dir = PathBuf::from("/tmp/FastVideo_ROCm_build");
5739
5740 let initial_cleanup = inst.cleanup();
5743 execute_native_command(
5744 &NativeCommand::from_shell_cmd(
5745 &initial_cleanup.program,
5746 &initial_cleanup.args,
5747 &initial_cleanup.env,
5748 ),
5749 sudo_pw,
5750 sender,
5751 &component.name,
5752 )?;
5753
5754 let mkdir_cmd = inst.mkdir_build_dir();
5756 execute_native_command(
5757 &NativeCommand::from_shell_cmd(&mkdir_cmd.program, &mkdir_cmd.args, &mkdir_cmd.env),
5758 sudo_pw,
5759 sender,
5760 &component.name,
5761 )?;
5762
5763 let clone_cmd = inst.git_clone();
5765 execute_native_command(
5766 &NativeCommand::from_shell_cmd_with_dir(
5767 &clone_cmd.program,
5768 &clone_cmd.args,
5769 &clone_cmd.env,
5770 clone_cmd.working_dir,
5771 ),
5772 sudo_pw,
5773 sender,
5774 &component.name,
5775 )?;
5776
5777 let checkout_cmd = inst.git_checkout();
5779 execute_native_command(
5780 &NativeCommand::from_shell_cmd_with_dir(
5781 &checkout_cmd.program,
5782 &checkout_cmd.args,
5783 &checkout_cmd.env,
5784 Some(build_dir.clone()),
5785 ),
5786 sudo_pw,
5787 sender,
5788 &component.name,
5789 )?;
5790
5791 let sudo_user = std::env::var("SUDO_USER").ok();
5792 let user = std::env::var("USER").ok();
5793 let run_user = resolve_install_user(sudo_user.as_deref(), user.as_deref())
5794 .context("FastVideo install user unavailable from SUDO_USER or USER")?;
5795
5796 for step in post_checkout_plan() {
5797 let step_sudo = if step.requires_sudo() { sudo_pw } else { None };
5798 if step == FastVideoInstallStep::KernelInstall {
5799 let snapshot = inst.distribution_snapshot(
5800 "before",
5801 crate::installers::common::nvidia_blocklist::NVIDIA_RUNTIME_PREFIXES,
5802 );
5803 execute_native_command(
5804 &NativeCommand::from_shell_cmd(
5805 &snapshot.program,
5806 &snapshot.args,
5807 &snapshot.env,
5808 ),
5809 None,
5810 sender,
5811 &component.name,
5812 )?;
5813 }
5814
5815 let command = match step {
5816 FastVideoInstallStep::Chown => {
5817 let cmd = inst.chown_build_dir(&run_user);
5818 NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env)
5819 }
5820 FastVideoInstallStep::ManagedDependenciesPreflight => {
5821 let cmd = inst.managed_dependencies_preflight();
5822 NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env)
5823 }
5824 FastVideoInstallStep::FlashAttentionPreflight => {
5825 let cmd = inst.flash_attention_preflight();
5826 NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env)
5827 }
5828 FastVideoInstallStep::SourcePolicyPreflight => {
5829 let cmd = inst.source_policy_preflight();
5830 NativeCommand::from_shell_cmd_with_dir(
5831 &cmd.program,
5832 &cmd.args,
5833 &cmd.env,
5834 cmd.working_dir,
5835 )
5836 }
5837 FastVideoInstallStep::KernelInstall => {
5838 let _ = sender.send(InstallerEvent::Log(
5839 fastvideo_build_warning(&gpu_archs),
5840 false,
5841 ));
5842 let cmd = inst.pip_install_kernel();
5843 NativeCommand::from_shell_cmd_with_dir(
5844 &cmd.program,
5845 &cmd.args,
5846 &cmd.env,
5847 cmd.working_dir,
5848 )
5849 }
5850 FastVideoInstallStep::PackageInstall => {
5851 let cmd = inst.pip_install_package();
5852 NativeCommand::from_shell_cmd_with_dir(
5853 &cmd.program,
5854 &cmd.args,
5855 &cmd.env,
5856 cmd.working_dir,
5857 )
5858 }
5859 };
5860 execute_native_command(&command, step_sudo, sender, &component.name)?;
5861 }
5862
5863 let snapshot = inst.distribution_snapshot(
5864 "after",
5865 crate::installers::common::nvidia_blocklist::NVIDIA_RUNTIME_PREFIXES,
5866 );
5867 execute_native_command(
5868 &NativeCommand::from_shell_cmd(&snapshot.program, &snapshot.args, &snapshot.env),
5869 None,
5870 sender,
5871 &component.name,
5872 )?;
5873
5874 let cleanup_cmd = inst.cleanup();
5876 execute_native_command(
5877 &NativeCommand::from_shell_cmd(
5878 &cleanup_cmd.program,
5879 &cleanup_cmd.args,
5880 &cleanup_cmd.env,
5881 ),
5882 sudo_pw,
5883 sender,
5884 &component.name,
5885 )?;
5886 }
5887 "llama-cpp" => {
5889 use crate::installers::components::llama_cpp::{LlamaCppConfig, LlamaCppInstaller};
5890 let gpu_arch = ctx
5891 .env_exports
5892 .get("GPU_ARCH")
5893 .map(|v| v.trim().to_string())
5894 .filter(|v| !v.is_empty())
5895 .unwrap_or_else(|| {
5896 detect_gpu_arch().unwrap_or_else(|| {
5897 eprintln!("Warning: no GPU arch detected, falling back to gfx1100");
5898 "gfx1100".to_string()
5899 })
5900 });
5901 let channel = ctx
5902 .env_exports
5903 .get("ROCM_CHANNEL")
5904 .map(|v| v.trim().to_string())
5905 .filter(|v| !v.is_empty())
5906 .unwrap_or_else(|| "latest".to_string());
5907 let inst = LlamaCppInstaller::new(LlamaCppConfig {
5908 gpu_arch: gpu_arch.clone(),
5909 channel: channel.clone(),
5910 ..Default::default()
5911 });
5912 let home = ctx.user_home;
5913 let manifest = inst.check_latest_release();
5914 let strategy = inst.resolve_install_strategy(home, manifest.as_ref());
5915 match &strategy {
5916 crate::installers::components::llama_cpp::InstallStrategy::Prebuilt(plan) => {
5917 let _ = sender.send(InstallerEvent::Log(
5918 format!(
5919 "llama-cpp install strategy selected: prebuilt ({})",
5920 plan.arch
5921 ),
5922 false,
5923 ));
5924 if let Err(err) = inst.download_prebuilt_binary(plan) {
5925 let _ = sender.send(InstallerEvent::Log(
5926 format!(
5927 "llama-cpp prebuilt download failed; falling back to source: {}",
5928 err
5929 ),
5930 false,
5931 ));
5932 inst.run_source_install(home).map_err(|source_err| {
5933 anyhow::anyhow!(
5934 "llama-cpp prebuilt install failed and source fallback failed: {}",
5935 source_err
5936 )
5937 })?;
5938 }
5939 }
5940 crate::installers::components::llama_cpp::InstallStrategy::Source(plan) => {
5941 let _ = sender.send(InstallerEvent::Log(
5942 format!(
5943 "llama-cpp install strategy selected: source ({})",
5944 plan.reason
5945 ),
5946 false,
5947 ));
5948 inst.run_source_install(home).map_err(|err| {
5949 anyhow::anyhow!("llama-cpp source install failed: {}", err)
5950 })?;
5951 }
5952 }
5953
5954 let clone_dir = inst.build_dir_path().to_string();
5958 let verification =
5959 crate::installers::components::llama_cpp::verify_installed_binary(home, &clone_dir);
5960 let _ = sender.send(InstallerEvent::Log(
5961 format!(
5962 "llama-cpp post-install verification: {}",
5963 verification.summary
5964 ),
5965 false,
5966 ));
5967 if verification.stronger_check_passed {
5968 let _ = sender.send(InstallerEvent::Log(
5969 format!(
5970 "llama-cpp llama-bench: prefill={:.0} t/s, decode={:.0} t/s",
5971 verification.bench_prefill_tps.unwrap_or(0.0),
5972 verification.bench_decode_tps.unwrap_or(0.0),
5973 ),
5974 false,
5975 ));
5976 }
5977 if verification.rdna3_is_valid.unwrap_or(false) {
5978 let _ = sender.send(InstallerEvent::Log(
5979 format!(
5980 "llama-cpp RDNA3 proof: device=RDNA3, WMMA={:.0} t/s",
5981 verification.rdna3_wmma_tps.unwrap_or(0.0),
5982 ),
5983 false,
5984 ));
5985 }
5986 }
5987
5988 _ => {
5990 bail!(
5991 "No native installer dispatch for component '{}'. \
5992 This should not happen — is_native_component() should return false for this ID.",
5993 component.id
5994 );
5995 }
5996 }
5997
5998 let _ = sender.send(InstallerEvent::Progress {
5999 component_id: component.id.clone(),
6000 progress: 0.9,
6001 message: format!("{} installation complete", component.name),
6002 });
6003
6004 Ok(())
6005}
6006
6007fn run_script(
6008 component: &Component,
6009 script_path: &str,
6010 sudo_password: Option<String>,
6011 batch_mode: bool,
6012 install_method: &str,
6013 sender: &Sender<InstallerEvent>,
6014 input_rx: Arc<Mutex<Receiver<String>>>,
6015) -> Result<()> {
6016 let user_home = resolve_mlstack_user_home();
6017 let user_name = std::env::var("USER").unwrap_or_else(|_| "user".into());
6018 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,ROCR_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";
6019
6020 let venv_bin = PathBuf::from(&user_home).join("rocm_venv").join("bin");
6021 let local_bin = PathBuf::from(&user_home).join(".local").join("bin");
6022 let mut path = std::env::var("PATH").unwrap_or_default();
6023
6024 if local_bin.exists() {
6026 let local_str = local_bin.to_string_lossy();
6027 if !path.split(':').any(|entry| entry == local_str) {
6028 if path.is_empty() {
6029 path = local_str.to_string();
6030 } else {
6031 path = format!("{}:{}", path, local_str);
6032 }
6033 }
6034 }
6035
6036 if venv_bin.exists() {
6037 let venv_str = venv_bin.to_string_lossy();
6038 if !path.split(':').any(|entry| entry == venv_str) {
6039 if path.is_empty() {
6040 path = venv_str.to_string();
6041 } else {
6042 path = format!("{}:{}", venv_str, path);
6043 }
6044 }
6045 }
6046
6047 let env_exports = load_mlstack_env_exports(&user_home);
6048 let persistent_python = env_exports
6049 .get("MLSTACK_PYTHON_BIN")
6050 .map(|v| v.trim().to_string())
6051 .filter(|v| !v.is_empty());
6052 let python_bin = std::env::var("MLSTACK_ENV_NAME")
6056 .ok()
6057 .map(|n| n.trim().to_string())
6058 .filter(|n| !n.is_empty())
6059 .map(|n| {
6060 crate::platform::environment::mlstack_env_dir(&n)
6061 .join("bin")
6062 .join("python")
6063 .to_string_lossy()
6064 .to_string()
6065 })
6066 .or(persistent_python)
6067 .unwrap_or_else(resolve_python_bin);
6068 let force_reinstall = is_force_reinstall();
6069 let force_env_value = if force_reinstall { "true" } else { "false" };
6070 let sudo_password_env = sudo_password.clone();
6071 let user_home_meta = fs::metadata(&user_home).ok();
6072 #[cfg(unix)]
6073 let (mlstack_target_uid, mlstack_target_gid) = {
6074 use std::os::unix::fs::MetadataExt;
6075 (
6076 user_home_meta.as_ref().map(|meta| meta.uid().to_string()),
6077 user_home_meta.as_ref().map(|meta| meta.gid().to_string()),
6078 )
6079 };
6080 #[cfg(not(unix))]
6081 let (mlstack_target_uid, mlstack_target_gid): (Option<String>, Option<String>) = (None, None);
6082
6083 let component_needs_sudo = component.needs_sudo && needs_sudo();
6085
6086 let mut command = if component_needs_sudo {
6087 let mut cmd = Command::new("sudo");
6088 cmd.arg("-S")
6089 .arg("-p")
6090 .arg("")
6091 .arg(format!("--preserve-env={}", preserve_env))
6092 .arg("bash")
6093 .arg(script_path);
6094 cmd
6095 } else {
6096 let mut cmd = Command::new("bash");
6097 cmd.arg(script_path);
6098 cmd
6099 };
6100
6101 let mut pythonpath = std::env::var("PYTHONPATH").unwrap_or_default();
6102 let rocm_lib = "/opt/rocm/lib";
6103 if !pythonpath.split(':').any(|entry| entry == rocm_lib) {
6104 if pythonpath.is_empty() {
6105 pythonpath = rocm_lib.to_string();
6106 } else {
6107 pythonpath = format!("{}:{}", rocm_lib, pythonpath);
6108 }
6109 }
6110
6111 pythonpath = pythonpath
6113 .split(':')
6114 .filter(|part| !part.starts_with("/tmp/"))
6115 .collect::<Vec<_>>()
6116 .join(":");
6117
6118 let rocm_version = env_exports
6119 .get("ROCM_VERSION")
6120 .map(|v| v.trim().to_string())
6121 .filter(|v| !v.is_empty())
6122 .unwrap_or_else(detect_rocm_version);
6123 let gpu_arch = env_exports
6124 .get("GPU_ARCH")
6125 .map(|v| v.trim().to_string())
6126 .filter(|v| !v.is_empty())
6127 .or_else(|| {
6128 env_exports
6129 .get("PYTORCH_ROCM_ARCH")
6130 .map(|v| v.split(';').next().unwrap_or("").trim().to_string())
6131 .filter(|v| !v.is_empty())
6132 })
6133 .unwrap_or_else(|| {
6134 detect_gpu_arch().unwrap_or_else(|| {
6135 eprintln!("Warning: no GPU arch detected, falling back to gfx1100");
6136 "gfx1100".to_string()
6137 })
6138 });
6139 let py_rocm_arch = env_exports
6140 .get("PYTORCH_ROCM_ARCH")
6141 .map(|v| v.trim().to_string())
6142 .filter(|v| !v.is_empty())
6143 .unwrap_or_else(|| gpu_arch.clone());
6144 let mut gpu_archs = env_exports
6145 .get("GPU_ARCHS")
6146 .map(|v| v.trim().to_string())
6147 .filter(|v| !v.is_empty())
6148 .unwrap_or_else(|| gpu_arch.clone());
6149 let mut gpu_list = env_exports
6150 .get("HIP_VISIBLE_DEVICES")
6151 .map(|v| v.trim().to_string())
6152 .filter(|v| !v.is_empty())
6153 .or_else(|| {
6154 env_exports
6155 .get("CUDA_VISIBLE_DEVICES")
6156 .map(|v| v.trim().to_string())
6157 .filter(|v| !v.is_empty())
6158 })
6159 .unwrap_or_else(detect_gpu_list);
6160 let hsa_override = env_exports
6161 .get("HSA_OVERRIDE_GFX_VERSION")
6162 .map(|v| v.trim().to_string())
6163 .filter(|v| !v.is_empty());
6164 let aiter_jit_dir = PathBuf::from(&user_home)
6165 .join(".mlstack")
6166 .join("aiter")
6167 .join("jit");
6168
6169 if component.id == "aiter" {
6170 gpu_archs = py_rocm_arch
6171 .split(';')
6172 .next()
6173 .map(|v| v.trim().to_string())
6174 .filter(|v| !v.is_empty())
6175 .unwrap_or_else(|| gpu_arch.clone());
6176 gpu_list = prioritize_gpu_list_for_arch(&gpu_list, &gpu_archs);
6177 let _ = fs::create_dir_all(&aiter_jit_dir);
6178 }
6179
6180 command
6181 .stdin(Stdio::piped())
6182 .stdout(Stdio::piped())
6183 .stderr(Stdio::piped())
6184 .env("HOME", &user_home)
6185 .env("USER", &user_name)
6186 .env("LOGNAME", &user_name)
6187 .env("PATH", &path)
6188 .env("PYTHONPATH", &pythonpath)
6189 .env("MLSTACK_USER_HOME", &user_home)
6190 .env("MLSTACK_PYTHON_BIN", &python_bin)
6191 .env("UV_PYTHON", &python_bin)
6192 .env("ROCM_HOME", "/opt/rocm")
6193 .env("ROCM_PATH", "/opt/rocm")
6194 .env("ROCM_VERSION", &rocm_version)
6195 .env("ROCM_CHANNEL", "latest")
6196 .env("GPU_ARCH", &gpu_arch)
6197 .env("GPU_ARCHS", &gpu_archs)
6198 .env("PYTORCH_ROCM_ARCH", &py_rocm_arch)
6199 .env("HIP_PATH", "/opt/rocm")
6200 .env("HIP_ROOT_DIR", "/opt/rocm")
6201 .env("ROCM_ROOT", "/opt/rocm")
6202 .env("HIPCC_BIN_DIR", "/opt/rocm/bin")
6203 .env("HIP_VISIBLE_DEVICES", &gpu_list)
6204 .env("CUDA_VISIBLE_DEVICES", &gpu_list)
6205 .env("ROCR_VISIBLE_DEVICES", &gpu_list)
6206 .env("MLSTACK_BATCH_MODE", if batch_mode { "1" } else { "0" })
6207 .env("MLSTACK_INSTALL_METHOD", install_method)
6208 .env("INSTALL_METHOD", install_method)
6209 .env("PIP_BREAK_SYSTEM_PACKAGES", "1")
6210 .env("PIP_ROOT_USER_ACTION", "ignore")
6211 .env("UV_PIP_BREAK_SYSTEM_PACKAGES", "1")
6212 .env("UV_SYSTEM_PYTHON", "1")
6213 .env("FORCE", force_env_value)
6214 .env("PYTORCH_REINSTALL", force_env_value)
6215 .env("MLSTACK_FORCE_REINSTALL", force_env_value)
6216 .env(
6217 "MLSTACK_SKIP_TORCH_INSTALL",
6218 if component.id == "pytorch" { "0" } else { "1" },
6219 )
6220 .env("RUSTY_STACK", "true");
6221
6222 if let Some(uid) = mlstack_target_uid.as_deref() {
6223 command.env("MLSTACK_TARGET_UID", uid);
6224 }
6225 if let Some(gid) = mlstack_target_gid.as_deref() {
6226 command.env("MLSTACK_TARGET_GID", gid);
6227 }
6228
6229 if let Some(password) = sudo_password_env {
6230 command.env("MLSTACK_SUDO_PASSWORD", password);
6231 }
6232
6233 if force_reinstall {
6234 if let Ok(script_contents) = fs::read_to_string(script_path) {
6235 let supports_force_flag = script_contents.contains("--force)")
6236 || script_contents.contains("--force ")
6237 || script_contents.contains("\"--force\"")
6238 || script_contents.contains("'--force'");
6239 if supports_force_flag {
6240 command.arg("--force");
6241 }
6242 }
6243 }
6244
6245 if let Some(value) = hsa_override {
6246 command.env("HSA_OVERRIDE_GFX_VERSION", value);
6247 }
6248 if component.id == "aiter" {
6249 command.env("AITER_JIT_DIR", aiter_jit_dir);
6250 }
6251
6252 let mut child = command.spawn().with_context(|| {
6253 format!(
6254 "Failed to spawn installer script '{}' for {}. \
6255 Ensure bash is installed and the script file is accessible.",
6256 script_path, component.name
6257 )
6258 })?;
6259
6260 let mut child_stdin = child.stdin.take().with_context(|| {
6261 format!(
6262 "Failed to open stdin for script '{}' — \
6263 this usually indicates a system resource issue.",
6264 script_path
6265 )
6266 })?;
6267 if let Some(password) = sudo_password {
6268 let _ = child_stdin.write_all(password.as_bytes());
6269 let _ = child_stdin.write_all(b"\n");
6270 let _ = child_stdin.flush();
6271 }
6272
6273 if let Some(default_input) = default_component_input(&component.id) {
6275 let _ = child_stdin.write_all(default_input.as_bytes());
6276 let _ = child_stdin.flush();
6277 }
6278
6279 let stop_input = Arc::new(AtomicBool::new(false));
6280
6281 let stop_input_clone = Arc::clone(&stop_input);
6283 let input_thread = thread::spawn(move || {
6284 while !stop_input_clone.load(Ordering::Relaxed) {
6285 let rx = match input_rx.lock() {
6286 Ok(guard) => guard,
6287 Err(_) => break, };
6289 match rx.recv_timeout(Duration::from_millis(100)) {
6290 Ok(input) => {
6291 let _ = child_stdin.write_all(input.as_bytes());
6292 let _ = child_stdin.flush();
6293 }
6294 Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {
6295 continue;
6296 }
6297 Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => break,
6298 }
6299 }
6300 });
6301
6302 let sender_stdout = sender.clone();
6303 let component_id = component.id.clone();
6304 let component_name = component_name(component);
6305 let stdout_stream = match child.stdout.take() {
6306 Some(s) => s,
6307 None => {
6308 stop_input.store(true, Ordering::Relaxed);
6309 let _ = input_thread.join();
6310 let _ = sender.send(InstallerEvent::Log(
6311 format!(
6312 "[ERROR] {} — failed to capture stdout from script process",
6313 component.name
6314 ),
6315 false,
6316 ));
6317 bail!(
6318 "{} — failed to capture stdout from script process. \
6319 This usually indicates a system resource issue.",
6320 component.name
6321 );
6322 }
6323 };
6324 let stdout_handle = thread::spawn(move || {
6325 let mut reader = BufReader::new(stdout_stream);
6326 let mut buffer = Vec::new();
6327 let mut chunk = [0u8; 1024];
6328
6329 while let Ok(n) = reader.read(&mut chunk) {
6330 if n == 0 {
6331 break;
6332 }
6333 for &b in chunk.iter().take(n) {
6334 if b == b'\n' || b == b'\r' {
6335 let is_transient = b == b'\r';
6336 if !buffer.is_empty() {
6337 let line = String::from_utf8_lossy(&buffer).to_string();
6338 let clean_line = strip_ansi_codes(&line);
6339 if !clean_line.trim().is_empty() {
6340 let _ = sender_stdout
6341 .send(InstallerEvent::Log(clean_line.clone(), is_transient));
6342 if let Some(progress) = parse_progress(&clean_line) {
6343 let _ = sender_stdout.send(InstallerEvent::Progress {
6344 component_id: component_id.clone(),
6345 progress,
6346 message: format!(
6347 "{}: {}%",
6348 component_name,
6349 (progress * 100.0) as i32
6350 ),
6351 });
6352 }
6353 }
6354 buffer.clear();
6355 }
6356 } else {
6357 buffer.push(b);
6358 if buffer.len() > 8192 {
6359 buffer.clear();
6360 }
6361 }
6362 }
6363 }
6364 });
6365
6366 let sender_stderr = sender.clone();
6367 let stderr_stream = match child.stderr.take() {
6368 Some(s) => s,
6369 None => {
6370 stop_input.store(true, Ordering::Relaxed);
6371 let _ = input_thread.join();
6372 let _ = stdout_handle.join();
6373 let _ = sender.send(InstallerEvent::Log(
6374 format!(
6375 "[ERROR] {} — failed to capture stderr from script process",
6376 component.name
6377 ),
6378 false,
6379 ));
6380 bail!(
6381 "{} — failed to capture stderr from script process. \
6382 This usually indicates a system resource issue.",
6383 component.name
6384 );
6385 }
6386 };
6387 let stderr_handle = thread::spawn(move || {
6388 let mut reader = BufReader::new(stderr_stream);
6389 let mut buffer = Vec::new();
6390 let mut chunk = [0u8; 1024];
6391
6392 while let Ok(n) = reader.read(&mut chunk) {
6393 if n == 0 {
6394 break;
6395 }
6396 for &b in chunk.iter().take(n) {
6397 if b == b'\n' || b == b'\r' {
6398 let is_transient = b == b'\r';
6399 if !buffer.is_empty() {
6400 let line = String::from_utf8_lossy(&buffer).to_string();
6401 let clean_line = strip_ansi_codes(&line);
6402 if !clean_line.trim().is_empty() {
6403 let _ =
6404 sender_stderr.send(InstallerEvent::Log(clean_line, is_transient));
6405 }
6406 buffer.clear();
6407 }
6408 } else {
6409 buffer.push(b);
6410 if buffer.len() > 8192 {
6411 buffer.clear();
6412 }
6413 }
6414 }
6415 }
6416 });
6417
6418 let status = child.wait().with_context(|| {
6419 format!(
6420 "{} — script process wait failed. \
6421 The process may have been killed or the system is under heavy load.",
6422 component.name
6423 )
6424 })?;
6425 stop_input.store(true, Ordering::Relaxed);
6426 let _ = input_thread.join();
6427 let _ = stdout_handle.join();
6428 let _ = stderr_handle.join();
6429
6430 if !status.success() {
6431 let code = status.code().unwrap_or(-1);
6432 let error_detail = match code {
6433 1 => "General error — check the script output above for details".to_string(),
6434 2 => "Misuse of shell builtins — check script arguments".to_string(),
6435 126 => "Script not executable — check file permissions (chmod +x)".to_string(),
6436 127 => "Command not found in script — ensure required tools are installed".to_string(),
6437 130 => "Script terminated by Ctrl+C (SIGINT)".to_string(),
6438 137 => "Script killed (SIGKILL) — possibly out of memory".to_string(),
6439 139 => "Segmentation fault in script subprocess".to_string(),
6440 n if n > 128 => format!(
6441 "Script process terminated by signal {} — check system logs",
6442 n - 128
6443 ),
6444 _ => format!(
6445 "Script exited with code {} — check the output above for details",
6446 code
6447 ),
6448 };
6449 let _ = sender.send(InstallerEvent::Log(
6450 format!(
6451 "[ERROR] {} failed: {} (exit code {} at {})",
6452 component.name,
6453 error_detail,
6454 code,
6455 Local::now().format("%H:%M:%S")
6456 ),
6457 false,
6458 ));
6459 bail!(
6460 "{} failed: {} (exit code {})",
6461 component.name,
6462 error_detail,
6463 code
6464 );
6465 }
6466
6467 Ok(())
6468}
6469
6470fn star_repo(config: &InstallerConfig, sender: &Sender<InstallerEvent>) -> Result<()> {
6471 if !config.star_repos {
6472 return Ok(());
6473 }
6474
6475 let repo_url = "https://github.com/scooter-lacroix/Stan-s-ML-Stack";
6476
6477 let gh_check = Command::new("gh")
6479 .arg("auth")
6480 .arg("status")
6481 .stdout(Stdio::null())
6482 .stderr(Stdio::null())
6483 .status();
6484
6485 if let Ok(status) = gh_check {
6486 if status.success() {
6487 let _ = sender.send(InstallerEvent::Log(
6488 "Starring ML Stack repository on GitHub...".into(),
6489 false,
6490 ));
6491 let star_status = Command::new("gh")
6492 .arg("repo")
6493 .arg("star")
6494 .arg(repo_url)
6495 .stdout(Stdio::null())
6496 .stderr(Stdio::null())
6497 .status();
6498
6499 if let Ok(s) = star_status {
6500 if s.success() {
6501 let _ = sender.send(InstallerEvent::Log(
6502 "Thank you for starring the repo! ✦".into(),
6503 false,
6504 ));
6505 }
6506 }
6507 }
6508 }
6509
6510 Ok(())
6511}
6512
6513fn component_name(comp: &Component) -> String {
6514 comp.name.clone()
6515}
6516
6517fn format_user_friendly_error(component_id: &str, raw_error: &str) -> String {
6526 let error_lower = raw_error.to_lowercase();
6527
6528 match component_id {
6531 "permanent-env" => {
6532 if error_lower.contains("rocminfo") {
6533 return format!(
6534 "{}: Failed to detect GPU — rocminfo command failed. \
6535 Ensure ROCm is properly installed and /opt/rocm/bin is in your PATH. \
6536 Original error: {}",
6537 component_id, raw_error
6538 );
6539 }
6540 return format!(
6541 "{}: Environment setup failed. \
6542 Ensure ROCm is installed and GPU is accessible. \
6543 Original error: {}",
6544 component_id, raw_error
6545 );
6546 }
6547 "rocm" if error_lower.contains("no such file") || error_lower.contains("not found") => {
6548 return format!(
6549 "{}: ROCm installation failed — required packages or repositories not found. \
6550 Ensure your system's package manager is configured correctly and ROCm repositories are added. \
6551 Original error: {}",
6552 component_id, raw_error
6553 );
6554 }
6555 "pytorch" => {
6556 if error_lower.contains("no such file") || error_lower.contains("not found") {
6557 return format!(
6558 "{}: PyTorch installation failed — required packages not found. \
6559 Ensure ROCm is installed first and Python 3.10+ is available. \
6560 Original error: {}",
6561 component_id, raw_error
6562 );
6563 }
6564 return format!(
6565 "{}: PyTorch installation failed. \
6566 Ensure ROCm is installed first, Python 3.10-3.13 is available, \
6567 and pip/uv can reach the PyTorch package index. \
6568 Original error: {}",
6569 component_id, raw_error
6570 );
6571 }
6572 "triton" | "flash-attn" | "flash-attn-triton" | "flash-attn-ck" | "vllm" | "aiter" => {
6573 return format!(
6574 "{}: Installation failed. \
6575 Ensure PyTorch and ROCm are properly installed first \
6576 (these are required dependencies). \
6577 Original error: {}",
6578 component_id, raw_error
6579 );
6580 }
6581 "megatron" => {
6582 return format!(
6583 "{}: Megatron-LM installation failed. \
6584 Ensure PyTorch and MPI4Py are installed first (required dependencies). \
6585 Original error: {}",
6586 component_id, raw_error
6587 );
6588 }
6589 "deepspeed" => {
6590 return format!(
6591 "{}: DeepSpeed installation failed. \
6592 Ensure PyTorch is installed first (required dependency). \
6593 Original error: {}",
6594 component_id, raw_error
6595 );
6596 }
6597 _ => {}
6598 }
6599
6600 if error_lower.contains("permission denied") || error_lower.contains("sudo") {
6603 return format!(
6604 "{}: Permission denied. Try running with elevated privileges or check file/directory permissions. \
6605 Original error: {}",
6606 component_id, raw_error
6607 );
6608 }
6609
6610 if error_lower.contains("network")
6611 || error_lower.contains("connection")
6612 || error_lower.contains("timeout")
6613 || error_lower.contains("resolve")
6614 {
6615 return format!(
6616 "{}: Network error — failed to download packages or clone repositories. \
6617 Check your internet connection and try again. \
6618 Original error: {}",
6619 component_id, raw_error
6620 );
6621 }
6622
6623 if error_lower.contains("out of memory")
6624 || error_lower.contains("cannot allocate")
6625 || error_lower.contains("killed")
6626 {
6627 return format!(
6628 "{}: Out of memory — the build process was killed due to insufficient memory. \
6629 Try closing other applications or increasing swap space. \
6630 Original error: {}",
6631 component_id, raw_error
6632 );
6633 }
6634
6635 if error_lower.contains("exit code") {
6636 return raw_error.to_string();
6638 }
6639
6640 if error_lower.contains("no such file") || error_lower.contains("not found") {
6641 return format!(
6642 "{}: Required file or command not found. \
6643 Ensure all dependencies are installed before this component. \
6644 Original error: {}",
6645 component_id, raw_error
6646 );
6647 }
6648
6649 format!(
6652 "{}: Installation failed. Check the log output above for details. \
6653 Original error: {}",
6654 component_id, raw_error
6655 )
6656}
6657
6658pub(crate) fn detect_gpu_arch() -> Option<String> {
6663 crate::gpu::detect_discrete_amd_gpus()
6664 .first()
6665 .and_then(|g| g.gfx_arch.clone())
6666}
6667
6668pub(crate) fn detect_gpu_list() -> String {
6691 let gpus = detect_discrete_gpus();
6692 if gpus.is_empty() {
6693 "0".to_string()
6694 } else {
6695 gpus.join(",")
6696 }
6697}
6698
6699pub(crate) fn detect_discrete_gpus() -> Vec<String> {
6707 crate::gpu::detect_discrete_amd_gpus()
6708 .iter()
6709 .enumerate()
6710 .map(|(i, _)| i.to_string())
6711 .collect()
6712}
6713
6714fn detect_gpu_arch_by_id() -> HashMap<String, String> {
6716 let mut mapping: HashMap<String, String> = HashMap::new();
6717
6718 let rocminfo_path = crate::installers::common::utils::resolve_rocminfo_path();
6719 let Ok(output) = Command::new(&rocminfo_path).output() else {
6720 return mapping;
6721 };
6722 let stdout = String::from_utf8_lossy(&output.stdout);
6723 let mut current_gpu_id: Option<String> = None;
6724
6725 for line in stdout.lines() {
6726 let line = line.trim();
6727 if line.contains("GPU ID:") {
6728 if let Some(id_part) = line.split(':').nth(1) {
6729 let id: String = id_part.chars().filter(|c| c.is_ascii_digit()).collect();
6730 if !id.is_empty() {
6731 current_gpu_id = Some(id);
6732 }
6733 }
6734 continue;
6735 }
6736
6737 if !line.contains("Name:") || !line.contains("gfx") {
6738 continue;
6739 }
6740
6741 if let Some(id) = current_gpu_id.clone() {
6742 for token in line.split_whitespace() {
6743 let cleaned = token.trim_matches(|c: char| !c.is_alphanumeric());
6744 if cleaned.starts_with("gfx") {
6745 mapping.entry(id).or_insert_with(|| cleaned.to_string());
6746 break;
6747 }
6748 }
6749 }
6750 }
6751
6752 mapping
6753}
6754
6755fn prioritize_gpu_list_for_arch(gpu_list: &str, target_arch: &str) -> String {
6758 let target_arch = target_arch.trim();
6759 if target_arch.is_empty() {
6760 return gpu_list.to_string();
6761 }
6762
6763 let arch_by_id = detect_gpu_arch_by_id();
6764 if arch_by_id.is_empty() {
6765 return gpu_list.to_string();
6766 }
6767
6768 let mut ids: Vec<String> = gpu_list
6769 .split(',')
6770 .map(|v| v.trim().to_string())
6771 .filter(|v| !v.is_empty())
6772 .collect();
6773 if ids.len() < 2 {
6774 return gpu_list.to_string();
6775 }
6776
6777 let preferred_pos = ids.iter().position(|id| {
6778 arch_by_id
6779 .get(id)
6780 .map(|v| v == target_arch)
6781 .unwrap_or(false)
6782 });
6783 if let Some(pos) = preferred_pos {
6784 if pos > 0 {
6785 let preferred = ids.remove(pos);
6786 ids.insert(0, preferred);
6787 return ids.join(",");
6788 }
6789 }
6790
6791 gpu_list.to_string()
6792}
6793
6794#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6795enum VerificationResult {
6796 Verified,
6797 Failed,
6798 Missing,
6799}
6800
6801impl VerificationResult {
6802 fn label(&self) -> &'static str {
6803 match self {
6804 VerificationResult::Verified => "Verified",
6805 VerificationResult::Failed => "Failed",
6806 VerificationResult::Missing => "Missing",
6807 }
6808 }
6809}
6810
6811fn resolve_python_bin() -> String {
6812 if let Ok(value) = std::env::var("MLSTACK_PYTHON_BIN") {
6813 let value = value.trim();
6814 if !value.is_empty() {
6815 return value.to_string();
6816 }
6817 }
6818
6819 if let Ok(value) = std::env::var("UV_PYTHON") {
6820 let value = value.trim();
6821 if !value.is_empty() {
6822 return value.to_string();
6823 }
6824 }
6825
6826 let global_py = crate::platform::environment::mlstack_global_python();
6831 if global_py.exists() {
6832 return global_py.to_string_lossy().to_string();
6833 }
6834
6835 let home = std::env::var("HOME").unwrap_or_default();
6836
6837 let uv_python_dir = PathBuf::from(&home).join(".local/share/uv/python");
6842 if let Ok(entries) = std::fs::read_dir(&uv_python_dir) {
6843 let mut uv_pythons: Vec<PathBuf> = entries
6844 .filter_map(|e| e.ok())
6845 .filter_map(|e| {
6846 let bin = e.path().join("bin/python3");
6847 if bin.exists() {
6848 Some(bin)
6849 } else {
6850 None
6851 }
6852 })
6853 .collect();
6854 uv_pythons.sort_by(|a, b| {
6856 let va = python_version_tuple(a);
6857 let vb = python_version_tuple(b);
6858 python_ml_priority(&va).cmp(&python_ml_priority(&vb))
6859 });
6860 for python in &uv_pythons {
6862 if python_has_rocm_torch(&python.to_string_lossy()) {
6863 return python.to_string_lossy().to_string();
6864 }
6865 }
6866 if let Some(first) = uv_pythons.first() {
6868 return first.to_string_lossy().to_string();
6869 }
6870 }
6871
6872 let local_bin_python = format!("{home}/.local/bin/python3");
6874 if Path::new(&local_bin_python).exists() {
6875 if python_has_rocm_torch(&local_bin_python) {
6876 return local_bin_python;
6877 }
6878 return local_bin_python;
6879 }
6880
6881 let candidates = [
6883 Path::new(&home)
6884 .join("rocm_venv")
6885 .join("bin")
6886 .join("python"),
6887 PathBuf::from("/usr/local/bin/python3"),
6888 PathBuf::from("/usr/bin/python3"),
6889 PathBuf::from("python3"),
6890 PathBuf::from("python"),
6891 ];
6892
6893 for candidate in &candidates {
6894 let candidate_str = candidate.to_string_lossy();
6895 if candidate_str.is_empty() {
6896 continue;
6897 }
6898 if candidate.is_absolute() && !candidate.exists() {
6899 continue;
6900 }
6901 if python_has_rocm_torch(&candidate_str) {
6902 return candidate_str.to_string();
6903 }
6904 }
6905
6906 for candidate in &candidates {
6907 let candidate_str = candidate.to_string_lossy();
6908 if candidate.is_absolute() && !candidate.exists() {
6909 continue;
6910 }
6911 return candidate_str.to_string();
6912 }
6913
6914 "python3".to_string()
6915}
6916
6917fn python_version_tuple(path: &Path) -> (u32, u32) {
6919 let name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
6920 let parent = path
6922 .parent()
6923 .and_then(|p| p.parent())
6924 .and_then(|p| p.file_name())
6925 .and_then(|n| n.to_str())
6926 .unwrap_or("");
6927 if let Some(v) = parse_cpython_version(parent) {
6929 return v;
6930 }
6931 if let Some(rest) = name.strip_prefix("python") {
6933 let parts: Vec<&str> = rest.split('.').collect();
6934 if parts.len() >= 2 {
6935 if let (Ok(major), Ok(minor)) = (parts[0].parse::<u32>(), parts[1].parse::<u32>()) {
6936 return (major, minor);
6937 }
6938 }
6939 }
6940 (0, 0)
6941}
6942
6943fn parse_cpython_version(dir_name: &str) -> Option<(u32, u32)> {
6945 let rest = dir_name.strip_prefix("cpython-")?;
6946 let parts: Vec<&str> = rest.split('.').collect();
6947 if parts.len() >= 2 {
6948 let major = parts[0].parse::<u32>().ok()?;
6949 let minor = parts[1].parse::<u32>().ok()?;
6950 return Some((major, minor));
6951 }
6952 None
6953}
6954
6955fn python_ml_priority(version: &(u32, u32)) -> (i32, u32) {
6957 match version {
6958 (3, 12) => (0, version.1), (3, 13) => (1, version.1),
6960 (3, 11) => (2, version.1),
6961 (3, 10) => (3, version.1),
6962 _ => (4, version.1), }
6964}
6965
6966fn python_has_rocm_torch(python: &str) -> bool {
6967 Command::new(python)
6968 .arg("-c")
6969 .arg(
6970 "import torch, sys; ok=hasattr(torch,'__version__') and hasattr(torch,'version') and getattr(torch.version,'hip', None); sys.exit(0 if ok else 1)",
6971 )
6972 .stdout(Stdio::null())
6973 .stderr(Stdio::null())
6974 .status()
6975 .map(|status| status.success())
6976 .unwrap_or(false)
6977}
6978
6979fn default_component_input(component_id: &str) -> Option<&'static str> {
6980 match component_id {
6981 "ml-stack-core" => Some("9\n0\n"),
6982 _ => None,
6983 }
6984}
6985
6986fn parse_progress(line: &str) -> Option<f32> {
6987 let lower = line.to_lowercase();
6988 if let Some(percent_idx) = lower.find('%') {
6989 let prefix = &lower[..percent_idx];
6990 let digits: String = prefix
6991 .chars()
6992 .rev()
6993 .take_while(|c| c.is_ascii_digit())
6994 .collect();
6995 if !digits.is_empty() {
6996 let value: String = digits.chars().rev().collect();
6997 if let Ok(percent) = value.parse::<f32>() {
6998 return Some((percent / 100.0).min(1.0));
6999 }
7000 }
7001 }
7002 if lower.contains("progress") && lower.contains("100") {
7003 return Some(1.0);
7004 }
7005 if let Some(bracket) = line.find('[') {
7009 if let Some(close) = line[bracket..].find(']') {
7010 let inside = &line[bracket + 1..bracket + close];
7011 if let Some(slash) = inside.find('/') {
7012 let (n_str, m_str) = inside.split_at(slash);
7013 if let (Ok(n), Ok(m)) = (
7014 n_str.trim().parse::<f32>(),
7015 m_str[1..].trim().parse::<f32>(),
7016 ) {
7017 if m > 0.0 && n <= m {
7018 return Some((n / m).min(1.0));
7019 }
7020 }
7021 }
7022 }
7023 }
7024 None
7025}
7026
7027fn needs_sudo() -> bool {
7028 #[cfg(unix)]
7029 {
7030 unsafe { libc::geteuid() != 0 }
7031 }
7032 #[cfg(not(unix))]
7033 {
7034 false
7035 }
7036}
7037
7038fn strip_ansi_codes(s: &str) -> String {
7039 let mut clean = String::with_capacity(s.len());
7040 let mut in_escape = false;
7041 for c in s.chars() {
7042 if c == '\x1b' {
7043 in_escape = true;
7044 } else if in_escape {
7045 if ('@'..='~').contains(&c) || c == 'm' {
7046 in_escape = false;
7047 }
7048 } else {
7049 clean.push(c);
7050 }
7051 }
7052 clean
7053}
7054
7055fn sanitize_output_line(line: &str) -> String {
7056 let stripped = strip_ansi_codes(line);
7057 let mut clean = String::with_capacity(stripped.len());
7058 for c in stripped.chars() {
7059 if c == '\t' {
7060 clean.push(' ');
7061 } else if !c.is_control() {
7062 clean.push(c);
7063 }
7064 }
7065 clean
7066}
7067
7068fn should_suppress_verification_line(line: &str, target_id: &str) -> bool {
7069 let trimmed = line.trim();
7070 if trimmed.is_empty() {
7071 return true;
7072 }
7073
7074 let lower = trimmed.to_lowercase();
7075 if is_exception_or_traceback_line(&lower) {
7076 return false;
7077 }
7078
7079 if lower.contains("amdgpu.ids")
7080 || lower.contains("tool-lib")
7081 || lower.contains("cannot load tool")
7082 || lower.contains("libdrm/amdgpu.ids: no such file or directory")
7083 {
7084 return true;
7085 }
7086
7087 if (target_id == "rocm" || target_id == "rocm-smi")
7088 && (lower.starts_with("======")
7089 || lower.starts_with("====")
7090 || lower.starts_with("---")
7091 || lower.contains("the info is not accurate")
7092 || lower.contains("================================")
7093 || lower.starts_with('{')
7094 || lower.starts_with('[')
7095 || lower.starts_with(']')
7096 || lower.starts_with('}')
7097 || lower.starts_with("\"card")
7098 || lower.starts_with("device ")
7099 || lower.starts_with("*****"))
7100 {
7101 return true;
7102 }
7103
7104 false
7105}
7106
7107fn is_exception_or_traceback_line(lowercase_line: &str) -> bool {
7108 lowercase_line.contains("traceback")
7109 || lowercase_line.contains("exception")
7110 || lowercase_line.contains("segmentation fault")
7111 || lowercase_line.contains("fatal:")
7112 || lowercase_line.contains("error:")
7113 || lowercase_line.starts_with("error ")
7114}
7115
7116fn run_native_benchmark(component_id: &str, sender: &Sender<InstallerEvent>) -> Result<()> {
7121 use crate::benchmark_runners;
7122
7123 let bench_name = match component_id {
7124 "mlperf-inference" => "all",
7125 "rocm-benchmarks" => "gpu-capability",
7126 "gpu-memory-bandwidth" => "memory-bandwidth",
7127 "rocm-smi-bench" => "gpu-capability",
7128 "pytorch-performance" => "pytorch",
7129 "vllm-performance" => "vllm",
7130 "deepspeed-performance" => "deepspeed",
7131 "megatron-performance" => "megatron",
7132 "onnx-performance" => "onnx",
7133 "rusty-llama-performance" => "llama-cpp",
7134 "flash-attention-ck-performance" => "flash-attention-ck",
7135 "all-benchmarks" => "all",
7136 _ => bail!("Unknown benchmark component: {}", component_id),
7137 };
7138
7139 let _ = sender.send(InstallerEvent::Log(
7140 format!("[native] Running {} benchmark...", component_id),
7141 false,
7142 ));
7143
7144 let _ = sender.send(InstallerEvent::Progress {
7145 component_id: component_id.to_string(),
7146 progress: 0.3,
7147 message: format!("Running {} benchmark", component_id),
7148 });
7149
7150 match benchmark_runners::run_benchmark(bench_name) {
7151 Ok(output) => {
7152 let _ = sender.send(InstallerEvent::Log(
7153 format!(
7154 "[native] {} benchmark completed: {} ({} ms)",
7155 component_id,
7156 if output.success { "SUCCESS" } else { "FAILED" },
7157 output.execution_time_ms
7158 ),
7159 false,
7160 ));
7161 persist_benchmark_result(bench_name, &output);
7166 if !output.errors.is_empty() {
7167 for err in &output.errors {
7168 let _ = sender.send(InstallerEvent::Log(
7169 format!("[native] {} error: {}", component_id, err),
7170 false,
7171 ));
7172 }
7173 }
7174 if !output.success {
7175 let error_details = output.errors.join("; ");
7176 bail!(
7177 "{} benchmark failed: {}",
7178 component_id,
7179 if error_details.is_empty() {
7180 "unknown error".to_string()
7181 } else {
7182 error_details
7183 }
7184 );
7185 }
7186 Ok(())
7187 }
7188 Err(err) => {
7189 let _ = sender.send(InstallerEvent::Log(
7190 format!("[native] {} benchmark error: {}", component_id, err),
7191 false,
7192 ));
7193 bail!("{} benchmark error: {}", component_id, err)
7194 }
7195 }
7196}
7197
7198fn persist_onnx_install_status(version: &str, providers: &[&str]) {
7211 let provider = providers
7212 .iter()
7213 .find(|p| matches!(**p, "MIGraphXExecutionProvider" | "ROCMExecutionProvider"))
7214 .copied()
7215 .unwrap_or("none");
7216 persist_benchmark_result(
7217 "onnx",
7218 &crate::benchmark_runners::BenchmarkOutput {
7219 name: "onnx".to_string(),
7220 success: true,
7221 execution_time_ms: 0,
7222 results: serde_json::json!({
7223 "ort_version": version,
7224 "provider": provider,
7225 "providers_available": providers,
7226 "provider_priority": ["MIGraphXExecutionProvider", "ROCMExecutionProvider"],
7227 "install_status": "provider_ready",
7228 "benchmark_status": "not_run",
7229 "graph_opt_level": "ORT_ENABLE_ALL",
7230 "session_create_ms": 0.0,
7231 "inference_latency_p50_ms": 0.0,
7232 "inference_latency_p95_ms": 0.0,
7233 "inference_latency_p99_ms": 0.0,
7234 "throughput_inf_per_sec": 0.0,
7235 "peak_rss_mb": 0.0,
7236 "quantized_supported": false,
7237 "inference_samples": []
7238 }),
7239 errors: Vec::new(),
7240 },
7241 );
7242}
7243
7244fn persist_benchmark_result(bench_name: &str, output: &crate::benchmark_runners::BenchmarkOutput) {
7245 let metrics = output
7246 .results
7247 .get("metrics")
7248 .cloned()
7249 .unwrap_or_else(|| output.results.clone());
7250 let payload = serde_json::json!({
7251 "name": output.name.clone(),
7252 "success": output.success,
7253 "execution_time_ms": output.execution_time_ms as u64,
7254 "results": metrics,
7255 "errors": serde_json::to_value(&output.errors).unwrap_or_default(),
7256 });
7257 let body = match serde_json::to_string_pretty(&payload) {
7258 Ok(s) => s,
7259 Err(_) => return,
7260 };
7261 let ts = std::time::SystemTime::now()
7262 .duration_since(std::time::UNIX_EPOCH)
7263 .map(|d| d.as_secs())
7264 .unwrap_or(0);
7265 let fname = format!("{bench_name}_benchmarks_{ts}.json");
7266
7267 let mut candidates: Vec<std::path::PathBuf> =
7268 crate::benchmark_logs::benchmark_log_directories();
7269 if let Some(home) = std::env::var("HOME").ok().filter(|s| !s.is_empty()) {
7270 candidates.push(
7271 std::path::Path::new(&home)
7272 .join(".rusty-stack")
7273 .join("logs"),
7274 );
7275 }
7276 let tmp = std::env::var("TMPDIR").unwrap_or_else(|_| "/tmp".to_string());
7277 candidates.push(std::path::Path::new(&tmp).join("rusty-stack").join("logs"));
7278
7279 for dir in candidates {
7280 let _ = std::fs::create_dir_all(&dir);
7281 if std::fs::write(dir.join(&fname), &body).is_ok() {
7282 return;
7283 }
7284 }
7285}
7286
7287#[cfg(test)]
7288mod tests {
7289 use super::*;
7290 use std::sync::Mutex;
7291
7292 static ENV_LOCK: Mutex<()> = Mutex::new(());
7293
7294 #[test]
7300 fn test_safe_rocm_path_prepends_rocm_bin() {
7301 if !Path::new("/opt/rocm/bin").exists() {
7306 eprintln!("skip: /opt/rocm/bin absent (CI host) — prepend behavior is pass-through");
7307 return;
7308 }
7309 let p = safe_rocm_path("/home/scooter/.mlstack/global/bin:/usr/bin:/bin");
7310 let rocm = p.find("/opt/rocm/bin").unwrap();
7311 let shim = p.find("/home/scooter/.mlstack/global/bin").unwrap();
7312 assert!(
7313 rocm < shim,
7314 "/opt/rocm/bin must precede the venv bin: got {p}"
7315 );
7316 assert!(p.contains("/opt/rocm/hip/bin"));
7318 assert!(p.contains("/usr/bin") && p.contains("/bin"));
7319 }
7320
7321 #[test]
7322 fn test_safe_rocm_path_dedupes_rocm_entries() {
7323 if !Path::new("/opt/rocm/bin").exists() {
7327 eprintln!("skip: /opt/rocm/bin absent (CI host)");
7328 return;
7329 }
7330 let p = safe_rocm_path("/opt/rocm/bin:/usr/bin:/opt/rocm/hip/bin");
7331 assert_eq!(
7332 p.matches("/opt/rocm/bin").count(),
7333 1,
7334 "rocm/bin must appear exactly once: got {p}"
7335 );
7336 assert_eq!(p.matches("/opt/rocm/hip/bin").count(), 1);
7337 }
7338
7339 #[test]
7340 fn test_safe_rocm_path_passes_through_when_rocm_absent() {
7341 let inherited = "/usr/local/bin:/usr/bin";
7347 if Path::new("/opt/rocm/bin").exists() {
7348 assert!(safe_rocm_path(inherited).starts_with("/opt/rocm/bin"));
7349 } else {
7350 assert_eq!(safe_rocm_path(inherited), inherited);
7351 }
7352 }
7353
7354 #[test]
7363 fn test_format_user_friendly_error_permission_denied() {
7364 let msg = format_user_friendly_error("rocm", "Permission denied: /opt/rocm");
7365 assert!(
7366 msg.contains("Permission denied"),
7367 "Should mention permission issue"
7368 );
7369 assert!(
7370 msg.contains("elevated privileges"),
7371 "Should suggest running with sudo"
7372 );
7373 }
7374
7375 #[test]
7376 fn test_native_maintenance_actions_use_native_verdict() {
7377 let repair = Component {
7378 id: "rccl-repair".into(),
7379 name: "Repair RCCL Multi-GPU".into(),
7380 description: String::new(),
7381 script: String::new(),
7382 category: Category::Maintenance,
7383 required: false,
7384 selected: true,
7385 installed: false,
7386 progress: 0.0,
7387 estimate: String::new(),
7388 needs_sudo: false,
7389 experimental: false,
7390 note: None,
7391 };
7392 let verify = Component {
7393 id: "verify-basic".into(),
7394 name: "Verify Installation".into(),
7395 description: String::new(),
7396 script: String::new(),
7397 category: Category::Maintenance,
7398 required: false,
7399 selected: true,
7400 installed: false,
7401 progress: 0.0,
7402 estimate: String::new(),
7403 needs_sudo: false,
7404 experimental: false,
7405 note: None,
7406 };
7407 let normal = Component {
7408 id: "pytorch".into(),
7409 name: "PyTorch with ROCm".into(),
7410 description: String::new(),
7411 script: String::new(),
7412 category: Category::Foundation,
7413 required: true,
7414 selected: true,
7415 installed: false,
7416 progress: 0.0,
7417 estimate: String::new(),
7418 needs_sudo: false,
7419 experimental: false,
7420 note: None,
7421 };
7422
7423 assert!(is_native_maintenance_action(&repair));
7424 assert!(!is_native_maintenance_action(&verify));
7425 assert!(!is_native_maintenance_action(&normal));
7426
7427 let success = native_maintenance_action_outcome(&repair, true);
7428 assert!(success.success);
7429 assert!(success.report_lines[0].contains("reported success"));
7430
7431 let failure = native_maintenance_action_outcome(&repair, false);
7432 assert!(!failure.success);
7433 assert!(failure.report_lines[0].contains("reported failure"));
7434 }
7435
7436 #[test]
7437 fn test_benchmark_patterns_match_native_output_filenames() {
7438 assert_eq!(
7439 benchmark_pattern_for_target("rocm-benchmarks"),
7440 Some("gpu-capability_benchmarks")
7441 );
7442 assert_eq!(
7443 benchmark_pattern_for_target("gpu-memory-bandwidth"),
7444 Some("memory-bandwidth_benchmarks")
7445 );
7446 assert_eq!(
7447 benchmark_pattern_for_target("all-benchmarks"),
7448 Some("all_benchmarks")
7449 );
7450 assert_eq!(
7451 benchmark_pattern_for_target("onnx-performance"),
7452 Some("onnx_benchmarks")
7453 );
7454 }
7455
7456 #[test]
7457 fn test_single_benchmark_results_render_metrics_and_errors() {
7458 let lines = benchmark_result_lines_from_json(&serde_json::json!({
7459 "name": "vllm",
7460 "success": false,
7461 "execution_time_ms": 248842,
7462 "results": {
7463 "mode": "degraded",
7464 "reason": "model timed out",
7465 "throughput_tokens_per_sec": 0.0,
7466 "visible_devices": "0,1"
7467 },
7468 "errors": ["model timed out after 240 seconds"]
7469 }));
7470 let joined = lines.join("\n");
7471
7472 assert!(joined.contains("MODE: degraded"));
7473 assert!(joined.contains("REASON: model timed out"));
7474 assert!(joined.contains("THROUGHPUT TOKENS PER SEC: 0"));
7475 assert!(joined.contains("ERROR: model timed out after 240 seconds"));
7476 }
7477
7478 #[test]
7479 fn test_persist_onnx_install_status_writes_clean_status_artifact() {
7480 let _global_env = crate::test_support::lock_env();
7481 let _guard = ENV_LOCK.lock().unwrap();
7482 let old_home = std::env::var("HOME").ok();
7483 let old_tmpdir = std::env::var("TMPDIR").ok();
7484 let old_log_dir = std::env::var("MLSTACK_LOG_DIR").ok();
7485 let root = std::env::temp_dir().join(format!(
7486 "rusty-stack-persist-onnx-status-{}",
7487 std::process::id()
7488 ));
7489 let _ = std::fs::remove_dir_all(&root);
7490 let logs = root.join("logs");
7491 let tmp = root.join("tmp");
7492 std::fs::create_dir_all(&logs).unwrap();
7493 std::fs::create_dir_all(&tmp).unwrap();
7494 std::env::set_var("HOME", &root);
7495 std::env::set_var("TMPDIR", &tmp);
7496 std::env::set_var("MLSTACK_LOG_DIR", &logs);
7497
7498 persist_onnx_install_status(
7499 "1.27.1",
7500 &["MIGraphXExecutionProvider", "CPUExecutionProvider"],
7501 );
7502
7503 let written = std::fs::read_dir(&logs)
7504 .unwrap()
7505 .find_map(|entry| {
7506 let path = entry.ok()?.path();
7507 (path
7508 .file_name()?
7509 .to_string_lossy()
7510 .starts_with("onnx_benchmarks_"))
7511 .then_some(path)
7512 })
7513 .expect("onnx status log should be written");
7514 let json: serde_json::Value =
7515 serde_json::from_str(&std::fs::read_to_string(written).unwrap()).unwrap();
7516
7517 if let Some(home) = old_home {
7518 std::env::set_var("HOME", home);
7519 } else {
7520 std::env::remove_var("HOME");
7521 }
7522 if let Some(tmpdir) = old_tmpdir {
7523 std::env::set_var("TMPDIR", tmpdir);
7524 } else {
7525 std::env::remove_var("TMPDIR");
7526 }
7527 if let Some(log_dir) = old_log_dir {
7528 std::env::set_var("MLSTACK_LOG_DIR", log_dir);
7529 } else {
7530 std::env::remove_var("MLSTACK_LOG_DIR");
7531 }
7532 let _ = std::fs::remove_dir_all(&root);
7533
7534 assert_eq!(json["success"], true);
7535 assert_eq!(json["errors"].as_array().unwrap().len(), 0);
7536 assert_eq!(json["results"]["provider"], "MIGraphXExecutionProvider");
7537 assert_eq!(json["results"]["benchmark_status"], "not_run");
7538 }
7539
7540 #[test]
7541 fn test_persist_full_suite_keeps_nested_results() {
7542 let _global_env = crate::test_support::lock_env();
7543 let _guard = ENV_LOCK.lock().unwrap();
7544 let old_home = std::env::var("HOME").ok();
7545 let old_tmpdir = std::env::var("TMPDIR").ok();
7546 let old_log_dir = std::env::var("MLSTACK_LOG_DIR").ok();
7547 let root = std::env::temp_dir().join(format!(
7548 "rusty-stack-persist-full-suite-{}",
7549 std::process::id()
7550 ));
7551 let _ = std::fs::remove_dir_all(&root);
7552 let logs = root.join("logs");
7553 let tmp = root.join("tmp");
7554 std::fs::create_dir_all(&logs).unwrap();
7555 std::fs::create_dir_all(&tmp).unwrap();
7556 std::env::set_var("HOME", &root);
7557 std::env::set_var("TMPDIR", &tmp);
7558 std::env::set_var("MLSTACK_LOG_DIR", &logs);
7559
7560 let output = crate::benchmark_runners::BenchmarkOutput {
7561 name: "all".into(),
7562 success: true,
7563 execution_time_ms: 1,
7564 results: serde_json::json!({
7565 "memory_bandwidth": {
7566 "success": true,
7567 "execution_time_ms": 1,
7568 "metrics": {"hbm_peak_gb_s": 123.0}
7569 }
7570 }),
7571 errors: Vec::new(),
7572 };
7573 persist_benchmark_result("all", &output);
7574
7575 let written = std::fs::read_dir(&logs)
7576 .unwrap()
7577 .find_map(|entry| {
7578 let path = entry.ok()?.path();
7579 (path
7580 .file_name()?
7581 .to_string_lossy()
7582 .starts_with("all_benchmarks_"))
7583 .then_some(path)
7584 })
7585 .expect("all benchmark log should be written");
7586 let json: serde_json::Value =
7587 serde_json::from_str(&std::fs::read_to_string(written).unwrap()).unwrap();
7588
7589 if let Some(home) = old_home {
7590 std::env::set_var("HOME", home);
7591 } else {
7592 std::env::remove_var("HOME");
7593 }
7594 if let Some(tmpdir) = old_tmpdir {
7595 std::env::set_var("TMPDIR", tmpdir);
7596 } else {
7597 std::env::remove_var("TMPDIR");
7598 }
7599 if let Some(log_dir) = old_log_dir {
7600 std::env::set_var("MLSTACK_LOG_DIR", log_dir);
7601 } else {
7602 std::env::remove_var("MLSTACK_LOG_DIR");
7603 }
7604 let _ = std::fs::remove_dir_all(&root);
7605
7606 assert_eq!(
7607 json["results"]["memory_bandwidth"]["metrics"]["hbm_peak_gb_s"],
7608 123.0
7609 );
7610 }
7611
7612 #[test]
7613 fn test_format_user_friendly_error_not_found() {
7614 let msg = format_user_friendly_error("pytorch", "No such file or directory: python3");
7615 assert!(msg.contains("not found"), "Should mention file not found");
7616 assert!(
7617 msg.contains("ROCm") || msg.contains("Python"),
7618 "Should mention ROCm or Python dependency, got: {}",
7619 msg
7620 );
7621 }
7622
7623 #[test]
7624 fn test_format_user_friendly_error_network() {
7625 let msg = format_user_friendly_error(
7627 "onnxruntime",
7628 "network error: Could not resolve host: github.com",
7629 );
7630 assert!(
7631 msg.contains("Network error"),
7632 "Should mention network issue, got: {}",
7633 msg
7634 );
7635 assert!(
7636 msg.contains("internet connection"),
7637 "Should suggest checking internet, got: {}",
7638 msg
7639 );
7640 }
7641
7642 #[test]
7643 fn test_format_user_friendly_error_out_of_memory() {
7644 let msg = format_user_friendly_error("onnxruntime", "out of memory during build");
7646 assert!(
7647 msg.contains("Out of memory"),
7648 "Should mention memory issue, got: {}",
7649 msg
7650 );
7651 assert!(
7652 msg.contains("swap space"),
7653 "Should suggest increasing swap, got: {}",
7654 msg
7655 );
7656 }
7657
7658 #[test]
7659 fn test_format_user_friendly_error_permanent_env_rocminfo() {
7660 let msg = format_user_friendly_error("permanent-env", "rocminfo command not found");
7661 assert!(
7662 msg.contains("rocminfo"),
7663 "Should mention rocminfo specifically"
7664 );
7665 assert!(
7666 msg.contains("ROCm is properly installed"),
7667 "Should suggest ROCm installation"
7668 );
7669 }
7670
7671 #[test]
7672 fn test_format_user_friendly_error_pytorch_specific() {
7673 let msg = format_user_friendly_error("pytorch", "pip install failed: build error");
7675 assert!(msg.contains("PyTorch"), "Should mention PyTorch by name");
7676 assert!(
7677 msg.contains("ROCm is installed first"),
7678 "Should mention ROCm dependency"
7679 );
7680 assert!(
7681 msg.contains("Python 3.10-3.13"),
7682 "Should mention Python version, got: {}",
7683 msg
7684 );
7685 }
7686
7687 #[test]
7688 fn test_format_user_friendly_error_triton_dependency() {
7689 let msg = format_user_friendly_error("triton", "git clone failed");
7690 assert!(
7691 msg.contains("PyTorch and ROCm"),
7692 "Should mention required dependencies"
7693 );
7694 }
7695
7696 #[test]
7697 fn test_format_user_friendly_error_megatron_dependency() {
7698 let msg = format_user_friendly_error("megatron", "pip install failed");
7699 assert!(
7700 msg.contains("PyTorch and MPI4Py"),
7701 "Should mention Megatron's specific dependencies"
7702 );
7703 }
7704
7705 #[test]
7706 fn test_format_user_friendly_error_deepspeed_dependency() {
7707 let msg = format_user_friendly_error("deepspeed", "build failed");
7708 assert!(
7709 msg.contains("PyTorch"),
7710 "Should mention DeepSpeed's dependency on PyTorch"
7711 );
7712 }
7713
7714 #[test]
7715 fn test_format_user_friendly_error_exit_code_passthrough() {
7716 let msg = format_user_friendly_error("rocm", "ROCm failed: General error (exit code 1)");
7717 assert!(
7719 msg.contains("exit code 1"),
7720 "Should preserve exit code context"
7721 );
7722 }
7723
7724 #[test]
7725 fn test_format_user_friendly_error_unknown_component() {
7726 let msg = format_user_friendly_error("some-unknown-component", "something went wrong");
7727 assert!(
7728 msg.contains("Installation failed"),
7729 "Should have generic failure message"
7730 );
7731 assert!(
7732 msg.contains("something went wrong"),
7733 "Should include original error"
7734 );
7735 }
7736
7737 #[test]
7738 fn test_format_user_friendly_error_rocm_not_found() {
7739 let msg = format_user_friendly_error("rocm", "package not found: amdgpu-dkms");
7740 assert!(msg.contains("ROCm"), "Should mention ROCm specifically");
7741 assert!(msg.contains("repositories"), "Should mention repositories");
7742 }
7743
7744 #[test]
7745 fn test_execute_native_command_nonexistent_program() {
7746 let (tx, rx) = std::sync::mpsc::channel();
7749 let cmd = NativeCommand::Shell {
7750 program: "/nonexistent/program/that/does/not/exist".to_string(),
7751 args: vec![],
7752 env: vec![],
7753 working_dir: None,
7754 };
7755 let result = execute_native_command(&cmd, None, &tx, "test-component");
7756 assert!(result.is_err(), "Should fail for non-existent program");
7757 let err_msg = result.unwrap_err().to_string();
7758 assert!(!err_msg.is_empty(), "Error message should not be empty");
7759
7760 drop(tx);
7762 let events: Vec<_> = rx.try_iter().collect();
7763 assert!(
7764 !events.is_empty(),
7765 "Should have sent at least one event before failing"
7766 );
7767 }
7768
7769 #[test]
7770 fn test_execute_native_command_captures_exit_code() {
7771 let (tx, rx) = std::sync::mpsc::channel();
7773 let cmd = NativeCommand::Shell {
7774 program: "bash".to_string(),
7775 args: vec!["-c".to_string(), "echo test_output; exit 42".to_string()],
7776 env: vec![],
7777 working_dir: None,
7778 };
7779 let result = execute_native_command(&cmd, None, &tx, "test-component");
7780 assert!(
7781 result.is_err(),
7782 "Should fail for non-zero exit code, got: {:?}",
7783 result
7784 );
7785
7786 let err_msg = result.unwrap_err().to_string();
7787 assert!(
7788 err_msg.contains("42"),
7789 "Error should include exit code 42, got: {}",
7790 err_msg
7791 );
7792
7793 drop(tx);
7795 let events: Vec<_> = rx.try_iter().collect();
7796 let log_messages: Vec<&str> = events
7797 .iter()
7798 .filter_map(|e| match e {
7799 InstallerEvent::Log(msg, _) => Some(msg.as_str()),
7800 _ => None,
7801 })
7802 .collect();
7803 assert!(
7804 log_messages.iter().any(|m| m.contains("test_output")),
7805 "Should have captured stdout output before failure"
7806 );
7807 assert!(
7808 log_messages
7809 .iter()
7810 .any(|m| m.contains("[ERROR]") || m.contains("exit code")),
7811 "Should have logged error with exit code"
7812 );
7813 }
7814
7815 #[test]
7816 fn test_execute_native_command_success_sends_events() {
7817 let (tx, rx) = std::sync::mpsc::channel();
7819 let cmd = NativeCommand::Shell {
7820 program: "bash".to_string(),
7821 args: vec![
7822 "-c".to_string(),
7823 "echo 'hello world'; echo 'second line'".to_string(),
7824 ],
7825 env: vec![],
7826 working_dir: None,
7827 };
7828 let result = execute_native_command(&cmd, None, &tx, "test-component");
7829 assert!(result.is_ok(), "Should succeed for zero exit code");
7830
7831 drop(tx);
7832 let events: Vec<_> = rx.try_iter().collect();
7833 let log_messages: Vec<&str> = events
7834 .iter()
7835 .filter_map(|e| match e {
7836 InstallerEvent::Log(msg, _) => Some(msg.as_str()),
7837 _ => None,
7838 })
7839 .collect();
7840 assert!(
7841 log_messages.iter().any(|m| m.contains("hello world")),
7842 "Should have captured first line of output"
7843 );
7844 assert!(
7845 log_messages.iter().any(|m| m.contains("second line")),
7846 "Should have captured second line of output"
7847 );
7848 }
7849
7850 #[test]
7855 fn test_no_sudo_wrap_when_password_is_none() {
7856 let (tx, rx) = std::sync::mpsc::channel();
7860 let cmd = NativeCommand::Shell {
7861 program: "echo".to_string(),
7862 args: vec!["hello".to_string()],
7863 env: vec![],
7864 working_dir: None,
7865 };
7866 let result = execute_native_command(&cmd, None, &tx, "test-component");
7867 assert!(result.is_ok(), "Should succeed without sudo: {:?}", result);
7868
7869 drop(tx);
7870 let events: Vec<_> = rx.try_iter().collect();
7871 let log_messages: Vec<String> = events
7872 .iter()
7873 .filter_map(|e| match e {
7874 InstallerEvent::Log(msg, _) => Some(msg.clone()),
7875 _ => None,
7876 })
7877 .collect();
7878
7879 let cmd_log = log_messages.iter().find(|m| m.contains("[native] $"));
7881 assert!(cmd_log.is_some(), "Should have logged the command");
7882 let cmd_log = cmd_log.unwrap();
7883 assert!(
7884 !cmd_log.contains("sudo"),
7885 "Command should NOT be wrapped with sudo when sudo_pw is None, got: {}",
7886 cmd_log
7887 );
7888 assert!(
7889 cmd_log.contains("echo hello"),
7890 "Command should be 'echo hello', got: {}",
7891 cmd_log
7892 );
7893 }
7894
7895 #[test]
7896 fn test_sudo_wrap_when_password_is_some() {
7897 let (tx, rx) = std::sync::mpsc::channel();
7901 let cmd = NativeCommand::Shell {
7902 program: "echo".to_string(),
7903 args: vec!["hello".to_string()],
7904 env: vec![],
7905 working_dir: None,
7906 };
7907 let _result = execute_native_command(&cmd, Some("dummy_pw"), &tx, "test-component");
7909 drop(tx);
7912 let events: Vec<_> = rx.try_iter().collect();
7913 let log_messages: Vec<String> = events
7914 .iter()
7915 .filter_map(|e| match e {
7916 InstallerEvent::Log(msg, _) => Some(msg.clone()),
7917 _ => None,
7918 })
7919 .collect();
7920
7921 let has_sudo_log = log_messages.iter().any(|m| m.contains("Running with sudo"));
7923 assert!(
7924 has_sudo_log,
7925 "Should have logged 'Running with sudo' when sudo_pw is Some"
7926 );
7927 }
7928
7929 #[test]
7930 fn test_git_clone_without_sudo_works() {
7931 let (tx, rx) = std::sync::mpsc::channel();
7936 let cmd = NativeCommand::Shell {
7937 program: "git".to_string(),
7938 args: vec!["--version".to_string()],
7939 env: vec![],
7940 working_dir: None,
7941 };
7942 let result = execute_native_command(&cmd, None, &tx, "textgen");
7943 assert!(
7944 result.is_ok(),
7945 "git --version should succeed without sudo wrapping: {:?}",
7946 result
7947 );
7948
7949 drop(tx);
7950 let events: Vec<_> = rx.try_iter().collect();
7951 let log_messages: Vec<String> = events
7952 .iter()
7953 .filter_map(|e| match e {
7954 InstallerEvent::Log(msg, _) => Some(msg.clone()),
7955 _ => None,
7956 })
7957 .collect();
7958
7959 assert!(
7961 log_messages.iter().any(|m| m.contains("git version")),
7962 "Should have captured 'git version' output, got: {:?}",
7963 log_messages
7964 );
7965
7966 let cmd_log = log_messages.iter().find(|m| m.contains("[native] $"));
7968 assert!(cmd_log.is_some(), "Should have logged the command");
7969 assert!(
7970 !cmd_log.unwrap().contains("sudo"),
7971 "git command should NOT be wrapped with sudo"
7972 );
7973 }
7974
7975 #[test]
7976 fn test_textgen_git_clone_args_not_empty() {
7977 use crate::installers::components::textgen::{TextgenConfig, TextgenInstaller};
7980 let inst = TextgenInstaller::new(TextgenConfig {
7981 install_dir: "/tmp/test-textgen".to_string(),
7982 python_bin: "python3".to_string(),
7983 use_uv: false,
7984 break_system_packages: true,
7985 dry_run: false,
7986 });
7987 let cmd = inst.build_git_clone_command();
7988
7989 for (i, arg) in cmd.args.iter().enumerate() {
7991 assert!(
7992 !arg.is_empty(),
7993 "Argument {} should not be empty in git clone command",
7994 i
7995 );
7996 }
7997
7998 assert_eq!(cmd.program, "git");
8000 assert_eq!(
8001 cmd.args.len(),
8002 3,
8003 "git clone should have 3 args: clone, URL, dir"
8004 );
8005 assert_eq!(cmd.args[0], "clone");
8006 assert_eq!(
8007 cmd.args[1],
8008 "https://github.com/oobabooga/text-generation-webui.git"
8009 );
8010 assert_eq!(cmd.args[2], "/tmp/test-textgen");
8011
8012 assert!(
8014 !cmd.args.contains(&"--".to_string()),
8015 "git clone args should not contain '--' separator (from sudo wrapping)"
8016 );
8017 }
8018
8019 #[test]
8024 fn test_verification_result_labels() {
8025 assert_eq!(VerificationResult::Verified.label(), "Verified");
8026 assert_eq!(VerificationResult::Failed.label(), "Failed");
8027 assert_eq!(VerificationResult::Missing.label(), "Missing");
8028 }
8029}