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 .rfind(|l| !l.is_empty() && !l.contains("[WARNING]") && !l.contains("[ERROR]"))
522 .unwrap_or("")
523 .to_string();
524 if !version.is_empty() {
525 return version;
526 }
527 }
528 }
529 }
530 "not installed".to_string()
531}
532
533fn get_version_git(info: &ComponentInfo, home: &Path) -> String {
536 let Some(ref clone_dir) = info.clone_dir else {
537 return "unknown".to_string();
538 };
539
540 let repo_path = home.join(clone_dir);
541
542 if !repo_path.join(".git").exists() {
543 return "not installed".to_string();
544 }
545
546 if let Ok(output) = Command::new("git")
548 .args(["-C", repo_path.to_str().unwrap_or(".")])
549 .args(["describe", "--tags", "--abbrev=0"])
550 .output()
551 {
552 if output.status.success() {
553 let stdout = String::from_utf8_lossy(&output.stdout);
554 let tag = stdout.trim().to_string();
555 if !tag.is_empty() {
556 return tag;
557 }
558 }
559 }
560
561 if let Ok(output) = Command::new("git")
563 .args(["-C", repo_path.to_str().unwrap_or(".")])
564 .args(["log", "-1", "--format=%h %s"])
565 .output()
566 {
567 if output.status.success() {
568 let stdout = String::from_utf8_lossy(&output.stdout);
569 let line = stdout.trim().to_string();
570 if !line.is_empty() {
571 return line;
572 }
573 }
574 }
575
576 "installed".to_string()
577}
578
579fn extract_semver(raw: &str) -> String {
591 let s = raw.trim().trim_start_matches('v').to_string();
593
594 let re = regex::Regex::new(r"(\d+\.\d+(?:\.\d+)?)").unwrap();
596 if let Some(caps) = re.captures(&s) {
597 return caps[1].to_string();
598 }
599
600 s
602}
603
604pub fn git_component_has_updates(id: &str) -> bool {
609 let Some(info) = get_component(id) else {
610 return false;
611 };
612 let Some(ref clone_dir) = info.clone_dir else {
613 return false;
614 };
615 let home = resolve_user_home();
616 let repo_path = home.join(clone_dir);
617
618 if !repo_path.join(".git").exists() {
619 return false;
620 }
621
622 let _ = Command::new("git")
624 .args(["-C", repo_path.to_str().unwrap_or(".")])
625 .args(["fetch", "--quiet"])
626 .output();
627
628 if let Ok(output) = Command::new("git")
630 .args(["-C", repo_path.to_str().unwrap_or(".")])
631 .args(["log", "HEAD..@{u}", "--oneline"])
632 .output()
633 {
634 if output.status.success() {
635 let stdout = String::from_utf8_lossy(&output.stdout);
636 return !stdout.trim().is_empty();
637 }
638 }
639
640 false
641}
642
643fn python_module_mappings() -> &'static [(String, String)] {
645 static MAPPINGS: std::sync::OnceLock<Vec<(String, String)>> = std::sync::OnceLock::new();
646 MAPPINGS.get_or_init(|| {
647 known_components()
648 .iter()
649 .filter(|c| c.detection_method == DetectionMethod::PythonModule)
650 .filter_map(|c| {
651 c.python_import
652 .as_ref()
653 .map(|imp| (c.id.clone(), imp.clone()))
654 })
655 .collect()
656 })
657}
658
659pub fn detect_python_modules() -> Vec<String> {
664 let interpreters = super::environment::python_interpreters();
665 detect_python_modules_with_interpreters(&interpreters)
666}
667
668pub fn detect_python_modules_with_interpreters(interpreters: &[PathBuf]) -> Vec<String> {
670 let mappings = python_module_mappings();
671 if mappings.is_empty() || interpreters.is_empty() {
672 return Vec::new();
673 }
674
675 let mut script = String::from("mods = [\n");
677 for (comp_id, import_name) in mappings {
678 script.push_str(&format!(" ({comp_id:?}, {import_name:?}),\n"));
679 }
680 script.push_str(
681 "]\nfor comp_id, import_name in mods:\n try:\n __import__(import_name)\n print(comp_id)\n except ImportError:\n pass\n",
682 );
683
684 for python in interpreters {
685 if let Ok(output) = Command::new(python).arg("-c").arg(&script).output() {
686 tracing::debug!(
687 "Python detection: Interpreter {} returned status: {}",
688 python.display(),
689 output.status
690 );
691 if output.status.success() {
692 let stdout = String::from_utf8_lossy(&output.stdout);
693 let detected: Vec<String> = stdout
694 .lines()
695 .map(|l| l.trim().to_string())
696 .filter(|l| !l.is_empty() && is_known_component(l))
697 .collect();
698 if !detected.is_empty() {
699 return detected;
700 }
701 }
702 }
703 }
704
705 Vec::new()
706}
707
708pub fn get_versions_batch(component_ids: &[&str]) -> Vec<VersionInfo> {
713 let interpreters = super::environment::python_interpreters();
714 get_versions_batch_with_interpreters(component_ids, &interpreters)
715}
716
717pub fn get_versions_batch_with_interpreters(
719 component_ids: &[&str],
720 interpreters: &[PathBuf],
721) -> Vec<VersionInfo> {
722 let mappings = python_module_mappings();
723
724 if interpreters.is_empty() {
725 return component_ids
726 .iter()
727 .map(|id| VersionInfo {
728 component_id: (*id).to_string(),
729 version: "unknown".to_string(),
730 })
731 .collect();
732 }
733
734 let mut script = String::from("import sys\nmapping = {\n");
736 for (comp_id, import_name) in mappings {
737 script.push_str(&format!(
738 " {:?}: ({:?}, '__version__'),\n",
739 comp_id, import_name
740 ));
741 }
742 script.push_str(
743 "}\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",
744 );
745
746 for python in interpreters {
747 if let Ok(output) = Command::new(python)
748 .arg("-c")
749 .arg(&script)
750 .args(component_ids)
751 .output()
752 {
753 if output.status.success() {
754 let stdout = String::from_utf8_lossy(&output.stdout);
755 return stdout
756 .lines()
757 .filter_map(|line| {
758 let line = line.trim();
759 if line.is_empty() {
760 return None;
761 }
762 let (id, version) = line.split_once('=')?;
763 if id.contains('[') || id.contains(' ') || id.contains(']') {
767 return None;
768 }
769 Some(VersionInfo {
770 component_id: id.to_string(),
771 version: version.to_string(),
772 })
773 })
774 .collect();
775 }
776 }
777 }
778
779 component_ids
781 .iter()
782 .map(|id| VersionInfo {
783 component_id: (*id).to_string(),
784 version: "unknown".to_string(),
785 })
786 .collect()
787}
788
789pub fn detect_all_installed() -> Vec<String> {
794 let home = resolve_user_home();
795 detect_all_installed_with_home(&home)
796}
797
798pub fn detect_all_installed_with_home(home: &Path) -> Vec<String> {
800 let mut installed = Vec::new();
801
802 if is_component_installed_with_home("rocm", home) {
804 installed.push("rocm".to_string());
805 }
806
807 for info in known_components() {
809 if info.detection_method == DetectionMethod::CommandBased
810 && is_component_installed(&info.id)
811 {
812 installed.push(info.id.clone());
813 }
814 }
815
816 let python_installed = detect_python_modules();
818 installed.extend(python_installed);
819
820 for info in known_components() {
822 if info.detection_method == DetectionMethod::GitBased && detect_git_based(info, home) {
823 installed.push(info.id.clone());
824 }
825 }
826
827 if is_component_installed_with_home("permanent-env", home) {
829 installed.push("permanent-env".to_string());
830 }
831
832 installed.sort();
833 installed.dedup();
834 installed
835}
836
837#[cfg(test)]
842mod tests {
843 use super::*;
844
845 #[test]
850 fn test_registry_has_exactly_19_components() {
851 let components = known_components();
852 assert_eq!(
853 components.len(),
854 19,
855 "Registry must contain exactly 19 known components, found {}",
856 components.len()
857 );
858 }
859
860 #[test]
861 fn test_registry_contains_all_expected_ids() {
862 let expected_ids = [
863 "rocm",
864 "pytorch",
865 "triton",
866 "megatron",
867 "deepspeed",
868 "vllm",
869 "aiter",
870 "onnx",
871 "bitsandbytes",
872 "migraphx",
873 "flash-attn",
874 "mpi4py",
875 "wandb",
876 "comfyui",
877 "vllm-studio",
878 "textgen",
879 "rocm-smi",
880 "permanent-env",
881 "llama-cpp",
882 ];
883 let components = known_components();
884 for expected in &expected_ids {
885 assert!(
886 components.iter().any(|c| c.id == *expected),
887 "Missing component: {expected}"
888 );
889 }
890 }
891
892 #[test]
893 fn test_is_known_component_accepts_all_19() {
894 let ids = [
895 "rocm",
896 "pytorch",
897 "triton",
898 "megatron",
899 "deepspeed",
900 "vllm",
901 "aiter",
902 "onnx",
903 "bitsandbytes",
904 "migraphx",
905 "flash-attn",
906 "mpi4py",
907 "wandb",
908 "comfyui",
909 "vllm-studio",
910 "textgen",
911 "rocm-smi",
912 "permanent-env",
913 "llama-cpp",
914 ];
915 for id in &ids {
916 assert!(
917 is_known_component(id),
918 "is_known_component should accept '{id}'"
919 );
920 }
921 }
922
923 #[test]
924 fn test_is_known_component_rejects_unknown() {
925 assert!(!is_known_component("nonexistent"));
926 assert!(!is_known_component("cuda"));
927 assert!(!is_known_component("tensorflow"));
928 assert!(!is_known_component(""));
929 assert!(!is_known_component("PyTorch")); }
931
932 #[test]
933 fn test_get_component_returns_info_for_known() {
934 let info = get_component("pytorch").unwrap();
935 assert_eq!(info.id, "pytorch");
936 assert_eq!(info.display_name, "PyTorch");
937 assert_eq!(info.detection_method, DetectionMethod::PythonModule);
938 assert_eq!(info.python_import.as_deref(), Some("torch"));
939 }
940
941 #[test]
942 fn test_get_component_returns_none_for_unknown() {
943 assert!(get_component("nonexistent").is_none());
944 }
945
946 #[test]
947 fn test_all_components_have_non_empty_ids() {
948 for c in known_components() {
949 assert!(!c.id.is_empty(), "Component has empty id");
950 }
951 }
952
953 #[test]
954 fn test_all_component_ids_are_unique() {
955 let ids: Vec<&str> = known_components().iter().map(|c| c.id.as_str()).collect();
956 let unique: std::collections::HashSet<&str> = ids.iter().copied().collect();
957 assert_eq!(ids.len(), unique.len(), "Component IDs must be unique");
958 }
959
960 #[test]
965 fn test_all_display_names_non_empty() {
966 for c in known_components() {
967 assert!(
968 !c.display_name.is_empty(),
969 "Display name for '{}' is empty",
970 c.id
971 );
972 }
973 }
974
975 #[test]
976 fn test_display_names_are_distinct() {
977 let names: Vec<&str> = known_components()
978 .iter()
979 .map(|c| c.display_name.as_str())
980 .collect();
981 let unique: std::collections::HashSet<&str> = names.iter().copied().collect();
982 assert_eq!(
983 names.len(),
984 unique.len(),
985 "Display names must all be distinct (bijective)"
986 );
987 }
988
989 #[test]
990 fn test_display_name_to_id_is_bijective() {
991 let map = display_name_to_id_map();
992 assert_eq!(
993 map.len(),
994 19,
995 "Display name to ID map must have exactly 19 entries"
996 );
997
998 for c in known_components() {
1000 let name = display_name(&c.id);
1001 let back_id = map.get(&name);
1002 assert_eq!(
1003 back_id,
1004 Some(&c.id),
1005 "Round-trip failed for '{}': display_name='{}'",
1006 c.id,
1007 name
1008 );
1009 }
1010 }
1011
1012 #[test]
1013 fn test_display_name_unknown_passthrough() {
1014 assert_eq!(display_name("nonexistent"), "nonexistent");
1015 assert_eq!(display_name("custom-tool"), "custom-tool");
1016 }
1017
1018 #[test]
1019 fn test_display_name_known_values() {
1020 assert_eq!(display_name("rocm"), "ROCm");
1021 assert_eq!(display_name("pytorch"), "PyTorch");
1022 assert_eq!(display_name("flash-attn"), "Flash Attention");
1023 assert_eq!(display_name("wandb"), "Weights & Biases");
1024 assert_eq!(display_name("comfyui"), "ComfyUI");
1025 assert_eq!(display_name("permanent-env"), "ML Stack Environment");
1026 }
1027
1028 #[test]
1033 fn test_rocm_detection_requires_both_version_file_and_rocminfo() {
1034 let result = is_component_installed("rocm");
1038 let _ = result;
1040 }
1041
1042 #[test]
1043 fn test_permanent_env_detection_checks_marker() {
1044 let dir = tempfile::tempdir().unwrap();
1045
1046 std::fs::write(dir.path().join(".mlstack_env"), "# empty\n").unwrap();
1048 assert!(!is_component_installed_with_home(
1049 "permanent-env",
1050 dir.path()
1051 ));
1052
1053 std::fs::write(
1055 dir.path().join(".mlstack_env"),
1056 "export MLSTACK_PYTHON_BIN=/usr/bin/python3\n",
1057 )
1058 .unwrap();
1059 assert!(is_component_installed_with_home(
1060 "permanent-env",
1061 dir.path()
1062 ));
1063 }
1064
1065 #[test]
1066 fn test_permanent_env_not_installed_when_file_missing() {
1067 let dir = tempfile::tempdir().unwrap();
1068 assert!(!is_component_installed_with_home(
1069 "permanent-env",
1070 dir.path()
1071 ));
1072 }
1073
1074 #[test]
1079 fn test_python_module_detection_for_pytorch() {
1080 let _ = is_component_installed("pytorch");
1084 }
1085
1086 #[test]
1087 fn test_detect_python_modules_with_interpreters_empty() {
1088 let result = detect_python_modules_with_interpreters(&[]);
1089 assert!(result.is_empty());
1090 }
1091
1092 #[test]
1093 fn test_detect_python_modules_uses_single_subprocess() {
1094 let result = detect_python_modules();
1097 for id in &result {
1099 assert!(is_known_component(id), "Unknown component ID: {id}");
1100 }
1101 }
1102
1103 #[test]
1108 fn test_git_based_detection_comfyui_with_git_dir() {
1109 let dir = tempfile::tempdir().unwrap();
1110
1111 std::fs::create_dir_all(dir.path().join("ComfyUI/.git")).unwrap();
1113 assert!(is_component_installed_with_home("comfyui", dir.path()));
1114 }
1115
1116 #[test]
1117 fn test_git_based_detection_comfyui_without_git_dir() {
1118 let dir = tempfile::tempdir().unwrap();
1119
1120 std::fs::create_dir_all(dir.path().join("ComfyUI")).unwrap();
1122 assert!(!is_component_installed_with_home("comfyui", dir.path()));
1123 }
1124
1125 #[test]
1126 fn test_git_based_detection_vllm_studio() {
1127 let dir = tempfile::tempdir().unwrap();
1128
1129 std::fs::create_dir_all(dir.path().join("vllm-studio/.git")).unwrap();
1130 assert!(is_component_installed_with_home("vllm-studio", dir.path()));
1131 }
1132
1133 #[test]
1134 fn test_git_based_detection_textgen() {
1135 let dir = tempfile::tempdir().unwrap();
1136
1137 std::fs::create_dir_all(dir.path().join("text-generation-webui/.git")).unwrap();
1138 assert!(is_component_installed_with_home("textgen", dir.path()));
1139 }
1140
1141 #[test]
1146 fn test_version_query_rocm_format() {
1147 let version = get_version("rocm");
1148 assert!(!version.is_empty());
1150 }
1151
1152 #[test]
1153 fn test_version_query_permanent_env() {
1154 let dir = tempfile::tempdir().unwrap();
1155
1156 let version = get_version_with_home("permanent-env", dir.path());
1158 assert_eq!(version, "not installed");
1159
1160 std::fs::write(
1162 dir.path().join(".mlstack_env"),
1163 "export MLSTACK_PYTHON_BIN=/usr/bin/python3\n",
1164 )
1165 .unwrap();
1166 let version = get_version_with_home("permanent-env", dir.path());
1167 assert_eq!(version, "installed");
1168 }
1169
1170 #[test]
1171 fn test_version_query_unknown_component() {
1172 assert_eq!(get_version("nonexistent"), "unknown");
1173 }
1174
1175 #[test]
1176 fn test_version_query_git_based_not_installed() {
1177 let dir = tempfile::tempdir().unwrap();
1178 let version = get_version_with_home("comfyui", dir.path());
1179 assert_eq!(version, "not installed");
1180 }
1181
1182 #[test]
1187 fn test_get_versions_batch_returns_results() {
1188 let ids = vec!["pytorch", "triton", "nonexistent-mod"];
1189 let results = get_versions_batch(&ids);
1190
1191 assert_eq!(results.len(), 3);
1192
1193 assert_eq!(results[0].component_id, "pytorch");
1195 assert_eq!(results[1].component_id, "triton");
1196 assert_eq!(results[2].component_id, "nonexistent-mod");
1197
1198 assert_eq!(results[2].version, "unknown");
1200 }
1201
1202 #[test]
1203 fn test_get_versions_batch_with_interpreters_empty() {
1204 let ids = vec!["pytorch"];
1205 let results = get_versions_batch_with_interpreters(&ids, &[]);
1206 assert_eq!(results.len(), 1);
1207 assert_eq!(results[0].version, "unknown");
1208 }
1209
1210 #[test]
1211 fn test_get_versions_batch_empty_ids() {
1212 let results = get_versions_batch(&[]);
1213 assert!(results.is_empty());
1214 }
1215
1216 #[test]
1221 fn test_detect_all_installed_returns_valid_ids() {
1222 let installed = detect_all_installed();
1223 for id in &installed {
1224 assert!(
1225 is_known_component(id),
1226 "detect_all_installed returned unknown component: {id}"
1227 );
1228 }
1229 }
1230
1231 #[test]
1232 fn test_detect_all_installed_no_duplicates() {
1233 let installed = detect_all_installed();
1234 let unique: std::collections::HashSet<&String> = installed.iter().collect();
1235 assert_eq!(
1236 installed.len(),
1237 unique.len(),
1238 "detect_all_installed should not return duplicates"
1239 );
1240 }
1241
1242 #[test]
1247 fn test_component_info_serde_roundtrip() {
1248 let info = get_component("pytorch").unwrap();
1249 let json = serde_json::to_string(info).unwrap();
1250 let back: ComponentInfo = serde_json::from_str(&json).unwrap();
1251 assert_eq!(info, &back);
1252 }
1253
1254 #[test]
1255 fn test_detection_method_serde_roundtrip() {
1256 let methods = [
1257 DetectionMethod::PathBased,
1258 DetectionMethod::PythonModule,
1259 DetectionMethod::GitBased,
1260 DetectionMethod::CommandBased,
1261 ];
1262 for method in &methods {
1263 let json = serde_json::to_string(method).unwrap();
1264 let back: DetectionMethod = serde_json::from_str(&json).unwrap();
1265 assert_eq!(*method, back);
1266 }
1267 }
1268
1269 #[test]
1270 fn test_version_info_serde_roundtrip() {
1271 let info = VersionInfo {
1272 component_id: "pytorch".to_string(),
1273 version: "2.4.0".to_string(),
1274 };
1275 let json = serde_json::to_string(&info).unwrap();
1276 let back: VersionInfo = serde_json::from_str(&json).unwrap();
1277 assert_eq!(info, back);
1278 }
1279
1280 #[test]
1285 fn test_python_components_have_import_names() {
1286 for c in known_components() {
1287 if c.detection_method == DetectionMethod::PythonModule {
1288 assert!(
1289 c.python_import.is_some(),
1290 "Python module component '{}' must have python_import",
1291 c.id
1292 );
1293 assert!(
1294 !c.python_import.as_ref().unwrap().is_empty(),
1295 "python_import for '{}' must be non-empty",
1296 c.id
1297 );
1298 }
1299 }
1300 }
1301
1302 #[test]
1303 fn test_git_components_have_clone_dirs() {
1304 for c in known_components() {
1305 if c.detection_method == DetectionMethod::GitBased {
1306 assert!(
1307 c.clone_dir.is_some(),
1308 "Git-based component '{}' must have clone_dir",
1309 c.id
1310 );
1311 assert!(
1312 !c.clone_dir.as_ref().unwrap().is_empty(),
1313 "clone_dir for '{}' must be non-empty",
1314 c.id
1315 );
1316 }
1317 }
1318 }
1319
1320 #[test]
1321 fn test_installer_scripts_for_installable_components() {
1322 let no_script = ["rocm-smi", "permanent-env", "llama-cpp"];
1324 for c in known_components() {
1325 if !no_script.contains(&c.id.as_str()) {
1326 assert!(
1327 !c.installer_script.is_empty(),
1328 "Component '{}' must have an installer script",
1329 c.id
1330 );
1331 }
1332 }
1333 }
1334
1335 #[test]
1340 fn test_real_system_detect_all() {
1341 let installed = detect_all_installed();
1342 println!("Installed components: {:?}", installed);
1345 }
1347
1348 #[test]
1349 fn test_real_system_version_queries() {
1350 for c in known_components() {
1352 let version = get_version(&c.id);
1353 println!("{}: {}", display_name(&c.id), version);
1354 }
1355 }
1357}