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 =
808 if crate::bootstrap::env_setup::rocm_version_is_therock(&detect_rocm_version())
809 || std::env::var("ROCM_CHANNEL")
810 .ok()
811 .and_then(|c| crate::installers::components::rocm::RocmChannel::parse(&c))
812 .map(|c| c.is_therock() && rocm_sdk_venv_active())
813 .unwrap_or(false)
814 {
815 String::new()
816 } else {
817 crate::bootstrap::env_setup::GpuArchInfo::from_arch(gpu_arch).hsa_override_gfx_version
818 };
819
820 if env_path.exists() && !force {
826 let contents = fs::read_to_string(&env_path).context("Failed to read .mlstack_env")?;
827 let rocm_version = detect_rocm_version();
828 let rocm_home = "/opt/rocm";
829 let rocm_lib = "/opt/rocm/lib";
830 let defaults = [
831 ("ROCM_VERSION", rocm_version.as_str()),
832 ("ROCM_CHANNEL", "latest"),
833 ("GPU_ARCH", gpu_arch),
834 ("PYTORCH_ROCM_ARCH", gpu_arch),
835 ("GPU_ARCHS", gpu_arch),
836 ("ROCM_HOME", rocm_home),
837 ("ROCM_PATH", rocm_home),
838 ("HIP_PATH", rocm_home),
839 ("HSA_OVERRIDE_GFX_VERSION", hsa_override.as_str()),
840 ("HIP_VISIBLE_DEVICES", discrete_indices.as_str()),
841 ("CUDA_VISIBLE_DEVICES", discrete_indices.as_str()),
842 ("ROCR_VISIBLE_DEVICES", discrete_indices.as_str()),
847 ("PYTORCH_ROCM_DEVICE", primary_gpu),
848 ("MLSTACK_PYTHON_BIN", python_bin.as_str()),
849 ("UV_PYTHON", python_bin.as_str()),
850 ("MLSTACK_INSTALL_METHOD", normalized_install_method),
851 ("INSTALL_METHOD", normalized_install_method),
852 ("PYTHONPATH", rocm_lib),
853 ("PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:512"),
859 ("PYTORCH_HIP_ALLOC_CONF", "max_split_size_mb:512"),
863 ("AMDGPU_ASIC_ID_TABLE_PATH", "/usr/share/libdrm/amdgpu.ids"),
865 ("AMDGPU_ASIC_ID_TABLE_PATHS", "/usr/share/libdrm"),
866 ];
867 let (sanitized, _changed) = sanitize_mlstack_env(
868 &contents,
869 &defaults,
870 discrete_indices.as_str(),
871 rocm_home,
872 python_bin.as_str(),
873 normalized_install_method,
874 hsa_override.as_str(),
875 );
876 let updated = ensure_gpu_metadata_exports(&sanitized, &gpu_pairs);
879 if updated != contents {
880 fs::write(&env_path, &updated).context("Failed to update .mlstack_env")?;
881 return Ok(EnvUpdate::Updated);
882 }
883 return Ok(EnvUpdate::Unchanged);
884 }
885
886 if let Some(parent) = env_path.parent() {
887 fs::create_dir_all(parent).context("Failed to create home directory")?;
888 }
889
890 let rocm_version = detect_rocm_version();
891 let rocm_home = "/opt/rocm";
892
893 let mut content = crate::bootstrap::env_setup::generate_env_file_content(
897 &discrete_indices,
898 rocm_home,
899 &rocm_version,
900 "latest",
901 gpu_arch,
902 &hsa_override,
903 &python_bin,
904 );
905 content.push_str(&format!(
906 "\n# Rusty-Stack install metadata\nexport MLSTACK_INSTALL_METHOD={}\nexport INSTALL_METHOD={}\n",
907 normalized_install_method, normalized_install_method,
908 ));
909 content = ensure_gpu_metadata_exports(&content, &gpu_pairs);
911
912 fs::write(&env_path, &content).context("Failed to create .mlstack_env")?;
913
914 let fish_dir = PathBuf::from(user_home).join(".config/fish/conf.d");
917 if fs::create_dir_all(&fish_dir).is_ok() {
918 let fish_content = crate::bootstrap::env_setup::generate_fish_env_file_content(
919 &discrete_indices,
920 rocm_home,
921 &rocm_version,
922 "latest",
923 gpu_arch,
924 &hsa_override,
925 &python_bin,
926 );
927 let fish_path = fish_dir.join("mlstack_env.fish");
928 fs::write(&fish_path, fish_content)
929 .context("Failed to write fish conf.d/mlstack_env.fish")?;
930 }
931
932 fix_libdrm_amdgpu_ids();
936
937 Ok(EnvUpdate::Created)
938}
939
940fn gpu_identity_from_rocminfo(dgpus: &[crate::gpu::AmdGpu]) -> (String, String) {
947 let count = dgpus.len();
948 let text =
949 match std::process::Command::new(crate::installers::common::utils::resolve_rocminfo_path())
950 .output()
951 {
952 Ok(o) if o.status.success() => String::from_utf8_lossy(&o.stdout).to_string(),
953 _ => return gpu_identity_fallback(dgpus),
954 };
955 let mut entries: Vec<(String, u32)> = Vec::new();
956 let mut name = String::new();
957 let mut cus: u32 = 0;
958 let mut is_gpu = false;
959 let mut started = false;
960 for line in text.lines() {
961 let t = line.trim();
962 if t.starts_with("Agent ") && t.contains(char::is_numeric) {
963 if started && is_gpu && !name.is_empty() && !crate::gpu::is_integrated_gpu_name(&name) {
964 entries.push((name.clone(), cus));
965 }
966 name.clear();
967 cus = 0;
968 is_gpu = false;
969 started = false;
970 } else if let Some(v) = t.strip_prefix("Marketing Name:") {
971 name = v.trim().to_string();
972 started = true;
973 } else if let Some(v) = t.strip_prefix("Device Type:") {
974 is_gpu = v.trim().contains("GPU");
975 started = true;
976 } else if let Some(v) = t.strip_prefix("Compute Unit:") {
977 cus = v.trim().parse().unwrap_or(0);
978 started = true;
979 }
980 }
981 if started && is_gpu && !name.is_empty() && !crate::gpu::is_integrated_gpu_name(&name) {
982 entries.push((name, cus));
983 }
984 entries.truncate(count);
985 if entries.len() == count {
986 let names = entries
987 .iter()
988 .map(|(n, _)| n.clone())
989 .collect::<Vec<_>>()
990 .join(",");
991 let cus = entries
992 .iter()
993 .map(|(_, c)| c.to_string())
994 .collect::<Vec<_>>()
995 .join(",");
996 (names, cus)
997 } else {
998 gpu_identity_fallback(dgpus)
999 }
1000}
1001
1002fn gpu_identity_fallback(dgpus: &[crate::gpu::AmdGpu]) -> (String, String) {
1005 let names = dgpus
1006 .iter()
1007 .map(|g| {
1008 g.marketing_name
1009 .clone()
1010 .unwrap_or_else(|| format!("AMD GPU {}", g.pci_device_id))
1011 })
1012 .collect::<Vec<_>>()
1013 .join(",");
1014 let cus = dgpus
1015 .iter()
1016 .map(|_| "0".to_string())
1017 .collect::<Vec<_>>()
1018 .join(",");
1019 (names, cus)
1020}
1021
1022fn ensure_gpu_metadata_exports(contents: &str, pairs: &[(&str, &str)]) -> String {
1030 let mut lines: Vec<String> = contents.lines().map(|l| l.to_string()).collect();
1031 for (key, value) in pairs {
1032 let quoted = format!("export {key}=\"{value}\"");
1033 let prefix = format!("export {key}=");
1034 match lines
1035 .iter()
1036 .position(|l| l.trim_start().starts_with(&prefix))
1037 {
1038 Some(idx) => lines[idx] = quoted,
1039 None => lines.push(quoted),
1040 }
1041 }
1042 let mut out = lines.join("\n");
1043 if !out.ends_with('\n') {
1044 out.push('\n');
1045 }
1046 out
1047}
1048
1049fn sanitize_mlstack_env(
1050 contents: &str,
1051 defaults: &[(&str, &str)],
1052 gpu_list: &str,
1053 rocm_home: &str,
1054 python_bin: &str,
1055 install_method: &str,
1056 hsa_override: &str,
1057) -> (String, bool) {
1058 let mut changed = false;
1059 let mut lines = Vec::new();
1060 let mut seen = std::collections::HashSet::new();
1061 let default_for = |key: &str| -> &str {
1066 defaults
1067 .iter()
1068 .find(|(k, _)| *k == key)
1069 .map(|(_, v)| *v)
1070 .unwrap_or("")
1071 };
1072
1073 for line in contents.lines() {
1074 let trimmed = line.trim_start();
1075 if let Some(key) = extract_env_key(trimmed) {
1076 seen.insert(key.to_string());
1077 }
1078 if hsa_override.is_empty() && trimmed.starts_with("export HSA_OVERRIDE_GFX_VERSION=") {
1088 changed = true;
1089 continue;
1090 }
1091 if trimmed.starts_with("export TORCH_CUDA_ARCH_LIST=")
1092 || trimmed.starts_with("export PYTORCH_ALLOC_CONF=")
1093 {
1094 changed = true;
1095 continue;
1096 }
1097 let (line, changed_hip) = normalize_env_value(line, "HIP_PATH", "/opt/rocm");
1098 let (line, changed_rocm) = normalize_env_value(&line, "ROCM_PATH", "/opt/rocm");
1099 let (line, changed_python) = normalize_env_value(&line, "MLSTACK_PYTHON_BIN", python_bin);
1100 let (line, changed_uv_python) = normalize_env_value(&line, "UV_PYTHON", python_bin);
1101 let (line, changed_method) =
1102 normalize_env_value(&line, "MLSTACK_INSTALL_METHOD", install_method);
1103 let (line, changed_method_alias) =
1104 normalize_env_value(&line, "INSTALL_METHOD", install_method);
1105 let (line, changed_hsa) = if hsa_override.is_empty() {
1106 (line, false)
1107 } else {
1108 normalize_env_value(&line, "HSA_OVERRIDE_GFX_VERSION", hsa_override)
1109 };
1110 let (line, changed_hip_vis) =
1111 normalize_visible_devices(&line, "HIP_VISIBLE_DEVICES", gpu_list);
1112 let (line, changed_cuda_vis) =
1113 normalize_visible_devices(&line, "CUDA_VISIBLE_DEVICES", gpu_list);
1114 let (line, changed_rocr_vis) =
1115 normalize_visible_devices(&line, "ROCR_VISIBLE_DEVICES", gpu_list);
1116 let (line, changed_rocm_ver) =
1120 normalize_env_value(&line, "ROCM_VERSION", default_for("ROCM_VERSION"));
1121 let (line, changed_gpu_arch) =
1122 normalize_env_value(&line, "GPU_ARCH", default_for("GPU_ARCH"));
1123 let (line, changed_pytorch_arch) =
1124 normalize_env_value(&line, "PYTORCH_ROCM_ARCH", default_for("PYTORCH_ROCM_ARCH"));
1125 let (line, changed_gpu_archs) =
1126 normalize_env_value(&line, "GPU_ARCHS", default_for("GPU_ARCHS"));
1127 let (line, changed_rocm_home) =
1128 normalize_env_value(&line, "ROCM_HOME", default_for("ROCM_HOME"));
1129 let (line, changed_pytorch_dev) = normalize_env_value(
1130 &line,
1131 "PYTORCH_ROCM_DEVICE",
1132 default_for("PYTORCH_ROCM_DEVICE"),
1133 );
1134 let (line, changed_pythonpath) = if line.starts_with("export PYTHONPATH=") {
1135 let desired = format!("export PYTHONPATH={}/lib:$PYTHONPATH", rocm_home);
1136 if line.trim() != desired {
1137 (desired, true)
1138 } else {
1139 (line, false)
1140 }
1141 } else {
1142 (line, false)
1143 };
1144
1145 let (line, changed_path) = if line.starts_with("export PATH=") {
1147 let desired = format!("export PATH=\"/usr/local/bin:/usr/bin:/bin:/usr/local/games:/usr/games:{}/bin:{}/hip/bin:$PATH\"", rocm_home, rocm_home);
1148 if line.trim() != desired {
1149 (desired, true)
1150 } else {
1151 (line, false)
1152 }
1153 } else {
1154 (line, false)
1155 };
1156
1157 let (line, changed_ld) = if line.starts_with("export LD_LIBRARY_PATH=") {
1158 let desired = format!(
1159 "export LD_LIBRARY_PATH=\"$HOME/.mlstack/libmpi-compat:$HOME/.mlstack/libmpi-compat-user-$(id -u):{}/lib:{}/hip/lib:{}/opencl/lib:$LD_LIBRARY_PATH\"",
1160 rocm_home, rocm_home, rocm_home
1161 );
1162 if line.trim() != desired {
1163 (desired, true)
1164 } else {
1165 (line, false)
1166 }
1167 } else {
1168 (line, false)
1169 };
1170
1171 changed |= changed_hip
1172 || changed_rocm
1173 || changed_python
1174 || changed_uv_python
1175 || changed_method
1176 || changed_method_alias
1177 || changed_hsa
1178 || changed_hip_vis
1179 || changed_cuda_vis
1180 || changed_rocr_vis
1181 || changed_rocm_ver
1182 || changed_gpu_arch
1183 || changed_pytorch_arch
1184 || changed_gpu_archs
1185 || changed_rocm_home
1186 || changed_pytorch_dev
1187 || changed_pythonpath
1188 || changed_path
1189 || changed_ld;
1190 lines.push(line);
1191 }
1192
1193 for (key, value) in defaults {
1194 if !seen.contains(*key) {
1195 lines.push(format!("export {}={}", key, value));
1196 changed = true;
1197 }
1198 }
1199
1200 let mut output = lines.join("\n");
1201 output.push('\n');
1202 (output, changed)
1203}
1204
1205fn extract_env_key(line: &str) -> Option<&str> {
1206 let export_idx = line.rfind("export ")?;
1208 let rest = &line[export_idx + 7..]; rest.split('=').next().map(str::trim)
1210}
1211
1212fn normalize_env_value(line: &str, key: &str, desired: &str) -> (String, bool) {
1213 let marker = format!("export {}=", key);
1214 if let Some(export_idx) = line.rfind(&marker) {
1216 let after_marker = &line[export_idx + marker.len()..];
1217 let current = after_marker.trim().trim_matches('"');
1218 let current = strip_shell_suffixes_from_value(current);
1220 if current != desired {
1221 return (format!("export {}={}", key, desired), true);
1223 }
1224 }
1225 (line.to_string(), false)
1226}
1227
1228fn normalize_visible_devices(line: &str, key: &str, desired: &str) -> (String, bool) {
1229 let desired = desired.trim();
1230 if desired.is_empty() {
1231 return (line.to_string(), false);
1232 }
1233
1234 let marker = format!("export {}=", key);
1235 if let Some(export_idx) = line.rfind(&marker) {
1237 let after_marker = &line[export_idx + marker.len()..];
1238 let current = after_marker.trim().trim_matches('"');
1239 let current = strip_shell_suffixes_from_value(current);
1241
1242 if current != desired {
1248 return (format!("export {}={}", key, desired), true);
1250 }
1251 }
1252 (line.to_string(), false)
1253}
1254
1255fn strip_hsa_override_for_therock(command: &mut Command, env_exports: &HashMap<String, String>) {
1259 let version = env_exports
1260 .get("ROCM_VERSION")
1261 .cloned()
1262 .unwrap_or_else(detect_rocm_version);
1263 if crate::bootstrap::env_setup::rocm_version_is_therock(&version)
1264 || crate::bootstrap::env_setup::rocm_version_is_therock(&detect_rocm_version())
1265 {
1266 command.env_remove("HSA_OVERRIDE_GFX_VERSION");
1267 }
1268}
1269
1270fn rocm_sdk_venv_python() -> Option<std::path::PathBuf> {
1273 let home = std::env::var("HOME").ok()?;
1274 let entries = match fs::read_dir(format!("{home}/.mlstack/rocm")) {
1275 Ok(d) => d,
1276 Err(_) => return None,
1277 };
1278 entries
1279 .flatten()
1280 .map(|e| e.path().join("bin/python"))
1281 .find(|p| p.exists())
1282}
1283
1284fn rocm_sdk_venv_active() -> bool {
1285 rocm_sdk_venv_python().is_some()
1286}
1287
1288fn detect_rocm_version() -> String {
1289 if let Ok(version) = fs::read_to_string("/opt/rocm/.info/version") {
1290 return version.trim().to_string();
1291 }
1292 if let Some(py) = rocm_sdk_venv_python() {
1294 if let Ok(out) = Command::new(py)
1295 .args(["-c", "import importlib.metadata as m; print(m.version('rocm'))"])
1296 .output()
1297 {
1298 let v = String::from_utf8_lossy(&out.stdout).trim().to_string();
1299 if out.status.success() && !v.is_empty() {
1300 return v;
1301 }
1302 }
1303 }
1304 "7.2.0".to_string()
1305}
1306
1307fn run_verification(
1308 component: &Component,
1309 python_candidates: &[String],
1310 sender: &Sender<InstallerEvent>,
1311 user_home: &str,
1312) -> VerificationOutcome {
1313 let steps = verification_commands(&component.id, python_candidates);
1314 let mut results: HashMap<String, VerificationCheckResult> = HashMap::new();
1315 let mut executed = 0usize;
1316 let mut failed = 0usize;
1317
1318 for step in &steps {
1319 if let Some((benchmark_ok, benchmark_output)) = verify_benchmark_target(&step.target_id) {
1320 executed += 1;
1321 if benchmark_ok {
1322 results.insert(
1323 step.target_id.clone(),
1324 VerificationCheckResult {
1325 status: VerificationResult::Verified,
1326 output: benchmark_output,
1327 },
1328 );
1329 } else {
1330 failed += 1;
1331 results.insert(
1332 step.target_id.clone(),
1333 VerificationCheckResult {
1334 status: VerificationResult::Failed,
1335 output: benchmark_output,
1336 },
1337 );
1338 let _ = sender.send(InstallerEvent::Log(
1339 format!("Verification failed for {}", step.label),
1340 false,
1341 ));
1342 }
1343 continue;
1344 }
1345
1346 if !is_component_installed_by_id(&step.target_id, python_candidates) {
1347 results.insert(
1348 step.target_id.clone(),
1349 VerificationCheckResult {
1350 status: VerificationResult::Missing,
1351 output: vec!["Component not found in path or site-packages".to_string()],
1352 },
1353 );
1354 let _ = sender.send(InstallerEvent::Log(
1355 format!("Skipping {} (not installed)", step.label),
1356 false,
1357 ));
1358 continue;
1359 }
1360 if !modules_available(&step.modules, python_candidates) {
1361 results.insert(
1362 step.target_id.clone(),
1363 VerificationCheckResult {
1364 status: VerificationResult::Missing,
1365 output: vec![format!(
1366 "Required Python modules missing: {}",
1367 step.modules.join(", ")
1368 )],
1369 },
1370 );
1371 let _ = sender.send(InstallerEvent::Log(
1372 format!("Skipping {} (module not available)", step.label),
1373 false,
1374 ));
1375 continue;
1376 }
1377 executed += 1;
1378 let _ = sender.send(InstallerEvent::Log(
1379 format!("Verifying {}...", step.label),
1380 false,
1381 ));
1382 match run_verification_command(step, sender) {
1383 Ok((success, output)) => {
1384 if success {
1385 results.insert(
1386 step.target_id.clone(),
1387 VerificationCheckResult {
1388 status: VerificationResult::Verified,
1389 output,
1390 },
1391 );
1392 } else {
1393 failed += 1;
1394 results.insert(
1395 step.target_id.clone(),
1396 VerificationCheckResult {
1397 status: VerificationResult::Failed,
1398 output,
1399 },
1400 );
1401 let _ = sender.send(InstallerEvent::Log(
1402 format!("Verification failed for {}", step.label),
1403 false,
1404 ));
1405 }
1406 }
1407 Err(err) => {
1408 failed += 1;
1409 results.insert(
1410 step.target_id.clone(),
1411 VerificationCheckResult {
1412 status: VerificationResult::Failed,
1413 output: vec![err.to_string()],
1414 },
1415 );
1416 let _ = sender.send(InstallerEvent::Log(
1417 format!("Failed to execute verification for {}: {}", step.label, err),
1418 false,
1419 ));
1420 }
1421 }
1422 }
1423
1424 let report_lines = build_verification_report(python_candidates, user_home, &results, &steps);
1425 let success = executed > 0 && failed == 0;
1426
1427 if executed == 0 {
1428 let _ = sender.send(InstallerEvent::Log(
1429 "No installed components found to verify".into(),
1430 false,
1431 ));
1432 }
1433
1434 if failed > 0 {
1435 let _ = sender.send(InstallerEvent::Log(
1436 format!("Verification completed with {} failed checks", failed),
1437 false,
1438 ));
1439 }
1440
1441 VerificationOutcome {
1442 success,
1443 report_lines,
1444 }
1445}
1446
1447fn build_verification_report(
1448 python_candidates: &[String],
1449 user_home: &str,
1450 results: &HashMap<String, VerificationCheckResult>,
1451 steps: &[VerificationCommand],
1452) -> Vec<String> {
1453 let mut lines = Vec::new();
1454 lines.push("Environment:".to_string());
1455 for (key, value) in collect_env_info(user_home) {
1456 lines.push(format!(" {}: {}", key, value));
1457 }
1458
1459 lines.push("".to_string());
1460 lines.push("Python Interpreters:".to_string());
1461 for py in python_candidates {
1462 lines.push(format!(" - {}", py));
1463 }
1464
1465 lines.push("".to_string());
1466 lines.push("Checks:".to_string());
1467 for step in steps {
1468 if let Some(res) = results.get(&step.target_id) {
1469 lines.push(format!(" - {}: {}", step.label, res.status.label()));
1470 for out_line in &res.output {
1471 lines.push(format!(" > {}", out_line));
1472 }
1473 } else {
1474 lines.push(format!(" - {}: Skipped", step.label));
1475 }
1476 }
1477
1478 lines
1479}
1480
1481fn run_component_verification(
1482 component: &Component,
1483 python_candidates: &[String],
1484 sender: &Sender<InstallerEvent>,
1485 user_home: &str,
1486) -> VerificationOutcome {
1487 let steps = component_verification_commands(&component.id, python_candidates);
1488 let mut results: HashMap<String, VerificationCheckResult> = HashMap::new();
1489 let mut executed = 0usize;
1490 let mut failed = 0usize;
1491
1492 for step in &steps {
1493 if let Some((benchmark_ok, benchmark_output)) = verify_benchmark_target(&step.target_id) {
1494 executed += 1;
1495 if benchmark_ok {
1496 results.insert(
1497 step.target_id.clone(),
1498 VerificationCheckResult {
1499 status: VerificationResult::Verified,
1500 output: benchmark_output,
1501 },
1502 );
1503 } else {
1504 failed += 1;
1505 results.insert(
1506 step.target_id.clone(),
1507 VerificationCheckResult {
1508 status: VerificationResult::Failed,
1509 output: benchmark_output,
1510 },
1511 );
1512 }
1513 continue;
1514 }
1515
1516 if !modules_available(&step.modules, python_candidates) {
1517 results.insert(
1518 step.target_id.clone(),
1519 VerificationCheckResult {
1520 status: VerificationResult::Missing,
1521 output: vec![format!(
1522 "Required Python modules missing: {}",
1523 step.modules.join(", ")
1524 )],
1525 },
1526 );
1527 continue;
1528 }
1529 executed += 1;
1530 match run_verification_command(step, sender) {
1531 Ok((success, output)) => {
1532 if success {
1533 results.insert(
1534 step.target_id.clone(),
1535 VerificationCheckResult {
1536 status: VerificationResult::Verified,
1537 output,
1538 },
1539 );
1540 } else {
1541 failed += 1;
1542 results.insert(
1543 step.target_id.clone(),
1544 VerificationCheckResult {
1545 status: VerificationResult::Failed,
1546 output,
1547 },
1548 );
1549 }
1550 }
1551 Err(err) => {
1552 failed += 1;
1553 results.insert(
1554 step.target_id.clone(),
1555 VerificationCheckResult {
1556 status: VerificationResult::Failed,
1557 output: vec![err.to_string()],
1558 },
1559 );
1560 }
1561 }
1562 }
1563
1564 let report_lines = build_component_report(component, user_home, &results, &steps);
1565 let success = executed > 0 && failed == 0;
1566
1567 VerificationOutcome {
1568 success,
1569 report_lines,
1570 }
1571}
1572
1573fn is_native_maintenance_action(component: &Component) -> bool {
1574 component.category == Category::Maintenance && !component.id.starts_with("verify-")
1575}
1576
1577fn native_maintenance_action_outcome(
1578 component: &Component,
1579 install_success: bool,
1580) -> VerificationOutcome {
1581 VerificationOutcome {
1582 success: install_success,
1583 report_lines: vec![format!(
1584 "{} native action {}",
1585 component.name,
1586 if install_success {
1587 "reported success"
1588 } else {
1589 "reported failure"
1590 }
1591 )],
1592 }
1593}
1594
1595fn build_component_report(
1596 component: &Component,
1597 user_home: &str,
1598 results: &HashMap<String, VerificationCheckResult>,
1599 steps: &[VerificationCommand],
1600) -> Vec<String> {
1601 let mut lines = Vec::new();
1602 lines.push(format!("Component: {}", component.name));
1603
1604 for (key, value) in collect_env_info(user_home) {
1605 lines.push(format!(" {}: {}", key, value));
1606 }
1607
1608 if let Some(benchmark_lines) = extract_benchmark_results(&component.id) {
1610 lines.push("".to_string());
1611 lines.push("Benchmark Results:".to_string());
1612 lines.extend(benchmark_lines);
1613 }
1614
1615 lines.push("".to_string());
1616 lines.push("Verification Checks:".to_string());
1617
1618 for step in steps {
1619 if let Some(res) = results.get(&step.target_id) {
1620 lines.push(format!(" - {}: {}", step.label, res.status.label()));
1621 for out_line in &res.output {
1622 lines.push(format!(" > {}", out_line));
1623 }
1624 } else {
1625 lines.push(format!(" - {}: Skipped", step.label));
1626 }
1627 }
1628
1629 lines
1630}
1631
1632fn extract_benchmark_results(component_id: &str) -> Option<Vec<String>> {
1633 let log_dirs = benchmark_log_directories();
1634 if log_dirs.is_empty() {
1635 return None;
1636 }
1637
1638 let pattern = match component_id {
1640 "mlperf-inference" => "all_benchmarks",
1641 "rocm-benchmarks" => "gpu-capability_benchmarks",
1642 "gpu-memory-bandwidth" => "memory-bandwidth_benchmarks",
1643 "rocm-smi-bench" => "gpu-capability_benchmarks",
1644 "pytorch-performance" => "pytorch_performance",
1645 "vllm-performance" => "vllm_benchmarks",
1646 "deepspeed-performance" => "deepspeed_benchmarks",
1647 "megatron-performance" => "megatron_benchmarks",
1648 "onnx-performance" => "onnx_benchmarks",
1649 "all-benchmarks" => "all_benchmarks",
1650 _ => return None,
1651 };
1652
1653 let log_file = find_latest_log_in_dirs(&log_dirs, pattern)?;
1654 let contents = std::fs::read_to_string(log_file).ok()?;
1655
1656 let Some(json) = extract_benchmark_json_value(&contents) else {
1657 return Some(vec![" No metrics found in log".to_string()]);
1658 };
1659
1660 let mut result_lines = benchmark_result_lines_from_json(&json);
1661 if result_lines.is_empty() {
1662 result_lines.push(" No metrics found in log".to_string());
1663 }
1664
1665 Some(result_lines)
1666}
1667
1668fn benchmark_result_lines_from_json(json: &serde_json::Value) -> Vec<String> {
1669 let mut result_lines = Vec::new();
1670
1671 if let Some(results) = json.get("results").and_then(|r| r.as_object()) {
1674 let nested_suite = results.values().any(|value| {
1675 value.as_object().is_some_and(|obj| {
1676 obj.contains_key("success")
1677 || obj.contains_key("metrics")
1678 || obj.contains_key("execution_time_ms")
1679 })
1680 });
1681 if nested_suite {
1682 for (bench_name, bench_data) in results {
1683 if bench_data.is_object() {
1684 let status = bench_data
1685 .get("success")
1686 .and_then(|v| v.as_bool())
1687 .map(|ok| if ok { "SUCCESS" } else { "FAILED" })
1688 .unwrap_or("UNKNOWN");
1689 let time_ms = bench_data
1690 .get("execution_time_ms")
1691 .and_then(|v| v.as_u64())
1692 .unwrap_or(0);
1693 result_lines.push(format!(
1694 " === {} ({}, {} ms) ===",
1695 bench_name.replace('_', " ").to_uppercase(),
1696 status,
1697 time_ms
1698 ));
1699 if let Some(metrics) = bench_data.get("metrics").and_then(|m| m.as_object()) {
1700 append_benchmark_metric_lines(&mut result_lines, metrics, " ");
1701 }
1702 append_benchmark_error_lines(&mut result_lines, bench_data, " ");
1703 result_lines.push(String::new());
1704 }
1705 }
1706 } else {
1707 append_benchmark_metric_lines(&mut result_lines, results, " ");
1708 append_benchmark_error_lines(&mut result_lines, json, " ");
1709 }
1710 } else if let Some(metrics) = json.get("metrics").and_then(|m| m.as_object()) {
1711 append_benchmark_metric_lines(&mut result_lines, metrics, " ");
1713 append_benchmark_error_lines(&mut result_lines, json, " ");
1714 }
1715
1716 result_lines
1717}
1718
1719fn append_benchmark_metric_lines(
1720 lines: &mut Vec<String>,
1721 metrics: &serde_json::Map<String, serde_json::Value>,
1722 indent: &str,
1723) {
1724 for (key, value) in metrics {
1725 let value_str = match value {
1726 serde_json::Value::Number(n) => n.to_string(),
1727 serde_json::Value::String(s) => s.clone(),
1728 serde_json::Value::Bool(b) => b.to_string(),
1729 serde_json::Value::Array(arr) => format!("{} values", arr.len()),
1730 _ => continue,
1731 };
1732 let formatted_key = key.replace('_', " ").to_uppercase();
1733 lines.push(format!("{}{}: {}", indent, formatted_key, value_str));
1734 }
1735}
1736
1737fn append_benchmark_error_lines(lines: &mut Vec<String>, value: &serde_json::Value, indent: &str) {
1738 if let Some(errors) = value.get("errors").and_then(|v| v.as_array()) {
1739 for err in errors.iter().filter_map(|v| v.as_str()).take(3) {
1740 lines.push(format!("{}ERROR: {}", indent, err));
1741 }
1742 }
1743}
1744
1745fn benchmark_pattern_for_target(target_id: &str) -> Option<&'static str> {
1746 match target_id {
1747 "mlperf-inference" => Some("all_benchmarks"),
1748 "rocm-benchmarks" => Some("gpu-capability_benchmarks"),
1749 "gpu-memory-bandwidth" => Some("memory-bandwidth_benchmarks"),
1750 "rocm-smi-bench" => Some("gpu-capability_benchmarks"),
1751 "pytorch-performance" => Some("pytorch_performance"),
1752 "vllm-performance" => Some("vllm_benchmarks"),
1753 "deepspeed-performance" => Some("deepspeed_benchmarks"),
1754 "megatron-performance" => Some("megatron_benchmarks"),
1755 "onnx-performance" => Some("onnx_benchmarks"),
1756 "all-benchmarks" => Some("all_benchmarks"),
1757 _ => None,
1758 }
1759}
1760
1761fn benchmark_success_and_errors(value: &serde_json::Value) -> (bool, Vec<String>) {
1762 if let Some(success) = value.get("success").and_then(|v| v.as_bool()) {
1763 let mut details = Vec::new();
1764 if !success {
1765 if let Some(reason) = value.get("reason").and_then(|v| v.as_str()) {
1766 let reason = reason.trim();
1767 if !reason.is_empty() {
1768 details.push(reason.to_string());
1769 }
1770 }
1771 if let Some(errors) = value.get("errors").and_then(|v| v.as_array()) {
1772 for err in errors.iter().filter_map(|entry| entry.as_str()).take(4) {
1773 let trimmed = err.trim();
1774 if !trimmed.is_empty() {
1775 details.push(trimmed.to_string());
1776 }
1777 }
1778 }
1779 if details.is_empty() {
1780 details.push("Benchmark JSON reported success=false".to_string());
1781 }
1782 }
1783 return (success, details);
1784 }
1785
1786 if let Some(results) = value.get("results").and_then(|v| v.as_object()) {
1787 let mut failed = Vec::new();
1788 let mut seen_result = false;
1789 for (name, item) in results {
1790 if let Some(item_success) = item.get("success").and_then(|v| v.as_bool()) {
1791 seen_result = true;
1792 if !item_success {
1793 let mut reason = item
1794 .get("reason")
1795 .and_then(|v| v.as_str())
1796 .map(|s| s.trim().to_string())
1797 .filter(|s| !s.is_empty())
1798 .unwrap_or_else(|| format!("{name}: success=false"));
1799 if reason == format!("{name}: success=false") {
1800 if let Some(errors) = item.get("errors").and_then(|v| v.as_array()) {
1801 if let Some(first) = errors.iter().filter_map(|v| v.as_str()).next() {
1802 let first = first.trim();
1803 if !first.is_empty() {
1804 reason = format!("{name}: {first}");
1805 }
1806 }
1807 }
1808 }
1809 failed.push(reason);
1810 }
1811 }
1812 }
1813 if seen_result {
1814 return (failed.is_empty(), failed);
1815 }
1816 }
1817
1818 (
1819 false,
1820 vec!["Benchmark JSON does not contain a recognizable success field".to_string()],
1821 )
1822}
1823
1824fn verify_benchmark_target(target_id: &str) -> Option<(bool, Vec<String>)> {
1825 let pattern = benchmark_pattern_for_target(target_id)?;
1826 let log_dirs = benchmark_log_directories();
1827 if log_dirs.is_empty() {
1828 return Some((
1829 false,
1830 vec!["No benchmark log directories found".to_string()],
1831 ));
1832 }
1833
1834 let Some(log_file) = find_latest_log_in_dirs(&log_dirs, pattern) else {
1835 return Some((
1836 false,
1837 vec![format!("No benchmark logs found for pattern '{pattern}'")],
1838 ));
1839 };
1840
1841 let mut output = vec![format!("Latest benchmark log: {}", log_file.display())];
1842 let contents = match std::fs::read_to_string(&log_file) {
1843 Ok(contents) => contents,
1844 Err(err) => {
1845 output.push(format!("Could not read benchmark log: {err}"));
1846 return Some((false, output));
1847 }
1848 };
1849
1850 let Some(json) = extract_benchmark_json_value(&contents) else {
1851 output.push("Could not parse benchmark JSON payload from log".to_string());
1852 return Some((false, output));
1853 };
1854
1855 let (success, mut details) = benchmark_success_and_errors(&json);
1856 output.append(&mut details);
1857 Some((success, output))
1858}
1859
1860fn collect_env_info(user_home: &str) -> Vec<(String, String)> {
1861 let mut values = load_mlstack_env_exports(user_home);
1862
1863 let keys = [
1864 "ROCM_VERSION",
1865 "ROCM_CHANNEL",
1866 "GPU_ARCH",
1867 "GPU_ARCHS",
1868 "PYTORCH_ROCM_ARCH",
1869 "HIP_VISIBLE_DEVICES",
1870 "CUDA_VISIBLE_DEVICES",
1871 "ROCR_VISIBLE_DEVICES",
1872 "PYTORCH_ROCM_DEVICE",
1873 "HSA_OVERRIDE_GFX_VERSION",
1874 "ROCM_HOME",
1875 "ROCM_PATH",
1876 "HIP_PATH",
1877 "MLSTACK_PYTHON_BIN",
1878 "PYTHONPATH",
1879 ];
1880
1881 for key in keys {
1882 if !values.contains_key(key) {
1883 if let Ok(value) = std::env::var(key) {
1884 values.insert(key.to_string(), value);
1885 }
1886 }
1887 }
1888
1889 keys.iter()
1890 .map(|key| {
1891 let default_val = if *key == "ROCM_HOME" || *key == "ROCM_PATH" || *key == "HIP_PATH" {
1892 "/opt/rocm".to_string()
1893 } else {
1894 "unknown".to_string()
1895 };
1896 (
1897 key.to_string(),
1898 values.get(*key).cloned().unwrap_or(default_val),
1899 )
1900 })
1901 .collect()
1902}
1903
1904fn load_mlstack_env_exports(user_home: &str) -> HashMap<String, String> {
1905 let mut values: HashMap<String, String> = HashMap::new();
1906 for home in candidate_mlstack_homes(user_home) {
1907 let env_path = PathBuf::from(home).join(".mlstack_env");
1908 if let Ok(contents) = fs::read_to_string(&env_path) {
1909 for line in contents.lines() {
1910 if let Some((key, value)) = parse_env_export(line) {
1911 values.entry(key).or_insert(value);
1912 }
1913 }
1914 }
1915 }
1916
1917 values
1918}
1919
1920fn parse_env_export(line: &str) -> Option<(String, String)> {
1921 let export_idx = line.rfind("export ")?;
1922 let rest = &line[export_idx + 7..];
1923 let mut parts = rest.splitn(2, '=');
1924 let key = parts.next()?.trim().to_string();
1925 let value = parts.next()?.trim().trim_matches('"').to_string();
1926 let value = strip_shell_suffixes_from_value(&value);
1928 Some((key, value))
1929}
1930
1931fn strip_shell_suffixes_from_value(value: &str) -> String {
1939 let mut v = value.to_string();
1940 if let Some(idx) = v.find("; fi") {
1942 v.truncate(idx);
1943 }
1944 if let Some(idx) = v.find("&& ") {
1946 v.truncate(idx);
1947 }
1948 v.trim().to_string()
1949}
1950
1951fn detect_onnx_rocm_provider(python: &str) -> Option<String> {
1952 let output = Command::new(python)
1953 .arg("-c")
1954 .arg(
1955 "import onnxruntime, pathlib; base=pathlib.Path(onnxruntime.__file__).parent; libs=list(base.rglob('libonnxruntime_providers_rocm.so')); print(libs[0] if libs else '')",
1956 )
1957 .output()
1958 .ok()?;
1959 if !output.status.success() {
1960 return None;
1961 }
1962 let value = String::from_utf8_lossy(&output.stdout).trim().to_string();
1963 if value.is_empty() {
1964 None
1965 } else {
1966 Some(value)
1967 }
1968}
1969
1970fn push_unique_env_path(paths: &mut Vec<String>, value: &str) {
1971 let candidate = value.trim();
1972 if candidate.is_empty() {
1973 return;
1974 }
1975 if paths.iter().any(|existing| existing == candidate) {
1976 return;
1977 }
1978 paths.push(candidate.to_string());
1979}
1980
1981fn user_home_from_passwd(user_name: &str) -> Option<String> {
1982 let target = user_name.trim();
1983 if target.is_empty() {
1984 return None;
1985 }
1986 let passwd = fs::read_to_string("/etc/passwd").ok()?;
1987 passwd.lines().find_map(|line| {
1988 if line.trim().is_empty() || line.starts_with('#') {
1989 return None;
1990 }
1991 let fields: Vec<&str> = line.split(':').collect();
1992 if fields.len() < 6 || fields[0] != target {
1993 return None;
1994 }
1995 let home = fields[5].trim();
1996 if home.is_empty() {
1997 None
1998 } else {
1999 Some(home.to_string())
2000 }
2001 })
2002}
2003
2004fn passwd_homes_with_mlstack() -> Vec<String> {
2005 let mut homes = Vec::new();
2006 let Ok(passwd) = fs::read_to_string("/etc/passwd") else {
2007 return homes;
2008 };
2009
2010 for line in passwd.lines() {
2011 if line.trim().is_empty() || line.starts_with('#') {
2012 continue;
2013 }
2014 let fields: Vec<&str> = line.split(':').collect();
2015 if fields.len() < 6 {
2016 continue;
2017 }
2018 let home = fields[5].trim();
2019 if home.is_empty() {
2020 continue;
2021 }
2022 if Path::new(home).join(".mlstack").is_dir()
2023 || Path::new(home).join(".mlstack_env").is_file()
2024 {
2025 push_unique_env_path(&mut homes, home);
2026 }
2027 }
2028
2029 homes
2030}
2031
2032fn candidate_mlstack_homes(preferred_home: &str) -> Vec<String> {
2033 let mut homes = Vec::new();
2034 push_unique_env_path(&mut homes, preferred_home);
2035
2036 if let Ok(value) = std::env::var("MLSTACK_USER_HOME") {
2037 push_unique_env_path(&mut homes, &value);
2038 }
2039 if let Ok(value) = std::env::var("HOME") {
2040 push_unique_env_path(&mut homes, &value);
2041 }
2042
2043 for key in ["SUDO_USER", "USER", "LOGNAME"] {
2044 if let Ok(user_name) = std::env::var(key) {
2045 if let Some(home) = user_home_from_passwd(&user_name) {
2046 push_unique_env_path(&mut homes, &home);
2047 }
2048 }
2049 }
2050
2051 for home in passwd_homes_with_mlstack() {
2052 push_unique_env_path(&mut homes, &home);
2053 }
2054
2055 homes
2056}
2057
2058fn resolve_mlstack_user_home() -> String {
2059 let fallback = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
2060 let candidates = candidate_mlstack_homes(&fallback);
2061
2062 for home in &candidates {
2063 if Path::new(home).join(".mlstack_env").is_file() {
2064 return home.clone();
2065 }
2066 }
2067 for home in &candidates {
2068 if Path::new(home).join(".mlstack").is_dir() {
2069 return home.clone();
2070 }
2071 }
2072
2073 candidates.into_iter().next().unwrap_or(fallback)
2074}
2075
2076fn mlstack_mpi_compat_dirs(user_home: &str) -> Vec<String> {
2077 let mut dirs = Vec::new();
2078 for home in candidate_mlstack_homes(user_home) {
2079 let mlstack_dir = PathBuf::from(home).join(".mlstack");
2080
2081 let primary = mlstack_dir.join("libmpi-compat");
2082 if primary.is_dir() {
2083 dirs.push(primary.to_string_lossy().to_string());
2084 }
2085
2086 if let Ok(entries) = fs::read_dir(&mlstack_dir) {
2087 for entry in entries.flatten() {
2088 let path = entry.path();
2089 if !path.is_dir() {
2090 continue;
2091 }
2092 let Some(name) = path.file_name().and_then(|value| value.to_str()) else {
2093 continue;
2094 };
2095 if name.starts_with("libmpi-compat-user-") {
2096 dirs.push(path.to_string_lossy().to_string());
2097 }
2098 }
2099 }
2100 }
2101
2102 dirs.sort();
2103 dirs.dedup();
2104 dirs
2105}
2106
2107fn build_verification_ld_library_path(user_home: &str) -> String {
2108 let mut paths: Vec<String> = Vec::new();
2109
2110 for compat_dir in mlstack_mpi_compat_dirs(user_home) {
2111 push_unique_env_path(&mut paths, &compat_dir);
2112 }
2113
2114 for rocm_path in ["/opt/rocm/lib", "/opt/rocm/hip/lib", "/opt/rocm/opencl/lib"] {
2115 if Path::new(rocm_path).exists() {
2116 push_unique_env_path(&mut paths, rocm_path);
2117 }
2118 }
2119
2120 if let Ok(existing) = std::env::var("LD_LIBRARY_PATH") {
2121 for part in existing.split(':') {
2122 push_unique_env_path(&mut paths, part);
2123 }
2124 }
2125
2126 paths.join(":")
2127}
2128
2129fn run_verification_command(
2130 step: &VerificationCommand,
2131 sender: &Sender<InstallerEvent>,
2132) -> Result<(bool, Vec<String>)> {
2133 let user_home = resolve_mlstack_user_home();
2134 let user_name = std::env::var("USER").unwrap_or_else(|_| "user".into());
2135 let env_exports = load_mlstack_env_exports(&user_home);
2136
2137 let mut command = Command::new(&step.program);
2138 command.args(&step.args);
2139
2140 let extra_paths = python_search_paths();
2141 let mut pythonpath = std::env::var("PYTHONPATH").unwrap_or_default();
2142 if !extra_paths.is_empty() {
2143 if !pythonpath.is_empty() {
2144 pythonpath.push(':');
2145 }
2146 pythonpath.push_str(&extra_paths.join(":"));
2147 }
2148
2149 pythonpath = pythonpath
2151 .split(':')
2152 .filter(|part| !part.starts_with("/tmp/"))
2153 .collect::<Vec<_>>()
2154 .join(":");
2155
2156 if !pythonpath.is_empty() {
2157 command.env("PYTHONPATH", pythonpath);
2158 }
2159
2160 let verification_ld = build_verification_ld_library_path(&user_home);
2161 if !verification_ld.is_empty() {
2162 command.env("LD_LIBRARY_PATH", verification_ld);
2163 }
2164
2165 if step.target_id == "onnx" {
2166 if let Some(lib) = detect_onnx_rocm_provider(&step.program) {
2167 command.env("ORT_ROCM_EP_PROVIDER_PATH", lib);
2168 }
2169 }
2170
2171 let selected_python = if !step.modules.is_empty() {
2172 step.program.clone()
2173 } else {
2174 resolve_python_bin()
2175 };
2176 let gpu_arch = env_exports
2177 .get("GPU_ARCH")
2178 .map(|v| v.trim().to_string())
2179 .filter(|v| !v.is_empty())
2180 .unwrap_or_else(|| {
2181 detect_gpu_arch().unwrap_or_else(|| {
2182 eprintln!("Warning: no GPU arch detected, falling back to gfx1100");
2183 "gfx1100".to_string()
2184 })
2185 });
2186 let py_rocm_arch = env_exports
2187 .get("PYTORCH_ROCM_ARCH")
2188 .map(|v| v.trim().to_string())
2189 .filter(|v| !v.is_empty())
2190 .unwrap_or_else(|| gpu_arch.clone());
2191 let mut gpu_archs = env_exports
2192 .get("GPU_ARCHS")
2193 .map(|v| v.trim().to_string())
2194 .filter(|v| !v.is_empty())
2195 .unwrap_or_else(|| gpu_arch.clone());
2196 if step.target_id == "aiter" {
2197 gpu_archs = py_rocm_arch
2198 .split(';')
2199 .next()
2200 .map(|v| v.trim().to_string())
2201 .filter(|v| !v.is_empty())
2202 .unwrap_or_else(|| gpu_arch.clone());
2203 }
2204 let mut hip_visible = env_exports
2205 .get("HIP_VISIBLE_DEVICES")
2206 .map(|v| v.trim().to_string())
2207 .filter(|v| !v.is_empty());
2208 let mut cuda_visible = env_exports
2209 .get("CUDA_VISIBLE_DEVICES")
2210 .map(|v| v.trim().to_string())
2211 .filter(|v| !v.is_empty())
2212 .or_else(|| hip_visible.clone());
2213 let mut rocr_visible = env_exports
2214 .get("ROCR_VISIBLE_DEVICES")
2215 .map(|v| v.trim().to_string())
2216 .filter(|v| !v.is_empty())
2217 .or_else(|| hip_visible.clone());
2218 if step.target_id == "aiter" {
2219 let base_list = hip_visible
2220 .clone()
2221 .or_else(|| cuda_visible.clone())
2222 .unwrap_or_else(detect_gpu_list);
2223 let ordered = prioritize_gpu_list_for_arch(&base_list, &gpu_archs);
2224 hip_visible = Some(ordered.clone());
2225 cuda_visible = Some(ordered.clone());
2226 rocr_visible = Some(ordered);
2227 }
2228 let hsa_override = env_exports
2229 .get("HSA_OVERRIDE_GFX_VERSION")
2230 .map(|v| v.trim().to_string())
2231 .filter(|v| !v.is_empty());
2232
2233 command
2234 .env("HOME", &user_home)
2235 .env("USER", &user_name)
2236 .env("LOGNAME", &user_name)
2237 .env("MLSTACK_USER_HOME", &user_home)
2238 .env("MLSTACK_PYTHON_BIN", &selected_python)
2239 .env("UV_PYTHON", &selected_python)
2240 .env("PYTHONWARNINGS", "ignore")
2241 .env("PYTHONUNBUFFERED", "1")
2242 .env("ROCM_HOME", "/opt/rocm")
2243 .env("HIP_PATH", "/opt/rocm")
2244 .env("GPU_ARCH", &gpu_arch)
2245 .env("PYTORCH_ROCM_ARCH", &py_rocm_arch)
2246 .env("GPU_ARCHS", &gpu_archs)
2247 .stdout(Stdio::piped())
2248 .stderr(Stdio::piped());
2249
2250 if let Some(value) = hip_visible {
2251 command.env("HIP_VISIBLE_DEVICES", value);
2252 }
2253 if let Some(value) = cuda_visible {
2254 command.env("CUDA_VISIBLE_DEVICES", value);
2255 }
2256 if let Some(value) = rocr_visible {
2257 command.env("ROCR_VISIBLE_DEVICES", value);
2258 }
2259 if let Some(value) = hsa_override {
2260 command.env("HSA_OVERRIDE_GFX_VERSION", value);
2261 }
2262 strip_hsa_override_for_therock(&mut command, &env_exports);
2263 if step.target_id == "aiter" {
2264 let aiter_jit_dir = PathBuf::from(&user_home)
2265 .join(".mlstack")
2266 .join("aiter")
2267 .join("jit");
2268 let _ = fs::create_dir_all(&aiter_jit_dir);
2269 command.env("AITER_JIT_DIR", aiter_jit_dir);
2270 }
2271
2272 if step.target_id == "mpi4py" {
2273 command
2275 .env("OMPI_MCA_btl", "^smcuda")
2276 .env("OMPI_MCA_opal_cuda_support", "0")
2277 .env("OMPI_MCA_opal_accelerator_rocm_enable", "0");
2278 }
2279
2280 let output = command
2281 .output()
2282 .with_context(|| format!("Failed to run verification for {}", step.label))?;
2283
2284 let mut captured_output = Vec::new();
2285 const MAX_CAPTURED_OUTPUT_LINES: usize = 24;
2286
2287 let stdout = String::from_utf8_lossy(&output.stdout);
2288 for raw in stdout.replace('\r', "\n").lines() {
2289 let clean_line = sanitize_output_line(raw);
2290 if should_suppress_verification_line(&clean_line, &step.target_id) {
2291 continue;
2292 }
2293 if !clean_line.trim().is_empty() && captured_output.len() < MAX_CAPTURED_OUTPUT_LINES {
2294 let line = format!("{}: {}", step.label, clean_line);
2295 let _ = sender.send(InstallerEvent::Log(line.clone(), false));
2296 captured_output.push(clean_line);
2297 }
2298 }
2299
2300 let stderr = String::from_utf8_lossy(&output.stderr);
2301 for raw in stderr.replace('\r', "\n").lines() {
2302 let clean_line = sanitize_output_line(raw);
2303 if should_suppress_verification_line(&clean_line, &step.target_id) {
2304 continue;
2305 }
2306 if !clean_line.trim().is_empty() && captured_output.len() < MAX_CAPTURED_OUTPUT_LINES {
2307 let line = format!("{}: {}", step.label, clean_line);
2308 let _ = sender.send(InstallerEvent::Log(line.clone(), false));
2309 captured_output.push(clean_line);
2310 }
2311 }
2312
2313 if output.status.success() && captured_output.len() == MAX_CAPTURED_OUTPUT_LINES {
2314 captured_output.push("Output truncated".to_string());
2315 }
2316
2317 Ok((output.status.success(), captured_output))
2318}
2319
2320fn sort_components_by_dependencies(components: Vec<Component>) -> Vec<Component> {
2333 use crate::installers::components::topological_sort;
2334
2335 let mut installer_comps: Vec<Component> = Vec::new();
2337 let mut other_comps: Vec<Component> = Vec::new();
2338
2339 for comp in components {
2340 if comp.category == Category::Maintenance || comp.category == Category::Performance {
2341 other_comps.push(comp);
2342 } else {
2343 installer_comps.push(comp);
2344 }
2345 }
2346
2347 let ids: Vec<String> = installer_comps.iter().map(|c| c.id.clone()).collect();
2349
2350 match topological_sort(&ids) {
2351 Ok(sorted_ids) => {
2352 let mut reordered = Vec::with_capacity(installer_comps.len());
2354 for id in &sorted_ids {
2355 if let Some(pos) = installer_comps.iter().position(|c| c.id == *id) {
2356 reordered.push(installer_comps.remove(pos));
2357 }
2358 }
2359 reordered.extend(installer_comps);
2361
2362 reordered.extend(other_comps);
2364 reordered
2365 }
2366 Err(_) => {
2367 installer_comps.extend(other_comps);
2370 installer_comps
2371 }
2372 }
2373}
2374
2375struct NativeInstallerContext<'a> {
2382 sudo_password: Option<String>,
2383 batch_mode: bool,
2384 install_method: &'a str,
2385 sender: &'a Sender<InstallerEvent>,
2386 user_home: &'a str,
2387 env_exports: &'a HashMap<String, String>,
2388 vllm_version_lag: u32,
2390 vllm_version_min_age_days: u32,
2391 target_version: Option<String>,
2393}
2394
2395enum NativeCommand {
2405 Shell {
2407 program: String,
2408 args: Vec<String>,
2409 env: Vec<(String, String)>,
2410 working_dir: Option<PathBuf>,
2411 },
2412 Pip { program: String, args: Vec<String> },
2414 System { program: String, args: Vec<String> },
2416 Package { program: String, args: Vec<String> },
2418}
2419
2420impl NativeCommand {
2421 fn from_shell_cmd(program: &str, args: &[String], env: &[(String, String)]) -> Self {
2425 NativeCommand::Shell {
2426 program: program.to_string(),
2427 args: args.to_vec(),
2428 env: env.to_vec(),
2429 working_dir: None,
2430 }
2431 }
2432
2433 fn from_shell_cmd_with_dir(
2435 program: &str,
2436 args: &[String],
2437 env: &[(String, String)],
2438 working_dir: Option<PathBuf>,
2439 ) -> Self {
2440 NativeCommand::Shell {
2441 program: program.to_string(),
2442 args: args.to_vec(),
2443 env: env.to_vec(),
2444 working_dir,
2445 }
2446 }
2447}
2448
2449fn is_up_to_date_output(output: &str) -> bool {
2469 let lower = output.to_lowercase();
2470 lower.contains("up to date -- skipping")
2471 || lower.contains("up to date, skipping")
2472 || lower.contains("up to date -- reinstalling")
2473 || lower.contains("there is nothing to do")
2474 || lower.contains("already installed")
2475 || (lower.contains("warning:") && lower.contains("up to date"))
2476}
2477
2478fn safe_rocm_path(inherited: &str) -> String {
2490 const ROCM_BINS: &[&str] = &["/opt/rocm/bin", "/opt/rocm/hip/bin"];
2491 if !Path::new("/opt/rocm/bin").exists() {
2492 return inherited.to_string();
2493 }
2494 let mut parts: Vec<&str> = ROCM_BINS.to_vec();
2495 for entry in inherited.split(':') {
2496 if !entry.is_empty() && !ROCM_BINS.contains(&entry) {
2497 parts.push(entry);
2498 }
2499 }
2500 parts.join(":")
2501}
2502
2503fn execute_native_command(
2509 cmd: &NativeCommand,
2510 sudo_pw: Option<&str>,
2511 sender: &Sender<InstallerEvent>,
2512 component_name: &str,
2513) -> Result<()> {
2514 let (program, args, envs, working_dir) = match cmd {
2515 NativeCommand::Shell {
2516 program,
2517 args,
2518 env,
2519 working_dir,
2520 } => (
2521 program.clone(),
2522 args.clone(),
2523 env.clone(),
2524 working_dir.clone(),
2525 ),
2526 NativeCommand::Pip { program, args } => (program.clone(), args.clone(), vec![], None),
2527 NativeCommand::System { program, args } => (program.clone(), args.clone(), vec![], None),
2528 NativeCommand::Package { program, args } => (program.clone(), args.clone(), vec![], None),
2529 };
2530
2531 let invokes_pip = program == "uv"
2540 || program.ends_with("pip")
2541 || program.ends_with("pip3")
2542 || args.iter().any(|a| a == "pip");
2543 let is_install = args.iter().any(|a| a == "install");
2544 if invokes_pip && is_install {
2545 let mut skip_next = false;
2546 for a in &args {
2547 if skip_next {
2548 skip_next = false;
2552 continue;
2553 }
2554 if a == "-r" || a == "--requirement" || a == "-c" || a == "--constraint" {
2555 skip_next = true;
2556 continue;
2557 }
2558 if a.starts_with('-') || a == "install" || a == "." || a == "-e" || a.is_empty() {
2559 continue;
2560 }
2561 if crate::installers::common::nvidia_blocklist::is_nvidia_cuda_runtime_package(a) {
2562 let _ = sender.send(InstallerEvent::Log(
2563 format!("[native] BLOCKED nvidia/cuda package in {component_name}: {a}"),
2564 true,
2565 ));
2566 anyhow::bail!(
2567 "NVIDIA/CUDA contamination blocked: component '{component_name}' attempted \
2568 to pip install '{a}' — violates the No-CUDA hard-prime tenet"
2569 );
2570 }
2571 if a.starts_with("http")
2572 && crate::installers::common::nvidia_blocklist::is_cuda_wheel_url(a)
2573 {
2574 let _ = sender.send(InstallerEvent::Log(
2575 format!("[native] BLOCKED cuda wheel URL in {component_name}: {a}"),
2576 true,
2577 ));
2578 anyhow::bail!("NVIDIA/CUDA wheel URL blocked in {component_name}: {a}");
2579 }
2580 }
2581 }
2582
2583 let already_sudo = program == "sudo";
2586 let mut command = if let Some(_pw) = sudo_pw {
2587 if already_sudo {
2588 let mut c = Command::new(&program);
2590 c.args(&args);
2591 c
2592 } else {
2593 let mut c = Command::new("sudo");
2594 c.arg("-S").arg("-p").arg("").arg("--").arg(&program);
2595 c.args(&args);
2596 let _ = sender.send(InstallerEvent::Log(
2598 format!("[native] Running with sudo: {} {}", program, args.join(" ")),
2599 false,
2600 ));
2601 c
2602 }
2603 } else {
2604 let mut c = Command::new(&program);
2605 c.args(&args);
2606 c
2607 };
2608
2609 for (key, value) in &envs {
2611 command.env(key, value);
2612 }
2613
2614 command.env(
2623 "PATH",
2624 safe_rocm_path(&std::env::var("PATH").unwrap_or_default()),
2625 );
2626
2627 if let Some(ref dir) = working_dir {
2629 command.current_dir(dir);
2630 }
2631
2632 command
2634 .stdin(Stdio::piped())
2635 .stdout(Stdio::piped())
2636 .stderr(Stdio::piped());
2637
2638 let cmd_str = if args.is_empty() {
2640 program.clone()
2641 } else {
2642 format!("{} {}", program, args.join(" "))
2643 };
2644 let _ = sender.send(InstallerEvent::Log(
2645 format!("[native] $ {}", cmd_str),
2646 false,
2647 ));
2648
2649 let mut child = command.spawn().with_context(|| {
2650 let mut hint = String::new();
2651 if cmd_str.starts_with("sudo") {
2653 hint = " — ensure sudo is installed and you have the correct permissions".to_string();
2654 } else if cmd_str.starts_with("pip")
2655 || cmd_str.starts_with("uv")
2656 || cmd_str.starts_with("python")
2657 {
2658 hint = " — ensure Python/pip is installed and accessible".to_string();
2659 } else if cmd_str.starts_with("git") {
2660 hint = " — ensure git is installed".to_string();
2661 } else if cmd_str.starts_with("cmake") || cmd_str.starts_with("make") {
2662 hint = " — ensure build tools (cmake, make, gcc) are installed".to_string();
2663 } else if cmd_str.starts_with("rocminfo") || cmd_str.starts_with("rocm-smi") {
2664 hint = " — ensure ROCm is installed and /opt/rocm/bin is in PATH".to_string();
2665 }
2666 format!(
2667 "Failed to spawn command '{}' for {}.{}",
2668 cmd_str, component_name, hint
2669 )
2670 })?;
2671
2672 if let Some(pw) = sudo_pw {
2675 if let Some(mut stdin) = child.stdin.take() {
2676 let _ = stdin.write_all(pw.as_bytes());
2677 let _ = stdin.write_all(b"\n");
2678 let _ = stdin.flush();
2679 drop(stdin);
2680 }
2681 } else {
2682 drop(child.stdin.take());
2685 }
2686
2687 let sender_stdout = sender.clone();
2690 let stdout_stream = match child.stdout.take() {
2691 Some(s) => s,
2692 None => {
2693 let _ = sender.send(InstallerEvent::Log(
2694 format!(
2695 "[ERROR] {} — failed to capture stdout from spawned process",
2696 component_name
2697 ),
2698 false,
2699 ));
2700 bail!(
2701 "{} — failed to capture stdout from spawned process. \
2702 This usually indicates a system resource issue.",
2703 component_name
2704 );
2705 }
2706 };
2707 let stdout_captured: Arc<Mutex<String>> = Arc::new(Mutex::new(String::new()));
2708 let stdout_captured_clone = stdout_captured.clone();
2709 let stdout_handle = thread::spawn(move || {
2710 let reader = BufReader::new(stdout_stream);
2711 for line in reader.lines() {
2712 match line {
2713 Ok(l) => {
2714 let clean = strip_ansi_codes(&l);
2715 if !clean.trim().is_empty() {
2716 let _ = sender_stdout.send(InstallerEvent::Log(clean.clone(), false));
2717 if let Ok(mut buf) = stdout_captured_clone.lock() {
2718 buf.push_str(&clean);
2719 buf.push('\n');
2720 }
2721 }
2722 }
2723 Err(_) => break,
2724 }
2725 }
2726 });
2727
2728 let sender_stderr = sender.clone();
2729 let stderr_stream = match child.stderr.take() {
2730 Some(s) => s,
2731 None => {
2732 let _ = sender.send(InstallerEvent::Log(
2733 format!(
2734 "[ERROR] {} — failed to capture stderr from spawned process",
2735 component_name
2736 ),
2737 false,
2738 ));
2739 bail!(
2740 "{} — failed to capture stderr from spawned process. \
2741 This usually indicates a system resource issue.",
2742 component_name
2743 );
2744 }
2745 };
2746 let stderr_captured: Arc<Mutex<String>> = Arc::new(Mutex::new(String::new()));
2747 let stderr_captured_clone = stderr_captured.clone();
2748 let stderr_handle = thread::spawn(move || {
2749 let reader = BufReader::new(stderr_stream);
2750 for line in reader.lines() {
2751 match line {
2752 Ok(l) => {
2753 let clean = strip_ansi_codes(&l);
2754 if !clean.trim().is_empty() {
2755 let _ = sender_stderr.send(InstallerEvent::Log(clean.clone(), false));
2756 if let Ok(mut buf) = stderr_captured_clone.lock() {
2757 buf.push_str(&clean);
2758 buf.push('\n');
2759 }
2760 }
2761 }
2762 Err(_) => break,
2763 }
2764 }
2765 });
2766
2767 let status = child.wait().with_context(|| {
2769 format!(
2770 "{} — command process wait failed. \
2771 The process may have been killed or the system is under heavy load.",
2772 component_name
2773 )
2774 })?;
2775 let _ = stdout_handle.join();
2776 let _ = stderr_handle.join();
2777
2778 if !status.success() {
2779 let combined_output = {
2784 let stdout_buf = stdout_captured.lock().unwrap_or_else(|e| e.into_inner());
2785 let stderr_buf = stderr_captured.lock().unwrap_or_else(|e| e.into_inner());
2786 format!("{}{}", *stdout_buf, *stderr_buf)
2787 };
2788
2789 if is_up_to_date_output(&combined_output) {
2790 let _ = sender.send(InstallerEvent::Log(
2791 format!(
2792 "[native] {} — package already up to date, treating as success",
2793 component_name
2794 ),
2795 false,
2796 ));
2797 return Ok(());
2798 }
2799
2800 let code = status.code().unwrap_or(-1);
2801 let error_detail = match code {
2802 1 => "General error — check the command output above for details".to_string(),
2803 2 => "Misuse of shell builtins — check command arguments".to_string(),
2804 126 => "Command not executable — check file permissions".to_string(),
2805 127 => "Command not found — ensure the required tool is installed".to_string(),
2806 130 => "Process terminated by Ctrl+C (SIGINT)".to_string(),
2807 137 => "Process killed (SIGKILL) — possibly out of memory".to_string(),
2808 139 => "Segmentation fault — this is a bug in the invoked program".to_string(),
2809 n if n > 128 => format!(
2810 "Process terminated by signal {} — check system logs",
2811 n - 128
2812 ),
2813 _ => format!(
2814 "Command exited with code {} — check the output above for details",
2815 code
2816 ),
2817 };
2818 let _ = sender.send(InstallerEvent::Log(
2819 format!(
2820 "[ERROR] {} failed: {} (exit code {} at {})",
2821 component_name,
2822 error_detail,
2823 code,
2824 Local::now().format("%H:%M:%S")
2825 ),
2826 false,
2827 ));
2828 bail!(
2829 "{} failed: {} (exit code {})",
2830 component_name,
2831 error_detail,
2832 code
2833 );
2834 }
2835
2836 Ok(())
2837}
2838
2839fn build_common_env(ctx: &NativeInstallerContext) -> Vec<(String, String)> {
2848 let mut env = Vec::new();
2849
2850 env.push(("HOME".into(), ctx.user_home.to_string()));
2852 if let Ok(user) = std::env::var("USER") {
2853 env.push(("USER".into(), user.clone()));
2854 env.push(("LOGNAME".into(), user));
2855 }
2856 env.push(("MLSTACK_USER_HOME".into(), ctx.user_home.to_string()));
2857 env.push((
2858 "MLSTACK_BATCH_MODE".into(),
2859 if ctx.batch_mode { "1" } else { "0" }.into(),
2860 ));
2861 env.push((
2862 "MLSTACK_INSTALL_METHOD".into(),
2863 ctx.install_method.to_string(),
2864 ));
2865 env.push(("INSTALL_METHOD".into(), ctx.install_method.to_string()));
2866 env.push(("RUSTY_STACK".into(), "true".into()));
2867
2868 for key in &[
2870 "ROCM_VERSION",
2871 "ROCM_CHANNEL",
2872 "GPU_ARCH",
2873 "GPU_ARCHS",
2874 "PYTORCH_ROCM_ARCH",
2875 "HIP_VISIBLE_DEVICES",
2876 "CUDA_VISIBLE_DEVICES",
2877 "ROCR_VISIBLE_DEVICES",
2878 "PYTORCH_ROCM_DEVICE",
2879 "HSA_OVERRIDE_GFX_VERSION",
2880 "MLSTACK_PYTHON_BIN",
2881 "UV_PYTHON",
2882 ] {
2883 if let Some(value) = ctx.env_exports.get(*key) {
2884 let trimmed = value.trim();
2885 if !trimmed.is_empty() {
2886 env.push((key.to_string(), trimmed.to_string()));
2887 }
2888 }
2889 }
2890
2891 env.push(("ROCM_HOME".into(), "/opt/rocm".into()));
2893 env.push(("ROCM_PATH".into(), "/opt/rocm".into()));
2894 env.push(("HIP_PATH".into(), "/opt/rocm".into()));
2895 env.push(("HIP_ROOT_DIR".into(), "/opt/rocm".into()));
2896 env.push(("ROCM_ROOT".into(), "/opt/rocm".into()));
2897 env.push(("HIPCC_BIN_DIR".into(), "/opt/rocm/bin".into()));
2898
2899 env.push(("PIP_BREAK_SYSTEM_PACKAGES".into(), "1".into()));
2901 env.push(("PIP_ROOT_USER_ACTION".into(), "ignore".into()));
2902 env.push(("UV_PIP_BREAK_SYSTEM_PACKAGES".into(), "1".into()));
2903 env.push(("UV_SYSTEM_PYTHON".into(), "1".into()));
2904 env.push(("PYTHONUNBUFFERED".into(), "1".into()));
2905
2906 env
2907}
2908
2909fn is_existing_git_repo(target_dir: &str) -> bool {
2919 let path = Path::new(target_dir);
2920 match std::fs::symlink_metadata(target_dir) {
2923 Ok(meta) if meta.is_dir() => {
2924 path.join(".git")
2926 .symlink_metadata()
2927 .map(|m| m.is_dir() || m.is_file())
2928 .unwrap_or(false)
2929 }
2930 _ => false,
2931 }
2932}
2933
2934fn fix_directory_ownership(dir: &str, sudo_pw: Option<&str>, sender: &Sender<InstallerEvent>) {
2951 let stat_output = std::process::Command::new("stat")
2953 .args(["-c", "%U", dir])
2954 .output();
2955
2956 let needs_fix = match stat_output {
2957 Ok(out) if out.status.success() => {
2958 let owner = String::from_utf8_lossy(&out.stdout).trim().to_string();
2959 owner == "root"
2960 }
2961 _ => true, };
2963
2964 if !needs_fix {
2965 return;
2966 }
2967
2968 let real_user = std::env::var("SUDO_USER")
2969 .or_else(|_| std::env::var("USER"))
2970 .unwrap_or_else(|_| "scooter".to_string());
2971
2972 let _ = sender.send(InstallerEvent::Log(
2973 format!("[native] Fixing ownership of {} -> {}", dir, real_user),
2974 false,
2975 ));
2976
2977 if sudo_pw.is_some() {
2978 let chown_cmd = NativeCommand::from_shell_cmd(
2979 "chown",
2980 &[
2981 "-R".to_string(),
2982 format!("{}:{}", real_user, real_user),
2983 dir.to_string(),
2984 ],
2985 &[],
2986 );
2987 let _ = execute_native_command(&chown_cmd, sudo_pw, sender, "ownership-fix");
2988 }
2989 }
2991
2992fn git_clone_or_pull(
2993 repo_url: &str,
2994 target_dir: &str,
2995 extra_clone_args: &[&str],
2996 sudo_pw: Option<&str>,
2997 sender: &Sender<InstallerEvent>,
2998 component_name: &str,
2999) -> Result<()> {
3000 if is_existing_git_repo(target_dir) {
3001 let _ = sender.send(InstallerEvent::Log(
3002 format!(
3003 "[native] {} — directory '{}' already exists, pulling updates instead of cloning",
3004 component_name, target_dir
3005 ),
3006 false,
3007 ));
3008
3009 let pull_cmd = NativeCommand::from_shell_cmd(
3011 "git",
3012 &[
3013 "-C".to_string(),
3014 target_dir.to_string(),
3015 "pull".to_string(),
3016 "--ff-only".to_string(),
3017 ],
3018 &[],
3019 );
3020 let pull_result = execute_native_command(&pull_cmd, sudo_pw, sender, component_name);
3021
3022 if pull_result.is_ok() {
3023 return Ok(());
3024 }
3025
3026 let _ = sender.send(InstallerEvent::Log(
3028 format!(
3029 "[native] {} — git pull --ff-only failed, trying fetch + reset --hard origin/HEAD",
3030 component_name
3031 ),
3032 false,
3033 ));
3034
3035 let fetch_cmd = NativeCommand::from_shell_cmd(
3036 "git",
3037 &[
3038 "-C".to_string(),
3039 target_dir.to_string(),
3040 "fetch".to_string(),
3041 "--all".to_string(),
3042 ],
3043 &[],
3044 );
3045 let fetch_ok = execute_native_command(&fetch_cmd, sudo_pw, sender, component_name).is_ok();
3046
3047 let reset_ok = if fetch_ok {
3048 let reset_cmd = NativeCommand::from_shell_cmd(
3049 "git",
3050 &[
3051 "-C".to_string(),
3052 target_dir.to_string(),
3053 "reset".to_string(),
3054 "--hard".to_string(),
3055 "origin/HEAD".to_string(),
3056 ],
3057 &[],
3058 );
3059 execute_native_command(&reset_cmd, sudo_pw, sender, component_name).is_ok()
3060 } else {
3061 false
3062 };
3063
3064 if fetch_ok && reset_ok {
3065 return Ok(());
3066 }
3067
3068 let _ = sender.send(InstallerEvent::Log(
3072 format!(
3073 "[native] {} — existing .git is corrupt/unusable (pull+fetch+reset failed); removing and re-cloning",
3074 component_name
3075 ),
3076 false,
3077 ));
3078 if std::fs::remove_dir_all(target_dir).is_err() {
3079 let rm_cmd = NativeCommand::from_shell_cmd(
3080 "rm",
3081 &["-rf".to_string(), target_dir.to_string()],
3082 &[],
3083 );
3084 let _ = execute_native_command(&rm_cmd, sudo_pw, sender, component_name);
3085 }
3086 if Path::new(target_dir).join(".git").exists() {
3093 bail!(
3094 "{} — could not remove corrupt repo at '{}' (.git still present); \
3095 aborting re-clone to avoid infinite recursion",
3096 component_name,
3097 target_dir
3098 );
3099 }
3100 git_clone_or_pull(
3101 repo_url,
3102 target_dir,
3103 extra_clone_args,
3104 sudo_pw,
3105 sender,
3106 component_name,
3107 )
3108 } else {
3109 let target_path = Path::new(target_dir);
3114 if target_path.is_dir() && !target_path.join(".git").exists() {
3115 let _ = sender.send(InstallerEvent::Log(
3116 format!(
3117 "[native] {} — directory '{}' exists but is not a git repo; removing incomplete clone",
3118 component_name, target_dir
3119 ),
3120 false,
3121 ));
3122 if std::fs::remove_dir_all(target_dir).is_err() {
3124 let rm_cmd = NativeCommand::from_shell_cmd(
3126 "rm",
3127 &["-rf".to_string(), target_dir.to_string()],
3128 &[],
3129 );
3130 let rm_result = execute_native_command(&rm_cmd, sudo_pw, sender, component_name);
3131 if let Err(e) = rm_result {
3132 let _ = sender.send(InstallerEvent::Log(
3133 format!(
3134 "[native] [WARN] {} — could not remove stale dir: {}. Clone may fail.",
3135 component_name, e
3136 ),
3137 false,
3138 ));
3139 }
3140 }
3141 }
3142
3143 let _ = sender.send(InstallerEvent::Log(
3144 format!(
3145 "[native] {} — cloning '{}' into '{}'",
3146 component_name, repo_url, target_dir
3147 ),
3148 false,
3149 ));
3150
3151 let mut args: Vec<String> = vec!["clone".to_string()];
3152 for extra in extra_clone_args {
3153 args.push(extra.to_string());
3154 }
3155 args.push(repo_url.to_string());
3156 args.push(target_dir.to_string());
3157
3158 let clone_cmd = NativeCommand::from_shell_cmd("git", &args, &[]);
3159 execute_native_command(&clone_cmd, sudo_pw, sender, component_name)
3160 }
3161}
3162
3163fn filter_cuda_requirements(req_path: &str) -> std::io::Result<String> {
3194 crate::installers::common::nvidia_blocklist::filter_requirements_file(req_path)
3195}
3196
3197fn is_force_reinstall() -> bool {
3202 std::env::var("MLSTACK_FORCE_REINSTALL")
3203 .or_else(|_| std::env::var("FORCE"))
3204 .map(|v| {
3205 matches!(
3206 v.trim().to_ascii_lowercase().as_str(),
3207 "1" | "true" | "yes" | "on"
3208 )
3209 })
3210 .unwrap_or(false)
3211}
3212
3213fn triton_installed_satisfies_target(python_bin: &str, target: Option<&str>) -> bool {
3228 let target = match target {
3229 Some(t) => t.trim(),
3230 None => return true, };
3232 if target.is_empty() || target == "latest" || target == "installed" {
3234 return true;
3235 }
3236 if is_force_reinstall() {
3238 return false;
3239 }
3240 let installed = std::process::Command::new(python_bin)
3242 .arg("-c")
3243 .arg(
3244 "try:\n import triton, importlib.metadata as im\n \
3245 print(getattr(triton, '__version__', None) or im.version('triton') \
3246 or im.version('triton-rocm') or im.version('pytorch-triton'))\n\
3247 except Exception:\n pass",
3248 )
3249 .output()
3250 .ok()
3251 .and_then(|o| {
3252 if o.status.success() {
3253 String::from_utf8(o.stdout).ok()
3254 } else {
3255 None
3256 }
3257 })
3258 .map(|s| s.trim().to_string())
3259 .filter(|s| !s.is_empty());
3260 match installed {
3261 Some(ver) => version_satisfies(&ver, target),
3262 None => false, }
3264}
3265
3266fn version_satisfies(installed: &str, target: &str) -> bool {
3272 fn parts(v: &str) -> Vec<u64> {
3273 v.split('-')
3274 .next()
3275 .unwrap_or("")
3276 .split('.')
3277 .filter_map(|s| s.parse::<u64>().ok())
3278 .collect()
3279 }
3280 let a = parts(installed.trim());
3281 let b = parts(target.trim());
3282 if a.is_empty() || b.is_empty() {
3283 return false;
3285 }
3286 for (ai, bi) in a.iter().zip(b.iter()) {
3287 match ai.cmp(bi) {
3288 std::cmp::Ordering::Greater => return true,
3289 std::cmp::Ordering::Less => return false,
3290 std::cmp::Ordering::Equal => continue,
3291 }
3292 }
3293 a.len() >= b.len()
3296}
3297
3298fn build_pip_uninstall(packages: &[&str]) -> NativeCommand {
3303 let use_uv = crate::installers::common::utils::command_exists("uv");
3304 if use_uv {
3305 let mut args = vec!["pip".to_string(), "uninstall".to_string()];
3306 for pkg in packages {
3307 args.push(pkg.to_string());
3308 }
3309 NativeCommand::Pip {
3310 program: "uv".to_string(),
3311 args,
3312 }
3313 } else {
3314 let mut args = vec!["uninstall".to_string(), "-y".to_string()];
3315 for pkg in packages {
3316 args.push(pkg.to_string());
3317 }
3318 NativeCommand::Pip {
3319 program: "pip3".to_string(),
3320 args,
3321 }
3322 }
3323}
3324
3325fn purge_component_packages(
3331 packages: &[&str],
3332 sender: &Sender<InstallerEvent>,
3333 component_name: &str,
3334) {
3335 let _ = sender.send(InstallerEvent::Log(
3336 format!(
3337 "[native] {} — purging packages: {}",
3338 component_name,
3339 packages.join(", ")
3340 ),
3341 false,
3342 ));
3343
3344 let uninstall_cmd = build_pip_uninstall(packages);
3346 let _ = execute_native_command(&uninstall_cmd, None, sender, component_name);
3347}
3348
3349fn purge_pytorch(sender: &Sender<InstallerEvent>, component_name: &str) {
3351 purge_component_packages(
3352 &[
3353 "torch",
3354 "torchvision",
3355 "torchaudio",
3356 "pytorch-triton-rocm",
3357 "pytorch-triton",
3358 ],
3359 sender,
3360 component_name,
3361 );
3362}
3363
3364fn purge_triton(sender: &Sender<InstallerEvent>, component_name: &str) {
3366 purge_component_packages(
3367 &["triton", "triton-rocm", "pytorch-triton-rocm"],
3368 sender,
3369 component_name,
3370 );
3371}
3372
3373fn purge_aiter(sender: &Sender<InstallerEvent>, component_name: &str) {
3375 purge_component_packages(&["aiter", "amd-aiter"], sender, component_name);
3376}
3377
3378fn purge_flash_attn(sender: &Sender<InstallerEvent>, component_name: &str) {
3380 purge_component_packages(
3381 &["flash-attn", "flash_attn", "flash_attention_amd"],
3382 sender,
3383 component_name,
3384 );
3385}
3386
3387fn purge_vllm(sender: &Sender<InstallerEvent>, component_name: &str) {
3389 purge_component_packages(&["vllm"], sender, component_name);
3390}
3391
3392fn unseal_core_requested() -> bool {
3395 std::env::var("MLSTACK_UNSEAL_CORE")
3396 .map(|v| {
3397 matches!(
3398 v.trim().to_ascii_lowercase().as_str(),
3399 "1" | "true" | "yes" | "on"
3400 )
3401 })
3402 .unwrap_or(false)
3403}
3404
3405fn registry_gate(component: &Component, sender: &Sender<InstallerEvent>) -> Result<bool> {
3414 let registry = crate::core::registry::InstalledComponentRegistry::load();
3415 if !registry.is_sealed(&component.id) {
3416 return Ok(true);
3417 }
3418 if is_force_reinstall() && !unseal_core_requested() {
3419 let _ = sender.send(InstallerEvent::Log(
3426 format!(
3427 "[native] {} is a sealed CORE component — force ignored, reusing the installed copy. \
3428 Set MLSTACK_UNSEAL_CORE=1 to intentionally override it.",
3429 component.name
3430 ),
3431 false,
3432 ));
3433 return Ok(false);
3434 }
3435 if !is_force_reinstall() {
3436 let _ = sender.send(InstallerEvent::Log(
3437 format!(
3438 "[native] {} already installed (sealed core) — reusing, NOT reinstalling",
3439 component.name
3440 ),
3441 false,
3442 ));
3443 return Ok(false);
3444 }
3445 Ok(true)
3446}
3447
3448fn registry_record(component: &Component) {
3454 use crate::core::registry::{InstalledComponent, InstalledComponentRegistry};
3455 let mut registry = InstalledComponentRegistry::load();
3456 let record = InstalledComponent {
3457 id: component.id.clone(),
3458 version: component_version(&component.id),
3459 source_index: component_source_index(&component.id),
3460 location: install_target_location(),
3461 sealed: crate::core::registry::should_seal_component(&component.id),
3462 installed_at: Some(chrono::Utc::now().to_rfc3339()),
3463 pip_packages: component_pip_packages(&component.id),
3464 };
3465 registry.mark_installed(record);
3466 if let Err(e) = registry.save() {
3467 eprintln!("[WARN] Failed to save installed component registry: {e}");
3471 }
3472}
3473
3474fn component_pip_packages(id: &str) -> Vec<String> {
3478 match id {
3479 "pytorch" => vec!["torch", "torchvision", "torchaudio", "pytorch-triton-rocm"],
3480 "triton" => vec!["triton", "pytorch-triton-rocm"],
3481 "vllm" | "vllm-multi" => vec!["vllm"],
3482 "deepspeed" => vec!["deepspeed", "einops"],
3483 "megatron" => vec!["megatron-core"],
3484 "bitsandbytes" => vec!["bitsandbytes"],
3485 "flash-attn" | "flash-attn-triton" | "flash-attn-ck" | "flash_attention" => {
3486 vec!["flash-attn", "flash_attn"]
3487 }
3488 "aiter" => vec!["aiter"],
3489 "onnx" | "onnxruntime" => vec!["onnxruntime", "onnxruntime-rocm"],
3490 "migraphx" => vec!["migraphx"],
3491 "mpi4py" => vec!["mpi4py"],
3492 "wandb" => vec!["wandb"],
3493 "fastvideo" => vec!["fastvideo"],
3494 _ => Vec::new(),
3495 }
3496 .into_iter()
3497 .map(String::from)
3498 .collect()
3499}
3500
3501fn install_target_location() -> Option<String> {
3505 if let Ok(name) = std::env::var("MLSTACK_ENV_NAME") {
3506 let name = name.trim();
3507 if !name.is_empty() {
3508 return Some(
3509 crate::platform::environment::mlstack_env_dir(name)
3510 .to_string_lossy()
3511 .into_owned(),
3512 );
3513 }
3514 }
3515 let global = crate::platform::environment::mlstack_global_dir();
3516 if global.exists() {
3517 return Some(global.to_string_lossy().into_owned());
3518 }
3519 std::env::var("MLSTACK_PYTHON_BIN")
3520 .ok()
3521 .map(|s| s.trim().to_string())
3522 .filter(|s| !s.is_empty())
3523}
3524
3525fn component_source_index(id: &str) -> Option<String> {
3531 match id {
3532 "pytorch" => Some("https://download.pytorch.org/whl/rocm".to_string()),
3533 "triton" => Some("triton-rocm (bundled with PyTorch)".to_string()),
3534 "rocm" => Some("/opt/rocm".to_string()),
3535 _ => None,
3536 }
3537}
3538
3539fn component_version(id: &str) -> Option<String> {
3543 match id {
3544 "rocm" => {
3545 let v = detect_rocm_version();
3546 if v.is_empty() {
3547 None
3548 } else {
3549 Some(v)
3550 }
3551 }
3552 _ => None,
3553 }
3554}
3555
3556fn ensure_user_in_groups(
3573 groups: &[&str],
3574 sudo_pw: Option<&str>,
3575 sender: &Sender<InstallerEvent>,
3576 component_name: &str,
3577) {
3578 use std::process::Command;
3579
3580 let user = std::env::var("SUDO_USER")
3583 .ok()
3584 .map(|u| u.trim().to_string())
3585 .filter(|u| !u.is_empty())
3586 .or_else(|| {
3587 std::env::var("USER")
3588 .ok()
3589 .map(|u| u.trim().to_string())
3590 .filter(|u| !u.is_empty())
3591 });
3592 let user = match user {
3593 Some(u) => u,
3594 None => {
3595 let _ = sender.send(InstallerEvent::Log(
3596 format!(
3597 "[native] {component_name} — could not determine target user; skipping \
3598 group check. Add yourself manually: sudo usermod -aG {} $USER",
3599 groups.join(",")
3600 ),
3601 false,
3602 ));
3603 return;
3604 }
3605 };
3606
3607 let current_groups: Vec<String> = Command::new("id")
3610 .args(["-nG", &user])
3611 .output()
3612 .ok()
3613 .filter(|o| o.status.success())
3614 .map(|o| {
3615 String::from_utf8_lossy(&o.stdout)
3616 .split_whitespace()
3617 .map(|s| s.to_string())
3618 .collect()
3619 })
3620 .unwrap_or_default();
3621
3622 let mut to_add: Vec<&str> = Vec::new();
3624 for g in groups {
3625 let exists = Command::new("getent")
3626 .args(["group", g])
3627 .output()
3628 .map(|o| o.status.success())
3629 .unwrap_or(false);
3630 if !exists {
3631 let _ = sender.send(InstallerEvent::Log(
3632 format!(
3633 "[native] {component_name} — group '{g}' not present on this system; skipping"
3634 ),
3635 false,
3636 ));
3637 continue;
3638 }
3639 if !current_groups.iter().any(|c| c == g) {
3640 to_add.push(*g);
3641 }
3642 }
3643
3644 if to_add.is_empty() {
3645 let _ = sender.send(InstallerEvent::Log(
3646 format!(
3647 "[native] {component_name} — user '{user}' already in all required group(s); nothing to add"
3648 ),
3649 false,
3650 ));
3651 return;
3652 }
3653
3654 let csv = to_add.join(",");
3655
3656 let askpass = sudo_pw.filter(|p| !p.is_empty()).and_then(|p| {
3658 match crate::installers::common::askpass::Askpass::new(p) {
3659 Ok(g) => Some(g),
3660 Err(e) => {
3661 let _ = sender.send(InstallerEvent::Log(
3662 format!(
3663 "[native] {component_name} — askpass setup failed ({e}); cannot add \
3664 groups non-interactively"
3665 ),
3666 true,
3667 ));
3668 None
3669 }
3670 }
3671 });
3672
3673 let added = if let Some(ap) = askpass {
3674 let args = vec![
3677 "-A".to_string(),
3678 "usermod".to_string(),
3679 "-aG".to_string(),
3680 csv.clone(),
3681 user.clone(),
3682 ];
3683 let envs = vec![("SUDO_ASKPASS".to_string(), ap.path().to_string())];
3684 let cmd = NativeCommand::from_shell_cmd("sudo", &args, &envs);
3685 execute_native_command(&cmd, None, sender, component_name).is_ok()
3686 } else {
3687 let _ = sender.send(InstallerEvent::Log(
3688 format!(
3689 "[native] {component_name} — no sudo password available; please run: \
3690 sudo usermod -aG {csv} {user}"
3691 ),
3692 true,
3693 ));
3694 return;
3695 };
3696
3697 if added {
3698 let _ = sender.send(InstallerEvent::Log(
3699 format!(
3700 "[native] {component_name} — added user '{user}' to group(s): {csv}. \
3701 A re-login (or reboot) is required for the new groups to take effect."
3702 ),
3703 false,
3704 ));
3705 }
3706}
3707
3708fn run_rccl_remediation(ctx: &NativeInstallerContext, component_name: &str) -> Result<()> {
3710 let sender = ctx.sender;
3711 let rccl_remediation = crate::installers::components::rccl::RcclInstaller::new(
3712 ctx.user_home,
3713 resolve_python_bin(),
3714 ctx.env_exports
3715 .get("ROCM_PATH")
3716 .cloned()
3717 .unwrap_or_else(|| "/opt/rocm".to_string()),
3718 ctx.env_exports
3719 .get("ROCM_VERSION")
3720 .cloned()
3721 .unwrap_or_else(|| "7.2.0".to_string()),
3722 ctx.env_exports
3723 .get("ROCM_CHANNEL")
3724 .cloned()
3725 .unwrap_or_else(|| "latest".to_string()),
3726 ctx.env_exports
3727 .get("ROCR_VISIBLE_DEVICES")
3728 .cloned()
3729 .or_else(|| ctx.env_exports.get("HIP_VISIBLE_DEVICES").cloned()),
3730 );
3731
3732 let _ = sender.send(InstallerEvent::Log(
3733 "[native] Probing RCCL multi-GPU collectives...".into(),
3734 false,
3735 ));
3736 if let Some(build) = rccl_remediation.prepare_remediation()? {
3737 let _ = sender.send(InstallerEvent::Log(
3738 "[native] Known RCCL collective failure detected; building sealed overlay...".into(),
3739 false,
3740 ));
3741 execute_native_command(
3742 &NativeCommand::from_shell_cmd(&build.program, &build.args, &build.env),
3743 None,
3744 sender,
3745 component_name,
3746 )?;
3747 rccl_remediation.activate_and_verify()?;
3748 let _ = sender.send(InstallerEvent::Log(
3749 "[native] RCCL sealed overlay activated and validated on both GPUs".into(),
3750 false,
3751 ));
3752 } else {
3753 let _ = sender.send(InstallerEvent::Log(
3754 "[native] RCCL probe healthy or existing sealed overlay revalidated; no repair needed"
3755 .into(),
3756 false,
3757 ));
3758 }
3759 Ok(())
3760}
3761
3762fn run_native_installer(component: &Component, ctx: &NativeInstallerContext) -> Result<()> {
3770 let sender = ctx.sender;
3771 let _ = sender.send(InstallerEvent::Log(
3772 format!("[native] Installing {} via Rust module...", component.name),
3773 false,
3774 ));
3775
3776 let _ = sender.send(InstallerEvent::Progress {
3777 component_id: component.id.clone(),
3778 progress: 0.1,
3779 message: format!("Starting {} installation", component.name),
3780 });
3781
3782 let component_needs_sudo = component.needs_sudo && needs_sudo();
3784 let cached_sudo =
3785 component_needs_sudo && ctx.sudo_password.is_none() && can_sudo_non_interactive();
3786 let sudo_pw: Option<&str> = if component_needs_sudo {
3787 ctx.sudo_password
3788 .as_deref()
3789 .or(if cached_sudo { Some("") } else { None })
3790 } else {
3791 None
3792 };
3793
3794 if component_needs_sudo && sudo_pw.is_none() {
3796 bail!("{} requires sudo but no password provided", component.name);
3797 }
3798
3799 let common_env = build_common_env(ctx);
3801
3802 match component.id.as_str() {
3806 "permanent-env" => {
3808 use crate::installers::components::permanent_env::PermanentEnvInstaller;
3809 match ensure_mlstack_env(ctx.user_home, ctx.install_method, is_force_reinstall()) {
3820 Ok(EnvUpdate::Created) | Ok(EnvUpdate::Updated) => {
3821 let _ = sender.send(InstallerEvent::Log(
3822 "[native] Permanent ROCm Env — regenerated ~/.mlstack_env from current state".into(),
3823 false,
3824 ));
3825 }
3826 Ok(EnvUpdate::Unchanged) => {
3827 let _ = sender.send(InstallerEvent::Log(
3828 "[native] Permanent ROCm Env — ~/.mlstack_env already up to date".into(),
3829 false,
3830 ));
3831 }
3832 Err(e) => {
3833 let _ = sender.send(InstallerEvent::Log(
3834 format!(
3835 "[native] Permanent ROCm Env — could not regenerate ~/.mlstack_env: {e}"
3836 ),
3837 true,
3838 ));
3839 }
3840 }
3841 crate::platform::environment::write_activate_snippets()?;
3843 crate::platform::environment::offer_global_prepend(true)?;
3844 let inst = PermanentEnvInstaller::default();
3845 let cmd = inst.build_mkdir_triton_dirs_command();
3846 execute_native_command(
3847 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
3848 sudo_pw,
3849 sender,
3850 &component.name,
3851 )?;
3852 }
3853
3854 "rocm" => {
3856 use crate::installers::components::rocm::{RocmConfig, RocmInstaller};
3857 let rocm_force = is_force_reinstall();
3858 use crate::installers::components::rocm::RocmChannel;
3859 let rocm_channel = std::env::var("ROCM_CHANNEL")
3860 .ok()
3861 .or_else(|| ctx.env_exports.get("ROCM_CHANNEL").cloned())
3862 .and_then(|c| RocmChannel::parse(&c))
3863 .unwrap_or(RocmChannel::Latest);
3864 let inst = RocmInstaller::new(RocmConfig {
3865 force_reinstall: rocm_force,
3866 channel: rocm_channel,
3867 gfx_arch: detect_gpu_arch().unwrap_or_else(|| "gfx1100".to_string()),
3868 legacy_fallback: std::env::var("ROCM_LEGACY_VERSION")
3869 .map(|v| v.trim() == "7.2.3")
3870 .unwrap_or(false),
3871 ..RocmConfig::default()
3872 });
3873 let distro = crate::installers::common::DistroFacade::detect();
3874 let is_arch = distro.is_arch_family();
3875 if rocm_channel.is_therock() && (is_arch || rocm_channel == RocmChannel::Preview) {
3876 let home = std::env::var("HOME").unwrap_or_default();
3879 let venv = format!("{home}/.mlstack/rocm/{}", inst.effective_version());
3880 let index = inst
3881 .pip_index_url()
3882 .ok_or_else(|| anyhow::anyhow!("TheRock channel has no pip index"))?
3883 .to_string();
3884 if !std::process::Command::new("uv")
3885 .arg("--version")
3886 .output()
3887 .map(|o| o.status.success())
3888 .unwrap_or(false)
3889 {
3890 bail!(
3891 "`uv` is required to install ROCm {} into a user-space environment; install it from https://docs.astral.sh/uv/ and retry",
3892 inst.effective_version()
3893 );
3894 }
3895 let venv_python =
3896 std::env::var("ROCM_VENV_PYTHON").unwrap_or_else(|_| "3.12".to_string());
3897 let req = inst.pip_requirement();
3898 let steps: Vec<(&str, Vec<String>)> = vec![
3899 (
3900 "uv",
3901 vec![
3902 "venv".into(),
3903 "--python".into(),
3904 venv_python.clone(),
3905 venv.clone(),
3906 ],
3907 ),
3908 (
3909 "uv",
3910 vec![
3911 "pip".into(),
3912 "install".into(),
3913 "--python".into(),
3914 format!("{venv}/bin/python"),
3915 "--index-url".into(),
3916 index,
3917 "--extra-index-url".into(),
3918 "https://pypi.org/simple".into(),
3919 "--index-strategy".into(),
3920 "unsafe-best-match".into(),
3921 req,
3922 ],
3923 ),
3924 ];
3925 for (prog, args) in steps {
3926 execute_native_command(
3927 &NativeCommand::Pip {
3928 program: prog.to_string(),
3929 args,
3930 },
3931 None,
3932 sender,
3933 &component.name,
3934 )?;
3935 }
3936 let _ = sender.send(InstallerEvent::Log(
3937 format!(
3938 "[native] ROCm {} installed to {venv}",
3939 inst.effective_version()
3940 ),
3941 false,
3942 ));
3943 } else if is_arch {
3944 let all_pkgs = inst.pacman_rocm_packages();
3950 let force_reinstall = is_force_reinstall();
3951 let pkgs_to_install = if force_reinstall {
3952 all_pkgs.clone()
3953 } else {
3954 let need_install = RocmInstaller::filter_already_installed_pacman(&all_pkgs);
3955 if need_install.is_empty() {
3956 let _ = sender.send(InstallerEvent::Log(
3957 format!(
3958 "[native] {} — all ROCm packages already installed, skipping",
3959 component.name
3960 ),
3961 false,
3962 ));
3963 return Ok(());
3964 }
3965 need_install
3966 };
3967
3968 let _ = sender.send(InstallerEvent::Log(
3969 format!(
3970 "[native] {} — installing {} of {} ROCm packages via yay \
3971 (user-space, non-interactive): {}",
3972 component.name,
3973 pkgs_to_install.len(),
3974 all_pkgs.len(),
3975 pkgs_to_install.join(", ")
3976 ),
3977 false,
3978 ));
3979
3980 let mut args = vec!["-S".to_string()];
3990 if !force_reinstall {
3991 args.push("--needed".to_string());
3992 }
3993 args.push("--noconfirm".to_string());
3994
3995 let askpass = sudo_pw.filter(|p| !p.is_empty()).and_then(|p| {
3997 match crate::installers::common::askpass::Askpass::new(p) {
3998 Ok(g) => Some(g),
3999 Err(e) => {
4000 let _ = sender.send(InstallerEvent::Log(
4001 format!(
4002 "[native] {} — askpass setup failed ({e}); yay's \
4003 privileged steps may not authenticate",
4004 component.name
4005 ),
4006 true,
4007 ));
4008 None
4009 }
4010 }
4011 });
4012 let mut envs: Vec<(String, String)> = Vec::new();
4013 if let Some(ref ap) = askpass {
4014 envs.push(("SUDO_ASKPASS".to_string(), ap.path().to_string()));
4015 args.push("--sudoflags=-A".to_string());
4016 } else {
4017 let _ = sender.send(InstallerEvent::Log(
4018 format!(
4019 "[native] {} — no sudo password supplied; yay will only \
4020 authenticate if sudo is cached / passwordless",
4021 component.name
4022 ),
4023 false,
4024 ));
4025 }
4026 args.extend(pkgs_to_install);
4027
4028 let yay_cmd = NativeCommand::Shell {
4029 program: "yay".to_string(),
4030 args,
4031 env: envs,
4032 working_dir: None,
4033 };
4034 execute_native_command(&yay_cmd, None, sender, &component.name)?;
4037 } else {
4038 let therock = rocm_channel.is_therock();
4039 let commands = if distro.uses_apt() {
4040 if therock {
4041 inst.therock_apt_commands(&distro)
4042 .map_err(|e| anyhow::anyhow!(e))?
4043 } else {
4044 inst.apt_install_commands()
4045 }
4046 } else if distro.uses_dnf() || distro.uses_yum() {
4047 if therock {
4048 inst.therock_rpm_commands(&distro)
4049 .map_err(|e| anyhow::anyhow!(e))?
4050 } else {
4051 let rhel_ver = distro.version().split('.').next().unwrap_or("9");
4052 inst.dnf_install_commands(rhel_ver)
4053 }
4054 } else if distro.uses_zypper() {
4055 if therock {
4056 inst.therock_rpm_commands(&distro)
4057 .map_err(|e| anyhow::anyhow!(e))?
4058 } else {
4059 inst.zypper_install_commands()
4060 }
4061 } else {
4062 bail!(
4063 "Unsupported distribution for ROCm installation: {} ({})",
4064 distro.name(),
4065 distro.id()
4066 );
4067 };
4068 for pkg_cmd in &commands {
4069 let native_cmd = NativeCommand::Package {
4070 program: pkg_cmd.program.clone(),
4071 args: pkg_cmd.args.clone(),
4072 };
4073 execute_native_command(&native_cmd, sudo_pw, sender, &component.name)?;
4076 }
4077 }
4078
4079 if is_force_reinstall() {
4081 let marker_path = std::path::Path::new("/tmp/mlstack-reboot-required");
4082 let _ = std::fs::write(
4083 marker_path,
4084 "ROCm packages were reinstalled. Reboot required for kernel module changes.\n",
4085 );
4086 let _ = sender.send(InstallerEvent::Log(
4087 "[native] ROCm — reboot recommended for kernel module changes".to_string(),
4088 false,
4089 ));
4090 }
4091
4092 ensure_user_in_groups(&["render", "video"], sudo_pw, sender, &component.name);
4097 }
4098
4099 "pytorch" => {
4101 use crate::installers::components::pytorch::{PyTorchConfig, PyTorchInstaller};
4102 let force = is_force_reinstall();
4103 let python_bin = resolve_python_bin();
4104 let config = PyTorchConfig {
4105 python_bin: python_bin.clone(),
4106 force_reinstall: force,
4107 ..PyTorchConfig::default()
4108 };
4109 let inst = PyTorchInstaller::new(config);
4110 let use_uv = crate::installers::common::utils::command_exists("uv");
4111 let rocm_mm = ctx
4112 .env_exports
4113 .get("ROCM_VERSION")
4114 .map(|v| {
4115 let parts: Vec<&str> = v.split('.').collect();
4116 if parts.len() >= 2 {
4117 format!("{}.{}", parts[0], parts[1])
4118 } else {
4119 v.clone()
4120 }
4121 })
4122 .unwrap_or_else(|| "7.2".to_string());
4123
4124 if force {
4126 purge_pytorch(sender, &component.name);
4127 }
4128
4129 let install_cmd =
4138 inst.build_install_command(&rocm_mm, use_uv, ctx.target_version.as_deref());
4139 execute_native_command(
4140 &NativeCommand::Pip {
4141 program: install_cmd.program.clone(),
4142 args: install_cmd.args.clone(),
4143 },
4144 None,
4145 sender,
4146 &component.name,
4147 )?;
4148
4149 let deps_cmd = inst.build_common_deps_command(use_uv);
4154 let _ = common_env.clone(); execute_native_command(
4156 &NativeCommand::Pip {
4157 program: deps_cmd.program.clone(),
4158 args: deps_cmd.args.clone(),
4159 },
4160 None, sender,
4162 &component.name,
4163 )?;
4164
4165 let managed_python = resolve_python_bin();
4170 let managed_install =
4171 crate::installers::common::managed_python_install_command(managed_python.clone());
4172 execute_native_command(
4173 &NativeCommand::Pip {
4174 program: managed_install.program,
4175 args: managed_install.args,
4176 },
4177 None,
4178 sender,
4179 &component.name,
4180 )?;
4181
4182 let managed_verify =
4183 crate::installers::common::managed_python_verify_command(managed_python);
4184 execute_native_command(
4185 &NativeCommand::from_shell_cmd(&managed_verify.program, &managed_verify.args, &[]),
4186 None,
4187 sender,
4188 &component.name,
4189 )?;
4190 }
4191
4192 "triton" => {
4194 use crate::installers::components::triton::{TritonConfig, TritonInstaller};
4195
4196 let python_bin = resolve_python_bin();
4213 let triton_present = std::process::Command::new(&python_bin)
4214 .arg("-c")
4215 .arg("import triton")
4216 .status()
4217 .map(|s| s.success())
4218 .unwrap_or(false);
4219 let triton_satisfies_target = triton_present
4220 && triton_installed_satisfies_target(&python_bin, ctx.target_version.as_deref());
4221 if triton_satisfies_target {
4222 let _ = sender.send(InstallerEvent::Log(
4223 "[native] Triton — already importable (triton-rocm bundled with PyTorch) \
4224 and satisfies the requested target; skipping the redundant from-source \
4225 ROCm/triton build"
4226 .into(),
4227 false,
4228 ));
4229 return Ok(());
4230 }
4231
4232 let triton_force = is_force_reinstall();
4233 let inst = TritonInstaller::new(TritonConfig {
4234 python_bin: resolve_python_bin(),
4235 ..Default::default()
4236 });
4237 let target_dir = format!("{}/.mlstack/triton", ctx.user_home);
4238
4239 if triton_force {
4241 purge_triton(sender, &component.name);
4242 }
4243
4244 git_clone_or_pull(
4248 inst.git_clone_url(),
4249 &target_dir,
4250 &[],
4251 sudo_pw,
4252 sender,
4253 &component.name,
4254 )?;
4255
4256 fix_directory_ownership(&target_dir, sudo_pw, sender);
4259
4260 let safe_dir_cmd = NativeCommand::from_shell_cmd(
4262 "git",
4263 &[
4264 "config".to_string(),
4265 "--global".to_string(),
4266 "--add".to_string(),
4267 "safe.directory".to_string(),
4268 target_dir.clone(),
4269 ],
4270 &[],
4271 );
4272 let _ = execute_native_command(&safe_dir_cmd, None, sender, &component.name);
4273
4274 let branch = inst.select_branch().to_string();
4276 let fetch_branch_cmd = NativeCommand::from_shell_cmd(
4277 "git",
4278 &[
4279 "-C".to_string(),
4280 target_dir.clone(),
4281 "fetch".to_string(),
4282 "origin".to_string(),
4283 branch.clone(),
4284 ],
4285 &[],
4286 );
4287 let _ = execute_native_command(&fetch_branch_cmd, None, sender, &component.name);
4288
4289 let checkout_cmd = inst.build_git_checkout_command();
4291 let mut checkout_args = vec!["-C".to_string(), target_dir.clone()];
4292 checkout_args.extend(checkout_cmd.args.iter().cloned());
4293 let checkout_result = execute_native_command(
4294 &NativeCommand::from_shell_cmd(
4295 &checkout_cmd.program,
4296 &checkout_args,
4297 &checkout_cmd.env,
4298 ),
4299 None,
4300 sender,
4301 &component.name,
4302 );
4303 if let Err(err) = checkout_result {
4304 let _ = sender.send(InstallerEvent::Log(
4305 format!(
4306 "[native] [WARN] {} git checkout failed (non-fatal): {}",
4307 component.name, err
4308 ),
4309 false,
4310 ));
4311 }
4312
4313 let cmd = inst.build_prerequisites_command();
4315 execute_native_command(
4316 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4317 None,
4318 sender,
4319 &component.name,
4320 )?;
4321
4322 let cmd = inst.build_pip_install_command(&target_dir);
4324 execute_native_command(
4325 &NativeCommand::from_shell_cmd_with_dir(
4326 &cmd.program,
4327 &cmd.args,
4328 &cmd.env,
4329 cmd.working_dir.clone(),
4330 ),
4331 None,
4332 sender,
4333 &component.name,
4334 )?;
4335 }
4336
4337 "mpi4py" => {
4339 use crate::installers::components::mpi4py::{Mpi4PyConfig, Mpi4PyInstaller};
4340 let force = is_force_reinstall();
4341 let inst = Mpi4PyInstaller::new(Mpi4PyConfig {
4342 python_bin: resolve_python_bin(),
4343 force_reinstall: force,
4344 ..Mpi4PyConfig::default()
4345 });
4346 let distro = crate::installers::common::DistroFacade::detect();
4347
4348 if force {
4350 purge_component_packages(&["mpi4py"], sender, &component.name);
4351 }
4352
4353 let cmd = inst.build_system_mpi_install_command(&distro);
4355 execute_native_command(
4356 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4357 sudo_pw,
4358 sender,
4359 &component.name,
4360 )?;
4361
4362 let mpi_path = std::path::PathBuf::from("/usr");
4364 let cmd = inst.build_pip_install_command(&mpi_path);
4365 execute_native_command(
4366 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4367 None,
4368 sender,
4369 &component.name,
4370 )?;
4371 }
4372
4373 "deepspeed" => {
4375 use crate::installers::components::deepspeed::{DeepSpeedConfig, DeepSpeedInstaller};
4376 let force = is_force_reinstall();
4377 let config = DeepSpeedConfig {
4378 python_bin: resolve_python_bin(),
4379 force_reinstall: force,
4380 ..DeepSpeedConfig::default()
4381 };
4382 let inst = DeepSpeedInstaller::new(config);
4383
4384 if force {
4386 purge_component_packages(&["deepspeed"], sender, &component.name);
4387 }
4388
4389 let deps_cmd = inst.build_deps_install_command();
4391 execute_native_command(
4392 &NativeCommand::Pip {
4393 program: deps_cmd.program.clone(),
4394 args: deps_cmd.args.clone(),
4395 },
4396 None,
4397 sender,
4398 &component.name,
4399 )?;
4400
4401 let install_cmd = inst.build_install_command();
4403 execute_native_command(
4404 &NativeCommand::Pip {
4405 program: install_cmd.program.clone(),
4406 args: install_cmd.args.clone(),
4407 },
4408 None,
4409 sender,
4410 &component.name,
4411 )?;
4412 }
4413
4414 "ml-stack-core" => {
4416 use crate::installers::components::ml_stack::{MlStackConfig, MlStackInstaller};
4417 let inst = MlStackInstaller::new(MlStackConfig {
4418 python_bin: resolve_python_bin(),
4419 ..Default::default()
4420 });
4421 let distro = crate::installers::common::DistroFacade::detect();
4422
4423 let migraphx_cmd = inst.build_migraphx_install_command(&distro);
4425 execute_native_command(
4426 &NativeCommand::System {
4427 program: migraphx_cmd.program.clone(),
4428 args: migraphx_cmd.args.clone(),
4429 },
4430 sudo_pw,
4431 sender,
4432 &component.name,
4433 )?;
4434
4435 let rccl_cmd = inst.build_rccl_install_command(&distro);
4436 execute_native_command(
4437 &NativeCommand::System {
4438 program: rccl_cmd.program.clone(),
4439 args: rccl_cmd.args.clone(),
4440 },
4441 sudo_pw,
4442 sender,
4443 &component.name,
4444 )?;
4445
4446 run_rccl_remediation(ctx, &component.name)?;
4449
4450 let mpi_cmd = inst.build_mpi_install_command(&distro);
4451 execute_native_command(
4452 &NativeCommand::System {
4453 program: mpi_cmd.program.clone(),
4454 args: mpi_cmd.args.clone(),
4455 },
4456 sudo_pw,
4457 sender,
4458 &component.name,
4459 )?;
4460
4461 let megatron_dir = format!("{}/.mlstack/src/Megatron-LM", ctx.user_home);
4466 let megatron_clone = inst.build_megatron_clone_command(&megatron_dir);
4467 execute_native_command(
4468 &NativeCommand::System {
4469 program: megatron_clone.program.clone(),
4470 args: megatron_clone.args.clone(),
4471 },
4472 sudo_pw,
4473 sender,
4474 &component.name,
4475 )?;
4476
4477 let megatron_install = inst.build_megatron_install_command();
4478 execute_native_command(
4479 &NativeCommand::from_shell_cmd_with_dir(
4480 &megatron_install.program,
4481 &megatron_install.args,
4482 &[],
4483 Some(PathBuf::from(&megatron_dir)),
4484 ),
4485 None,
4486 sender,
4487 &component.name,
4488 )?;
4489 }
4490
4491 "flash-attn" | "flash-attn-triton" | "flash-attn-ck" => {
4501 let force = is_force_reinstall();
4502 let is_triton = matches!(component.id.as_str(), "flash-attn" | "flash-attn-triton");
4505 let backend_name = if is_triton { "triton" } else { "ck" };
4506
4507 if force {
4509 purge_flash_attn(sender, &component.name);
4510 }
4511
4512 let target_dir = format!("{}/.mlstack/flash-attention", ctx.user_home);
4513
4514 git_clone_or_pull(
4515 "https://github.com/ROCm/flash-attention.git",
4516 &target_dir,
4517 &[],
4518 sudo_pw,
4519 sender,
4520 &component.name,
4521 )?;
4522
4523 fix_directory_ownership(&target_dir, sudo_pw, sender);
4525
4526 let safe_dir_cmd = NativeCommand::from_shell_cmd(
4528 "git",
4529 &[
4530 "config".to_string(),
4531 "--global".to_string(),
4532 "--add".to_string(),
4533 "safe.directory".to_string(),
4534 target_dir.clone(),
4535 ],
4536 &[],
4537 );
4538 let _ = execute_native_command(&safe_dir_cmd, None, sender, &component.name);
4539
4540 let gpu_arch = ctx
4542 .env_exports
4543 .get("GPU_ARCH")
4544 .cloned()
4545 .unwrap_or_else(|| "gfx1100".to_string());
4546 let hsa_version = ctx
4547 .env_exports
4548 .get("HSA_OVERRIDE_GFX_VERSION")
4549 .cloned()
4550 .unwrap_or_else(|| "11.0.0".to_string());
4551
4552 let fa_deps_args = vec![
4558 "-m".to_string(),
4559 "pip".to_string(),
4560 "install".to_string(),
4561 "--break-system-packages".to_string(),
4562 "einops".to_string(),
4563 "ninja".to_string(),
4564 "packaging".to_string(),
4565 "psutil".to_string(),
4566 ];
4567 execute_native_command(
4568 &NativeCommand::Shell {
4569 program: resolve_python_bin(),
4570 args: fa_deps_args,
4571 env: vec![],
4572 working_dir: None,
4573 },
4574 None, sender,
4576 &component.name,
4577 )?;
4578
4579 let mut pip_args = vec![
4586 "-m".to_string(),
4587 "pip".to_string(),
4588 "install".to_string(),
4589 "-v".to_string(),
4590 "--no-build-isolation".to_string(),
4591 "--no-deps".to_string(),
4592 "--break-system-packages".to_string(),
4593 ".".to_string(),
4594 ];
4595 if force {
4596 pip_args.push("--force-reinstall".to_string());
4597 }
4598
4599 let mut pip_env = vec![
4600 ("GPU_ARCH".to_string(), gpu_arch.clone()),
4601 ("PYTORCH_ROCM_ARCH".to_string(), gpu_arch),
4602 ("ROCM_PATH".to_string(), "/opt/rocm".to_string()),
4603 ("HSA_OVERRIDE_GFX_VERSION".to_string(), hsa_version),
4604 ];
4605 if is_triton {
4606 pip_env.push((
4608 "FLASH_ATTENTION_TRITON_AMD_ENABLE".to_string(),
4609 "TRUE".to_string(),
4610 ));
4611 } else {
4612 let jobs = std::thread::available_parallelism()
4616 .map(|n| n.get())
4617 .unwrap_or(4)
4618 .min(8);
4619 pip_env.push(("MAX_JOBS".to_string(), jobs.to_string()));
4620 }
4621
4622 execute_native_command(
4623 &NativeCommand::Shell {
4624 program: resolve_python_bin(),
4627 args: pip_args,
4628 env: pip_env,
4629 working_dir: Some(PathBuf::from(&target_dir)),
4630 },
4631 None, sender,
4633 &component.name,
4634 )?;
4635
4636 let marker = PathBuf::from(&target_dir).join(".backend");
4642 std::fs::write(&marker, backend_name).map_err(|err| {
4643 anyhow::anyhow!(
4644 "{}: failed to write Flash Attention backend marker {}: {} \
4645 — package installed but untrackable; failing install",
4646 component.name,
4647 marker.display(),
4648 err
4649 )
4650 })?;
4651 }
4652
4653 "repair-stack" => {
4657 use crate::installers::components::repair::RepairInstaller;
4658 let repair_ids = RepairInstaller::native_component_ids();
4659
4660 let _ = sender.send(InstallerEvent::Log(
4661 "[native] Starting ML Stack repair (8 steps via native installers)...".to_string(),
4662 false,
4663 ));
4664
4665 for step_id in &repair_ids {
4666 let sub_component = Component {
4669 id: (*step_id).to_string(),
4670 name: (*step_id).to_string(),
4671 description: String::new(),
4672 script: String::new(),
4673 category: Category::Core,
4674 required: true,
4675 selected: true,
4676 installed: false,
4677 progress: 0.0,
4678 estimate: String::new(),
4679 needs_sudo: component.needs_sudo,
4680 experimental: false,
4681 note: None,
4682 };
4683
4684 let _ = sender.send(InstallerEvent::Progress {
4685 component_id: component.id.clone(),
4686 progress: 0.5,
4687 message: format!("Repairing {}...", step_id),
4688 });
4689
4690 if let Err(err) = run_native_installer(&sub_component, ctx) {
4691 let _ = sender.send(InstallerEvent::Log(
4693 format!("[native] Repair step '{}' failed: {}", step_id, err),
4694 false,
4695 ));
4696 }
4697 }
4698 }
4699
4700 "rccl-repair" => {
4701 run_rccl_remediation(ctx, &component.name)?;
4702 }
4703
4704 "megatron" => {
4706 use crate::installers::components::megatron::{MegatronConfig, MegatronInstaller};
4707 let force = is_force_reinstall();
4708 let inst = MegatronInstaller::new(MegatronConfig {
4709 python_bin: resolve_python_bin(),
4710 force_reinstall: force,
4711 ..MegatronConfig::default()
4712 });
4713
4714 if force {
4716 purge_component_packages(&["megatron-core", "megatron"], sender, &component.name);
4717 }
4718
4719 for dep_pkg in &["ninja", "packaging", "wheel"] {
4721 let cmd = inst.build_dep_install_command(dep_pkg);
4722 execute_native_command(
4723 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4724 None,
4725 sender,
4726 &component.name,
4727 )?;
4728 }
4729
4730 let clone_cmd = inst.build_git_clone_command();
4732 let clone_target_str = clone_cmd.args.last().cloned().unwrap_or_default();
4733 git_clone_or_pull(
4734 "https://github.com/NVIDIA/Megatron-LM.git",
4735 &clone_target_str,
4736 &[],
4737 sudo_pw,
4738 sender,
4739 &component.name,
4740 )?;
4741
4742 let run_user = std::env::var("SUDO_USER")
4748 .or_else(|_| std::env::var("USER"))
4749 .unwrap_or_default();
4750 if !run_user.is_empty() {
4751 let chown_cmd = NativeCommand::from_shell_cmd(
4752 "chown",
4753 &[
4754 "-R".to_string(),
4755 format!("{run_user}:{run_user}"),
4756 clone_target_str.clone(),
4757 ],
4758 &[],
4759 );
4760 let _ = execute_native_command(&chown_cmd, sudo_pw, sender, &component.name);
4761 }
4762
4763 let clone_target_path = PathBuf::from(&clone_target_str);
4767 if let Ok(entries) = std::fs::read_dir(&clone_target_path) {
4768 for entry in entries.flatten() {
4769 let name = entry.file_name();
4770 let name_str = name.to_string_lossy();
4771 if name_str.ends_with(".egg-info") {
4772 let egg_path = entry.path();
4773 if std::fs::remove_dir_all(&egg_path).is_err() {
4775 let rm_cmd = NativeCommand::from_shell_cmd(
4777 "rm",
4778 &["-rf".to_string(), egg_path.to_string_lossy().to_string()],
4779 &[],
4780 );
4781 let _ =
4782 execute_native_command(&rm_cmd, sudo_pw, sender, &component.name);
4783 }
4784 }
4785 }
4786 }
4787
4788 let cmd = inst.build_pip_install_command();
4790 execute_native_command(
4791 &NativeCommand::from_shell_cmd_with_dir(
4792 &cmd.program,
4793 &cmd.args,
4794 &cmd.env,
4795 cmd.working_dir.clone(),
4796 ),
4797 None,
4798 sender,
4799 &component.name,
4800 )?;
4801
4802 let requirements_candidates = [
4804 clone_target_path
4805 .join("requirements")
4806 .join("requirements.txt"),
4807 clone_target_path.join("requirements.txt"),
4808 clone_target_path
4809 .join("requirements")
4810 .join("requirements_amd.txt"),
4811 ];
4812 for req_path in requirements_candidates.iter().filter(|p| p.exists()) {
4813 let filtered = match filter_cuda_requirements(&req_path.to_string_lossy()) {
4814 Ok(f) => f,
4815 Err(e) => {
4816 let _ = sender.send(InstallerEvent::Log(
4817 format!(
4818 "[native] WARN: could not read Megatron requirements {}: {e} — skipping requirements install",
4819 req_path.display()
4820 ),
4821 true,
4822 ));
4823 anyhow::bail!(
4824 "failed to read Megatron requirements {}: {e}",
4825 req_path.display()
4826 );
4827 }
4828 };
4829 if filtered.is_empty() {
4830 continue;
4831 }
4832 let tmp_dir = std::env::temp_dir().join("rusty-stack-megatron");
4833 let _ = std::fs::create_dir_all(&tmp_dir);
4834 let filtered_path = tmp_dir.join("requirements_filtered.txt");
4835 std::fs::write(&filtered_path, &filtered).with_context(|| {
4836 format!(
4837 "Failed to write filtered Megatron requirements to {:?}",
4838 filtered_path
4839 )
4840 })?;
4841 let filtered_path_str = filtered_path.to_string_lossy().to_string();
4842 let filtered_cmd = inst.build_requirements_install_command(&filtered_path_str);
4843 execute_native_command(
4844 &NativeCommand::from_shell_cmd(
4845 &filtered_cmd.program,
4846 &filtered_cmd.args,
4847 &filtered_cmd.env,
4848 ),
4849 None,
4850 sender,
4851 &component.name,
4852 )?;
4853 let _ = std::fs::remove_file(&filtered_path);
4854 break;
4855 }
4856 }
4857
4858 "vllm" => {
4860 use crate::installers::components::vllm_multi::{VllmConfig, VllmInstaller};
4861 let vllm_force = is_force_reinstall();
4862 let inst = VllmInstaller::new(VllmConfig {
4863 python_bin: resolve_python_bin(),
4864 ..Default::default()
4865 });
4866
4867 if vllm_force {
4869 purge_vllm(sender, &component.name);
4870 }
4871
4872 let cmd = inst.build_tools_install_command();
4875 execute_native_command(
4876 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4877 None,
4878 sender,
4879 &component.name,
4880 )?;
4881
4882 let vllm_version = inst
4887 .resolve_version_with_policy(ctx.vllm_version_lag, ctx.vllm_version_min_age_days);
4888 let _ = sender.send(InstallerEvent::Log(
4889 format!(
4890 "[native] vLLM: selected {vllm_version} (gate: lag={}, min_age_days={}) — \
4891 building from source (ROCm, env torch)",
4892 ctx.vllm_version_lag, ctx.vllm_version_min_age_days
4893 ),
4894 false,
4895 ));
4896 let cmd = inst.build_source_install_command(&vllm_version);
4901 execute_native_command(
4902 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4903 None,
4904 sender,
4905 &component.name,
4906 )?;
4907
4908 let cmd = inst.build_versioned_deps_command();
4912 execute_native_command(
4913 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4914 None,
4915 sender,
4916 &component.name,
4917 )?;
4918
4919 let cmd = inst.build_deps_install_command();
4921 execute_native_command(
4922 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4923 None,
4924 sender,
4925 &component.name,
4926 )?;
4927
4928 let cmd = inst.build_integrity_check_command();
4932 execute_native_command(
4933 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4934 None,
4935 sender,
4936 &component.name,
4937 )?;
4938 }
4939
4940 "aiter" => {
4942 use crate::installers::components::aiter::{AiterConfig, AiterInstaller};
4943 let detected_gpu_arch = ctx
4946 .env_exports
4947 .get("GPU_ARCH")
4948 .map(|v| v.trim().to_string())
4949 .filter(|v| !v.is_empty())
4950 .unwrap_or_else(|| detect_gpu_arch().unwrap_or_else(|| "gfx1100".to_string()));
4951 let inst = AiterInstaller::new(AiterConfig {
4952 python_bin: resolve_python_bin(),
4953 gpu_arch: detected_gpu_arch,
4954 ..Default::default()
4955 });
4956 let target_dir = format!("{}/.mlstack/aiter", ctx.user_home);
4957
4958 if is_force_reinstall() {
4960 purge_aiter(sender, &component.name);
4961 }
4962
4963 git_clone_or_pull(
4965 inst.git_clone_url(),
4966 &target_dir,
4967 &["--recursive"],
4968 sudo_pw,
4969 sender,
4970 &component.name,
4971 )?;
4972
4973 let run_user = std::env::var("SUDO_USER")
4977 .or_else(|_| std::env::var("USER"))
4978 .unwrap_or_default();
4979 if !run_user.is_empty() {
4980 let chown_cmd = NativeCommand::from_shell_cmd(
4981 "chown",
4982 &[
4983 "-R".to_string(),
4984 format!("{run_user}:{run_user}"),
4985 target_dir.clone(),
4986 ],
4987 &[],
4988 );
4989 let _ = execute_native_command(&chown_cmd, sudo_pw, sender, &component.name);
4990 }
4991
4992 let cmd = inst.build_deps_install_command();
4994 execute_native_command(
4995 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
4996 sudo_pw,
4997 sender,
4998 &component.name,
4999 )?;
5000
5001 let cmd = inst.build_pip_install_command(&target_dir);
5003 execute_native_command(
5004 &NativeCommand::from_shell_cmd_with_dir(
5005 &cmd.program,
5006 &cmd.args,
5007 &cmd.env,
5008 cmd.working_dir.clone(),
5009 ),
5010 None,
5011 sender,
5012 &component.name,
5013 )?;
5014
5015 let run_user = std::env::var("SUDO_USER")
5019 .or_else(|_| std::env::var("USER"))
5020 .unwrap_or_default();
5021 if !run_user.is_empty() {
5022 let user_home = &ctx.user_home;
5023 for aiter_dir in [
5024 format!("{user_home}/.aiter"),
5025 format!("{user_home}/.mlstack/aiter"),
5026 ] {
5027 let dir_path = PathBuf::from(&aiter_dir);
5028 if dir_path.exists() {
5029 let needs_fix = std::fs::read_dir(&dir_path)
5031 .map(|mut entries| entries.next().is_some())
5032 .unwrap_or(false)
5033 && std::fs::write(dir_path.join(".permission_check"), b"").is_err();
5034 if needs_fix {
5035 let chown_cmd = NativeCommand::from_shell_cmd(
5036 "chown",
5037 &[
5038 "-R".to_string(),
5039 format!("{run_user}:{run_user}"),
5040 aiter_dir.clone(),
5041 ],
5042 &[],
5043 );
5044 let _ = execute_native_command(
5045 &chown_cmd,
5046 sudo_pw,
5047 sender,
5048 &component.name,
5049 );
5050 } else {
5051 let _ = std::fs::remove_file(dir_path.join(".permission_check"));
5053 }
5054 }
5055 }
5056 }
5057 }
5058 "vllm-studio" => {
5059 use crate::installers::components::vllm_studio::{
5060 VllmStudioConfig, VllmStudioInstaller,
5061 };
5062 let inst = VllmStudioInstaller::new(VllmStudioConfig {
5063 python_bin: resolve_python_bin(),
5064 ..Default::default()
5065 });
5066
5067 let clone_cmd = inst.build_git_clone_command();
5069 let target_dir = clone_cmd.args.last().cloned().unwrap_or_default();
5070 git_clone_or_pull(
5071 inst.repo_url(),
5072 &target_dir,
5073 &[],
5074 None,
5075 sender,
5076 &component.name,
5077 )?;
5078
5079 let pkg_mgr = if crate::installers::common::utils::command_exists("npm") {
5081 "npm"
5082 } else {
5083 "yarn"
5084 };
5085 let cmd = inst.build_pkg_install_command(pkg_mgr);
5086 execute_native_command(
5087 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5088 None,
5089 sender,
5090 &component.name,
5091 )?;
5092
5093 let frontend_dir = PathBuf::from(&target_dir).join("frontend");
5096 let cmd = inst.build_frontend_build_command(pkg_mgr);
5097 let build_result = execute_native_command(
5098 &NativeCommand::from_shell_cmd_with_dir(
5099 &cmd.program,
5100 &cmd.args,
5101 &cmd.env,
5102 Some(frontend_dir),
5103 ),
5104 None,
5105 sender,
5106 &component.name,
5107 );
5108 if let Err(err) = build_result {
5109 let _ = sender.send(InstallerEvent::Log(
5110 format!(
5111 "[native] [WARN] {} frontend build failed (non-fatal): {}",
5112 component.name, err
5113 ),
5114 false,
5115 ));
5116 }
5117 }
5118
5119 "comfyui" => {
5121 use crate::installers::components::comfyui::{ComfyuiConfig, ComfyuiInstaller};
5122 let inst = ComfyuiInstaller::new(ComfyuiConfig {
5123 python_bin: resolve_python_bin(),
5124 ..Default::default()
5125 });
5126
5127 let cmd = inst.build_pytorch_check_command();
5129 execute_native_command(
5130 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5131 None,
5132 sender,
5133 &component.name,
5134 )?;
5135
5136 let clone_cmd = inst.build_git_clone_command();
5138 let target_dir = clone_cmd.args.last().cloned().unwrap_or_default();
5139 git_clone_or_pull(
5140 inst.repo_url(),
5141 &target_dir,
5142 &[],
5143 None,
5144 sender,
5145 &component.name,
5146 )?;
5147
5148 let req_path = format!("{}/requirements.txt", target_dir);
5154 if std::path::Path::new(&req_path).exists() {
5155 let filtered = match filter_cuda_requirements(&req_path) {
5156 Ok(f) => f,
5157 Err(e) => {
5158 let _ = sender.send(InstallerEvent::Log(
5159 format!(
5160 "[native] WARN: could not read ComfyUI requirements {req_path}: {e} — skipping requirements install"
5161 ),
5162 true,
5163 ));
5164 anyhow::bail!("failed to read ComfyUI requirements {req_path}: {e}");
5165 }
5166 };
5167 if !filtered.is_empty() {
5168 let tmp_dir = std::env::temp_dir().join("rusty-stack-comfyui");
5169 let _ = std::fs::create_dir_all(&tmp_dir);
5170 let filtered_path = tmp_dir.join("requirements_filtered.txt");
5171 std::fs::write(&filtered_path, &filtered).with_context(|| {
5172 format!(
5173 "Failed to write filtered requirements to {:?}",
5174 filtered_path
5175 )
5176 })?;
5177 let filtered_path_str = filtered_path.to_string_lossy().to_string();
5178 let _ = sender.send(InstallerEvent::Log(
5179 format!(
5180 "[native] Installing filtered requirements (CUDA deps excluded) from {}",
5181 req_path
5182 ),
5183 false,
5184 ));
5185 let cmd = inst.build_pip_install_command(&filtered_path_str);
5186 execute_native_command(
5187 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5188 None,
5189 sender,
5190 &component.name,
5191 )?;
5192 let _ = std::fs::remove_file(&filtered_path);
5193 } else {
5194 let _ = sender.send(InstallerEvent::Log(
5195 "[native] [WARN] All ComfyUI requirements were filtered out, nothing to install".to_string(),
5196 false,
5197 ));
5198 }
5199 } else {
5200 let _ = sender.send(InstallerEvent::Log(
5201 "[native] [WARN] ComfyUI requirements.txt not found, skipping dependency installation".to_string(),
5202 false,
5203 ));
5204 }
5205 }
5206
5207 "textgen" => {
5209 use crate::installers::components::textgen::{TextgenConfig, TextgenInstaller};
5210 let textgen_python = resolve_python_bin();
5213 let inst = TextgenInstaller::new(TextgenConfig {
5218 python_bin: textgen_python.clone(),
5219 use_uv: false,
5220 break_system_packages: true,
5221 ..Default::default()
5222 });
5223
5224 let clone_cmd = inst.build_git_clone_command();
5226 let target_dir = clone_cmd.args.last().cloned().unwrap_or_default();
5227 git_clone_or_pull(
5228 inst.repo_url(),
5229 &target_dir,
5230 &[],
5231 None,
5232 sender,
5233 &component.name,
5234 )?;
5235
5236 let amd_req_full = format!("{}/requirements/full/requirements_amd.txt", target_dir);
5241 let amd_req_subdir = format!("{}/requirements/requirements_amd.txt", target_dir);
5242 let amd_req_path = if std::path::Path::new(&amd_req_full).exists() {
5243 amd_req_full
5244 } else if std::path::Path::new(&amd_req_subdir).exists() {
5245 amd_req_subdir
5246 } else {
5247 String::new()
5248 };
5249 if !amd_req_path.is_empty() {
5250 let _ = sender.send(InstallerEvent::Log(
5251 "[native] Installing AMD-specific requirements first (ROCm wheels)".to_string(),
5252 false,
5253 ));
5254 let cmd = inst.build_amd_requirements_command(&amd_req_path);
5255 execute_native_command(
5256 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5257 None,
5258 sender,
5259 &component.name,
5260 )?;
5261 }
5262
5263 let req_path_root = format!("{}/requirements.txt", target_dir);
5270 let req_path_full = format!("{}/requirements/full/requirements.txt", target_dir);
5271 let req_path_subdir = format!("{}/requirements/requirements.txt", target_dir);
5272 let raw_req_path = if std::path::Path::new(&req_path_full).exists() {
5273 req_path_full
5274 } else if std::path::Path::new(&req_path_root).exists() {
5275 req_path_root
5276 } else if std::path::Path::new(&req_path_subdir).exists() {
5277 req_path_subdir
5278 } else {
5279 String::new()
5280 };
5281
5282 if !raw_req_path.is_empty() {
5283 let filtered = match filter_cuda_requirements(&raw_req_path) {
5285 Ok(f) => f,
5286 Err(e) => {
5287 let _ = sender.send(InstallerEvent::Log(
5288 format!(
5289 "[native] WARN: could not read textgen requirements {raw_req_path}: {e} — skipping requirements install"
5290 ),
5291 true,
5292 ));
5293 anyhow::bail!("failed to read textgen requirements {raw_req_path}: {e}");
5294 }
5295 };
5296 if !filtered.is_empty() {
5297 let tmp_dir = std::env::temp_dir().join("rusty-stack-textgen");
5299 let _ = std::fs::create_dir_all(&tmp_dir);
5300 let filtered_path = tmp_dir.join("requirements_filtered.txt");
5301 std::fs::write(&filtered_path, &filtered).with_context(|| {
5302 format!(
5303 "Failed to write filtered requirements to {:?}",
5304 filtered_path
5305 )
5306 })?;
5307 let filtered_path_str = filtered_path.to_string_lossy().to_string();
5308 let _ = sender.send(InstallerEvent::Log(
5309 format!(
5310 "[native] Installing filtered requirements (CUDA deps excluded) from {}",
5311 raw_req_path
5312 ),
5313 false,
5314 ));
5315 let cmd = inst.build_pip_install_command(&filtered_path_str);
5316 execute_native_command(
5317 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5318 None,
5319 sender,
5320 &component.name,
5321 )?;
5322 let _ = std::fs::remove_file(&filtered_path);
5323 } else {
5324 let _ = sender.send(InstallerEvent::Log(
5325 "[native] [WARN] All requirements were filtered out, nothing to install"
5326 .to_string(),
5327 false,
5328 ));
5329 }
5330 } else {
5331 let _ = sender.send(InstallerEvent::Log(
5332 "[native] [WARN] requirements.txt not found, skipping dependency installation"
5333 .to_string(),
5334 false,
5335 ));
5336 }
5337 }
5338
5339 "onnx" => {
5341 use crate::installers::components::onnxruntime::{
5342 install_method_from_env, prebuilt_version_from_env, OnnxInstallMethod,
5343 OnnxRuntimeConfig, OnnxRuntimeInstaller,
5344 };
5345
5346 let detected_rocm_version = detect_rocm_version();
5347 let detected_gpu_arch = detect_gpu_arch().unwrap_or_else(|| {
5348 eprintln!("Warning: no GPU arch detected for ONNX, falling back to gfx1100");
5349 "gfx1100".to_string()
5350 });
5351
5352 let install_method = install_method_from_env();
5359 let prebuilt_version = prebuilt_version_from_env();
5360
5361 let config = OnnxRuntimeConfig {
5363 rocm_version: Some(detected_rocm_version.clone()),
5364 rocm_release: Some({
5365 let parts: Vec<&str> = detected_rocm_version.split('.').collect();
5368 if parts.len() >= 3 {
5369 detected_rocm_version.clone()
5370 } else if parts.len() >= 2 && parts[0] == "7" && parts[1] == "2" {
5371 "7.2.4".to_string()
5372 } else if parts.len() >= 2 {
5373 format!("{}.{}.3", parts[0], parts[1])
5374 } else {
5375 "7.2.4".to_string()
5376 }
5377 }),
5378 gpu_arch: Some(detected_gpu_arch),
5379 python_bin: resolve_python_bin(),
5380 install_method,
5381 runtime_version: ctx.target_version.clone(),
5382 prebuilt_version,
5383 ..Default::default()
5384 };
5385 let inst = OnnxRuntimeInstaller::new(config);
5386
5387 let cmd = inst.build_uninstall_command();
5389 let _ = execute_native_command(
5390 &NativeCommand::from_shell_cmd_with_dir(
5391 &cmd.program,
5392 &cmd.args,
5393 &cmd.env,
5394 cmd.working_dir.clone(),
5395 ),
5396 None,
5397 sender,
5398 &component.name,
5399 );
5400
5401 match install_method {
5403 OnnxInstallMethod::MigraphxWheel => {
5404 let _ = sender.send(InstallerEvent::Log(
5408 format!(
5409 "[native] {} — installing onnxruntime-migraphx=={} from PyPI (ROCm {})",
5410 component.name,
5411 inst.prebuilt_version(),
5412 detected_rocm_version
5413 ),
5414 false,
5415 ));
5416 let cmd = inst.build_migraphx_install_command();
5417 execute_native_command(
5418 &NativeCommand::from_shell_cmd_with_dir(
5419 &cmd.program,
5420 &cmd.args,
5421 &cmd.env,
5422 cmd.working_dir.clone(),
5423 ),
5424 None,
5425 sender,
5426 &component.name,
5427 )?;
5428 }
5429 OnnxInstallMethod::PrebuiltWheel => {
5430 let _ = sender.send(InstallerEvent::Log(
5433 format!(
5434 "[native] {} — installing prebuilt onnxruntime-rocm from PyPI \
5435 (legacy; may be ABI-incompatible with ROCm {})",
5436 component.name, detected_rocm_version
5437 ),
5438 false,
5439 ));
5440 let cmd = inst.build_prebuilt_install_command();
5441 execute_native_command(
5442 &NativeCommand::from_shell_cmd_with_dir(
5443 &cmd.program,
5444 &cmd.args,
5445 &cmd.env,
5446 cmd.working_dir.clone(),
5447 ),
5448 None,
5449 sender,
5450 &component.name,
5451 )?;
5452 }
5453 OnnxInstallMethod::SourceBuild => {
5454 let _ = sender.send(InstallerEvent::Log(
5456 format!(
5457 "[native] {} — building from source (target {}, ROCm {}). Heavy build.",
5458 component.name,
5459 inst.runtime_version(),
5460 detected_rocm_version
5461 ),
5462 false,
5463 ));
5464 let rocm_env = crate::installers::common::RocmEnv::from_known(
5465 Some(std::path::PathBuf::from("/opt/rocm")),
5466 detected_rocm_version.clone(),
5467 );
5468 for cmd in [
5469 inst.build_git_clone_command(),
5470 inst.build_git_checkout_command(),
5471 inst.build_git_submodule_command(),
5472 inst.build_build_command(&rocm_env),
5473 inst.build_wheel_install_command(),
5474 ] {
5475 execute_native_command(
5476 &NativeCommand::from_shell_cmd_with_dir(
5477 &cmd.program,
5478 &cmd.args,
5479 &cmd.env,
5480 cmd.working_dir.clone(),
5481 ),
5482 None,
5483 sender,
5484 &component.name,
5485 )?;
5486 }
5487 }
5488 }
5489
5490 let cmd = inst.build_provider_validation_command();
5492 execute_native_command(
5493 &NativeCommand::from_shell_cmd_with_dir(
5494 &cmd.program,
5495 &cmd.args,
5496 &cmd.env,
5497 cmd.working_dir.clone(),
5498 ),
5499 None,
5500 sender,
5501 &component.name,
5502 )?;
5503
5504 let model_dirs = [
5506 dirs::home_dir().map(|h| h.join(".mlstack/models")),
5507 dirs::home_dir().map(|h| h.join(".local/share/models")),
5508 ];
5509 for model_dir in model_dirs.iter().flatten() {
5510 if model_dir.exists() {
5511 if let Ok(entries) = std::fs::read_dir(model_dir) {
5512 for entry in entries.flatten() {
5513 let path = entry.path();
5514 if path.extension().map(|e| e == "onnx").unwrap_or(false) {
5515 let model_path = path.to_string_lossy().to_string();
5516 if model_path.ends_with(".onnx.optimized") {
5518 continue;
5519 }
5520 let _ = sender.send(InstallerEvent::Log(
5521 format!(
5522 "[native] {} — optimizing model: {}",
5523 component.name, model_path
5524 ),
5525 false,
5526 ));
5527 let cmd = inst.build_model_optimizer_command(&model_path);
5528 let _ = execute_native_command(
5529 &NativeCommand::from_shell_cmd_with_dir(
5530 &cmd.program,
5531 &cmd.args,
5532 &cmd.env,
5533 cmd.working_dir.clone(),
5534 ),
5535 None,
5536 sender,
5537 &component.name,
5538 );
5539 }
5540 }
5541 }
5542 }
5543 }
5544
5545 persist_onnx_install_status(
5546 inst.prebuilt_version(),
5547 &["MIGraphXExecutionProvider", "CPUExecutionProvider"],
5548 );
5549
5550 let _ = sender.send(InstallerEvent::Log(
5551 format!("[native] {} — ONNX Runtime (MIGraphX) installed successfully. \
5552 Source build available via OnnxInstallMethod::SourceBuild for ROCMExecutionProvider.",
5553 component.name),
5554 false,
5555 ));
5556 }
5557
5558 "bitsandbytes" => {
5560 use crate::installers::components::bitsandbytes_multi::{
5561 BitsAndBytesConfig, BitsAndBytesInstaller,
5562 };
5563 let inst = BitsAndBytesInstaller::new(BitsAndBytesConfig {
5564 python_bin: resolve_python_bin(),
5565 ..Default::default()
5566 });
5567 let target_dir = format!("{}/.mlstack/bitsandbytes", ctx.user_home);
5568
5569 if is_force_reinstall() {
5571 purge_component_packages(&["bitsandbytes"], sender, &component.name);
5572 }
5573
5574 git_clone_or_pull(
5576 "https://github.com/ROCm/bitsandbytes.git",
5577 &target_dir,
5578 &["--recursive"],
5579 sudo_pw,
5580 sender,
5581 &component.name,
5582 )?;
5583
5584 let cmd = inst.build_pypi_install_command();
5586 execute_native_command(
5587 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5588 None,
5589 sender,
5590 &component.name,
5591 )?;
5592 }
5593
5594 "rocm-smi" => {
5596 use crate::installers::components::rocm_smi::RocmSmiInstaller;
5597 let inst = RocmSmiInstaller::new();
5598 let distro = crate::installers::common::DistroFacade::detect();
5599 let target_dir = format!("{}/.mlstack/rocm_smi", ctx.user_home);
5600
5601 git_clone_or_pull(
5603 inst.git_clone_url(),
5604 &target_dir,
5605 &[],
5606 sudo_pw,
5607 sender,
5608 &component.name,
5609 )?;
5610
5611 if !(distro.family() == crate::platform::detection::DistroFamily::Arch
5613 && crate::installers::common::utils::command_exists("rocm-smi"))
5614 {
5615 let cmd = inst.build_package_install_command(&distro);
5616 execute_native_command(
5617 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5618 sudo_pw,
5619 sender,
5620 &component.name,
5621 )?;
5622 } else {
5623 let _ = sender.send(InstallerEvent::Log(
5624 format!(
5625 "[native] {} — rocm-smi already on PATH, skipping package install",
5626 component.name
5627 ),
5628 false,
5629 ));
5630 }
5631 }
5632
5633 "migraphx" => {
5635 use crate::installers::components::migraphx_multi::{
5636 MigraphxConfig, MigraphxInstaller,
5637 };
5638 let inst = MigraphxInstaller::new(MigraphxConfig {
5639 python_bin: resolve_python_bin(),
5640 ..Default::default()
5641 });
5642 let distro = crate::installers::common::DistroFacade::detect();
5643 let packages = inst.required_packages(&distro);
5644
5645 let cmd = inst.build_package_install_command(&distro, &packages);
5647 execute_native_command(
5648 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5649 sudo_pw,
5650 sender,
5651 &component.name,
5652 )?;
5653
5654 if inst.should_skip_pip_install(&distro) {
5656 if let Some(msg) = inst.build_limitation_message(&distro) {
5657 let _ = sender.send(InstallerEvent::Log(
5658 format!("[native] [INFO] {}", msg),
5659 false,
5660 ));
5661 }
5662 } else {
5663 let cmd = inst.build_pip_install_command();
5664 execute_native_command(
5665 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5666 None,
5667 sender,
5668 &component.name,
5669 )?;
5670 }
5671
5672 match crate::installers::components::migraphx_multi::migraphx_driver_status() {
5677 Ok(()) => {}
5678 Err(missing) => {
5679 let detail = if missing.is_empty() {
5680 "migraphx-driver would not run".to_string()
5681 } else {
5682 format!("missing shared libraries: {}", missing.join(", "))
5683 };
5684 let _ = sender.send(InstallerEvent::Log(
5685 format!(
5686 "[native] [WARN] {} installed but NOT functional — {detail}. \
5687 This is usually a library version desync from a partial upgrade. \
5688 Run `sudo pacman -Syu` (Arch/CachyOS) or the equivalent system \
5689 update to sync dependencies, then re-verify.",
5690 component.name
5691 ),
5692 true,
5693 ));
5694 }
5695 }
5696 }
5697
5698 "pytorch-profiler" => {
5700 use crate::installers::components::pytorch_profiler::{
5701 PytorchProfilerConfig, PytorchProfilerInstaller,
5702 };
5703 let inst = PytorchProfilerInstaller::new(PytorchProfilerConfig {
5704 python_bin: resolve_python_bin(),
5705 ..Default::default()
5706 });
5707
5708 let cmd = inst.build_install_command();
5710 execute_native_command(
5711 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5712 None,
5713 sender,
5714 &component.name,
5715 )?;
5716 }
5717
5718 "wandb" => {
5720 use crate::installers::components::wandb::{WandbConfig, WandbInstaller};
5721 let inst = WandbInstaller::new(WandbConfig {
5722 python_bin: resolve_python_bin(),
5723 ..Default::default()
5724 });
5725
5726 let cmd = inst.build_install_command();
5728 execute_native_command(
5729 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5730 None,
5731 sender,
5732 &component.name,
5733 )?;
5734 }
5735
5736 "amdgpu-drivers" => {
5738 use crate::installers::components::amdgpu_drivers::{AmdgpuConfig, AmdgpuInstaller};
5739 let inst = AmdgpuInstaller::new(AmdgpuConfig::default());
5740 let distro = crate::installers::common::DistroFacade::detect();
5741 let packages = inst.required_packages(&distro);
5742
5743 let cmd = inst.build_package_install_command(&distro, &packages);
5745 execute_native_command(
5746 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5747 sudo_pw,
5748 sender,
5749 &component.name,
5750 )?;
5751
5752 let cmd = inst.build_verify_command();
5760 execute_native_command(
5761 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5762 None,
5763 sender,
5764 &component.name,
5765 )?;
5766
5767 ensure_user_in_groups(&["render", "video"], sudo_pw, sender, &component.name);
5770 }
5771
5772 "migraphx-python" => {
5774 use crate::installers::components::migraphx_python::{
5775 build_python_source_requested, source_branch, MigraphxPythonConfig,
5776 MigraphxPythonInstaller,
5777 };
5778 let inst = MigraphxPythonInstaller::new(MigraphxPythonConfig {
5779 python_bin: resolve_python_bin(),
5780 ..Default::default()
5781 });
5782
5783 let distro = crate::installers::common::DistroFacade::detect();
5785 if !inst.is_available_on_distro(&distro) {
5786 if build_python_source_requested() {
5787 let _ = sender.send(InstallerEvent::Log(
5790 "[native] MLSTACK_MIGRAPHX_BUILD_PYTHON=1 — building Python \
5791 bindings from source (AMDMIGraphX). Heavy build (~20-40 min)."
5792 .into(),
5793 false,
5794 ));
5795 let workdir = std::env::temp_dir()
5796 .join("rusty-migraphx-python")
5797 .to_string_lossy()
5798 .to_string();
5799 let _ = std::fs::remove_dir_all(&workdir);
5800 let gpu_arch =
5801 std::env::var("GPU_ARCH").unwrap_or_else(|_| "gfx1100".to_string());
5802 let cmds = inst.build_source_commands(
5803 &workdir,
5804 "/opt/rocm",
5805 &gpu_arch,
5806 &source_branch(),
5807 );
5808 for cmd in cmds {
5809 execute_native_command(
5810 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5811 None,
5812 sender,
5813 &component.name,
5814 )?;
5815 }
5816 } else if let Some(msg) = inst.build_unavailable_message(&distro) {
5817 let _ = sender.send(InstallerEvent::Log(
5818 format!("[native] [SKIP] {}", msg),
5819 false,
5820 ));
5821 }
5822 } else {
5823 let cmd = inst.build_install_command();
5825 execute_native_command(
5826 &NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env),
5827 None,
5828 sender,
5829 &component.name,
5830 )?;
5831 }
5832 }
5833
5834 "enhanced-env" => {
5836 let _ = sender.send(InstallerEvent::Log(
5840 "[native] Enhanced environment setup handled during pre-install phase".into(),
5841 false,
5842 ));
5843 }
5844
5845 "mlperf-inference"
5847 | "rocm-benchmarks"
5848 | "gpu-memory-bandwidth"
5849 | "rocm-smi-bench"
5850 | "pytorch-performance"
5851 | "vllm-performance"
5852 | "deepspeed-performance"
5853 | "megatron-performance"
5854 | "onnx-performance"
5855 | "rusty-llama-performance"
5856 | "flash-attention-ck-performance"
5857 | "all-benchmarks" => {
5858 run_native_benchmark(&component.id, sender)?;
5859 }
5860
5861 "fastvideo" => {
5863 use crate::installers::components::fastvideo::{
5864 detect_discrete_gpu_archs_from_sysfs, fastvideo_build_warning, post_checkout_plan,
5865 resolve_gpu_archs, resolve_install_user, select_visible_gpu_archs, FastVideoConfig,
5866 FastVideoInstallStep, FastVideoInstaller,
5867 };
5868 let detected_archs =
5869 detect_discrete_gpu_archs_from_sysfs(PathBuf::from("/sys/class/drm").as_path());
5870 let selected_archs = select_visible_gpu_archs(
5871 &detected_archs,
5872 ctx.env_exports
5873 .get("ROCR_VISIBLE_DEVICES")
5874 .map(String::as_str),
5875 ctx.env_exports
5876 .get("HIP_VISIBLE_DEVICES")
5877 .map(String::as_str),
5878 ctx.env_exports
5879 .get("CUDA_VISIBLE_DEVICES")
5880 .map(String::as_str),
5881 );
5882 let gpu_archs = resolve_gpu_archs(
5883 &selected_archs,
5884 ctx.env_exports.get("GPU_ARCHS").map(String::as_str),
5885 ctx.env_exports.get("GPU_ARCH").map(String::as_str),
5886 None,
5887 );
5888 let python_bin = ctx
5889 .env_exports
5890 .get("MLSTACK_PYTHON_BIN")
5891 .map(|v| v.trim().to_string())
5892 .filter(|v| !v.is_empty())
5893 .unwrap_or_else(resolve_python_bin);
5894 let inst = FastVideoInstaller::new(FastVideoConfig {
5895 gpu_archs: gpu_archs.clone(),
5896 python_bin: python_bin.clone(),
5897 });
5898 let build_dir = PathBuf::from("/tmp/FastVideo_ROCm_build");
5899
5900 let initial_cleanup = inst.cleanup();
5903 execute_native_command(
5904 &NativeCommand::from_shell_cmd(
5905 &initial_cleanup.program,
5906 &initial_cleanup.args,
5907 &initial_cleanup.env,
5908 ),
5909 sudo_pw,
5910 sender,
5911 &component.name,
5912 )?;
5913
5914 let mkdir_cmd = inst.mkdir_build_dir();
5916 execute_native_command(
5917 &NativeCommand::from_shell_cmd(&mkdir_cmd.program, &mkdir_cmd.args, &mkdir_cmd.env),
5918 sudo_pw,
5919 sender,
5920 &component.name,
5921 )?;
5922
5923 let clone_cmd = inst.git_clone();
5925 execute_native_command(
5926 &NativeCommand::from_shell_cmd_with_dir(
5927 &clone_cmd.program,
5928 &clone_cmd.args,
5929 &clone_cmd.env,
5930 clone_cmd.working_dir,
5931 ),
5932 sudo_pw,
5933 sender,
5934 &component.name,
5935 )?;
5936
5937 let checkout_cmd = inst.git_checkout();
5939 execute_native_command(
5940 &NativeCommand::from_shell_cmd_with_dir(
5941 &checkout_cmd.program,
5942 &checkout_cmd.args,
5943 &checkout_cmd.env,
5944 Some(build_dir.clone()),
5945 ),
5946 sudo_pw,
5947 sender,
5948 &component.name,
5949 )?;
5950
5951 let sudo_user = std::env::var("SUDO_USER").ok();
5952 let user = std::env::var("USER").ok();
5953 let run_user = resolve_install_user(sudo_user.as_deref(), user.as_deref())
5954 .context("FastVideo install user unavailable from SUDO_USER or USER")?;
5955
5956 for step in post_checkout_plan() {
5957 let step_sudo = if step.requires_sudo() { sudo_pw } else { None };
5958 if step == FastVideoInstallStep::KernelInstall {
5959 let snapshot = inst.distribution_snapshot(
5960 "before",
5961 crate::installers::common::nvidia_blocklist::NVIDIA_RUNTIME_PREFIXES,
5962 );
5963 execute_native_command(
5964 &NativeCommand::from_shell_cmd(
5965 &snapshot.program,
5966 &snapshot.args,
5967 &snapshot.env,
5968 ),
5969 None,
5970 sender,
5971 &component.name,
5972 )?;
5973 }
5974
5975 let command = match step {
5976 FastVideoInstallStep::Chown => {
5977 let cmd = inst.chown_build_dir(&run_user);
5978 NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env)
5979 }
5980 FastVideoInstallStep::ManagedDependenciesPreflight => {
5981 let cmd = inst.managed_dependencies_preflight();
5982 NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env)
5983 }
5984 FastVideoInstallStep::FlashAttentionPreflight => {
5985 let cmd = inst.flash_attention_preflight();
5986 NativeCommand::from_shell_cmd(&cmd.program, &cmd.args, &cmd.env)
5987 }
5988 FastVideoInstallStep::SourcePolicyPreflight => {
5989 let cmd = inst.source_policy_preflight();
5990 NativeCommand::from_shell_cmd_with_dir(
5991 &cmd.program,
5992 &cmd.args,
5993 &cmd.env,
5994 cmd.working_dir,
5995 )
5996 }
5997 FastVideoInstallStep::KernelInstall => {
5998 let _ = sender.send(InstallerEvent::Log(
5999 fastvideo_build_warning(&gpu_archs),
6000 false,
6001 ));
6002 let cmd = inst.pip_install_kernel();
6003 NativeCommand::from_shell_cmd_with_dir(
6004 &cmd.program,
6005 &cmd.args,
6006 &cmd.env,
6007 cmd.working_dir,
6008 )
6009 }
6010 FastVideoInstallStep::PackageInstall => {
6011 let cmd = inst.pip_install_package();
6012 NativeCommand::from_shell_cmd_with_dir(
6013 &cmd.program,
6014 &cmd.args,
6015 &cmd.env,
6016 cmd.working_dir,
6017 )
6018 }
6019 };
6020 execute_native_command(&command, step_sudo, sender, &component.name)?;
6021 }
6022
6023 let snapshot = inst.distribution_snapshot(
6024 "after",
6025 crate::installers::common::nvidia_blocklist::NVIDIA_RUNTIME_PREFIXES,
6026 );
6027 execute_native_command(
6028 &NativeCommand::from_shell_cmd(&snapshot.program, &snapshot.args, &snapshot.env),
6029 None,
6030 sender,
6031 &component.name,
6032 )?;
6033
6034 let cleanup_cmd = inst.cleanup();
6036 execute_native_command(
6037 &NativeCommand::from_shell_cmd(
6038 &cleanup_cmd.program,
6039 &cleanup_cmd.args,
6040 &cleanup_cmd.env,
6041 ),
6042 sudo_pw,
6043 sender,
6044 &component.name,
6045 )?;
6046 }
6047 "llama-cpp" => {
6049 use crate::installers::components::llama_cpp::{LlamaCppConfig, LlamaCppInstaller};
6050 let gpu_arch = ctx
6051 .env_exports
6052 .get("GPU_ARCH")
6053 .map(|v| v.trim().to_string())
6054 .filter(|v| !v.is_empty())
6055 .unwrap_or_else(|| {
6056 detect_gpu_arch().unwrap_or_else(|| {
6057 eprintln!("Warning: no GPU arch detected, falling back to gfx1100");
6058 "gfx1100".to_string()
6059 })
6060 });
6061 let channel = ctx
6062 .env_exports
6063 .get("ROCM_CHANNEL")
6064 .map(|v| v.trim().to_string())
6065 .filter(|v| !v.is_empty())
6066 .unwrap_or_else(|| "latest".to_string());
6067 let inst = LlamaCppInstaller::new(LlamaCppConfig {
6068 gpu_arch: gpu_arch.clone(),
6069 channel: channel.clone(),
6070 ..Default::default()
6071 });
6072 let home = ctx.user_home;
6073 let manifest = inst.check_latest_release();
6074 let strategy = inst.resolve_install_strategy(home, manifest.as_ref());
6075 match &strategy {
6076 crate::installers::components::llama_cpp::InstallStrategy::Prebuilt(plan) => {
6077 let _ = sender.send(InstallerEvent::Log(
6078 format!(
6079 "llama-cpp install strategy selected: prebuilt ({})",
6080 plan.arch
6081 ),
6082 false,
6083 ));
6084 if let Err(err) = inst.download_prebuilt_binary(plan) {
6085 let _ = sender.send(InstallerEvent::Log(
6086 format!(
6087 "llama-cpp prebuilt download failed; falling back to source: {}",
6088 err
6089 ),
6090 false,
6091 ));
6092 inst.run_source_install(home).map_err(|source_err| {
6093 anyhow::anyhow!(
6094 "llama-cpp prebuilt install failed and source fallback failed: {}",
6095 source_err
6096 )
6097 })?;
6098 }
6099 }
6100 crate::installers::components::llama_cpp::InstallStrategy::Source(plan) => {
6101 let _ = sender.send(InstallerEvent::Log(
6102 format!(
6103 "llama-cpp install strategy selected: source ({})",
6104 plan.reason
6105 ),
6106 false,
6107 ));
6108 inst.run_source_install(home).map_err(|err| {
6109 anyhow::anyhow!("llama-cpp source install failed: {}", err)
6110 })?;
6111 }
6112 }
6113
6114 let clone_dir = inst.build_dir_path().to_string();
6118 let verification =
6119 crate::installers::components::llama_cpp::verify_installed_binary(home, &clone_dir);
6120 let _ = sender.send(InstallerEvent::Log(
6121 format!(
6122 "llama-cpp post-install verification: {}",
6123 verification.summary
6124 ),
6125 false,
6126 ));
6127 if verification.stronger_check_passed {
6128 let _ = sender.send(InstallerEvent::Log(
6129 format!(
6130 "llama-cpp llama-bench: prefill={:.0} t/s, decode={:.0} t/s",
6131 verification.bench_prefill_tps.unwrap_or(0.0),
6132 verification.bench_decode_tps.unwrap_or(0.0),
6133 ),
6134 false,
6135 ));
6136 }
6137 if verification.rdna3_is_valid.unwrap_or(false) {
6138 let _ = sender.send(InstallerEvent::Log(
6139 format!(
6140 "llama-cpp RDNA3 proof: device=RDNA3, WMMA={:.0} t/s",
6141 verification.rdna3_wmma_tps.unwrap_or(0.0),
6142 ),
6143 false,
6144 ));
6145 }
6146 }
6147
6148 _ => {
6150 bail!(
6151 "No native installer dispatch for component '{}'. \
6152 This should not happen — is_native_component() should return false for this ID.",
6153 component.id
6154 );
6155 }
6156 }
6157
6158 let _ = sender.send(InstallerEvent::Progress {
6159 component_id: component.id.clone(),
6160 progress: 0.9,
6161 message: format!("{} installation complete", component.name),
6162 });
6163
6164 Ok(())
6165}
6166
6167fn run_script(
6168 component: &Component,
6169 script_path: &str,
6170 sudo_password: Option<String>,
6171 batch_mode: bool,
6172 install_method: &str,
6173 sender: &Sender<InstallerEvent>,
6174 input_rx: Arc<Mutex<Receiver<String>>>,
6175) -> Result<()> {
6176 let user_home = resolve_mlstack_user_home();
6177 let user_name = std::env::var("USER").unwrap_or_else(|_| "user".into());
6178 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";
6179
6180 let venv_bin = PathBuf::from(&user_home).join("rocm_venv").join("bin");
6181 let local_bin = PathBuf::from(&user_home).join(".local").join("bin");
6182 let mut path = std::env::var("PATH").unwrap_or_default();
6183
6184 if local_bin.exists() {
6186 let local_str = local_bin.to_string_lossy();
6187 if !path.split(':').any(|entry| entry == local_str) {
6188 if path.is_empty() {
6189 path = local_str.to_string();
6190 } else {
6191 path = format!("{}:{}", path, local_str);
6192 }
6193 }
6194 }
6195
6196 if venv_bin.exists() {
6197 let venv_str = venv_bin.to_string_lossy();
6198 if !path.split(':').any(|entry| entry == venv_str) {
6199 if path.is_empty() {
6200 path = venv_str.to_string();
6201 } else {
6202 path = format!("{}:{}", venv_str, path);
6203 }
6204 }
6205 }
6206
6207 let env_exports = load_mlstack_env_exports(&user_home);
6208 let persistent_python = env_exports
6209 .get("MLSTACK_PYTHON_BIN")
6210 .map(|v| v.trim().to_string())
6211 .filter(|v| !v.is_empty());
6212 let python_bin = std::env::var("MLSTACK_ENV_NAME")
6216 .ok()
6217 .map(|n| n.trim().to_string())
6218 .filter(|n| !n.is_empty())
6219 .map(|n| {
6220 crate::platform::environment::mlstack_env_dir(&n)
6221 .join("bin")
6222 .join("python")
6223 .to_string_lossy()
6224 .to_string()
6225 })
6226 .or(persistent_python)
6227 .unwrap_or_else(resolve_python_bin);
6228 let force_reinstall = is_force_reinstall();
6229 let force_env_value = if force_reinstall { "true" } else { "false" };
6230 let sudo_password_env = sudo_password.clone();
6231 let user_home_meta = fs::metadata(&user_home).ok();
6232 #[cfg(unix)]
6233 let (mlstack_target_uid, mlstack_target_gid) = {
6234 use std::os::unix::fs::MetadataExt;
6235 (
6236 user_home_meta.as_ref().map(|meta| meta.uid().to_string()),
6237 user_home_meta.as_ref().map(|meta| meta.gid().to_string()),
6238 )
6239 };
6240 #[cfg(not(unix))]
6241 let (mlstack_target_uid, mlstack_target_gid): (Option<String>, Option<String>) = (None, None);
6242
6243 let component_needs_sudo = component.needs_sudo && needs_sudo();
6245
6246 let mut command = if component_needs_sudo {
6247 let mut cmd = Command::new("sudo");
6248 cmd.arg("-S")
6249 .arg("-p")
6250 .arg("")
6251 .arg(format!("--preserve-env={}", preserve_env))
6252 .arg("bash")
6253 .arg(script_path);
6254 cmd
6255 } else {
6256 let mut cmd = Command::new("bash");
6257 cmd.arg(script_path);
6258 cmd
6259 };
6260
6261 let mut pythonpath = std::env::var("PYTHONPATH").unwrap_or_default();
6262 let rocm_lib = "/opt/rocm/lib";
6263 if !pythonpath.split(':').any(|entry| entry == rocm_lib) {
6264 if pythonpath.is_empty() {
6265 pythonpath = rocm_lib.to_string();
6266 } else {
6267 pythonpath = format!("{}:{}", rocm_lib, pythonpath);
6268 }
6269 }
6270
6271 pythonpath = pythonpath
6273 .split(':')
6274 .filter(|part| !part.starts_with("/tmp/"))
6275 .collect::<Vec<_>>()
6276 .join(":");
6277
6278 let rocm_version = env_exports
6279 .get("ROCM_VERSION")
6280 .map(|v| v.trim().to_string())
6281 .filter(|v| !v.is_empty())
6282 .unwrap_or_else(detect_rocm_version);
6283 let gpu_arch = env_exports
6284 .get("GPU_ARCH")
6285 .map(|v| v.trim().to_string())
6286 .filter(|v| !v.is_empty())
6287 .or_else(|| {
6288 env_exports
6289 .get("PYTORCH_ROCM_ARCH")
6290 .map(|v| v.split(';').next().unwrap_or("").trim().to_string())
6291 .filter(|v| !v.is_empty())
6292 })
6293 .unwrap_or_else(|| {
6294 detect_gpu_arch().unwrap_or_else(|| {
6295 eprintln!("Warning: no GPU arch detected, falling back to gfx1100");
6296 "gfx1100".to_string()
6297 })
6298 });
6299 let py_rocm_arch = env_exports
6300 .get("PYTORCH_ROCM_ARCH")
6301 .map(|v| v.trim().to_string())
6302 .filter(|v| !v.is_empty())
6303 .unwrap_or_else(|| gpu_arch.clone());
6304 let mut gpu_archs = env_exports
6305 .get("GPU_ARCHS")
6306 .map(|v| v.trim().to_string())
6307 .filter(|v| !v.is_empty())
6308 .unwrap_or_else(|| gpu_arch.clone());
6309 let mut gpu_list = env_exports
6310 .get("HIP_VISIBLE_DEVICES")
6311 .map(|v| v.trim().to_string())
6312 .filter(|v| !v.is_empty())
6313 .or_else(|| {
6314 env_exports
6315 .get("CUDA_VISIBLE_DEVICES")
6316 .map(|v| v.trim().to_string())
6317 .filter(|v| !v.is_empty())
6318 })
6319 .unwrap_or_else(detect_gpu_list);
6320 let hsa_override = env_exports
6321 .get("HSA_OVERRIDE_GFX_VERSION")
6322 .map(|v| v.trim().to_string())
6323 .filter(|v| !v.is_empty());
6324 let aiter_jit_dir = PathBuf::from(&user_home)
6325 .join(".mlstack")
6326 .join("aiter")
6327 .join("jit");
6328
6329 if component.id == "aiter" {
6330 gpu_archs = py_rocm_arch
6331 .split(';')
6332 .next()
6333 .map(|v| v.trim().to_string())
6334 .filter(|v| !v.is_empty())
6335 .unwrap_or_else(|| gpu_arch.clone());
6336 gpu_list = prioritize_gpu_list_for_arch(&gpu_list, &gpu_archs);
6337 let _ = fs::create_dir_all(&aiter_jit_dir);
6338 }
6339
6340 command
6341 .stdin(Stdio::piped())
6342 .stdout(Stdio::piped())
6343 .stderr(Stdio::piped())
6344 .env("HOME", &user_home)
6345 .env("USER", &user_name)
6346 .env("LOGNAME", &user_name)
6347 .env("PATH", &path)
6348 .env("PYTHONPATH", &pythonpath)
6349 .env("MLSTACK_USER_HOME", &user_home)
6350 .env("MLSTACK_PYTHON_BIN", &python_bin)
6351 .env("UV_PYTHON", &python_bin)
6352 .env("ROCM_HOME", "/opt/rocm")
6353 .env("ROCM_PATH", "/opt/rocm")
6354 .env("ROCM_VERSION", &rocm_version)
6355 .env("ROCM_CHANNEL", "latest")
6356 .env("GPU_ARCH", &gpu_arch)
6357 .env("GPU_ARCHS", &gpu_archs)
6358 .env("PYTORCH_ROCM_ARCH", &py_rocm_arch)
6359 .env("HIP_PATH", "/opt/rocm")
6360 .env("HIP_ROOT_DIR", "/opt/rocm")
6361 .env("ROCM_ROOT", "/opt/rocm")
6362 .env("HIPCC_BIN_DIR", "/opt/rocm/bin")
6363 .env("HIP_VISIBLE_DEVICES", &gpu_list)
6364 .env("CUDA_VISIBLE_DEVICES", &gpu_list)
6365 .env("ROCR_VISIBLE_DEVICES", &gpu_list)
6366 .env("MLSTACK_BATCH_MODE", if batch_mode { "1" } else { "0" })
6367 .env("MLSTACK_INSTALL_METHOD", install_method)
6368 .env("INSTALL_METHOD", install_method)
6369 .env("PIP_BREAK_SYSTEM_PACKAGES", "1")
6370 .env("PIP_ROOT_USER_ACTION", "ignore")
6371 .env("UV_PIP_BREAK_SYSTEM_PACKAGES", "1")
6372 .env("UV_SYSTEM_PYTHON", "1")
6373 .env("FORCE", force_env_value)
6374 .env("PYTORCH_REINSTALL", force_env_value)
6375 .env("MLSTACK_FORCE_REINSTALL", force_env_value)
6376 .env(
6377 "MLSTACK_SKIP_TORCH_INSTALL",
6378 if component.id == "pytorch" { "0" } else { "1" },
6379 )
6380 .env("RUSTY_STACK", "true");
6381
6382 if let Some(uid) = mlstack_target_uid.as_deref() {
6383 command.env("MLSTACK_TARGET_UID", uid);
6384 }
6385 if let Some(gid) = mlstack_target_gid.as_deref() {
6386 command.env("MLSTACK_TARGET_GID", gid);
6387 }
6388
6389 if let Some(password) = sudo_password_env {
6390 command.env("MLSTACK_SUDO_PASSWORD", password);
6391 }
6392
6393 if force_reinstall {
6394 if let Ok(script_contents) = fs::read_to_string(script_path) {
6395 let supports_force_flag = script_contents.contains("--force)")
6396 || script_contents.contains("--force ")
6397 || script_contents.contains("\"--force\"")
6398 || script_contents.contains("'--force'");
6399 if supports_force_flag {
6400 command.arg("--force");
6401 }
6402 }
6403 }
6404
6405 if let Some(value) = hsa_override {
6406 command.env("HSA_OVERRIDE_GFX_VERSION", value);
6407 }
6408 strip_hsa_override_for_therock(&mut command, &env_exports);
6409 if component.id == "aiter" {
6410 command.env("AITER_JIT_DIR", aiter_jit_dir);
6411 }
6412
6413 let mut child = command.spawn().with_context(|| {
6414 format!(
6415 "Failed to spawn installer script '{}' for {}. \
6416 Ensure bash is installed and the script file is accessible.",
6417 script_path, component.name
6418 )
6419 })?;
6420
6421 let mut child_stdin = child.stdin.take().with_context(|| {
6422 format!(
6423 "Failed to open stdin for script '{}' — \
6424 this usually indicates a system resource issue.",
6425 script_path
6426 )
6427 })?;
6428 if let Some(password) = sudo_password {
6429 let _ = child_stdin.write_all(password.as_bytes());
6430 let _ = child_stdin.write_all(b"\n");
6431 let _ = child_stdin.flush();
6432 }
6433
6434 if let Some(default_input) = default_component_input(&component.id) {
6436 let _ = child_stdin.write_all(default_input.as_bytes());
6437 let _ = child_stdin.flush();
6438 }
6439
6440 let stop_input = Arc::new(AtomicBool::new(false));
6441
6442 let stop_input_clone = Arc::clone(&stop_input);
6444 let input_thread = thread::spawn(move || {
6445 while !stop_input_clone.load(Ordering::Relaxed) {
6446 let rx = match input_rx.lock() {
6447 Ok(guard) => guard,
6448 Err(_) => break, };
6450 match rx.recv_timeout(Duration::from_millis(100)) {
6451 Ok(input) => {
6452 let _ = child_stdin.write_all(input.as_bytes());
6453 let _ = child_stdin.flush();
6454 }
6455 Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {
6456 continue;
6457 }
6458 Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => break,
6459 }
6460 }
6461 });
6462
6463 let sender_stdout = sender.clone();
6464 let component_id = component.id.clone();
6465 let component_name = component_name(component);
6466 let stdout_stream = match child.stdout.take() {
6467 Some(s) => s,
6468 None => {
6469 stop_input.store(true, Ordering::Relaxed);
6470 let _ = input_thread.join();
6471 let _ = sender.send(InstallerEvent::Log(
6472 format!(
6473 "[ERROR] {} — failed to capture stdout from script process",
6474 component.name
6475 ),
6476 false,
6477 ));
6478 bail!(
6479 "{} — failed to capture stdout from script process. \
6480 This usually indicates a system resource issue.",
6481 component.name
6482 );
6483 }
6484 };
6485 let stdout_handle = thread::spawn(move || {
6486 let mut reader = BufReader::new(stdout_stream);
6487 let mut buffer = Vec::new();
6488 let mut chunk = [0u8; 1024];
6489
6490 while let Ok(n) = reader.read(&mut chunk) {
6491 if n == 0 {
6492 break;
6493 }
6494 for &b in chunk.iter().take(n) {
6495 if b == b'\n' || b == b'\r' {
6496 let is_transient = b == b'\r';
6497 if !buffer.is_empty() {
6498 let line = String::from_utf8_lossy(&buffer).to_string();
6499 let clean_line = strip_ansi_codes(&line);
6500 if !clean_line.trim().is_empty() {
6501 let _ = sender_stdout
6502 .send(InstallerEvent::Log(clean_line.clone(), is_transient));
6503 if let Some(progress) = parse_progress(&clean_line) {
6504 let _ = sender_stdout.send(InstallerEvent::Progress {
6505 component_id: component_id.clone(),
6506 progress,
6507 message: format!(
6508 "{}: {}%",
6509 component_name,
6510 (progress * 100.0) as i32
6511 ),
6512 });
6513 }
6514 }
6515 buffer.clear();
6516 }
6517 } else {
6518 buffer.push(b);
6519 if buffer.len() > 8192 {
6520 buffer.clear();
6521 }
6522 }
6523 }
6524 }
6525 });
6526
6527 let sender_stderr = sender.clone();
6528 let stderr_stream = match child.stderr.take() {
6529 Some(s) => s,
6530 None => {
6531 stop_input.store(true, Ordering::Relaxed);
6532 let _ = input_thread.join();
6533 let _ = stdout_handle.join();
6534 let _ = sender.send(InstallerEvent::Log(
6535 format!(
6536 "[ERROR] {} — failed to capture stderr from script process",
6537 component.name
6538 ),
6539 false,
6540 ));
6541 bail!(
6542 "{} — failed to capture stderr from script process. \
6543 This usually indicates a system resource issue.",
6544 component.name
6545 );
6546 }
6547 };
6548 let stderr_handle = thread::spawn(move || {
6549 let mut reader = BufReader::new(stderr_stream);
6550 let mut buffer = Vec::new();
6551 let mut chunk = [0u8; 1024];
6552
6553 while let Ok(n) = reader.read(&mut chunk) {
6554 if n == 0 {
6555 break;
6556 }
6557 for &b in chunk.iter().take(n) {
6558 if b == b'\n' || b == b'\r' {
6559 let is_transient = b == b'\r';
6560 if !buffer.is_empty() {
6561 let line = String::from_utf8_lossy(&buffer).to_string();
6562 let clean_line = strip_ansi_codes(&line);
6563 if !clean_line.trim().is_empty() {
6564 let _ =
6565 sender_stderr.send(InstallerEvent::Log(clean_line, is_transient));
6566 }
6567 buffer.clear();
6568 }
6569 } else {
6570 buffer.push(b);
6571 if buffer.len() > 8192 {
6572 buffer.clear();
6573 }
6574 }
6575 }
6576 }
6577 });
6578
6579 let status = child.wait().with_context(|| {
6580 format!(
6581 "{} — script process wait failed. \
6582 The process may have been killed or the system is under heavy load.",
6583 component.name
6584 )
6585 })?;
6586 stop_input.store(true, Ordering::Relaxed);
6587 let _ = input_thread.join();
6588 let _ = stdout_handle.join();
6589 let _ = stderr_handle.join();
6590
6591 if !status.success() {
6592 let code = status.code().unwrap_or(-1);
6593 let error_detail = match code {
6594 1 => "General error — check the script output above for details".to_string(),
6595 2 => "Misuse of shell builtins — check script arguments".to_string(),
6596 126 => "Script not executable — check file permissions (chmod +x)".to_string(),
6597 127 => "Command not found in script — ensure required tools are installed".to_string(),
6598 130 => "Script terminated by Ctrl+C (SIGINT)".to_string(),
6599 137 => "Script killed (SIGKILL) — possibly out of memory".to_string(),
6600 139 => "Segmentation fault in script subprocess".to_string(),
6601 n if n > 128 => format!(
6602 "Script process terminated by signal {} — check system logs",
6603 n - 128
6604 ),
6605 _ => format!(
6606 "Script exited with code {} — check the output above for details",
6607 code
6608 ),
6609 };
6610 let _ = sender.send(InstallerEvent::Log(
6611 format!(
6612 "[ERROR] {} failed: {} (exit code {} at {})",
6613 component.name,
6614 error_detail,
6615 code,
6616 Local::now().format("%H:%M:%S")
6617 ),
6618 false,
6619 ));
6620 bail!(
6621 "{} failed: {} (exit code {})",
6622 component.name,
6623 error_detail,
6624 code
6625 );
6626 }
6627
6628 Ok(())
6629}
6630
6631fn star_repo(config: &InstallerConfig, sender: &Sender<InstallerEvent>) -> Result<()> {
6632 if !config.star_repos {
6633 return Ok(());
6634 }
6635
6636 let repo_url = "https://github.com/scooter-lacroix/Stan-s-ML-Stack";
6637
6638 let gh_check = Command::new("gh")
6640 .arg("auth")
6641 .arg("status")
6642 .stdout(Stdio::null())
6643 .stderr(Stdio::null())
6644 .status();
6645
6646 if let Ok(status) = gh_check {
6647 if status.success() {
6648 let _ = sender.send(InstallerEvent::Log(
6649 "Starring ML Stack repository on GitHub...".into(),
6650 false,
6651 ));
6652 let star_status = Command::new("gh")
6653 .arg("repo")
6654 .arg("star")
6655 .arg(repo_url)
6656 .stdout(Stdio::null())
6657 .stderr(Stdio::null())
6658 .status();
6659
6660 if let Ok(s) = star_status {
6661 if s.success() {
6662 let _ = sender.send(InstallerEvent::Log(
6663 "Thank you for starring the repo! ✦".into(),
6664 false,
6665 ));
6666 }
6667 }
6668 }
6669 }
6670
6671 Ok(())
6672}
6673
6674fn component_name(comp: &Component) -> String {
6675 comp.name.clone()
6676}
6677
6678fn format_user_friendly_error(component_id: &str, raw_error: &str) -> String {
6687 let error_lower = raw_error.to_lowercase();
6688
6689 match component_id {
6692 "permanent-env" => {
6693 if error_lower.contains("rocminfo") {
6694 return format!(
6695 "{}: Failed to detect GPU — rocminfo command failed. \
6696 Ensure ROCm is properly installed and /opt/rocm/bin is in your PATH. \
6697 Original error: {}",
6698 component_id, raw_error
6699 );
6700 }
6701 return format!(
6702 "{}: Environment setup failed. \
6703 Ensure ROCm is installed and GPU is accessible. \
6704 Original error: {}",
6705 component_id, raw_error
6706 );
6707 }
6708 "rocm" if error_lower.contains("no such file") || error_lower.contains("not found") => {
6709 return format!(
6710 "{}: ROCm installation failed — required packages or repositories not found. \
6711 Ensure your system's package manager is configured correctly and ROCm repositories are added. \
6712 Original error: {}",
6713 component_id, raw_error
6714 );
6715 }
6716 "pytorch" => {
6717 if error_lower.contains("no such file") || error_lower.contains("not found") {
6718 return format!(
6719 "{}: PyTorch installation failed — required packages not found. \
6720 Ensure ROCm is installed first and Python 3.10+ is available. \
6721 Original error: {}",
6722 component_id, raw_error
6723 );
6724 }
6725 return format!(
6726 "{}: PyTorch installation failed. \
6727 Ensure ROCm is installed first, Python 3.10-3.13 is available, \
6728 and pip/uv can reach the PyTorch package index. \
6729 Original error: {}",
6730 component_id, raw_error
6731 );
6732 }
6733 "triton" | "flash-attn" | "flash-attn-triton" | "flash-attn-ck" | "vllm" | "aiter" => {
6734 return format!(
6735 "{}: Installation failed. \
6736 Ensure PyTorch and ROCm are properly installed first \
6737 (these are required dependencies). \
6738 Original error: {}",
6739 component_id, raw_error
6740 );
6741 }
6742 "megatron" => {
6743 return format!(
6744 "{}: Megatron-LM installation failed. \
6745 Ensure PyTorch and MPI4Py are installed first (required dependencies). \
6746 Original error: {}",
6747 component_id, raw_error
6748 );
6749 }
6750 "deepspeed" => {
6751 return format!(
6752 "{}: DeepSpeed installation failed. \
6753 Ensure PyTorch is installed first (required dependency). \
6754 Original error: {}",
6755 component_id, raw_error
6756 );
6757 }
6758 _ => {}
6759 }
6760
6761 if error_lower.contains("permission denied") || error_lower.contains("sudo") {
6764 return format!(
6765 "{}: Permission denied. Try running with elevated privileges or check file/directory permissions. \
6766 Original error: {}",
6767 component_id, raw_error
6768 );
6769 }
6770
6771 if error_lower.contains("network")
6772 || error_lower.contains("connection")
6773 || error_lower.contains("timeout")
6774 || error_lower.contains("resolve")
6775 {
6776 return format!(
6777 "{}: Network error — failed to download packages or clone repositories. \
6778 Check your internet connection and try again. \
6779 Original error: {}",
6780 component_id, raw_error
6781 );
6782 }
6783
6784 if error_lower.contains("out of memory")
6785 || error_lower.contains("cannot allocate")
6786 || error_lower.contains("killed")
6787 {
6788 return format!(
6789 "{}: Out of memory — the build process was killed due to insufficient memory. \
6790 Try closing other applications or increasing swap space. \
6791 Original error: {}",
6792 component_id, raw_error
6793 );
6794 }
6795
6796 if error_lower.contains("exit code") {
6797 return raw_error.to_string();
6799 }
6800
6801 if error_lower.contains("no such file") || error_lower.contains("not found") {
6802 return format!(
6803 "{}: Required file or command not found. \
6804 Ensure all dependencies are installed before this component. \
6805 Original error: {}",
6806 component_id, raw_error
6807 );
6808 }
6809
6810 format!(
6813 "{}: Installation failed. Check the log output above for details. \
6814 Original error: {}",
6815 component_id, raw_error
6816 )
6817}
6818
6819pub(crate) fn detect_gpu_arch() -> Option<String> {
6824 crate::gpu::detect_discrete_amd_gpus()
6825 .first()
6826 .and_then(|g| g.gfx_arch.clone())
6827}
6828
6829pub(crate) fn detect_gpu_list() -> String {
6852 let gpus = detect_discrete_gpus();
6853 if gpus.is_empty() {
6854 "0".to_string()
6855 } else {
6856 gpus.join(",")
6857 }
6858}
6859
6860pub(crate) fn detect_discrete_gpus() -> Vec<String> {
6868 crate::gpu::detect_discrete_amd_gpus()
6869 .iter()
6870 .enumerate()
6871 .map(|(i, _)| i.to_string())
6872 .collect()
6873}
6874
6875fn detect_gpu_arch_by_id() -> HashMap<String, String> {
6877 let mut mapping: HashMap<String, String> = HashMap::new();
6878
6879 let rocminfo_path = crate::installers::common::utils::resolve_rocminfo_path();
6880 let Ok(output) = Command::new(&rocminfo_path).output() else {
6881 return mapping;
6882 };
6883 let stdout = String::from_utf8_lossy(&output.stdout);
6884 let mut current_gpu_id: Option<String> = None;
6885
6886 for line in stdout.lines() {
6887 let line = line.trim();
6888 if line.contains("GPU ID:") {
6889 if let Some(id_part) = line.split(':').nth(1) {
6890 let id: String = id_part.chars().filter(|c| c.is_ascii_digit()).collect();
6891 if !id.is_empty() {
6892 current_gpu_id = Some(id);
6893 }
6894 }
6895 continue;
6896 }
6897
6898 if !line.contains("Name:") || !line.contains("gfx") {
6899 continue;
6900 }
6901
6902 if let Some(id) = current_gpu_id.clone() {
6903 for token in line.split_whitespace() {
6904 let cleaned = token.trim_matches(|c: char| !c.is_alphanumeric());
6905 if cleaned.starts_with("gfx") {
6906 mapping.entry(id).or_insert_with(|| cleaned.to_string());
6907 break;
6908 }
6909 }
6910 }
6911 }
6912
6913 mapping
6914}
6915
6916fn prioritize_gpu_list_for_arch(gpu_list: &str, target_arch: &str) -> String {
6919 let target_arch = target_arch.trim();
6920 if target_arch.is_empty() {
6921 return gpu_list.to_string();
6922 }
6923
6924 let arch_by_id = detect_gpu_arch_by_id();
6925 if arch_by_id.is_empty() {
6926 return gpu_list.to_string();
6927 }
6928
6929 let mut ids: Vec<String> = gpu_list
6930 .split(',')
6931 .map(|v| v.trim().to_string())
6932 .filter(|v| !v.is_empty())
6933 .collect();
6934 if ids.len() < 2 {
6935 return gpu_list.to_string();
6936 }
6937
6938 let preferred_pos = ids.iter().position(|id| {
6939 arch_by_id
6940 .get(id)
6941 .map(|v| v == target_arch)
6942 .unwrap_or(false)
6943 });
6944 if let Some(pos) = preferred_pos {
6945 if pos > 0 {
6946 let preferred = ids.remove(pos);
6947 ids.insert(0, preferred);
6948 return ids.join(",");
6949 }
6950 }
6951
6952 gpu_list.to_string()
6953}
6954
6955#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6956enum VerificationResult {
6957 Verified,
6958 Failed,
6959 Missing,
6960}
6961
6962impl VerificationResult {
6963 fn label(&self) -> &'static str {
6964 match self {
6965 VerificationResult::Verified => "Verified",
6966 VerificationResult::Failed => "Failed",
6967 VerificationResult::Missing => "Missing",
6968 }
6969 }
6970}
6971
6972fn resolve_python_bin() -> String {
6973 if let Ok(value) = std::env::var("MLSTACK_PYTHON_BIN") {
6974 let value = value.trim();
6975 if !value.is_empty() {
6976 return value.to_string();
6977 }
6978 }
6979
6980 if let Ok(value) = std::env::var("UV_PYTHON") {
6981 let value = value.trim();
6982 if !value.is_empty() {
6983 return value.to_string();
6984 }
6985 }
6986
6987 let global_py = crate::platform::environment::mlstack_global_python();
6992 if global_py.exists() {
6993 return global_py.to_string_lossy().to_string();
6994 }
6995
6996 let home = std::env::var("HOME").unwrap_or_default();
6997
6998 let uv_python_dir = PathBuf::from(&home).join(".local/share/uv/python");
7003 if let Ok(entries) = std::fs::read_dir(&uv_python_dir) {
7004 let mut uv_pythons: Vec<PathBuf> = entries
7005 .filter_map(|e| e.ok())
7006 .filter_map(|e| {
7007 let bin = e.path().join("bin/python3");
7008 if bin.exists() {
7009 Some(bin)
7010 } else {
7011 None
7012 }
7013 })
7014 .collect();
7015 uv_pythons.sort_by(|a, b| {
7017 let va = python_version_tuple(a);
7018 let vb = python_version_tuple(b);
7019 python_ml_priority(&va).cmp(&python_ml_priority(&vb))
7020 });
7021 for python in &uv_pythons {
7023 if python_has_rocm_torch(&python.to_string_lossy()) {
7024 return python.to_string_lossy().to_string();
7025 }
7026 }
7027 if let Some(first) = uv_pythons.first() {
7029 return first.to_string_lossy().to_string();
7030 }
7031 }
7032
7033 let local_bin_python = format!("{home}/.local/bin/python3");
7035 if Path::new(&local_bin_python).exists() {
7036 if python_has_rocm_torch(&local_bin_python) {
7037 return local_bin_python;
7038 }
7039 return local_bin_python;
7040 }
7041
7042 let candidates = [
7044 Path::new(&home)
7045 .join("rocm_venv")
7046 .join("bin")
7047 .join("python"),
7048 PathBuf::from("/usr/local/bin/python3"),
7049 PathBuf::from("/usr/bin/python3"),
7050 PathBuf::from("python3"),
7051 PathBuf::from("python"),
7052 ];
7053
7054 for candidate in &candidates {
7055 let candidate_str = candidate.to_string_lossy();
7056 if candidate_str.is_empty() {
7057 continue;
7058 }
7059 if candidate.is_absolute() && !candidate.exists() {
7060 continue;
7061 }
7062 if python_has_rocm_torch(&candidate_str) {
7063 return candidate_str.to_string();
7064 }
7065 }
7066
7067 for candidate in &candidates {
7068 let candidate_str = candidate.to_string_lossy();
7069 if candidate.is_absolute() && !candidate.exists() {
7070 continue;
7071 }
7072 return candidate_str.to_string();
7073 }
7074
7075 "python3".to_string()
7076}
7077
7078fn python_version_tuple(path: &Path) -> (u32, u32) {
7080 let name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
7081 let parent = path
7083 .parent()
7084 .and_then(|p| p.parent())
7085 .and_then(|p| p.file_name())
7086 .and_then(|n| n.to_str())
7087 .unwrap_or("");
7088 if let Some(v) = parse_cpython_version(parent) {
7090 return v;
7091 }
7092 if let Some(rest) = name.strip_prefix("python") {
7094 let parts: Vec<&str> = rest.split('.').collect();
7095 if parts.len() >= 2 {
7096 if let (Ok(major), Ok(minor)) = (parts[0].parse::<u32>(), parts[1].parse::<u32>()) {
7097 return (major, minor);
7098 }
7099 }
7100 }
7101 (0, 0)
7102}
7103
7104fn parse_cpython_version(dir_name: &str) -> Option<(u32, u32)> {
7106 let rest = dir_name.strip_prefix("cpython-")?;
7107 let parts: Vec<&str> = rest.split('.').collect();
7108 if parts.len() >= 2 {
7109 let major = parts[0].parse::<u32>().ok()?;
7110 let minor = parts[1].parse::<u32>().ok()?;
7111 return Some((major, minor));
7112 }
7113 None
7114}
7115
7116fn python_ml_priority(version: &(u32, u32)) -> (i32, u32) {
7118 match version {
7119 (3, 12) => (0, version.1), (3, 13) => (1, version.1),
7121 (3, 11) => (2, version.1),
7122 (3, 10) => (3, version.1),
7123 _ => (4, version.1), }
7125}
7126
7127fn python_has_rocm_torch(python: &str) -> bool {
7128 Command::new(python)
7129 .arg("-c")
7130 .arg(
7131 "import torch, sys; ok=hasattr(torch,'__version__') and hasattr(torch,'version') and getattr(torch.version,'hip', None); sys.exit(0 if ok else 1)",
7132 )
7133 .stdout(Stdio::null())
7134 .stderr(Stdio::null())
7135 .status()
7136 .map(|status| status.success())
7137 .unwrap_or(false)
7138}
7139
7140fn default_component_input(component_id: &str) -> Option<&'static str> {
7141 match component_id {
7142 "ml-stack-core" => Some("9\n0\n"),
7143 _ => None,
7144 }
7145}
7146
7147fn parse_progress(line: &str) -> Option<f32> {
7148 let lower = line.to_lowercase();
7149 if let Some(percent_idx) = lower.find('%') {
7150 let prefix = &lower[..percent_idx];
7151 let digits: String = prefix
7152 .chars()
7153 .rev()
7154 .take_while(|c| c.is_ascii_digit())
7155 .collect();
7156 if !digits.is_empty() {
7157 let value: String = digits.chars().rev().collect();
7158 if let Ok(percent) = value.parse::<f32>() {
7159 return Some((percent / 100.0).min(1.0));
7160 }
7161 }
7162 }
7163 if lower.contains("progress") && lower.contains("100") {
7164 return Some(1.0);
7165 }
7166 if let Some(bracket) = line.find('[') {
7170 if let Some(close) = line[bracket..].find(']') {
7171 let inside = &line[bracket + 1..bracket + close];
7172 if let Some(slash) = inside.find('/') {
7173 let (n_str, m_str) = inside.split_at(slash);
7174 if let (Ok(n), Ok(m)) = (
7175 n_str.trim().parse::<f32>(),
7176 m_str[1..].trim().parse::<f32>(),
7177 ) {
7178 if m > 0.0 && n <= m {
7179 return Some((n / m).min(1.0));
7180 }
7181 }
7182 }
7183 }
7184 }
7185 None
7186}
7187
7188fn needs_sudo() -> bool {
7189 #[cfg(unix)]
7190 {
7191 unsafe { libc::geteuid() != 0 }
7192 }
7193 #[cfg(not(unix))]
7194 {
7195 false
7196 }
7197}
7198
7199fn strip_ansi_codes(s: &str) -> String {
7200 let mut clean = String::with_capacity(s.len());
7201 let mut in_escape = false;
7202 for c in s.chars() {
7203 if c == '\x1b' {
7204 in_escape = true;
7205 } else if in_escape {
7206 if ('@'..='~').contains(&c) || c == 'm' {
7207 in_escape = false;
7208 }
7209 } else {
7210 clean.push(c);
7211 }
7212 }
7213 clean
7214}
7215
7216fn sanitize_output_line(line: &str) -> String {
7217 let stripped = strip_ansi_codes(line);
7218 let mut clean = String::with_capacity(stripped.len());
7219 for c in stripped.chars() {
7220 if c == '\t' {
7221 clean.push(' ');
7222 } else if !c.is_control() {
7223 clean.push(c);
7224 }
7225 }
7226 clean
7227}
7228
7229fn should_suppress_verification_line(line: &str, target_id: &str) -> bool {
7230 let trimmed = line.trim();
7231 if trimmed.is_empty() {
7232 return true;
7233 }
7234
7235 let lower = trimmed.to_lowercase();
7236 if is_exception_or_traceback_line(&lower) {
7237 return false;
7238 }
7239
7240 if lower.contains("amdgpu.ids")
7241 || lower.contains("tool-lib")
7242 || lower.contains("cannot load tool")
7243 || lower.contains("libdrm/amdgpu.ids: no such file or directory")
7244 {
7245 return true;
7246 }
7247
7248 if (target_id == "rocm" || target_id == "rocm-smi")
7249 && (lower.starts_with("======")
7250 || lower.starts_with("====")
7251 || lower.starts_with("---")
7252 || lower.contains("the info is not accurate")
7253 || lower.contains("================================")
7254 || lower.starts_with('{')
7255 || lower.starts_with('[')
7256 || lower.starts_with(']')
7257 || lower.starts_with('}')
7258 || lower.starts_with("\"card")
7259 || lower.starts_with("device ")
7260 || lower.starts_with("*****"))
7261 {
7262 return true;
7263 }
7264
7265 false
7266}
7267
7268fn is_exception_or_traceback_line(lowercase_line: &str) -> bool {
7269 lowercase_line.contains("traceback")
7270 || lowercase_line.contains("exception")
7271 || lowercase_line.contains("segmentation fault")
7272 || lowercase_line.contains("fatal:")
7273 || lowercase_line.contains("error:")
7274 || lowercase_line.starts_with("error ")
7275}
7276
7277fn run_native_benchmark(component_id: &str, sender: &Sender<InstallerEvent>) -> Result<()> {
7282 use crate::benchmark_runners;
7283
7284 let bench_name = match component_id {
7285 "mlperf-inference" => "all",
7286 "rocm-benchmarks" => "gpu-capability",
7287 "gpu-memory-bandwidth" => "memory-bandwidth",
7288 "rocm-smi-bench" => "gpu-capability",
7289 "pytorch-performance" => "pytorch",
7290 "vllm-performance" => "vllm",
7291 "deepspeed-performance" => "deepspeed",
7292 "megatron-performance" => "megatron",
7293 "onnx-performance" => "onnx",
7294 "rusty-llama-performance" => "llama-cpp",
7295 "flash-attention-ck-performance" => "flash-attention-ck",
7296 "all-benchmarks" => "all",
7297 _ => bail!("Unknown benchmark component: {}", component_id),
7298 };
7299
7300 let _ = sender.send(InstallerEvent::Log(
7301 format!("[native] Running {} benchmark...", component_id),
7302 false,
7303 ));
7304
7305 let _ = sender.send(InstallerEvent::Progress {
7306 component_id: component_id.to_string(),
7307 progress: 0.3,
7308 message: format!("Running {} benchmark", component_id),
7309 });
7310
7311 match benchmark_runners::run_benchmark(bench_name) {
7312 Ok(output) => {
7313 let _ = sender.send(InstallerEvent::Log(
7314 format!(
7315 "[native] {} benchmark completed: {} ({} ms)",
7316 component_id,
7317 if output.success { "SUCCESS" } else { "FAILED" },
7318 output.execution_time_ms
7319 ),
7320 false,
7321 ));
7322 persist_benchmark_result(bench_name, &output);
7327 if !output.errors.is_empty() {
7328 for err in &output.errors {
7329 let _ = sender.send(InstallerEvent::Log(
7330 format!("[native] {} error: {}", component_id, err),
7331 false,
7332 ));
7333 }
7334 }
7335 if !output.success {
7336 let error_details = output.errors.join("; ");
7337 bail!(
7338 "{} benchmark failed: {}",
7339 component_id,
7340 if error_details.is_empty() {
7341 "unknown error".to_string()
7342 } else {
7343 error_details
7344 }
7345 );
7346 }
7347 Ok(())
7348 }
7349 Err(err) => {
7350 let _ = sender.send(InstallerEvent::Log(
7351 format!("[native] {} benchmark error: {}", component_id, err),
7352 false,
7353 ));
7354 bail!("{} benchmark error: {}", component_id, err)
7355 }
7356 }
7357}
7358
7359fn persist_onnx_install_status(version: &str, providers: &[&str]) {
7372 let provider = providers
7373 .iter()
7374 .find(|p| matches!(**p, "MIGraphXExecutionProvider" | "ROCMExecutionProvider"))
7375 .copied()
7376 .unwrap_or("none");
7377 persist_benchmark_result(
7378 "onnx",
7379 &crate::benchmark_runners::BenchmarkOutput {
7380 name: "onnx".to_string(),
7381 success: true,
7382 execution_time_ms: 0,
7383 results: serde_json::json!({
7384 "ort_version": version,
7385 "provider": provider,
7386 "providers_available": providers,
7387 "provider_priority": ["MIGraphXExecutionProvider", "ROCMExecutionProvider"],
7388 "install_status": "provider_ready",
7389 "benchmark_status": "not_run",
7390 "graph_opt_level": "ORT_ENABLE_ALL",
7391 "session_create_ms": 0.0,
7392 "inference_latency_p50_ms": 0.0,
7393 "inference_latency_p95_ms": 0.0,
7394 "inference_latency_p99_ms": 0.0,
7395 "throughput_inf_per_sec": 0.0,
7396 "peak_rss_mb": 0.0,
7397 "quantized_supported": false,
7398 "inference_samples": []
7399 }),
7400 errors: Vec::new(),
7401 },
7402 );
7403}
7404
7405fn persist_benchmark_result(bench_name: &str, output: &crate::benchmark_runners::BenchmarkOutput) {
7406 let metrics = output
7407 .results
7408 .get("metrics")
7409 .cloned()
7410 .unwrap_or_else(|| output.results.clone());
7411 let payload = serde_json::json!({
7412 "name": output.name.clone(),
7413 "success": output.success,
7414 "execution_time_ms": output.execution_time_ms as u64,
7415 "results": metrics,
7416 "errors": serde_json::to_value(&output.errors).unwrap_or_default(),
7417 });
7418 let body = match serde_json::to_string_pretty(&payload) {
7419 Ok(s) => s,
7420 Err(_) => return,
7421 };
7422 let ts = std::time::SystemTime::now()
7423 .duration_since(std::time::UNIX_EPOCH)
7424 .map(|d| d.as_secs())
7425 .unwrap_or(0);
7426 let fname = format!("{bench_name}_benchmarks_{ts}.json");
7427
7428 let mut candidates: Vec<std::path::PathBuf> =
7429 crate::benchmark_logs::benchmark_log_directories();
7430 if let Some(home) = std::env::var("HOME").ok().filter(|s| !s.is_empty()) {
7431 candidates.push(
7432 std::path::Path::new(&home)
7433 .join(".rusty-stack")
7434 .join("logs"),
7435 );
7436 }
7437 let tmp = std::env::var("TMPDIR").unwrap_or_else(|_| "/tmp".to_string());
7438 candidates.push(std::path::Path::new(&tmp).join("rusty-stack").join("logs"));
7439
7440 for dir in candidates {
7441 let _ = std::fs::create_dir_all(&dir);
7442 if std::fs::write(dir.join(&fname), &body).is_ok() {
7443 return;
7444 }
7445 }
7446}
7447
7448#[cfg(test)]
7449mod tests {
7450 use super::*;
7451 use std::sync::Mutex;
7452
7453 static ENV_LOCK: Mutex<()> = Mutex::new(());
7454
7455 #[test]
7461 fn test_safe_rocm_path_prepends_rocm_bin() {
7462 if !Path::new("/opt/rocm/bin").exists() {
7467 eprintln!("skip: /opt/rocm/bin absent (CI host) — prepend behavior is pass-through");
7468 return;
7469 }
7470 let p = safe_rocm_path("/home/scooter/.mlstack/global/bin:/usr/bin:/bin");
7471 let rocm = p.find("/opt/rocm/bin").unwrap();
7472 let shim = p.find("/home/scooter/.mlstack/global/bin").unwrap();
7473 assert!(
7474 rocm < shim,
7475 "/opt/rocm/bin must precede the venv bin: got {p}"
7476 );
7477 assert!(p.contains("/opt/rocm/hip/bin"));
7479 assert!(p.contains("/usr/bin") && p.contains("/bin"));
7480 }
7481
7482 #[test]
7483 fn test_safe_rocm_path_dedupes_rocm_entries() {
7484 if !Path::new("/opt/rocm/bin").exists() {
7488 eprintln!("skip: /opt/rocm/bin absent (CI host)");
7489 return;
7490 }
7491 let p = safe_rocm_path("/opt/rocm/bin:/usr/bin:/opt/rocm/hip/bin");
7492 assert_eq!(
7493 p.matches("/opt/rocm/bin").count(),
7494 1,
7495 "rocm/bin must appear exactly once: got {p}"
7496 );
7497 assert_eq!(p.matches("/opt/rocm/hip/bin").count(), 1);
7498 }
7499
7500 #[test]
7501 fn test_safe_rocm_path_passes_through_when_rocm_absent() {
7502 let inherited = "/usr/local/bin:/usr/bin";
7508 if Path::new("/opt/rocm/bin").exists() {
7509 assert!(safe_rocm_path(inherited).starts_with("/opt/rocm/bin"));
7510 } else {
7511 assert_eq!(safe_rocm_path(inherited), inherited);
7512 }
7513 }
7514
7515 #[test]
7524 fn test_format_user_friendly_error_permission_denied() {
7525 let msg = format_user_friendly_error("rocm", "Permission denied: /opt/rocm");
7526 assert!(
7527 msg.contains("Permission denied"),
7528 "Should mention permission issue"
7529 );
7530 assert!(
7531 msg.contains("elevated privileges"),
7532 "Should suggest running with sudo"
7533 );
7534 }
7535
7536 #[test]
7537 fn test_native_maintenance_actions_use_native_verdict() {
7538 let repair = Component {
7539 id: "rccl-repair".into(),
7540 name: "Repair RCCL Multi-GPU".into(),
7541 description: String::new(),
7542 script: String::new(),
7543 category: Category::Maintenance,
7544 required: false,
7545 selected: true,
7546 installed: false,
7547 progress: 0.0,
7548 estimate: String::new(),
7549 needs_sudo: false,
7550 experimental: false,
7551 note: None,
7552 };
7553 let verify = Component {
7554 id: "verify-basic".into(),
7555 name: "Verify Installation".into(),
7556 description: String::new(),
7557 script: String::new(),
7558 category: Category::Maintenance,
7559 required: false,
7560 selected: true,
7561 installed: false,
7562 progress: 0.0,
7563 estimate: String::new(),
7564 needs_sudo: false,
7565 experimental: false,
7566 note: None,
7567 };
7568 let normal = Component {
7569 id: "pytorch".into(),
7570 name: "PyTorch with ROCm".into(),
7571 description: String::new(),
7572 script: String::new(),
7573 category: Category::Foundation,
7574 required: true,
7575 selected: true,
7576 installed: false,
7577 progress: 0.0,
7578 estimate: String::new(),
7579 needs_sudo: false,
7580 experimental: false,
7581 note: None,
7582 };
7583
7584 assert!(is_native_maintenance_action(&repair));
7585 assert!(!is_native_maintenance_action(&verify));
7586 assert!(!is_native_maintenance_action(&normal));
7587
7588 let success = native_maintenance_action_outcome(&repair, true);
7589 assert!(success.success);
7590 assert!(success.report_lines[0].contains("reported success"));
7591
7592 let failure = native_maintenance_action_outcome(&repair, false);
7593 assert!(!failure.success);
7594 assert!(failure.report_lines[0].contains("reported failure"));
7595 }
7596
7597 #[test]
7598 fn test_benchmark_patterns_match_native_output_filenames() {
7599 assert_eq!(
7600 benchmark_pattern_for_target("rocm-benchmarks"),
7601 Some("gpu-capability_benchmarks")
7602 );
7603 assert_eq!(
7604 benchmark_pattern_for_target("gpu-memory-bandwidth"),
7605 Some("memory-bandwidth_benchmarks")
7606 );
7607 assert_eq!(
7608 benchmark_pattern_for_target("all-benchmarks"),
7609 Some("all_benchmarks")
7610 );
7611 assert_eq!(
7612 benchmark_pattern_for_target("onnx-performance"),
7613 Some("onnx_benchmarks")
7614 );
7615 }
7616
7617 #[test]
7618 fn test_single_benchmark_results_render_metrics_and_errors() {
7619 let lines = benchmark_result_lines_from_json(&serde_json::json!({
7620 "name": "vllm",
7621 "success": false,
7622 "execution_time_ms": 248842,
7623 "results": {
7624 "mode": "degraded",
7625 "reason": "model timed out",
7626 "throughput_tokens_per_sec": 0.0,
7627 "visible_devices": "0,1"
7628 },
7629 "errors": ["model timed out after 240 seconds"]
7630 }));
7631 let joined = lines.join("\n");
7632
7633 assert!(joined.contains("MODE: degraded"));
7634 assert!(joined.contains("REASON: model timed out"));
7635 assert!(joined.contains("THROUGHPUT TOKENS PER SEC: 0"));
7636 assert!(joined.contains("ERROR: model timed out after 240 seconds"));
7637 }
7638
7639 #[test]
7640 fn test_persist_onnx_install_status_writes_clean_status_artifact() {
7641 let _global_env = crate::test_support::lock_env();
7642 let _guard = ENV_LOCK.lock().unwrap();
7643 let old_home = std::env::var("HOME").ok();
7644 let old_tmpdir = std::env::var("TMPDIR").ok();
7645 let old_log_dir = std::env::var("MLSTACK_LOG_DIR").ok();
7646 let root = std::env::temp_dir().join(format!(
7647 "rusty-stack-persist-onnx-status-{}",
7648 std::process::id()
7649 ));
7650 let _ = std::fs::remove_dir_all(&root);
7651 let logs = root.join("logs");
7652 let tmp = root.join("tmp");
7653 std::fs::create_dir_all(&logs).unwrap();
7654 std::fs::create_dir_all(&tmp).unwrap();
7655 std::env::set_var("HOME", &root);
7656 std::env::set_var("TMPDIR", &tmp);
7657 std::env::set_var("MLSTACK_LOG_DIR", &logs);
7658
7659 persist_onnx_install_status(
7660 "1.27.1",
7661 &["MIGraphXExecutionProvider", "CPUExecutionProvider"],
7662 );
7663
7664 let written = std::fs::read_dir(&logs)
7665 .unwrap()
7666 .find_map(|entry| {
7667 let path = entry.ok()?.path();
7668 (path
7669 .file_name()?
7670 .to_string_lossy()
7671 .starts_with("onnx_benchmarks_"))
7672 .then_some(path)
7673 })
7674 .expect("onnx status log should be written");
7675 let json: serde_json::Value =
7676 serde_json::from_str(&std::fs::read_to_string(written).unwrap()).unwrap();
7677
7678 if let Some(home) = old_home {
7679 std::env::set_var("HOME", home);
7680 } else {
7681 std::env::remove_var("HOME");
7682 }
7683 if let Some(tmpdir) = old_tmpdir {
7684 std::env::set_var("TMPDIR", tmpdir);
7685 } else {
7686 std::env::remove_var("TMPDIR");
7687 }
7688 if let Some(log_dir) = old_log_dir {
7689 std::env::set_var("MLSTACK_LOG_DIR", log_dir);
7690 } else {
7691 std::env::remove_var("MLSTACK_LOG_DIR");
7692 }
7693 let _ = std::fs::remove_dir_all(&root);
7694
7695 assert_eq!(json["success"], true);
7696 assert_eq!(json["errors"].as_array().unwrap().len(), 0);
7697 assert_eq!(json["results"]["provider"], "MIGraphXExecutionProvider");
7698 assert_eq!(json["results"]["benchmark_status"], "not_run");
7699 }
7700
7701 #[test]
7702 fn test_persist_full_suite_keeps_nested_results() {
7703 let _global_env = crate::test_support::lock_env();
7704 let _guard = ENV_LOCK.lock().unwrap();
7705 let old_home = std::env::var("HOME").ok();
7706 let old_tmpdir = std::env::var("TMPDIR").ok();
7707 let old_log_dir = std::env::var("MLSTACK_LOG_DIR").ok();
7708 let root = std::env::temp_dir().join(format!(
7709 "rusty-stack-persist-full-suite-{}",
7710 std::process::id()
7711 ));
7712 let _ = std::fs::remove_dir_all(&root);
7713 let logs = root.join("logs");
7714 let tmp = root.join("tmp");
7715 std::fs::create_dir_all(&logs).unwrap();
7716 std::fs::create_dir_all(&tmp).unwrap();
7717 std::env::set_var("HOME", &root);
7718 std::env::set_var("TMPDIR", &tmp);
7719 std::env::set_var("MLSTACK_LOG_DIR", &logs);
7720
7721 let output = crate::benchmark_runners::BenchmarkOutput {
7722 name: "all".into(),
7723 success: true,
7724 execution_time_ms: 1,
7725 results: serde_json::json!({
7726 "memory_bandwidth": {
7727 "success": true,
7728 "execution_time_ms": 1,
7729 "metrics": {"hbm_peak_gb_s": 123.0}
7730 }
7731 }),
7732 errors: Vec::new(),
7733 };
7734 persist_benchmark_result("all", &output);
7735
7736 let written = std::fs::read_dir(&logs)
7737 .unwrap()
7738 .find_map(|entry| {
7739 let path = entry.ok()?.path();
7740 (path
7741 .file_name()?
7742 .to_string_lossy()
7743 .starts_with("all_benchmarks_"))
7744 .then_some(path)
7745 })
7746 .expect("all benchmark log should be written");
7747 let json: serde_json::Value =
7748 serde_json::from_str(&std::fs::read_to_string(written).unwrap()).unwrap();
7749
7750 if let Some(home) = old_home {
7751 std::env::set_var("HOME", home);
7752 } else {
7753 std::env::remove_var("HOME");
7754 }
7755 if let Some(tmpdir) = old_tmpdir {
7756 std::env::set_var("TMPDIR", tmpdir);
7757 } else {
7758 std::env::remove_var("TMPDIR");
7759 }
7760 if let Some(log_dir) = old_log_dir {
7761 std::env::set_var("MLSTACK_LOG_DIR", log_dir);
7762 } else {
7763 std::env::remove_var("MLSTACK_LOG_DIR");
7764 }
7765 let _ = std::fs::remove_dir_all(&root);
7766
7767 assert_eq!(
7768 json["results"]["memory_bandwidth"]["metrics"]["hbm_peak_gb_s"],
7769 123.0
7770 );
7771 }
7772
7773 #[test]
7774 fn test_format_user_friendly_error_not_found() {
7775 let msg = format_user_friendly_error("pytorch", "No such file or directory: python3");
7776 assert!(msg.contains("not found"), "Should mention file not found");
7777 assert!(
7778 msg.contains("ROCm") || msg.contains("Python"),
7779 "Should mention ROCm or Python dependency, got: {}",
7780 msg
7781 );
7782 }
7783
7784 #[test]
7785 fn test_format_user_friendly_error_network() {
7786 let msg = format_user_friendly_error(
7788 "onnxruntime",
7789 "network error: Could not resolve host: github.com",
7790 );
7791 assert!(
7792 msg.contains("Network error"),
7793 "Should mention network issue, got: {}",
7794 msg
7795 );
7796 assert!(
7797 msg.contains("internet connection"),
7798 "Should suggest checking internet, got: {}",
7799 msg
7800 );
7801 }
7802
7803 #[test]
7804 fn test_format_user_friendly_error_out_of_memory() {
7805 let msg = format_user_friendly_error("onnxruntime", "out of memory during build");
7807 assert!(
7808 msg.contains("Out of memory"),
7809 "Should mention memory issue, got: {}",
7810 msg
7811 );
7812 assert!(
7813 msg.contains("swap space"),
7814 "Should suggest increasing swap, got: {}",
7815 msg
7816 );
7817 }
7818
7819 #[test]
7820 fn test_format_user_friendly_error_permanent_env_rocminfo() {
7821 let msg = format_user_friendly_error("permanent-env", "rocminfo command not found");
7822 assert!(
7823 msg.contains("rocminfo"),
7824 "Should mention rocminfo specifically"
7825 );
7826 assert!(
7827 msg.contains("ROCm is properly installed"),
7828 "Should suggest ROCm installation"
7829 );
7830 }
7831
7832 #[test]
7833 fn test_format_user_friendly_error_pytorch_specific() {
7834 let msg = format_user_friendly_error("pytorch", "pip install failed: build error");
7836 assert!(msg.contains("PyTorch"), "Should mention PyTorch by name");
7837 assert!(
7838 msg.contains("ROCm is installed first"),
7839 "Should mention ROCm dependency"
7840 );
7841 assert!(
7842 msg.contains("Python 3.10-3.13"),
7843 "Should mention Python version, got: {}",
7844 msg
7845 );
7846 }
7847
7848 #[test]
7849 fn test_format_user_friendly_error_triton_dependency() {
7850 let msg = format_user_friendly_error("triton", "git clone failed");
7851 assert!(
7852 msg.contains("PyTorch and ROCm"),
7853 "Should mention required dependencies"
7854 );
7855 }
7856
7857 #[test]
7858 fn test_format_user_friendly_error_megatron_dependency() {
7859 let msg = format_user_friendly_error("megatron", "pip install failed");
7860 assert!(
7861 msg.contains("PyTorch and MPI4Py"),
7862 "Should mention Megatron's specific dependencies"
7863 );
7864 }
7865
7866 #[test]
7867 fn test_format_user_friendly_error_deepspeed_dependency() {
7868 let msg = format_user_friendly_error("deepspeed", "build failed");
7869 assert!(
7870 msg.contains("PyTorch"),
7871 "Should mention DeepSpeed's dependency on PyTorch"
7872 );
7873 }
7874
7875 #[test]
7876 fn test_format_user_friendly_error_exit_code_passthrough() {
7877 let msg = format_user_friendly_error("rocm", "ROCm failed: General error (exit code 1)");
7878 assert!(
7880 msg.contains("exit code 1"),
7881 "Should preserve exit code context"
7882 );
7883 }
7884
7885 #[test]
7886 fn test_format_user_friendly_error_unknown_component() {
7887 let msg = format_user_friendly_error("some-unknown-component", "something went wrong");
7888 assert!(
7889 msg.contains("Installation failed"),
7890 "Should have generic failure message"
7891 );
7892 assert!(
7893 msg.contains("something went wrong"),
7894 "Should include original error"
7895 );
7896 }
7897
7898 #[test]
7899 fn test_format_user_friendly_error_rocm_not_found() {
7900 let msg = format_user_friendly_error("rocm", "package not found: amdgpu-dkms");
7901 assert!(msg.contains("ROCm"), "Should mention ROCm specifically");
7902 assert!(msg.contains("repositories"), "Should mention repositories");
7903 }
7904
7905 #[test]
7906 fn test_execute_native_command_nonexistent_program() {
7907 let (tx, rx) = std::sync::mpsc::channel();
7910 let cmd = NativeCommand::Shell {
7911 program: "/nonexistent/program/that/does/not/exist".to_string(),
7912 args: vec![],
7913 env: vec![],
7914 working_dir: None,
7915 };
7916 let result = execute_native_command(&cmd, None, &tx, "test-component");
7917 assert!(result.is_err(), "Should fail for non-existent program");
7918 let err_msg = result.unwrap_err().to_string();
7919 assert!(!err_msg.is_empty(), "Error message should not be empty");
7920
7921 drop(tx);
7923 let events: Vec<_> = rx.try_iter().collect();
7924 assert!(
7925 !events.is_empty(),
7926 "Should have sent at least one event before failing"
7927 );
7928 }
7929
7930 #[test]
7931 fn test_execute_native_command_captures_exit_code() {
7932 let (tx, rx) = std::sync::mpsc::channel();
7934 let cmd = NativeCommand::Shell {
7935 program: "bash".to_string(),
7936 args: vec!["-c".to_string(), "echo test_output; exit 42".to_string()],
7937 env: vec![],
7938 working_dir: None,
7939 };
7940 let result = execute_native_command(&cmd, None, &tx, "test-component");
7941 assert!(
7942 result.is_err(),
7943 "Should fail for non-zero exit code, got: {:?}",
7944 result
7945 );
7946
7947 let err_msg = result.unwrap_err().to_string();
7948 assert!(
7949 err_msg.contains("42"),
7950 "Error should include exit code 42, got: {}",
7951 err_msg
7952 );
7953
7954 drop(tx);
7956 let events: Vec<_> = rx.try_iter().collect();
7957 let log_messages: Vec<&str> = events
7958 .iter()
7959 .filter_map(|e| match e {
7960 InstallerEvent::Log(msg, _) => Some(msg.as_str()),
7961 _ => None,
7962 })
7963 .collect();
7964 assert!(
7965 log_messages.iter().any(|m| m.contains("test_output")),
7966 "Should have captured stdout output before failure"
7967 );
7968 assert!(
7969 log_messages
7970 .iter()
7971 .any(|m| m.contains("[ERROR]") || m.contains("exit code")),
7972 "Should have logged error with exit code"
7973 );
7974 }
7975
7976 #[test]
7977 fn test_execute_native_command_success_sends_events() {
7978 let (tx, rx) = std::sync::mpsc::channel();
7980 let cmd = NativeCommand::Shell {
7981 program: "bash".to_string(),
7982 args: vec![
7983 "-c".to_string(),
7984 "echo 'hello world'; echo 'second line'".to_string(),
7985 ],
7986 env: vec![],
7987 working_dir: None,
7988 };
7989 let result = execute_native_command(&cmd, None, &tx, "test-component");
7990 assert!(result.is_ok(), "Should succeed for zero exit code");
7991
7992 drop(tx);
7993 let events: Vec<_> = rx.try_iter().collect();
7994 let log_messages: Vec<&str> = events
7995 .iter()
7996 .filter_map(|e| match e {
7997 InstallerEvent::Log(msg, _) => Some(msg.as_str()),
7998 _ => None,
7999 })
8000 .collect();
8001 assert!(
8002 log_messages.iter().any(|m| m.contains("hello world")),
8003 "Should have captured first line of output"
8004 );
8005 assert!(
8006 log_messages.iter().any(|m| m.contains("second line")),
8007 "Should have captured second line of output"
8008 );
8009 }
8010
8011 #[test]
8016 fn test_no_sudo_wrap_when_password_is_none() {
8017 let (tx, rx) = std::sync::mpsc::channel();
8021 let cmd = NativeCommand::Shell {
8022 program: "echo".to_string(),
8023 args: vec!["hello".to_string()],
8024 env: vec![],
8025 working_dir: None,
8026 };
8027 let result = execute_native_command(&cmd, None, &tx, "test-component");
8028 assert!(result.is_ok(), "Should succeed without sudo: {:?}", result);
8029
8030 drop(tx);
8031 let events: Vec<_> = rx.try_iter().collect();
8032 let log_messages: Vec<String> = events
8033 .iter()
8034 .filter_map(|e| match e {
8035 InstallerEvent::Log(msg, _) => Some(msg.clone()),
8036 _ => None,
8037 })
8038 .collect();
8039
8040 let cmd_log = log_messages.iter().find(|m| m.contains("[native] $"));
8042 assert!(cmd_log.is_some(), "Should have logged the command");
8043 let cmd_log = cmd_log.unwrap();
8044 assert!(
8045 !cmd_log.contains("sudo"),
8046 "Command should NOT be wrapped with sudo when sudo_pw is None, got: {}",
8047 cmd_log
8048 );
8049 assert!(
8050 cmd_log.contains("echo hello"),
8051 "Command should be 'echo hello', got: {}",
8052 cmd_log
8053 );
8054 }
8055
8056 #[test]
8057 fn test_sudo_wrap_when_password_is_some() {
8058 let (tx, rx) = std::sync::mpsc::channel();
8062 let cmd = NativeCommand::Shell {
8063 program: "echo".to_string(),
8064 args: vec!["hello".to_string()],
8065 env: vec![],
8066 working_dir: None,
8067 };
8068 let _result = execute_native_command(&cmd, Some("dummy_pw"), &tx, "test-component");
8070 drop(tx);
8073 let events: Vec<_> = rx.try_iter().collect();
8074 let log_messages: Vec<String> = events
8075 .iter()
8076 .filter_map(|e| match e {
8077 InstallerEvent::Log(msg, _) => Some(msg.clone()),
8078 _ => None,
8079 })
8080 .collect();
8081
8082 let has_sudo_log = log_messages.iter().any(|m| m.contains("Running with sudo"));
8084 assert!(
8085 has_sudo_log,
8086 "Should have logged 'Running with sudo' when sudo_pw is Some"
8087 );
8088 }
8089
8090 #[test]
8091 fn test_git_clone_without_sudo_works() {
8092 let (tx, rx) = std::sync::mpsc::channel();
8097 let cmd = NativeCommand::Shell {
8098 program: "git".to_string(),
8099 args: vec!["--version".to_string()],
8100 env: vec![],
8101 working_dir: None,
8102 };
8103 let result = execute_native_command(&cmd, None, &tx, "textgen");
8104 assert!(
8105 result.is_ok(),
8106 "git --version should succeed without sudo wrapping: {:?}",
8107 result
8108 );
8109
8110 drop(tx);
8111 let events: Vec<_> = rx.try_iter().collect();
8112 let log_messages: Vec<String> = events
8113 .iter()
8114 .filter_map(|e| match e {
8115 InstallerEvent::Log(msg, _) => Some(msg.clone()),
8116 _ => None,
8117 })
8118 .collect();
8119
8120 assert!(
8122 log_messages.iter().any(|m| m.contains("git version")),
8123 "Should have captured 'git version' output, got: {:?}",
8124 log_messages
8125 );
8126
8127 let cmd_log = log_messages.iter().find(|m| m.contains("[native] $"));
8129 assert!(cmd_log.is_some(), "Should have logged the command");
8130 assert!(
8131 !cmd_log.unwrap().contains("sudo"),
8132 "git command should NOT be wrapped with sudo"
8133 );
8134 }
8135
8136 #[test]
8137 fn test_textgen_git_clone_args_not_empty() {
8138 use crate::installers::components::textgen::{TextgenConfig, TextgenInstaller};
8141 let inst = TextgenInstaller::new(TextgenConfig {
8142 install_dir: "/tmp/test-textgen".to_string(),
8143 python_bin: "python3".to_string(),
8144 use_uv: false,
8145 break_system_packages: true,
8146 dry_run: false,
8147 });
8148 let cmd = inst.build_git_clone_command();
8149
8150 for (i, arg) in cmd.args.iter().enumerate() {
8152 assert!(
8153 !arg.is_empty(),
8154 "Argument {} should not be empty in git clone command",
8155 i
8156 );
8157 }
8158
8159 assert_eq!(cmd.program, "git");
8161 assert_eq!(
8162 cmd.args.len(),
8163 3,
8164 "git clone should have 3 args: clone, URL, dir"
8165 );
8166 assert_eq!(cmd.args[0], "clone");
8167 assert_eq!(
8168 cmd.args[1],
8169 "https://github.com/oobabooga/text-generation-webui.git"
8170 );
8171 assert_eq!(cmd.args[2], "/tmp/test-textgen");
8172
8173 assert!(
8175 !cmd.args.contains(&"--".to_string()),
8176 "git clone args should not contain '--' separator (from sudo wrapping)"
8177 );
8178 }
8179
8180 #[test]
8185 fn test_verification_result_labels() {
8186 assert_eq!(VerificationResult::Verified.label(), "Verified");
8187 assert_eq!(VerificationResult::Failed.label(), "Failed");
8188 assert_eq!(VerificationResult::Missing.label(), "Missing");
8189 }
8190}