1use std::collections::HashMap;
14use std::path::{Path, PathBuf};
15use std::process::Command;
16
17use serde::{Deserialize, Serialize};
18
19use super::environment::resolve_user_home;
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
27pub enum DetectionMethod {
28 PathBased,
30 PythonModule,
32 GitBased,
34 CommandBased,
36}
37
38#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
40pub struct ComponentInfo {
41 pub id: String,
43 pub display_name: String,
45 pub detection_method: DetectionMethod,
47 pub installer_script: String,
49 pub python_import: Option<String>,
51 pub clone_dir: Option<String>,
53}
54
55#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
57pub struct VersionInfo {
58 pub component_id: String,
60 pub version: String,
62}
63
64pub fn known_components() -> &'static [ComponentInfo] {
70 static COMPONENTS: std::sync::OnceLock<Vec<ComponentInfo>> = std::sync::OnceLock::new();
72 COMPONENTS.get_or_init(|| {
73 vec![
74 ComponentInfo {
76 id: "rocm".into(),
77 display_name: "ROCm".into(),
78 detection_method: DetectionMethod::PathBased,
79 installer_script: "install_rocm.sh".into(),
80 python_import: None,
81 clone_dir: None,
82 },
83 ComponentInfo {
84 id: "rocm-smi".into(),
85 display_name: "ROCm SMI".into(),
86 detection_method: DetectionMethod::CommandBased,
87 installer_script: String::new(), python_import: None,
89 clone_dir: None,
90 },
91 ComponentInfo {
93 id: "pytorch".into(),
94 display_name: "PyTorch".into(),
95 detection_method: DetectionMethod::PythonModule,
96 installer_script: "install_pytorch_rocm.sh".into(),
97 python_import: Some("torch".into()),
98 clone_dir: None,
99 },
100 ComponentInfo {
101 id: "triton".into(),
102 display_name: "Triton".into(),
103 detection_method: DetectionMethod::PythonModule,
104 installer_script: "install_triton_multi.sh".into(),
105 python_import: Some("triton".into()),
106 clone_dir: None,
107 },
108 ComponentInfo {
109 id: "onnx".into(),
110 display_name: "ONNX Runtime".into(),
111 detection_method: DetectionMethod::PythonModule,
112 installer_script: "build_onnxruntime_multi.sh".into(),
113 python_import: Some("onnxruntime".into()),
114 clone_dir: None,
115 },
116 ComponentInfo {
117 id: "migraphx".into(),
118 display_name: "MIGraphX".into(),
119 detection_method: DetectionMethod::PythonModule,
120 installer_script: "install_migraphx_multi.sh".into(),
121 python_import: Some("migraphx".into()),
122 clone_dir: None,
123 },
124 ComponentInfo {
125 id: "flash-attn".into(),
126 display_name: "Flash Attention".into(),
127 detection_method: DetectionMethod::PythonModule,
128 installer_script: "install_flash_attention_ck.sh".into(),
129 python_import: Some("flash_attn".into()),
130 clone_dir: None,
131 },
132 ComponentInfo {
133 id: "mpi4py".into(),
134 display_name: "MPI4Py".into(),
135 detection_method: DetectionMethod::PythonModule,
136 installer_script: "install_mpi4py.sh".into(),
137 python_import: Some("mpi4py".into()),
138 clone_dir: None,
139 },
140 ComponentInfo {
141 id: "megatron".into(),
142 display_name: "Megatron-LM".into(),
143 detection_method: DetectionMethod::GitBased,
144 installer_script: "install_megatron.sh".into(),
145 python_import: None,
146 clone_dir: Some("Megatron-LM".into()),
147 },
148 ComponentInfo {
149 id: "aiter".into(),
150 display_name: "AITER".into(),
151 detection_method: DetectionMethod::PythonModule,
152 installer_script: "install_aiter.sh".into(),
153 python_import: Some("aiter".into()),
154 clone_dir: None,
155 },
156 ComponentInfo {
158 id: "deepspeed".into(),
159 display_name: "DeepSpeed".into(),
160 detection_method: DetectionMethod::PythonModule,
161 installer_script: "install_deepspeed.sh".into(),
162 python_import: Some("deepspeed".into()),
163 clone_dir: None,
164 },
165 ComponentInfo {
166 id: "vllm".into(),
167 display_name: "vLLM".into(),
168 detection_method: DetectionMethod::PythonModule,
169 installer_script: "install_vllm_multi.sh".into(),
170 python_import: Some("vllm".into()),
171 clone_dir: None,
172 },
173 ComponentInfo {
174 id: "bitsandbytes".into(),
175 display_name: "bitsandbytes".into(),
176 detection_method: DetectionMethod::PythonModule,
177 installer_script: "install_bitsandbytes_multi.sh".into(),
178 python_import: Some("bitsandbytes".into()),
179 clone_dir: None,
180 },
181 ComponentInfo {
182 id: "wandb".into(),
183 display_name: "Weights & Biases".into(),
184 detection_method: DetectionMethod::PythonModule,
185 installer_script: "install_wandb.sh".into(),
186 python_import: Some("wandb".into()),
187 clone_dir: None,
188 },
189 ComponentInfo {
191 id: "llama-cpp".into(),
192 display_name: "llama.cpp (HIP)".into(),
193 detection_method: DetectionMethod::CommandBased,
194 installer_script: String::new(), python_import: None,
196 clone_dir: None,
197 },
198 ComponentInfo {
200 id: "comfyui".into(),
201 display_name: "ComfyUI".into(),
202 detection_method: DetectionMethod::GitBased,
203 installer_script: "install_comfyui.sh".into(),
204 python_import: None,
205 clone_dir: Some("ComfyUI".into()),
206 },
207 ComponentInfo {
208 id: "vllm-studio".into(),
209 display_name: "vLLM Studio".into(),
210 detection_method: DetectionMethod::GitBased,
211 installer_script: "install_vllm_studio.sh".into(),
212 python_import: None,
213 clone_dir: Some("vllm-studio".into()),
214 },
215 ComponentInfo {
216 id: "textgen".into(),
217 display_name: "text-generation-webui".into(),
218 detection_method: DetectionMethod::GitBased,
219 installer_script: "install_textgen.sh".into(),
220 python_import: None,
221 clone_dir: Some("text-generation-webui".into()),
222 },
223 ComponentInfo {
225 id: "permanent-env".into(),
226 display_name: "ML Stack Environment".into(),
227 detection_method: DetectionMethod::PathBased,
228 installer_script: String::new(), python_import: None,
230 clone_dir: None,
231 },
232 ]
233 })
234}
235
236pub fn is_known_component(id: &str) -> bool {
238 known_components().iter().any(|c| c.id == id)
239}
240
241pub fn get_component(id: &str) -> Option<&'static ComponentInfo> {
243 known_components().iter().find(|c| c.id == id)
244}
245
246pub fn display_name(id: &str) -> String {
251 get_component(id)
252 .map(|c| c.display_name.clone())
253 .unwrap_or_else(|| id.to_string())
254}
255
256pub fn display_name_to_id_map() -> HashMap<String, String> {
260 known_components()
261 .iter()
262 .map(|c| (c.display_name.clone(), c.id.clone()))
263 .collect()
264}
265
266pub fn is_component_installed(id: &str) -> bool {
280 is_component_installed_with_home(id, &resolve_user_home())
281}
282
283pub fn is_component_installed_with_home(id: &str, home: &Path) -> bool {
287 let Some(info) = get_component(id) else {
288 return false;
289 };
290
291 match info.detection_method {
292 DetectionMethod::PathBased => detect_path_based(info, home),
293 DetectionMethod::CommandBased => detect_command_based(info, home),
294 DetectionMethod::PythonModule => detect_python_module_single(info),
295 DetectionMethod::GitBased => detect_git_based(info, home),
296 }
297}
298
299fn detect_path_based(info: &ComponentInfo, home: &Path) -> bool {
305 match info.id.as_str() {
306 "rocm" => {
307 let version_file_exists = Path::new("/opt/rocm/.info/version").exists();
309 let rocminfo_works = Command::new("rocminfo")
310 .arg("--version")
311 .output()
312 .map(|o| o.status.success())
313 .unwrap_or(false);
314 version_file_exists && rocminfo_works
315 }
316 "permanent-env" => {
317 let env_file = home.join(".mlstack_env");
319 if !env_file.exists() {
320 return false;
321 }
322 std::fs::read_to_string(&env_file)
323 .map(|content| content.contains("MLSTACK_PYTHON_BIN"))
324 .unwrap_or(false)
325 }
326 _ => false,
327 }
328}
329
330fn detect_command_based(info: &ComponentInfo, home: &Path) -> bool {
336 match info.id.as_str() {
337 "rocm-smi" => Command::new("rocm-smi")
338 .arg("--version")
339 .output()
340 .map(|o| o.status.success())
341 .unwrap_or(false),
342 "llama-cpp" => {
343 let home_str = home.to_string_lossy();
346 crate::installers::components::llama_cpp::is_llama_cli_functional(&home_str)
347 || Command::new("llama-cli")
348 .arg("--help")
349 .output()
350 .map(|o| o.status.success())
351 .unwrap_or(false)
352 }
353 _ => false,
354 }
355}
356
357fn detect_python_module_single(info: &ComponentInfo) -> bool {
363 let Some(ref import_name) = info.python_import else {
364 return false;
365 };
366
367 let interpreters = super::environment::python_interpreters();
368 for python in &interpreters {
369 if try_python_import(python, import_name) {
370 return true;
371 }
372 }
373 false
374}
375
376fn try_python_import(python: &Path, module: &str) -> bool {
378 let script = format!(
379 "import importlib.util; spec = importlib.util.find_spec('{module}'); exit(0 if spec else 1)"
380 );
381 Command::new(python)
382 .arg("-c")
383 .arg(&script)
384 .output()
385 .map(|o| o.status.success())
386 .unwrap_or(false)
387}
388
389fn detect_git_based(info: &ComponentInfo, home: &Path) -> bool {
395 let Some(ref clone_dir) = info.clone_dir else {
396 return false;
397 };
398
399 let git_dir = home.join(clone_dir).join(".git");
400 git_dir.exists()
401}
402
403pub fn get_version(id: &str) -> String {
415 get_version_with_home(id, &resolve_user_home())
416}
417
418pub fn get_version_with_home(id: &str, home: &Path) -> String {
420 let Some(info) = get_component(id) else {
421 return "unknown".to_string();
422 };
423
424 match info.detection_method {
425 DetectionMethod::PathBased => get_version_path_based(info, home),
426 DetectionMethod::CommandBased => get_version_command_based(info, home),
427 DetectionMethod::PythonModule => get_version_python_single(info),
428 DetectionMethod::GitBased => get_version_git(info, home),
429 }
430}
431
432fn get_version_path_based(info: &ComponentInfo, home: &Path) -> String {
434 match info.id.as_str() {
435 "rocm" => {
436 if let Ok(content) = std::fs::read_to_string("/opt/rocm/.info/version") {
438 let v = content.trim().to_string();
439 if !v.is_empty() {
440 return v;
441 }
442 }
443 if let Ok(output) = Command::new("rocminfo").arg("--version").output() {
444 if output.status.success() {
445 let stdout = String::from_utf8_lossy(&output.stdout);
446 if let Some(line) = stdout.lines().next() {
447 return line.trim().to_string();
448 }
449 }
450 }
451 "unknown".to_string()
452 }
453 "permanent-env" => {
454 if is_component_installed_with_home("permanent-env", home) {
455 "installed".to_string()
456 } else {
457 "not installed".to_string()
458 }
459 }
460 _ => "unknown".to_string(),
461 }
462}
463
464fn get_version_command_based(info: &ComponentInfo, home: &Path) -> String {
466 match info.id.as_str() {
467 "rocm-smi" => {
468 if let Ok(output) = Command::new("rocm-smi").arg("--version").output() {
469 if output.status.success() {
470 let stdout = String::from_utf8_lossy(&output.stdout);
471 if let Some(line) = stdout.lines().next() {
472 return extract_semver(line.trim());
473 }
474 }
475 }
476 "not installed".to_string()
477 }
478 "llama-cpp" => {
479 if let Ok(output) = Command::new("llama-cli").arg("--version").output() {
480 if output.status.success() {
481 let stdout = String::from_utf8_lossy(&output.stdout);
482 let stderr = String::from_utf8_lossy(&output.stderr);
483 let combined = format!("{stdout}\n{stderr}");
484 let parsed = extract_semver(combined.trim());
485 if parsed != "unknown" {
486 return parsed;
487 }
488 return "installed".to_string();
489 }
490 }
491 if detect_command_based(info, home) {
492 return "installed".to_string();
493 }
494 "not installed".to_string()
495 }
496 _ => "unknown".to_string(),
497 }
498}
499
500fn get_version_python_single(info: &ComponentInfo) -> String {
502 let Some(ref import_name) = info.python_import else {
503 return "unknown".to_string();
504 };
505
506 let interpreters = super::environment::python_interpreters();
507 for python in &interpreters {
508 let script = format!(
509 "try:\n import {import_name}\n ver = getattr({import_name}, '__version__', None)\n if ver is None:\n import importlib.metadata as im\n for _n in ['{import_name}', '{import_name}'.replace('-', '_'), 'amd_' + '{import_name}'.replace('-', '_'), 'amd-' + '{import_name}']:\n try:\n ver = im.version(_n)\n break\n except Exception:\n pass\n if ver is not None:\n print(ver)\nexcept (ImportError, AttributeError):\n pass",
510 import_name = import_name
511 );
512 if let Ok(output) = Command::new(python).arg("-c").arg(&script).output() {
513 if output.status.success() {
514 let stdout = String::from_utf8_lossy(&output.stdout);
515 let version = stdout
519 .lines()
520 .map(|l| l.trim())
521 .filter(|l| !l.is_empty() && !l.contains("[WARNING]") && !l.contains("[ERROR]"))
522 .last()
523 .unwrap_or("")
524 .to_string();
525 if !version.is_empty() {
526 return version;
527 }
528 }
529 }
530 }
531 "not installed".to_string()
532}
533
534fn get_version_git(info: &ComponentInfo, home: &Path) -> String {
537 let Some(ref clone_dir) = info.clone_dir else {
538 return "unknown".to_string();
539 };
540
541 let repo_path = home.join(clone_dir);
542
543 if !repo_path.join(".git").exists() {
544 return "not installed".to_string();
545 }
546
547 if let Ok(output) = Command::new("git")
549 .args(["-C", repo_path.to_str().unwrap_or(".")])
550 .args(["describe", "--tags", "--abbrev=0"])
551 .output()
552 {
553 if output.status.success() {
554 let stdout = String::from_utf8_lossy(&output.stdout);
555 let tag = stdout.trim().to_string();
556 if !tag.is_empty() {
557 return tag;
558 }
559 }
560 }
561
562 if let Ok(output) = Command::new("git")
564 .args(["-C", repo_path.to_str().unwrap_or(".")])
565 .args(["log", "-1", "--format=%h %s"])
566 .output()
567 {
568 if output.status.success() {
569 let stdout = String::from_utf8_lossy(&output.stdout);
570 let line = stdout.trim().to_string();
571 if !line.is_empty() {
572 return line;
573 }
574 }
575 }
576
577 "installed".to_string()
578}
579
580fn extract_semver(raw: &str) -> String {
592 let s = raw.trim().trim_start_matches('v').to_string();
594
595 let re = regex::Regex::new(r"(\d+\.\d+(?:\.\d+)?)").unwrap();
597 if let Some(caps) = re.captures(&s) {
598 return caps[1].to_string();
599 }
600
601 s
603}
604
605pub fn git_component_has_updates(id: &str) -> bool {
610 let Some(info) = get_component(id) else {
611 return false;
612 };
613 let Some(ref clone_dir) = info.clone_dir else {
614 return false;
615 };
616 let home = resolve_user_home();
617 let repo_path = home.join(clone_dir);
618
619 if !repo_path.join(".git").exists() {
620 return false;
621 }
622
623 let _ = Command::new("git")
625 .args(["-C", repo_path.to_str().unwrap_or(".")])
626 .args(["fetch", "--quiet"])
627 .output();
628
629 if let Ok(output) = Command::new("git")
631 .args(["-C", repo_path.to_str().unwrap_or(".")])
632 .args(["log", "HEAD..@{u}", "--oneline"])
633 .output()
634 {
635 if output.status.success() {
636 let stdout = String::from_utf8_lossy(&output.stdout);
637 return !stdout.trim().is_empty();
638 }
639 }
640
641 false
642}
643
644fn python_module_mappings() -> &'static [(String, String)] {
646 static MAPPINGS: std::sync::OnceLock<Vec<(String, String)>> = std::sync::OnceLock::new();
647 MAPPINGS.get_or_init(|| {
648 known_components()
649 .iter()
650 .filter(|c| c.detection_method == DetectionMethod::PythonModule)
651 .filter_map(|c| {
652 c.python_import
653 .as_ref()
654 .map(|imp| (c.id.clone(), imp.clone()))
655 })
656 .collect()
657 })
658}
659
660pub fn detect_python_modules() -> Vec<String> {
665 let interpreters = super::environment::python_interpreters();
666 detect_python_modules_with_interpreters(&interpreters)
667}
668
669pub fn detect_python_modules_with_interpreters(interpreters: &[PathBuf]) -> Vec<String> {
671 let mappings = python_module_mappings();
672 if mappings.is_empty() || interpreters.is_empty() {
673 return Vec::new();
674 }
675
676 let mut script = String::from("mods = [\n");
678 for (comp_id, import_name) in mappings {
679 script.push_str(&format!(" ({comp_id:?}, {import_name:?}),\n"));
680 }
681 script.push_str(
682 "]\nfor comp_id, import_name in mods:\n try:\n __import__(import_name)\n print(comp_id)\n except ImportError:\n pass\n",
683 );
684
685 for python in interpreters {
686 if let Ok(output) = Command::new(python).arg("-c").arg(&script).output() {
687 tracing::debug!(
688 "Python detection: Interpreter {} returned status: {}",
689 python.display(),
690 output.status
691 );
692 if output.status.success() {
693 let stdout = String::from_utf8_lossy(&output.stdout);
694 let detected: Vec<String> = stdout
695 .lines()
696 .map(|l| l.trim().to_string())
697 .filter(|l| !l.is_empty() && is_known_component(&l))
698 .collect();
699 if !detected.is_empty() {
700 return detected;
701 }
702 }
703 }
704 }
705
706 Vec::new()
707}
708
709pub fn get_versions_batch(component_ids: &[&str]) -> Vec<VersionInfo> {
714 let interpreters = super::environment::python_interpreters();
715 get_versions_batch_with_interpreters(component_ids, &interpreters)
716}
717
718pub fn get_versions_batch_with_interpreters(
720 component_ids: &[&str],
721 interpreters: &[PathBuf],
722) -> Vec<VersionInfo> {
723 let mappings = python_module_mappings();
724
725 if interpreters.is_empty() {
726 return component_ids
727 .iter()
728 .map(|id| VersionInfo {
729 component_id: (*id).to_string(),
730 version: "unknown".to_string(),
731 })
732 .collect();
733 }
734
735 let mut script = String::from("import sys\nmapping = {\n");
737 for (comp_id, import_name) in mappings {
738 script.push_str(&format!(
739 " {:?}: ({:?}, '__version__'),\n",
740 comp_id, import_name
741 ));
742 }
743 script.push_str(
744 "}\nfor comp_id in sys.argv[1:]:\n if comp_id not in mapping:\n print(f'{comp_id}=unknown')\n continue\n mod_name, ver_attr = mapping[comp_id]\n try:\n mod = __import__(mod_name)\n ver = getattr(mod, ver_attr, 'unknown')\n if ver == 'unknown':\n try:\n import importlib.metadata as im\n for _n in [mod_name, mod_name.replace('-', '_'), 'amd_' + mod_name.replace('-', '_'), 'amd-' + mod_name]:\n try:\n ver = im.version(_n)\n break\n except Exception:\n pass\n except Exception:\n pass\n print(f'{comp_id}={ver}')\n except ImportError:\n print(f'{comp_id}=not installed')\n",
745 );
746
747 for python in interpreters {
748 if let Ok(output) = Command::new(python)
749 .arg("-c")
750 .arg(&script)
751 .args(component_ids)
752 .output()
753 {
754 if output.status.success() {
755 let stdout = String::from_utf8_lossy(&output.stdout);
756 return stdout
757 .lines()
758 .filter_map(|line| {
759 let line = line.trim();
760 if line.is_empty() {
761 return None;
762 }
763 let (id, version) = line.split_once('=')?;
764 if id.contains('[') || id.contains(' ') || id.contains(']') {
768 return None;
769 }
770 Some(VersionInfo {
771 component_id: id.to_string(),
772 version: version.to_string(),
773 })
774 })
775 .collect();
776 }
777 }
778 }
779
780 component_ids
782 .iter()
783 .map(|id| VersionInfo {
784 component_id: (*id).to_string(),
785 version: "unknown".to_string(),
786 })
787 .collect()
788}
789
790pub fn detect_all_installed() -> Vec<String> {
795 let home = resolve_user_home();
796 detect_all_installed_with_home(&home)
797}
798
799pub fn detect_all_installed_with_home(home: &Path) -> Vec<String> {
801 let mut installed = Vec::new();
802
803 if is_component_installed_with_home("rocm", home) {
805 installed.push("rocm".to_string());
806 }
807
808 for info in known_components() {
810 if info.detection_method == DetectionMethod::CommandBased
811 && is_component_installed(&info.id)
812 {
813 installed.push(info.id.clone());
814 }
815 }
816
817 let python_installed = detect_python_modules();
819 installed.extend(python_installed);
820
821 for info in known_components() {
823 if info.detection_method == DetectionMethod::GitBased && detect_git_based(info, home) {
824 installed.push(info.id.clone());
825 }
826 }
827
828 if is_component_installed_with_home("permanent-env", home) {
830 installed.push("permanent-env".to_string());
831 }
832
833 installed.sort();
834 installed.dedup();
835 installed
836}
837
838#[cfg(test)]
843mod tests {
844 use super::*;
845
846 #[test]
851 fn test_registry_has_exactly_19_components() {
852 let components = known_components();
853 assert_eq!(
854 components.len(),
855 19,
856 "Registry must contain exactly 19 known components, found {}",
857 components.len()
858 );
859 }
860
861 #[test]
862 fn test_registry_contains_all_expected_ids() {
863 let expected_ids = [
864 "rocm",
865 "pytorch",
866 "triton",
867 "megatron",
868 "deepspeed",
869 "vllm",
870 "aiter",
871 "onnx",
872 "bitsandbytes",
873 "migraphx",
874 "flash-attn",
875 "mpi4py",
876 "wandb",
877 "comfyui",
878 "vllm-studio",
879 "textgen",
880 "rocm-smi",
881 "permanent-env",
882 "llama-cpp",
883 ];
884 let components = known_components();
885 for expected in &expected_ids {
886 assert!(
887 components.iter().any(|c| c.id == *expected),
888 "Missing component: {expected}"
889 );
890 }
891 }
892
893 #[test]
894 fn test_is_known_component_accepts_all_19() {
895 let ids = [
896 "rocm",
897 "pytorch",
898 "triton",
899 "megatron",
900 "deepspeed",
901 "vllm",
902 "aiter",
903 "onnx",
904 "bitsandbytes",
905 "migraphx",
906 "flash-attn",
907 "mpi4py",
908 "wandb",
909 "comfyui",
910 "vllm-studio",
911 "textgen",
912 "rocm-smi",
913 "permanent-env",
914 "llama-cpp",
915 ];
916 for id in &ids {
917 assert!(
918 is_known_component(id),
919 "is_known_component should accept '{id}'"
920 );
921 }
922 }
923
924 #[test]
925 fn test_is_known_component_rejects_unknown() {
926 assert!(!is_known_component("nonexistent"));
927 assert!(!is_known_component("cuda"));
928 assert!(!is_known_component("tensorflow"));
929 assert!(!is_known_component(""));
930 assert!(!is_known_component("PyTorch")); }
932
933 #[test]
934 fn test_get_component_returns_info_for_known() {
935 let info = get_component("pytorch").unwrap();
936 assert_eq!(info.id, "pytorch");
937 assert_eq!(info.display_name, "PyTorch");
938 assert_eq!(info.detection_method, DetectionMethod::PythonModule);
939 assert_eq!(info.python_import.as_deref(), Some("torch"));
940 }
941
942 #[test]
943 fn test_get_component_returns_none_for_unknown() {
944 assert!(get_component("nonexistent").is_none());
945 }
946
947 #[test]
948 fn test_all_components_have_non_empty_ids() {
949 for c in known_components() {
950 assert!(!c.id.is_empty(), "Component has empty id");
951 }
952 }
953
954 #[test]
955 fn test_all_component_ids_are_unique() {
956 let ids: Vec<&str> = known_components().iter().map(|c| c.id.as_str()).collect();
957 let unique: std::collections::HashSet<&str> = ids.iter().copied().collect();
958 assert_eq!(ids.len(), unique.len(), "Component IDs must be unique");
959 }
960
961 #[test]
966 fn test_all_display_names_non_empty() {
967 for c in known_components() {
968 assert!(
969 !c.display_name.is_empty(),
970 "Display name for '{}' is empty",
971 c.id
972 );
973 }
974 }
975
976 #[test]
977 fn test_display_names_are_distinct() {
978 let names: Vec<&str> = known_components()
979 .iter()
980 .map(|c| c.display_name.as_str())
981 .collect();
982 let unique: std::collections::HashSet<&str> = names.iter().copied().collect();
983 assert_eq!(
984 names.len(),
985 unique.len(),
986 "Display names must all be distinct (bijective)"
987 );
988 }
989
990 #[test]
991 fn test_display_name_to_id_is_bijective() {
992 let map = display_name_to_id_map();
993 assert_eq!(
994 map.len(),
995 19,
996 "Display name to ID map must have exactly 19 entries"
997 );
998
999 for c in known_components() {
1001 let name = display_name(&c.id);
1002 let back_id = map.get(&name);
1003 assert_eq!(
1004 back_id,
1005 Some(&c.id),
1006 "Round-trip failed for '{}': display_name='{}'",
1007 c.id,
1008 name
1009 );
1010 }
1011 }
1012
1013 #[test]
1014 fn test_display_name_unknown_passthrough() {
1015 assert_eq!(display_name("nonexistent"), "nonexistent");
1016 assert_eq!(display_name("custom-tool"), "custom-tool");
1017 }
1018
1019 #[test]
1020 fn test_display_name_known_values() {
1021 assert_eq!(display_name("rocm"), "ROCm");
1022 assert_eq!(display_name("pytorch"), "PyTorch");
1023 assert_eq!(display_name("flash-attn"), "Flash Attention");
1024 assert_eq!(display_name("wandb"), "Weights & Biases");
1025 assert_eq!(display_name("comfyui"), "ComfyUI");
1026 assert_eq!(display_name("permanent-env"), "ML Stack Environment");
1027 }
1028
1029 #[test]
1034 fn test_rocm_detection_requires_both_version_file_and_rocminfo() {
1035 let result = is_component_installed("rocm");
1039 let _ = result;
1041 }
1042
1043 #[test]
1044 fn test_permanent_env_detection_checks_marker() {
1045 let dir = tempfile::tempdir().unwrap();
1046
1047 std::fs::write(dir.path().join(".mlstack_env"), "# empty\n").unwrap();
1049 assert!(!is_component_installed_with_home(
1050 "permanent-env",
1051 dir.path()
1052 ));
1053
1054 std::fs::write(
1056 dir.path().join(".mlstack_env"),
1057 "export MLSTACK_PYTHON_BIN=/usr/bin/python3\n",
1058 )
1059 .unwrap();
1060 assert!(is_component_installed_with_home(
1061 "permanent-env",
1062 dir.path()
1063 ));
1064 }
1065
1066 #[test]
1067 fn test_permanent_env_not_installed_when_file_missing() {
1068 let dir = tempfile::tempdir().unwrap();
1069 assert!(!is_component_installed_with_home(
1070 "permanent-env",
1071 dir.path()
1072 ));
1073 }
1074
1075 #[test]
1080 fn test_python_module_detection_for_pytorch() {
1081 let _ = is_component_installed("pytorch");
1085 }
1086
1087 #[test]
1088 fn test_detect_python_modules_with_interpreters_empty() {
1089 let result = detect_python_modules_with_interpreters(&[]);
1090 assert!(result.is_empty());
1091 }
1092
1093 #[test]
1094 fn test_detect_python_modules_uses_single_subprocess() {
1095 let result = detect_python_modules();
1098 for id in &result {
1100 assert!(is_known_component(id), "Unknown component ID: {id}");
1101 }
1102 }
1103
1104 #[test]
1109 fn test_git_based_detection_comfyui_with_git_dir() {
1110 let dir = tempfile::tempdir().unwrap();
1111
1112 std::fs::create_dir_all(dir.path().join("ComfyUI/.git")).unwrap();
1114 assert!(is_component_installed_with_home("comfyui", dir.path()));
1115 }
1116
1117 #[test]
1118 fn test_git_based_detection_comfyui_without_git_dir() {
1119 let dir = tempfile::tempdir().unwrap();
1120
1121 std::fs::create_dir_all(dir.path().join("ComfyUI")).unwrap();
1123 assert!(!is_component_installed_with_home("comfyui", dir.path()));
1124 }
1125
1126 #[test]
1127 fn test_git_based_detection_vllm_studio() {
1128 let dir = tempfile::tempdir().unwrap();
1129
1130 std::fs::create_dir_all(dir.path().join("vllm-studio/.git")).unwrap();
1131 assert!(is_component_installed_with_home("vllm-studio", dir.path()));
1132 }
1133
1134 #[test]
1135 fn test_git_based_detection_textgen() {
1136 let dir = tempfile::tempdir().unwrap();
1137
1138 std::fs::create_dir_all(dir.path().join("text-generation-webui/.git")).unwrap();
1139 assert!(is_component_installed_with_home("textgen", dir.path()));
1140 }
1141
1142 #[test]
1147 fn test_version_query_rocm_format() {
1148 let version = get_version("rocm");
1149 assert!(!version.is_empty());
1151 }
1152
1153 #[test]
1154 fn test_version_query_permanent_env() {
1155 let dir = tempfile::tempdir().unwrap();
1156
1157 let version = get_version_with_home("permanent-env", dir.path());
1159 assert_eq!(version, "not installed");
1160
1161 std::fs::write(
1163 dir.path().join(".mlstack_env"),
1164 "export MLSTACK_PYTHON_BIN=/usr/bin/python3\n",
1165 )
1166 .unwrap();
1167 let version = get_version_with_home("permanent-env", dir.path());
1168 assert_eq!(version, "installed");
1169 }
1170
1171 #[test]
1172 fn test_version_query_unknown_component() {
1173 assert_eq!(get_version("nonexistent"), "unknown");
1174 }
1175
1176 #[test]
1177 fn test_version_query_git_based_not_installed() {
1178 let dir = tempfile::tempdir().unwrap();
1179 let version = get_version_with_home("comfyui", dir.path());
1180 assert_eq!(version, "not installed");
1181 }
1182
1183 #[test]
1188 fn test_get_versions_batch_returns_results() {
1189 let ids = vec!["pytorch", "triton", "nonexistent-mod"];
1190 let results = get_versions_batch(&ids);
1191
1192 assert_eq!(results.len(), 3);
1193
1194 assert_eq!(results[0].component_id, "pytorch");
1196 assert_eq!(results[1].component_id, "triton");
1197 assert_eq!(results[2].component_id, "nonexistent-mod");
1198
1199 assert_eq!(results[2].version, "unknown");
1201 }
1202
1203 #[test]
1204 fn test_get_versions_batch_with_interpreters_empty() {
1205 let ids = vec!["pytorch"];
1206 let results = get_versions_batch_with_interpreters(&ids, &[]);
1207 assert_eq!(results.len(), 1);
1208 assert_eq!(results[0].version, "unknown");
1209 }
1210
1211 #[test]
1212 fn test_get_versions_batch_empty_ids() {
1213 let results = get_versions_batch(&[]);
1214 assert!(results.is_empty());
1215 }
1216
1217 #[test]
1222 fn test_detect_all_installed_returns_valid_ids() {
1223 let installed = detect_all_installed();
1224 for id in &installed {
1225 assert!(
1226 is_known_component(id),
1227 "detect_all_installed returned unknown component: {id}"
1228 );
1229 }
1230 }
1231
1232 #[test]
1233 fn test_detect_all_installed_no_duplicates() {
1234 let installed = detect_all_installed();
1235 let unique: std::collections::HashSet<&String> = installed.iter().collect();
1236 assert_eq!(
1237 installed.len(),
1238 unique.len(),
1239 "detect_all_installed should not return duplicates"
1240 );
1241 }
1242
1243 #[test]
1248 fn test_component_info_serde_roundtrip() {
1249 let info = get_component("pytorch").unwrap();
1250 let json = serde_json::to_string(info).unwrap();
1251 let back: ComponentInfo = serde_json::from_str(&json).unwrap();
1252 assert_eq!(info, &back);
1253 }
1254
1255 #[test]
1256 fn test_detection_method_serde_roundtrip() {
1257 let methods = [
1258 DetectionMethod::PathBased,
1259 DetectionMethod::PythonModule,
1260 DetectionMethod::GitBased,
1261 DetectionMethod::CommandBased,
1262 ];
1263 for method in &methods {
1264 let json = serde_json::to_string(method).unwrap();
1265 let back: DetectionMethod = serde_json::from_str(&json).unwrap();
1266 assert_eq!(*method, back);
1267 }
1268 }
1269
1270 #[test]
1271 fn test_version_info_serde_roundtrip() {
1272 let info = VersionInfo {
1273 component_id: "pytorch".to_string(),
1274 version: "2.4.0".to_string(),
1275 };
1276 let json = serde_json::to_string(&info).unwrap();
1277 let back: VersionInfo = serde_json::from_str(&json).unwrap();
1278 assert_eq!(info, back);
1279 }
1280
1281 #[test]
1286 fn test_python_components_have_import_names() {
1287 for c in known_components() {
1288 if c.detection_method == DetectionMethod::PythonModule {
1289 assert!(
1290 c.python_import.is_some(),
1291 "Python module component '{}' must have python_import",
1292 c.id
1293 );
1294 assert!(
1295 !c.python_import.as_ref().unwrap().is_empty(),
1296 "python_import for '{}' must be non-empty",
1297 c.id
1298 );
1299 }
1300 }
1301 }
1302
1303 #[test]
1304 fn test_git_components_have_clone_dirs() {
1305 for c in known_components() {
1306 if c.detection_method == DetectionMethod::GitBased {
1307 assert!(
1308 c.clone_dir.is_some(),
1309 "Git-based component '{}' must have clone_dir",
1310 c.id
1311 );
1312 assert!(
1313 !c.clone_dir.as_ref().unwrap().is_empty(),
1314 "clone_dir for '{}' must be non-empty",
1315 c.id
1316 );
1317 }
1318 }
1319 }
1320
1321 #[test]
1322 fn test_installer_scripts_for_installable_components() {
1323 let no_script = ["rocm-smi", "permanent-env", "llama-cpp"];
1325 for c in known_components() {
1326 if !no_script.contains(&c.id.as_str()) {
1327 assert!(
1328 !c.installer_script.is_empty(),
1329 "Component '{}' must have an installer script",
1330 c.id
1331 );
1332 }
1333 }
1334 }
1335
1336 #[test]
1341 fn test_real_system_detect_all() {
1342 let installed = detect_all_installed();
1343 println!("Installed components: {:?}", installed);
1346 }
1348
1349 #[test]
1350 fn test_real_system_version_queries() {
1351 for c in known_components() {
1353 let version = get_version(&c.id);
1354 println!("{}: {}", display_name(&c.id), version);
1355 }
1356 }
1358}