1use std::collections::HashSet;
33use std::env;
34use std::fs;
35use std::io::IsTerminal;
36use std::path::{Path, PathBuf};
37
38#[derive(Debug, Clone, PartialEq, Eq)]
44pub enum EnvUpdateResult {
45 Created,
47 Updated,
49 Unchanged,
51}
52
53pub fn resolve_user_home() -> PathBuf {
65 if let Ok(home) = env::var("MLSTACK_USER_HOME") {
67 let home = home.trim().to_string();
68 if !home.is_empty() && Path::new(&home).is_dir() {
69 return PathBuf::from(home);
70 }
71 }
72
73 if let Ok(home) = env::var("HOME") {
75 let home = home.trim().to_string();
76 if !home.is_empty() {
77 let running_sudo = env::var("SUDO_USER").is_ok();
79 if !(running_sudo && home == "/root") {
80 return PathBuf::from(home);
81 }
82 }
83 }
84
85 for key in ["SUDO_USER", "USER", "LOGNAME"] {
87 if let Ok(user_name) = env::var(key) {
88 let user_name = user_name.trim().to_string();
89 if !user_name.is_empty() && user_name != "root" {
90 if let Some(home) = lookup_home_from_passwd(&user_name) {
91 return home;
92 }
93 }
94 }
95 }
96
97 PathBuf::from("/tmp")
99}
100
101pub fn lookup_home_from_passwd(username: &str) -> Option<PathBuf> {
103 let content = fs::read_to_string("/etc/passwd").ok()?;
104 for line in content.lines() {
105 let fields: Vec<&str> = line.split(':').collect();
106 if fields.len() >= 6 && fields[0] == username {
107 let home = fields[5].trim();
108 if !home.is_empty() {
109 return Some(PathBuf::from(home));
110 }
111 }
112 }
113 None
114}
115
116pub fn mlstack_root() -> PathBuf {
126 resolve_user_home().join(".mlstack")
127}
128
129pub fn mlstack_global_dir() -> PathBuf {
131 mlstack_root().join("global")
132}
133
134pub fn mlstack_global_python() -> PathBuf {
136 mlstack_global_dir().join("bin").join("python")
137}
138
139pub fn mlstack_envs_dir() -> PathBuf {
141 mlstack_root().join("envs")
142}
143
144pub fn mlstack_env_dir(name: &str) -> PathBuf {
146 mlstack_envs_dir().join(name)
147}
148
149pub fn mlstack_cache_dir() -> PathBuf {
151 mlstack_root().join("cache")
152}
153
154pub fn mlstack_logs_dir() -> PathBuf {
156 mlstack_root().join("logs")
157}
158
159pub fn mlstack_triton_dir() -> PathBuf {
161 mlstack_root().join("triton")
162}
163
164pub fn ensure_mlstack_dirs() -> std::io::Result<()> {
166 for dir in [
167 mlstack_root(),
168 mlstack_global_dir(),
169 mlstack_envs_dir(),
170 mlstack_cache_dir(),
171 mlstack_logs_dir(),
172 mlstack_triton_dir(),
173 ] {
174 std::fs::create_dir_all(&dir)?;
175 }
176 Ok(())
177}
178
179pub fn ensure_global_venv(bootstrap_python: &str) -> anyhow::Result<PathBuf> {
187 let global_python = mlstack_global_python();
188 if !global_python.exists() {
189 ensure_mlstack_dirs()?;
190 let global_dir = mlstack_global_dir();
191 create_venv(&global_dir, bootstrap_python)?;
192 if !global_python.exists() {
193 anyhow::bail!(
194 "global venv creation reported success but {} is missing",
195 global_python.display()
196 );
197 }
198 }
199 ensure_pip_in_venv(&global_python)?;
202 Ok(global_python)
203}
204
205pub fn ensure_named_venv(name: &str, bootstrap_python: &str) -> anyhow::Result<PathBuf> {
211 let name = name.trim();
212 if name.is_empty() {
213 anyhow::bail!("named env name must not be empty");
214 }
215 if name.contains(std::path::MAIN_SEPARATOR)
216 || name.contains('/')
217 || name.contains('\\')
218 || name.contains("..")
219 {
220 anyhow::bail!("invalid env name '{name}'");
221 }
222 ensure_mlstack_dirs()?;
223 let env_dir = mlstack_env_dir(name);
224 let env_python = env_dir.join("bin").join("python");
225 if !env_python.exists() {
226 create_venv(&env_dir, bootstrap_python)?;
227 if !env_python.exists() {
228 anyhow::bail!(
229 "named venv creation reported success but {} is missing",
230 env_python.display()
231 );
232 }
233 }
234 ensure_pip_in_venv(&env_python)?;
236 Ok(env_python)
237}
238
239fn create_venv(dir: &Path, bootstrap_python: &str) -> anyhow::Result<()> {
241 if command_on_path("uv") {
242 let uv_status = std::process::Command::new("uv")
249 .arg("venv")
250 .arg("--seed")
251 .arg("--python")
252 .arg(bootstrap_python)
253 .arg(dir)
254 .status();
255 if let Ok(status) = uv_status {
256 if status.success() {
257 return Ok(());
258 }
259 }
260 }
261 let status = std::process::Command::new(bootstrap_python)
262 .arg("-m")
263 .arg("venv")
264 .arg(dir)
265 .status()?;
266 if !status.success() {
267 anyhow::bail!(
268 "failed to create venv at {} via uv/venv (bootstrap={})",
269 dir.display(),
270 bootstrap_python
271 );
272 }
273 Ok(())
274}
275
276fn ensure_pip_in_venv(venv_python: &Path) -> anyhow::Result<()> {
282 let has_pip = std::process::Command::new(venv_python)
284 .args(["-m", "pip", "--version"])
285 .output()
286 .map(|o| o.status.success())
287 .unwrap_or(false);
288 if has_pip {
289 return Ok(());
290 }
291 if command_on_path("uv") {
293 let s = std::process::Command::new("uv")
294 .args(["pip", "install", "--python"])
295 .arg(venv_python)
296 .arg("pip")
297 .status();
298 if let Ok(status) = s {
299 if status.success() {
300 return Ok(());
301 }
302 }
303 }
304 let status = std::process::Command::new(venv_python)
306 .args(["-m", "ensurepip", "--upgrade"])
307 .status()?;
308 if !status.success() {
309 anyhow::bail!(
310 "could not bootstrap pip into {} (tried `uv pip install pip` and `ensurepip`)",
311 venv_python.display()
312 );
313 }
314 Ok(())
315}
316
317pub fn python_interpreters() -> Vec<PathBuf> {
331 let home = resolve_user_home();
332 python_interpreters_for_home(&home)
333}
334
335pub fn resolve_canonical_python_bin() -> String {
350 if let Ok(name) = env::var("MLSTACK_ENV_NAME") {
352 let name = name.trim();
353 if !name.is_empty() {
354 let p = mlstack_env_dir(name).join("bin").join("python");
355 if p.exists() {
356 return p.to_string_lossy().to_string();
357 }
358 }
359 }
360 for key in ["MLSTACK_PYTHON_BIN", "UV_PYTHON"] {
362 if let Ok(val) = env::var(key) {
363 let val = val.trim();
364 if !val.is_empty() && Path::new(val).exists() {
365 return val.to_string();
366 }
367 }
368 }
369 let global = mlstack_global_python();
371 if global.exists() {
372 return global.to_string_lossy().to_string();
373 }
374 let interpreters = python_interpreters();
376 if let Some(first) = interpreters.first() {
377 return first.to_string_lossy().to_string();
378 }
379 "python3".to_string()
380}
381
382pub fn python_interpreters_for_home(home: &Path) -> Vec<PathBuf> {
384 let mut paths = Vec::new();
385 let mut seen = HashSet::new();
386
387 macro_rules! push_if_exists {
388 ($path:expr) => {
389 let p = $path;
390 if p.exists() && seen.insert(p.clone()) {
391 paths.push(p);
392 }
393 };
394 }
395
396 for key in ["MLSTACK_PYTHON_BIN", "UV_PYTHON"] {
399 if let Ok(val) = env::var(key) {
400 let val = val.trim().to_string();
401 if !val.is_empty() {
402 push_if_exists!(PathBuf::from(&val));
403 }
404 }
405 }
406
407 push_if_exists!(home
411 .join(".mlstack")
412 .join("global")
413 .join("bin")
414 .join("python"));
415 push_if_exists!(home
416 .join(".mlstack")
417 .join("global")
418 .join("bin")
419 .join("python3"));
420
421 if let Ok(venv) = env::var("VIRTUAL_ENV") {
423 let venv = venv.trim().to_string();
424 if !venv.is_empty() {
425 push_if_exists!(PathBuf::from(&venv).join("bin/python"));
426 push_if_exists!(PathBuf::from(&venv).join("bin/python3"));
427 }
428 }
429
430 if let Ok(conda_prefix) = env::var("CONDA_PREFIX") {
432 let conda_prefix = conda_prefix.trim().to_string();
433 if !conda_prefix.is_empty() {
434 push_if_exists!(PathBuf::from(&conda_prefix).join("bin/python"));
435 push_if_exists!(PathBuf::from(&conda_prefix).join("bin/python3"));
436 }
437 }
438
439 push_if_exists!(home.join(".local/bin/python3"));
443 push_if_exists!(home.join(".local/bin/python"));
444
445 let uv_python_dir = home.join(".local/share/uv/python");
447 if let Ok(entries) = fs::read_dir(&uv_python_dir) {
448 let mut uv_pythons: Vec<PathBuf> = entries
449 .filter_map(|e| e.ok())
450 .filter_map(|e| {
451 let bin = e.path().join("bin/python3");
452 if bin.exists() {
453 Some(bin)
454 } else {
455 None
456 }
457 })
458 .collect();
459 uv_pythons.sort_by(|a, b| {
460 let va = extract_python_version(a);
461 let vb = extract_python_version(b);
462 let priority = |v: (u32, u32)| -> (i32, u32) {
465 match v {
466 (3, 12) => (0, v.1), (3, 13) => (1, v.1),
468 (3, 11) => (2, v.1),
469 (3, 10) => (3, v.1),
470 _ => (4, v.1), }
472 };
473 priority(va).cmp(&priority(vb))
474 });
475 for p in uv_pythons {
476 push_if_exists!(p);
477 }
478 }
479
480 let component_venvs = ["rocm_venv", "pytorch", "ml_stack"];
482 for venv_name in &component_venvs {
483 push_if_exists!(home.join(venv_name).join("bin/python"));
484 push_if_exists!(home.join(venv_name).join("bin/python3"));
485 }
486
487 let system_pythons = [
489 "/usr/local/bin/python3",
490 "/usr/bin/python3",
491 "/usr/bin/python",
492 ];
493 for sys_path in &system_pythons {
494 push_if_exists!(PathBuf::from(sys_path));
495 }
496
497 if let Ok(entries) = fs::read_dir("/usr/bin") {
499 let mut versioned: Vec<PathBuf> = entries
500 .filter_map(|e| e.ok())
501 .filter_map(|e| {
502 let name = e.file_name().to_string_lossy().to_string();
503 if name.starts_with("python3.") && e.path().exists() {
504 Some(e.path())
505 } else {
506 None
507 }
508 })
509 .collect();
510 versioned.sort_by(|a, b| {
511 let va = extract_python_version(a);
513 let vb = extract_python_version(b);
514 vb.cmp(&va)
515 });
516 for p in versioned {
517 push_if_exists!(p);
518 }
519 }
520
521 let local_python3 = home.join(".local/bin/python3");
525 if local_python3.exists() && seen.insert(local_python3.clone()) {
526 paths.push(local_python3);
527 }
528
529 paths
530}
531
532fn extract_python_version(path: &Path) -> (u32, u32) {
534 let name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
535 if let Some(rest) = name.strip_prefix("python") {
537 let parts: Vec<&str> = rest.split('.').collect();
538 if parts.len() >= 2 {
539 let major = parts[0].parse::<u32>().unwrap_or(0);
540 let minor = parts[1].parse::<u32>().unwrap_or(0);
541 return (major, minor);
542 }
543 if let Ok(major) = rest.parse::<u32>() {
544 return (major, 0);
545 }
546 }
547 (0, 0)
548}
549
550pub fn resolve_managed_python(backend: &crate::config::PythonBackend) -> anyhow::Result<PathBuf> {
563 match backend {
564 crate::config::PythonBackend::Uv => {
565 if !command_on_path("uv") {
566 anyhow::bail!("uv backend selected but uv not on PATH");
567 }
568
569 let output = std::process::Command::new("uv")
571 .args(["python", "find", "3.12"])
572 .output();
573 if let Ok(out) = output {
574 if out.status.success() {
575 let stdout = String::from_utf8_lossy(&out.stdout);
576 if let Some(line) = stdout.lines().next() {
577 let path = line.trim();
578 if !path.is_empty() {
579 return Ok(PathBuf::from(path));
580 }
581 }
582 }
583 }
584
585 let output = std::process::Command::new("uv")
587 .args(["python", "find", ">=3.10,<3.14"])
588 .output()?;
589 if output.status.success() {
590 let stdout = String::from_utf8_lossy(&output.stdout);
591 if let Some(line) = stdout.lines().next() {
592 let path = line.trim();
593 if !path.is_empty() {
594 return Ok(PathBuf::from(path));
595 }
596 }
597 }
598
599 anyhow::bail!("uv python find failed to locate a Python ≤3.13");
600 }
601 crate::config::PythonBackend::Python => {
602 let candidates = ["python3.12", "python3.11", "python3.10", "python3.13"];
604 for name in &candidates {
605 if let Some(path) = check_python_version_cap(name) {
606 return Ok(path);
607 }
608 }
609
610 anyhow::bail!(
611 "No suitable Python ≤3.13 found on PATH (tried: {})",
612 candidates.join(", ")
613 );
614 }
615 }
616}
617
618fn check_python_version_cap(python_name: &str) -> Option<PathBuf> {
621 let output = std::process::Command::new(python_name)
622 .arg("--version")
623 .output()
624 .ok()?;
625
626 if !output.status.success() {
627 return None;
628 }
629
630 let version_str = String::from_utf8_lossy(&output.stdout);
631 let (major, minor) = parse_python_version_cap(&version_str)?;
633
634 if major > 3 || (major == 3 && minor > 13) {
636 return None;
637 }
638
639 for dir in env::split_paths(&env::var("PATH").unwrap_or_default()) {
641 let candidate = dir.join(python_name);
642 if candidate.is_file() {
643 return Some(candidate);
644 }
645 }
646 None
647}
648
649fn parse_python_version_cap(version_str: &str) -> Option<(u32, u32)> {
651 let rest = version_str.strip_prefix("Python ")?;
653 let parts: Vec<&str> = rest.split('.').collect();
654 if parts.len() >= 2 {
655 let major = parts[0].parse::<u32>().ok()?;
656 let minor = parts[1].parse::<u32>().ok()?;
657 Some((major, minor))
658 } else {
659 None
660 }
661}
662
663pub fn create_global_env(backend: &crate::config::PythonBackend) -> anyhow::Result<PathBuf> {
670 let bootstrap_python = resolve_managed_python(backend)?;
671 let global_dir = mlstack_global_dir();
672
673 if global_dir.exists() {
674 let global_python = mlstack_global_python();
676 if global_python.exists() {
677 ensure_pip_in_venv(&global_python)?;
681 return Ok(global_python);
682 }
683 std::fs::remove_dir_all(&global_dir)?;
685 }
686
687 std::fs::create_dir_all(&global_dir)?;
688
689 match backend {
690 crate::config::PythonBackend::Uv => {
691 let status = std::process::Command::new("uv")
693 .arg("venv")
694 .arg("--seed")
695 .arg(&global_dir)
696 .arg("--python")
697 .arg(&bootstrap_python)
698 .status()?;
699 if !status.success() {
700 anyhow::bail!("uv venv failed for global env");
701 }
702 }
703 crate::config::PythonBackend::Python => {
704 let status = std::process::Command::new(&bootstrap_python)
705 .arg("-m")
706 .arg("venv")
707 .arg(&global_dir)
708 .status()?;
709 if !status.success() {
710 anyhow::bail!("python -m venv failed for global env");
711 }
712 }
713 }
714
715 let global_python = mlstack_global_python();
716 if !global_python.exists() {
717 anyhow::bail!(
718 "global env creation reported success but {} is missing",
719 global_python.display()
720 );
721 }
722 ensure_pip_in_venv(&global_python)?;
725
726 Ok(global_python)
727}
728
729pub fn write_activate_snippets() -> anyhow::Result<(PathBuf, PathBuf)> {
737 let global_dir = mlstack_global_dir();
738 std::fs::create_dir_all(&global_dir)?;
739
740 let sh_path = global_dir.join("activate-global.sh");
741 let fish_path = global_dir.join("activate-global.fish");
742
743 let sh_content = r#"# mlstack-global-python activation — generated by Rusty-Stack
745# Do NOT edit manually — uninstall strips this marker line
746export PATH="$HOME/.mlstack/global/bin:$PATH"
747"#;
748 std::fs::write(&sh_path, sh_content)?;
749
750 let fish_content = r#"# mlstack-global-python activation — generated by Rusty-Stack
752# Do NOT edit manually — uninstall strips this marker line
753fish_add_path $HOME/.mlstack/global/bin
754"#;
755 std::fs::write(&fish_path, fish_content)?;
756
757 Ok((sh_path, fish_path))
758}
759
760pub fn offer_global_prepend(yes_default: bool) -> anyhow::Result<bool> {
772 let shell = std::env::var("SHELL").unwrap_or_default();
773 let (rc_path, _snippet_path) = if shell.contains("fish") {
774 let rc = resolve_user_home().join(".config/fish/config.fish");
775 let snippet = mlstack_global_dir().join("activate-global.fish");
776 (rc, snippet)
777 } else if shell.contains("zsh") {
778 let rc = resolve_user_home().join(".zshrc");
779 let snippet = mlstack_global_dir().join("activate-global.sh");
780 (rc, snippet)
781 } else {
782 let rc = resolve_user_home().join(".bashrc");
784 let snippet = mlstack_global_dir().join("activate-global.sh");
785 (rc, snippet)
786 };
787
788 let marker = "# mlstack-global-python";
789
790 if rc_path.exists() {
792 let content = std::fs::read_to_string(&rc_path).unwrap_or_default();
793 if content.contains(marker) {
794 return Ok(false); }
796 }
797
798 let act = if yes_default {
800 true
801 } else {
802 let is_tty = std::io::stdin().is_terminal();
804 if !is_tty {
805 true } else {
807 eprint!(
809 "Make ~/.mlstack/global the default Python in your {} rc? [Y/n] ",
810 if shell.contains("fish") {
811 "fish"
812 } else if shell.contains("zsh") {
813 "zsh"
814 } else {
815 "bash"
816 }
817 );
818 let mut input = String::new();
819 std::io::stdin().read_line(&mut input).ok();
820 let trimmed = input.trim().to_lowercase();
821 trimmed.is_empty() || trimmed == "y" || trimmed == "yes"
822 }
823 };
824
825 if !act {
826 return Ok(false);
827 }
828
829 if let Some(parent) = rc_path.parent() {
831 std::fs::create_dir_all(parent)?;
832 }
833 if !rc_path.exists() {
834 std::fs::write(&rc_path, "")?;
835 }
836
837 let source_line = if shell.contains("fish") {
839 format!(
840 "source \"$HOME/.mlstack/global/activate-global.fish\" {}",
841 marker
842 )
843 } else {
844 format!(
845 "source \"$HOME/.mlstack/global/activate-global.sh\" {}",
846 marker
847 )
848 };
849
850 let mut content = std::fs::read_to_string(&rc_path).unwrap_or_default();
851 if !content.ends_with('\n') {
852 content.push('\n');
853 }
854 content.push_str(&source_line);
855 content.push('\n');
856
857 std::fs::write(&rc_path, content)?;
858 Ok(true)
859}
860
861pub fn offer_mlstack_env_source() -> anyhow::Result<()> {
872 const MARKER: &str = "# mlstack-rocm-env";
873 let source_line = format!(
874 "[ -f \"$HOME/.mlstack_env\" ] && source \"$HOME/.mlstack_env\" {}",
875 MARKER
876 );
877 for name in ["bashrc", "zshrc"] {
878 let rc_path = resolve_user_home().join(format!(".{}", name));
879 let mut content = std::fs::read_to_string(&rc_path).unwrap_or_default();
880 if content.contains(".mlstack_env") {
881 continue; }
883 if !content.is_empty() && !content.ends_with('\n') {
884 content.push('\n');
885 }
886 content.push_str("# mlstack ROCm env — auto-sourced on Global install\n");
887 content.push_str(&source_line);
888 content.push('\n');
889 if let Some(parent) = rc_path.parent() {
890 std::fs::create_dir_all(parent)?;
891 }
892 std::fs::write(&rc_path, content)?;
893 }
894 Ok(())
895}
896
897pub fn normalize_mlstack_env(
914 user_home: &Path,
915 python_bin: &str,
916 rocm_path: &Path,
917) -> anyhow::Result<EnvUpdateResult> {
918 let env_path = user_home.join(".mlstack_env");
919
920 let rocm_str = rocm_path.to_string_lossy();
921 let rocm_lib_path = rocm_path.join("lib");
922 let rocm_lib = rocm_lib_path.to_string_lossy();
923 let user_home_str = user_home.to_string_lossy().to_string();
928
929 if env_path.exists() {
930 let contents = fs::read_to_string(&env_path)?;
931
932 let (normalized, changed) =
933 normalize_env_contents(&contents, python_bin, &rocm_str, &rocm_lib, &user_home_str);
934
935 if changed {
936 fs::write(&env_path, normalized)?;
937 return Ok(EnvUpdateResult::Updated);
938 }
939 return Ok(EnvUpdateResult::Unchanged);
940 }
941
942 if let Some(parent) = env_path.parent() {
944 fs::create_dir_all(parent)?;
945 }
946
947 let content = generate_env_file(python_bin, &rocm_str, &rocm_lib, &user_home_str);
948 fs::write(&env_path, content)?;
949 Ok(EnvUpdateResult::Created)
950}
951
952pub fn normalize_env_contents(
956 contents: &str,
957 python_bin: &str,
958 rocm_home: &str,
959 rocm_lib: &str,
960 user_home: &str,
961) -> (String, bool) {
962 let mut changed = false;
963 let mut lines: Vec<String> = Vec::new();
964 let mut seen_keys: HashSet<String> = HashSet::new();
965
966 for line in contents.lines() {
967 let trimmed = line.trim();
968 if trimmed.is_empty() || trimmed.starts_with('#') {
969 lines.push(line.to_string());
970 continue;
971 }
972
973 if let Some(key) = extract_export_key(trimmed) {
974 seen_keys.insert(key.clone());
975
976 let (normalized_line, line_changed) =
977 normalize_env_line(line, &key, python_bin, rocm_home, rocm_lib, user_home);
978 changed |= line_changed;
979 lines.push(normalized_line);
980 } else {
981 lines.push(line.to_string());
982 }
983 }
984
985 let required = [
987 (
988 "MLSTACK_PYTHON_BIN",
989 format!("export MLSTACK_PYTHON_BIN={}", python_bin),
990 ),
991 ("ROCM_PATH", format!("export ROCM_PATH={}", rocm_home)),
992 (
993 "ORT_MIGRAPHX_FP16_ENABLE",
994 "export ORT_MIGRAPHX_FP16_ENABLE=0".to_string(),
995 ),
996 ];
997
998 for (key, line) in &required {
999 if !seen_keys.contains(*key) {
1000 lines.push(line.clone());
1001 changed = true;
1002 }
1003 }
1004
1005 (lines.join("\n"), changed)
1006}
1007
1008pub fn generate_env_file(
1014 python_bin: &str,
1015 rocm_home: &str,
1016 rocm_lib: &str,
1017 user_home: &str,
1018) -> String {
1019 format!(
1020 "# ML Stack Environment File (generated by Rusty-Stack)\n\
1021export MLSTACK_PYTHON_BIN={python_bin}\n\
1022export UV_PYTHON={python_bin}\n\
1023export ROCM_HOME={rocm_home}\n\
1024export ROCM_PATH={rocm_home}\n\
1025export HIP_PATH={rocm_home}\n\
1026export PYTHONPATH={rocm_lib}:$PYTHONPATH\n\
1027export PATH=\"{rocm_home}/bin:{rocm_home}/hip/bin:$PATH\"\n\
1028export LD_LIBRARY_PATH=\"{rocm_home}/lib:{rocm_home}/hip/lib:{rocm_home}/opencl/lib:{user_home}/.mlstack/lib-compat:$LD_LIBRARY_PATH\"\n\
1029export ORT_MIGRAPHX_FP16_ENABLE=0\n"
1030 )
1031}
1032
1033fn extract_export_key(line: &str) -> Option<String> {
1039 let line = line.trim();
1040 let after_export = line.strip_prefix("export ")?;
1041 let (key, _) = after_export.split_once('=')?;
1042 Some(key.trim().to_string())
1043}
1044
1045fn normalize_env_line(
1047 line: &str,
1048 key: &str,
1049 python_bin: &str,
1050 rocm_home: &str,
1051 rocm_lib: &str,
1052 user_home: &str,
1053) -> (String, bool) {
1054 match key {
1055 "MLSTACK_PYTHON_BIN" | "UV_PYTHON" => {
1056 let desired = format!("export {}={}", key, python_bin);
1057 let trimmed = line.trim();
1058 if trimmed != desired {
1059 (desired, true)
1060 } else {
1061 (line.to_string(), false)
1062 }
1063 }
1064 "ROCM_PATH" | "ROCM_HOME" | "HIP_PATH" => {
1065 let desired = format!("export {}={}", key, rocm_home);
1066 let trimmed = line.trim();
1067 if trimmed != desired {
1068 (desired, true)
1069 } else {
1070 (line.to_string(), false)
1071 }
1072 }
1073 "PYTHONPATH" => {
1074 let desired = format!("export PYTHONPATH={}:$PYTHONPATH", rocm_lib);
1075 let trimmed = line.trim();
1076 if trimmed != desired {
1077 (desired, true)
1078 } else {
1079 (line.to_string(), false)
1080 }
1081 }
1082 "PATH" => {
1083 let desired = format!(
1084 "export PATH=\"{}/bin:{}/hip/bin:$PATH\"",
1085 rocm_home, rocm_home
1086 );
1087 let trimmed = line.trim();
1088 if trimmed != desired {
1089 (desired, true)
1090 } else {
1091 (line.to_string(), false)
1092 }
1093 }
1094 "LD_LIBRARY_PATH" => {
1095 let desired = format!(
1096 "export LD_LIBRARY_PATH=\"{}/lib:{}/hip/lib:{}/opencl/lib:{}/.mlstack/lib-compat:$LD_LIBRARY_PATH\"",
1097 rocm_home, rocm_home, rocm_home, user_home
1098 );
1099 let trimmed = line.trim();
1100 if trimmed != desired {
1101 (desired, true)
1102 } else {
1103 (line.to_string(), false)
1104 }
1105 }
1106 _ => (line.to_string(), false),
1107 }
1108}
1109
1110pub fn dedup_path_var(path_var: &str) -> String {
1112 let mut seen = HashSet::new();
1113 let mut result = Vec::new();
1114 for entry in path_var.split(':') {
1115 if !entry.is_empty() && seen.insert(entry.to_string()) {
1116 result.push(entry.to_string());
1117 }
1118 }
1119 result.join(":")
1120}
1121
1122pub fn command_on_path(cmd: &str) -> bool {
1124 let path = match env::var_os("PATH") {
1125 Some(p) => p,
1126 None => return false,
1127 };
1128 for dir in env::split_paths(&path) {
1129 let candidate = dir.join(cmd);
1130 if candidate.is_file() {
1131 return true;
1132 }
1133 }
1134 false
1135}
1136
1137#[cfg(test)]
1142mod tests {
1143 use super::*;
1144
1145 #[test]
1150 fn test_resolve_user_home_returns_valid_path() {
1151 let home = resolve_user_home();
1152 assert!(
1153 !home.as_os_str().is_empty(),
1154 "Home directory should not be empty"
1155 );
1156 }
1157
1158 #[test]
1159 fn test_resolve_user_home_env_override() {
1160 let _env = crate::test_support::lock_env();
1161 let saved = std::env::var("MLSTACK_USER_HOME").ok();
1162 let dir = tempfile::tempdir().unwrap();
1163 let dir_path = dir.path().to_path_buf();
1164
1165 std::env::set_var("MLSTACK_USER_HOME", dir_path.to_string_lossy().to_string());
1166 let result = resolve_user_home();
1167 assert_eq!(result, dir_path);
1168
1169 match saved {
1171 Some(v) => std::env::set_var("MLSTACK_USER_HOME", v),
1172 None => std::env::remove_var("MLSTACK_USER_HOME"),
1173 }
1174 }
1175
1176 #[test]
1177 fn test_resolve_user_home_from_passwd() {
1178 if let Ok(user) = env::var("USER") {
1180 if let Some(home) = lookup_home_from_passwd(&user) {
1181 assert!(
1182 home.is_dir(),
1183 "Home directory from passwd should exist: {:?}",
1184 home
1185 );
1186 assert!(
1187 !home.as_os_str().is_empty(),
1188 "Home directory should not be empty"
1189 );
1190 }
1191 }
1192 }
1193
1194 #[test]
1195 fn test_lookup_home_from_passwd_root() {
1196 let home = lookup_home_from_passwd("root");
1198 assert!(home.is_some(), "root should be in /etc/passwd");
1199 assert_eq!(home.unwrap(), PathBuf::from("/root"));
1200 }
1201
1202 #[test]
1203 fn test_lookup_home_from_passwd_nonexistent() {
1204 let home = lookup_home_from_passwd("nonexistent_user_xyz_12345");
1205 assert!(home.is_none());
1206 }
1207
1208 #[test]
1213 fn test_python_interpreters_returns_existing_paths() {
1214 let interpreters = python_interpreters();
1215 let valid: Vec<&PathBuf> = interpreters.iter().filter(|p| p.exists()).collect();
1219 assert!(
1221 !valid.is_empty(),
1222 "Should find at least one existing Python interpreter"
1223 );
1224 }
1225
1226 #[test]
1227 fn test_python_interpreters_deduplication() {
1228 let interpreters = python_interpreters();
1229 let unique: HashSet<PathBuf> = interpreters.clone().into_iter().collect();
1230 assert_eq!(
1231 interpreters.len(),
1232 unique.len(),
1233 "Python interpreters should be deduplicated"
1234 );
1235 }
1236
1237 #[test]
1238 fn test_python_interpreters_env_var_override() {
1239 let _env = crate::test_support::lock_env();
1240 let saved_py = std::env::var("MLSTACK_PYTHON_BIN").ok();
1243 let saved_home = std::env::var("MLSTACK_USER_HOME").ok();
1244
1245 let dir = tempfile::tempdir().unwrap();
1247 let fake_python = dir.path().join("bin/python3");
1248 fs::create_dir_all(dir.path().join("bin")).unwrap();
1249 fs::write(&fake_python, "#!/bin/sh").unwrap();
1250
1251 std::env::set_var(
1252 "MLSTACK_USER_HOME",
1253 dir.path().to_string_lossy().to_string(),
1254 );
1255 std::env::set_var(
1256 "MLSTACK_PYTHON_BIN",
1257 fake_python.to_string_lossy().to_string(),
1258 );
1259 let interpreters = python_interpreters();
1260 assert!(
1261 interpreters[0] == fake_python,
1262 "First interpreter should be the env var override: {:?} vs {:?}",
1263 interpreters[0],
1264 fake_python
1265 );
1266
1267 match saved_py {
1269 Some(v) => std::env::set_var("MLSTACK_PYTHON_BIN", v),
1270 None => std::env::remove_var("MLSTACK_PYTHON_BIN"),
1271 }
1272 match saved_home {
1273 Some(v) => std::env::set_var("MLSTACK_USER_HOME", v),
1274 None => std::env::remove_var("MLSTACK_USER_HOME"),
1275 }
1276 }
1277
1278 #[test]
1279 fn test_python_interpreters_for_home_with_component_venvs() {
1280 let dir = tempfile::tempdir().unwrap();
1281 let home = dir.path();
1282
1283 for venv_name in &["rocm_venv", "pytorch"] {
1285 let bin = home.join(venv_name).join("bin");
1286 fs::create_dir_all(&bin).unwrap();
1287 fs::write(bin.join("python3"), "#!/bin/sh").unwrap();
1288 }
1289
1290 let interpreters = python_interpreters_for_home(home);
1291 assert!(
1292 interpreters.len() >= 2,
1293 "Should find at least 2 component venv interpreters, found {}",
1294 interpreters.len()
1295 );
1296
1297 let rocm_venv_pos = interpreters
1299 .iter()
1300 .position(|p| p.to_string_lossy().contains("rocm_venv"));
1301 let system_pos = interpreters
1302 .iter()
1303 .position(|p| p.to_string_lossy().starts_with("/usr/bin"));
1304 if let (Some(rv), Some(sp)) = (rocm_venv_pos, system_pos) {
1305 assert!(rv < sp, "Component venvs should come before system Pythons");
1306 }
1307 }
1308
1309 #[test]
1310 fn test_extract_python_version() {
1311 assert_eq!(
1312 extract_python_version(Path::new("/usr/bin/python3.13")),
1313 (3, 13)
1314 );
1315 assert_eq!(
1316 extract_python_version(Path::new("/usr/bin/python3.12")),
1317 (3, 12)
1318 );
1319 assert_eq!(
1320 extract_python_version(Path::new("/usr/bin/python3")),
1321 (3, 0)
1322 );
1323 assert_eq!(extract_python_version(Path::new("/usr/bin/python")), (0, 0));
1324 }
1325
1326 #[test]
1331 fn test_normalize_mlstack_env_creates_new_file() {
1332 let dir = tempfile::tempdir().unwrap();
1333 let python_bin = "/usr/bin/python3";
1334 let rocm_path = Path::new("/opt/rocm");
1335
1336 let result = normalize_mlstack_env(dir.path(), python_bin, rocm_path).unwrap();
1337 assert_eq!(result, EnvUpdateResult::Created);
1338
1339 let env_path = dir.path().join(".mlstack_env");
1340 assert!(env_path.exists());
1341
1342 let contents = fs::read_to_string(&env_path).unwrap();
1343 assert!(contents.contains("export MLSTACK_PYTHON_BIN=/usr/bin/python3"));
1344 assert!(contents.contains("export ROCM_PATH=/opt/rocm"));
1345 assert!(contents.contains("export PYTHONPATH=/opt/rocm/lib:$PYTHONPATH"));
1346 }
1347
1348 #[test]
1349 fn test_normalize_mlstack_env_idempotent() {
1350 let dir = tempfile::tempdir().unwrap();
1351 let python_bin = "/usr/bin/python3";
1352 let rocm_path = Path::new("/opt/rocm");
1353
1354 normalize_mlstack_env(dir.path(), python_bin, rocm_path).unwrap();
1356 let first_contents = fs::read_to_string(dir.path().join(".mlstack_env")).unwrap();
1357
1358 let result = normalize_mlstack_env(dir.path(), python_bin, rocm_path).unwrap();
1360 assert_eq!(result, EnvUpdateResult::Unchanged);
1361
1362 let second_contents = fs::read_to_string(dir.path().join(".mlstack_env")).unwrap();
1363 assert_eq!(
1364 first_contents, second_contents,
1365 "Idempotent: second run should produce identical file"
1366 );
1367 }
1368
1369 #[test]
1370 fn test_normalize_mlstack_env_updates_stale_values() {
1371 let dir = tempfile::tempdir().unwrap();
1372 let env_path = dir.path().join(".mlstack_env");
1373
1374 fs::write(
1376 &env_path,
1377 "# ML Stack Environment File\n\
1378export MLSTACK_PYTHON_BIN=/usr/bin/python3.11\n\
1379export ROCM_PATH=/opt/rocm-6.0\n\
1380export PYTHONPATH=/opt/rocm-6.0/lib:$PYTHONPATH\n\
1381export PATH=\"/opt/rocm-6.0/bin:/opt/rocm-6.0/hip/bin:$PATH\"\n\
1382export LD_LIBRARY_PATH=\"/opt/rocm-6.0/lib:/opt/rocm-6.0/hip/lib:/opt/rocm-6.0/opencl/lib:$LD_LIBRARY_PATH\"\n",
1383 )
1384 .unwrap();
1385
1386 let result =
1387 normalize_mlstack_env(dir.path(), "/usr/bin/python3.13", Path::new("/opt/rocm"))
1388 .unwrap();
1389
1390 assert_eq!(result, EnvUpdateResult::Updated);
1391
1392 let contents = fs::read_to_string(&env_path).unwrap();
1393 assert!(
1394 contents.contains("export MLSTACK_PYTHON_BIN=/usr/bin/python3.13"),
1395 "Should update python bin"
1396 );
1397 assert!(
1398 contents.contains("export ROCM_PATH=/opt/rocm\n")
1399 || contents.contains("export ROCM_PATH=/opt/rocm\r"),
1400 "Should update ROCM_PATH"
1401 );
1402 }
1403
1404 #[test]
1405 fn test_normalize_mlstack_env_adds_missing_keys() {
1406 let dir = tempfile::tempdir().unwrap();
1407 let env_path = dir.path().join(".mlstack_env");
1408
1409 fs::write(&env_path, "# Minimal env\nexport PATH=/usr/bin:$PATH\n").unwrap();
1411
1412 let result =
1413 normalize_mlstack_env(dir.path(), "/usr/bin/python3", Path::new("/opt/rocm")).unwrap();
1414
1415 assert_eq!(result, EnvUpdateResult::Updated);
1416
1417 let contents = fs::read_to_string(&env_path).unwrap();
1418 assert!(contents.contains("export MLSTACK_PYTHON_BIN=/usr/bin/python3"));
1419 assert!(contents.contains("export ROCM_PATH=/opt/rocm"));
1420 }
1421
1422 #[test]
1427 fn test_extract_export_key() {
1428 assert_eq!(
1429 extract_export_key("export MLSTACK_PYTHON_BIN=/usr/bin/python3"),
1430 Some("MLSTACK_PYTHON_BIN".to_string())
1431 );
1432 assert_eq!(
1433 extract_export_key("export ROCM_PATH=/opt/rocm"),
1434 Some("ROCM_PATH".to_string())
1435 );
1436 assert_eq!(extract_export_key("# comment"), None);
1437 assert_eq!(extract_export_key(""), None);
1438 assert_eq!(extract_export_key("not an export"), None);
1439 }
1440
1441 #[test]
1442 fn test_dedup_path_var() {
1443 assert_eq!(
1444 dedup_path_var("/usr/bin:/usr/local/bin:/usr/bin"),
1445 "/usr/bin:/usr/local/bin"
1446 );
1447 assert_eq!(dedup_path_var("/usr/bin"), "/usr/bin");
1448 assert_eq!(dedup_path_var(""), "");
1449 assert_eq!(
1450 dedup_path_var("/usr/bin::/usr/local/bin"),
1451 "/usr/bin:/usr/local/bin"
1452 );
1453 }
1454
1455 #[test]
1456 fn test_generate_env_file() {
1457 let content = generate_env_file(
1458 "/usr/bin/python3",
1459 "/opt/rocm",
1460 "/opt/rocm/lib",
1461 "/home/test",
1462 );
1463 assert!(content.contains("export MLSTACK_PYTHON_BIN=/usr/bin/python3"));
1464 assert!(content.contains("export UV_PYTHON=/usr/bin/python3"));
1465 assert!(content.contains("export ROCM_HOME=/opt/rocm"));
1466 assert!(content.contains("export ROCM_PATH=/opt/rocm"));
1467 assert!(content.contains("export HIP_PATH=/opt/rocm"));
1468 assert!(content.contains("export PYTHONPATH=/opt/rocm/lib:$PYTHONPATH"));
1469 assert!(content.contains("/opt/rocm/bin"));
1470 assert!(content.contains(".mlstack/lib-compat"));
1471 assert!(
1474 content.contains("/home/test/.mlstack/lib-compat"),
1475 "LD_LIBRARY_PATH must use the passed user_home: {content}"
1476 );
1477 assert!(content.contains("export ORT_MIGRAPHX_FP16_ENABLE=0"));
1478 }
1479
1480 #[test]
1481 fn test_normalize_env_contents_no_change() {
1482 let contents =
1483 "# ML Stack Environment File\nexport MLSTACK_PYTHON_BIN=/usr/bin/python3\nexport ROCM_PATH=/opt/rocm\nexport ORT_MIGRAPHX_FP16_ENABLE=0\n";
1484 let (result, changed) = normalize_env_contents(
1485 contents,
1486 "/usr/bin/python3",
1487 "/opt/rocm",
1488 "/opt/rocm/lib",
1489 "/home/test",
1490 );
1491 assert!(
1492 !changed,
1493 "Should not report changes when content is already normalized"
1494 );
1495 assert!(result.contains("export MLSTACK_PYTHON_BIN=/usr/bin/python3"));
1496 assert!(result.contains("export ROCM_PATH=/opt/rocm"));
1497 assert!(result.contains("export ORT_MIGRAPHX_FP16_ENABLE=0"));
1498 }
1499
1500 #[test]
1501 fn test_normalize_env_contents_with_change() {
1502 let contents =
1503 "# ML Stack Environment File\nexport MLSTACK_PYTHON_BIN=/usr/bin/python3.11\n"
1504 .to_string();
1505 let (result, changed) = normalize_env_contents(
1506 &contents,
1507 "/usr/bin/python3.13",
1508 "/opt/rocm",
1509 "/opt/rocm/lib",
1510 "/home/test",
1511 );
1512 assert!(changed);
1513 assert!(result.contains("export MLSTACK_PYTHON_BIN=/usr/bin/python3.13"));
1514 assert!(result.contains("export ROCM_PATH=/opt/rocm"));
1516 }
1517
1518 #[test]
1523 fn test_parse_python_version_cap_3_12() {
1524 let (major, minor) = parse_python_version_cap("Python 3.12.12").unwrap();
1525 assert_eq!(major, 3);
1526 assert_eq!(minor, 12);
1527 }
1528
1529 #[test]
1530 fn test_parse_python_version_cap_3_14_rejected() {
1531 let parsed = parse_python_version_cap("Python 3.14.5");
1532 assert!(
1533 parsed.is_some(),
1534 "3.14.5 parses to (3, 14) but cap check rejects later"
1535 );
1536 let (major, minor) = parsed.unwrap();
1537 assert_eq!(major, 3);
1538 assert_eq!(minor, 14);
1539 }
1540
1541 #[test]
1542 fn test_parse_python_version_cap_invalid() {
1543 assert!(parse_python_version_cap("Not a version string").is_none());
1544 }
1545
1546 #[test]
1547 fn test_check_python_version_cap_rejects_3_14() {
1548 let (major, minor) = parse_python_version_cap("Python 3.14.0").unwrap();
1552 assert!(
1553 major > 3 || (major == 3 && minor > 13),
1554 "3.14.0 exceeds cap (≤3.13)"
1555 );
1556 }
1557
1558 #[test]
1559 fn test_check_python_version_cap_accepts_3_12() {
1560 let (major, minor) = parse_python_version_cap("Python 3.12.12").unwrap();
1561 assert!(major <= 3 && minor <= 13, "3.12.12 is within cap (≤3.13)");
1562 }
1563
1564 #[test]
1565 fn test_write_activate_snippets_creates_files() {
1566 let dir = tempfile::tempdir().unwrap();
1567 let _env = crate::test_support::lock_env();
1568 let saved_home = std::env::var("HOME").ok();
1569
1570 std::env::set_var("HOME", dir.path().to_string_lossy().to_string());
1572
1573 let (sh_path, fish_path) = write_activate_snippets().unwrap();
1574
1575 assert!(sh_path.exists());
1576 assert!(fish_path.exists());
1577
1578 let sh_content = std::fs::read_to_string(&sh_path).unwrap();
1579 assert!(sh_content.contains("export PATH=\"$HOME/.mlstack/global/bin:$PATH\""));
1580 assert!(sh_content.contains("# mlstack-global-python"));
1581
1582 let fish_content = std::fs::read_to_string(&fish_path).unwrap();
1583 assert!(fish_content.contains("fish_add_path $HOME/.mlstack/global/bin"));
1584 assert!(fish_content.contains("# mlstack-global-python"));
1585
1586 match saved_home {
1588 Some(v) => std::env::set_var("HOME", v),
1589 None => std::env::remove_var("HOME"),
1590 }
1591 }
1592
1593 #[test]
1594 fn test_offer_global_prepend_idempotent_marker() {
1595 let dir = tempfile::tempdir().unwrap();
1596 let _env = crate::test_support::lock_env();
1597 let saved_home = std::env::var("HOME").ok();
1598 let saved_shell = std::env::var("SHELL").ok();
1599
1600 std::env::set_var("HOME", dir.path().to_string_lossy().to_string());
1602 std::env::set_var("SHELL", "/bin/bash");
1603
1604 let rc_path = dir.path().join(".bashrc");
1605 std::fs::write(&rc_path, "# Existing rc\n").unwrap();
1606
1607 let act1 = offer_global_prepend(true).unwrap();
1609 assert!(act1, "First call should append marker line");
1610
1611 let content1 = std::fs::read_to_string(&rc_path).unwrap();
1612 assert!(
1613 content1.contains("# mlstack-global-python"),
1614 "Marker line should be present"
1615 );
1616
1617 let act2 = offer_global_prepend(true).unwrap();
1619 assert!(!act2, "Second call should detect existing marker and skip");
1620
1621 let content2 = std::fs::read_to_string(&rc_path).unwrap();
1622 let marker_count = content2.matches("# mlstack-global-python").count();
1623 assert_eq!(
1624 marker_count, 1,
1625 "Marker should appear exactly once (idempotent)"
1626 );
1627
1628 match saved_home {
1630 Some(v) => std::env::set_var("HOME", v),
1631 None => std::env::remove_var("HOME"),
1632 }
1633 match saved_shell {
1634 Some(v) => std::env::set_var("SHELL", v),
1635 None => std::env::remove_var("SHELL"),
1636 }
1637 }
1638
1639 #[test]
1640 fn test_offer_global_prepend_detects_shell() {
1641 let dir = tempfile::tempdir().unwrap();
1642 let _env = crate::test_support::lock_env();
1643 let saved_home = std::env::var("HOME").ok();
1644 let saved_shell = std::env::var("SHELL").ok();
1645
1646 std::env::set_var("HOME", dir.path().to_string_lossy().to_string());
1648 std::env::set_var("SHELL", "/usr/bin/fish");
1649
1650 let fish_config = dir.path().join(".config/fish/config.fish");
1651 std::fs::create_dir_all(fish_config.parent().unwrap()).unwrap();
1652 std::fs::write(&fish_config, "# Fish rc\n").unwrap();
1653
1654 let act_fish = offer_global_prepend(true).unwrap();
1655 assert!(act_fish);
1656
1657 let fish_content = std::fs::read_to_string(&fish_config).unwrap();
1658 assert!(fish_content.contains("source \"$HOME/.mlstack/global/activate-global.fish\""));
1659 assert!(fish_content.contains("# mlstack-global-python"));
1660
1661 match saved_home {
1663 Some(v) => std::env::set_var("HOME", v),
1664 None => std::env::remove_var("HOME"),
1665 }
1666 match saved_shell {
1667 Some(v) => std::env::set_var("SHELL", v),
1668 None => std::env::remove_var("SHELL"),
1669 }
1670 }
1671}