use std::collections::HashSet;
use std::env;
use std::fs;
use std::io::IsTerminal;
use std::path::{Path, PathBuf};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum EnvUpdateResult {
Created,
Updated,
Unchanged,
}
pub fn resolve_user_home() -> PathBuf {
if let Ok(home) = env::var("MLSTACK_USER_HOME") {
let home = home.trim().to_string();
if !home.is_empty() && Path::new(&home).is_dir() {
return PathBuf::from(home);
}
}
if let Ok(home) = env::var("HOME") {
let home = home.trim().to_string();
if !home.is_empty() {
let running_sudo = env::var("SUDO_USER").is_ok();
if !(running_sudo && home == "/root") {
return PathBuf::from(home);
}
}
}
for key in ["SUDO_USER", "USER", "LOGNAME"] {
if let Ok(user_name) = env::var(key) {
let user_name = user_name.trim().to_string();
if !user_name.is_empty() && user_name != "root" {
if let Some(home) = lookup_home_from_passwd(&user_name) {
return home;
}
}
}
}
PathBuf::from("/tmp")
}
pub fn lookup_home_from_passwd(username: &str) -> Option<PathBuf> {
let content = fs::read_to_string("/etc/passwd").ok()?;
for line in content.lines() {
let fields: Vec<&str> = line.split(':').collect();
if fields.len() >= 6 && fields[0] == username {
let home = fields[5].trim();
if !home.is_empty() {
return Some(PathBuf::from(home));
}
}
}
None
}
pub fn mlstack_root() -> PathBuf {
resolve_user_home().join(".mlstack")
}
pub fn mlstack_global_dir() -> PathBuf {
mlstack_root().join("global")
}
pub fn mlstack_global_python() -> PathBuf {
mlstack_global_dir().join("bin").join("python")
}
pub fn mlstack_envs_dir() -> PathBuf {
mlstack_root().join("envs")
}
pub fn mlstack_env_dir(name: &str) -> PathBuf {
mlstack_envs_dir().join(name)
}
pub fn mlstack_cache_dir() -> PathBuf {
mlstack_root().join("cache")
}
pub fn mlstack_logs_dir() -> PathBuf {
mlstack_root().join("logs")
}
pub fn mlstack_triton_dir() -> PathBuf {
mlstack_root().join("triton")
}
pub fn ensure_mlstack_dirs() -> std::io::Result<()> {
for dir in [
mlstack_root(),
mlstack_global_dir(),
mlstack_envs_dir(),
mlstack_cache_dir(),
mlstack_logs_dir(),
mlstack_triton_dir(),
] {
std::fs::create_dir_all(&dir)?;
}
Ok(())
}
pub fn ensure_global_venv(bootstrap_python: &str) -> anyhow::Result<PathBuf> {
let global_python = mlstack_global_python();
if !global_python.exists() {
ensure_mlstack_dirs()?;
let global_dir = mlstack_global_dir();
create_venv(&global_dir, bootstrap_python)?;
if !global_python.exists() {
anyhow::bail!(
"global venv creation reported success but {} is missing",
global_python.display()
);
}
}
ensure_pip_in_venv(&global_python)?;
Ok(global_python)
}
pub fn ensure_named_venv(name: &str, bootstrap_python: &str) -> anyhow::Result<PathBuf> {
let name = name.trim();
if name.is_empty() {
anyhow::bail!("named env name must not be empty");
}
if name.contains(std::path::MAIN_SEPARATOR)
|| name.contains('/')
|| name.contains('\\')
|| name.contains("..")
{
anyhow::bail!("invalid env name '{name}'");
}
ensure_mlstack_dirs()?;
let env_dir = mlstack_env_dir(name);
let env_python = env_dir.join("bin").join("python");
if !env_python.exists() {
create_venv(&env_dir, bootstrap_python)?;
if !env_python.exists() {
anyhow::bail!(
"named venv creation reported success but {} is missing",
env_python.display()
);
}
}
ensure_pip_in_venv(&env_python)?;
Ok(env_python)
}
fn create_venv(dir: &Path, bootstrap_python: &str) -> anyhow::Result<()> {
if command_on_path("uv") {
let uv_status = std::process::Command::new("uv")
.arg("venv")
.arg("--seed")
.arg("--python")
.arg(bootstrap_python)
.arg(dir)
.status();
if let Ok(status) = uv_status {
if status.success() {
return Ok(());
}
}
}
let status = std::process::Command::new(bootstrap_python)
.arg("-m")
.arg("venv")
.arg(dir)
.status()?;
if !status.success() {
anyhow::bail!(
"failed to create venv at {} via uv/venv (bootstrap={})",
dir.display(),
bootstrap_python
);
}
Ok(())
}
fn ensure_pip_in_venv(venv_python: &Path) -> anyhow::Result<()> {
let has_pip = std::process::Command::new(venv_python)
.args(["-m", "pip", "--version"])
.output()
.map(|o| o.status.success())
.unwrap_or(false);
if has_pip {
return Ok(());
}
if command_on_path("uv") {
let s = std::process::Command::new("uv")
.args(["pip", "install", "--python"])
.arg(venv_python)
.arg("pip")
.status();
if let Ok(status) = s {
if status.success() {
return Ok(());
}
}
}
let status = std::process::Command::new(venv_python)
.args(["-m", "ensurepip", "--upgrade"])
.status()?;
if !status.success() {
anyhow::bail!(
"could not bootstrap pip into {} (tried `uv pip install pip` and `ensurepip`)",
venv_python.display()
);
}
Ok(())
}
pub fn python_interpreters() -> Vec<PathBuf> {
let home = resolve_user_home();
python_interpreters_for_home(&home)
}
pub fn resolve_canonical_python_bin() -> String {
if let Ok(name) = env::var("MLSTACK_ENV_NAME") {
let name = name.trim();
if !name.is_empty() {
let p = mlstack_env_dir(name).join("bin").join("python");
if p.exists() {
return p.to_string_lossy().to_string();
}
}
}
for key in ["MLSTACK_PYTHON_BIN", "UV_PYTHON"] {
if let Ok(val) = env::var(key) {
let val = val.trim();
if !val.is_empty() && Path::new(val).exists() {
return val.to_string();
}
}
}
let global = mlstack_global_python();
if global.exists() {
return global.to_string_lossy().to_string();
}
let interpreters = python_interpreters();
if let Some(first) = interpreters.first() {
return first.to_string_lossy().to_string();
}
"python3".to_string()
}
pub fn python_interpreters_for_home(home: &Path) -> Vec<PathBuf> {
let mut paths = Vec::new();
let mut seen = HashSet::new();
macro_rules! push_if_exists {
($path:expr) => {
let p = $path;
if p.exists() && seen.insert(p.clone()) {
paths.push(p);
}
};
}
for key in ["MLSTACK_PYTHON_BIN", "UV_PYTHON"] {
if let Ok(val) = env::var(key) {
let val = val.trim().to_string();
if !val.is_empty() {
push_if_exists!(PathBuf::from(&val));
}
}
}
push_if_exists!(home
.join(".mlstack")
.join("global")
.join("bin")
.join("python"));
push_if_exists!(home
.join(".mlstack")
.join("global")
.join("bin")
.join("python3"));
if let Ok(venv) = env::var("VIRTUAL_ENV") {
let venv = venv.trim().to_string();
if !venv.is_empty() {
push_if_exists!(PathBuf::from(&venv).join("bin/python"));
push_if_exists!(PathBuf::from(&venv).join("bin/python3"));
}
}
if let Ok(conda_prefix) = env::var("CONDA_PREFIX") {
let conda_prefix = conda_prefix.trim().to_string();
if !conda_prefix.is_empty() {
push_if_exists!(PathBuf::from(&conda_prefix).join("bin/python"));
push_if_exists!(PathBuf::from(&conda_prefix).join("bin/python3"));
}
}
push_if_exists!(home.join(".local/bin/python3"));
push_if_exists!(home.join(".local/bin/python"));
let uv_python_dir = home.join(".local/share/uv/python");
if let Ok(entries) = fs::read_dir(&uv_python_dir) {
let mut uv_pythons: Vec<PathBuf> = entries
.filter_map(|e| e.ok())
.filter_map(|e| {
let bin = e.path().join("bin/python3");
if bin.exists() {
Some(bin)
} else {
None
}
})
.collect();
uv_pythons.sort_by(|a, b| {
let va = extract_python_version(a);
let vb = extract_python_version(b);
let priority = |v: (u32, u32)| -> (i32, u32) {
match v {
(3, 12) => (0, v.1), (3, 13) => (1, v.1),
(3, 11) => (2, v.1),
(3, 10) => (3, v.1),
_ => (4, v.1), }
};
priority(va).cmp(&priority(vb))
});
for p in uv_pythons {
push_if_exists!(p);
}
}
let component_venvs = ["rocm_venv", "pytorch", "ml_stack"];
for venv_name in &component_venvs {
push_if_exists!(home.join(venv_name).join("bin/python"));
push_if_exists!(home.join(venv_name).join("bin/python3"));
}
let system_pythons = [
"/usr/local/bin/python3",
"/usr/bin/python3",
"/usr/bin/python",
];
for sys_path in &system_pythons {
push_if_exists!(PathBuf::from(sys_path));
}
if let Ok(entries) = fs::read_dir("/usr/bin") {
let mut versioned: Vec<PathBuf> = entries
.filter_map(|e| e.ok())
.filter_map(|e| {
let name = e.file_name().to_string_lossy().to_string();
if name.starts_with("python3.") && e.path().exists() {
Some(e.path())
} else {
None
}
})
.collect();
versioned.sort_by(|a, b| {
let va = extract_python_version(a);
let vb = extract_python_version(b);
vb.cmp(&va)
});
for p in versioned {
push_if_exists!(p);
}
}
let local_python3 = home.join(".local/bin/python3");
if local_python3.exists() && seen.insert(local_python3.clone()) {
paths.push(local_python3);
}
paths
}
fn extract_python_version(path: &Path) -> (u32, u32) {
let name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
if let Some(rest) = name.strip_prefix("python") {
let parts: Vec<&str> = rest.split('.').collect();
if parts.len() >= 2 {
let major = parts[0].parse::<u32>().unwrap_or(0);
let minor = parts[1].parse::<u32>().unwrap_or(0);
return (major, minor);
}
if let Ok(major) = rest.parse::<u32>() {
return (major, 0);
}
}
(0, 0)
}
pub fn resolve_managed_python(backend: &crate::config::PythonBackend) -> anyhow::Result<PathBuf> {
match backend {
crate::config::PythonBackend::Uv => {
if !command_on_path("uv") {
anyhow::bail!("uv backend selected but uv not on PATH");
}
let output = std::process::Command::new("uv")
.args(["python", "find", "3.12"])
.output();
if let Ok(out) = output {
if out.status.success() {
let stdout = String::from_utf8_lossy(&out.stdout);
if let Some(line) = stdout.lines().next() {
let path = line.trim();
if !path.is_empty() {
return Ok(PathBuf::from(path));
}
}
}
}
let output = std::process::Command::new("uv")
.args(["python", "find", ">=3.10,<3.14"])
.output()?;
if output.status.success() {
let stdout = String::from_utf8_lossy(&output.stdout);
if let Some(line) = stdout.lines().next() {
let path = line.trim();
if !path.is_empty() {
return Ok(PathBuf::from(path));
}
}
}
anyhow::bail!("uv python find failed to locate a Python ≤3.13");
}
crate::config::PythonBackend::Python => {
let candidates = ["python3.12", "python3.11", "python3.10", "python3.13"];
for name in &candidates {
if let Some(path) = check_python_version_cap(name) {
return Ok(path);
}
}
anyhow::bail!(
"No suitable Python ≤3.13 found on PATH (tried: {})",
candidates.join(", ")
);
}
}
}
fn check_python_version_cap(python_name: &str) -> Option<PathBuf> {
let output = std::process::Command::new(python_name)
.arg("--version")
.output()
.ok()?;
if !output.status.success() {
return None;
}
let version_str = String::from_utf8_lossy(&output.stdout);
let (major, minor) = parse_python_version_cap(&version_str)?;
if major > 3 || (major == 3 && minor > 13) {
return None;
}
for dir in env::split_paths(&env::var("PATH").unwrap_or_default()) {
let candidate = dir.join(python_name);
if candidate.is_file() {
return Some(candidate);
}
}
None
}
fn parse_python_version_cap(version_str: &str) -> Option<(u32, u32)> {
let rest = version_str.strip_prefix("Python ")?;
let parts: Vec<&str> = rest.split('.').collect();
if parts.len() >= 2 {
let major = parts[0].parse::<u32>().ok()?;
let minor = parts[1].parse::<u32>().ok()?;
Some((major, minor))
} else {
None
}
}
pub fn create_global_env(backend: &crate::config::PythonBackend) -> anyhow::Result<PathBuf> {
let bootstrap_python = resolve_managed_python(backend)?;
let global_dir = mlstack_global_dir();
if global_dir.exists() {
let global_python = mlstack_global_python();
if global_python.exists() {
ensure_pip_in_venv(&global_python)?;
return Ok(global_python);
}
std::fs::remove_dir_all(&global_dir)?;
}
std::fs::create_dir_all(&global_dir)?;
match backend {
crate::config::PythonBackend::Uv => {
let status = std::process::Command::new("uv")
.arg("venv")
.arg("--seed")
.arg(&global_dir)
.arg("--python")
.arg(&bootstrap_python)
.status()?;
if !status.success() {
anyhow::bail!("uv venv failed for global env");
}
}
crate::config::PythonBackend::Python => {
let status = std::process::Command::new(&bootstrap_python)
.arg("-m")
.arg("venv")
.arg(&global_dir)
.status()?;
if !status.success() {
anyhow::bail!("python -m venv failed for global env");
}
}
}
let global_python = mlstack_global_python();
if !global_python.exists() {
anyhow::bail!(
"global env creation reported success but {} is missing",
global_python.display()
);
}
ensure_pip_in_venv(&global_python)?;
Ok(global_python)
}
pub fn write_activate_snippets() -> anyhow::Result<(PathBuf, PathBuf)> {
let global_dir = mlstack_global_dir();
std::fs::create_dir_all(&global_dir)?;
let sh_path = global_dir.join("activate-global.sh");
let fish_path = global_dir.join("activate-global.fish");
let sh_content = r#"# mlstack-global-python activation — generated by Rusty-Stack
# Do NOT edit manually — uninstall strips this marker line
export PATH="$HOME/.mlstack/global/bin:$PATH"
"#;
std::fs::write(&sh_path, sh_content)?;
let fish_content = r#"# mlstack-global-python activation — generated by Rusty-Stack
# Do NOT edit manually — uninstall strips this marker line
fish_add_path $HOME/.mlstack/global/bin
"#;
std::fs::write(&fish_path, fish_content)?;
Ok((sh_path, fish_path))
}
pub fn offer_global_prepend(yes_default: bool) -> anyhow::Result<bool> {
let shell = std::env::var("SHELL").unwrap_or_default();
let (rc_path, _snippet_path) = if shell.contains("fish") {
let rc = resolve_user_home().join(".config/fish/config.fish");
let snippet = mlstack_global_dir().join("activate-global.fish");
(rc, snippet)
} else if shell.contains("zsh") {
let rc = resolve_user_home().join(".zshrc");
let snippet = mlstack_global_dir().join("activate-global.sh");
(rc, snippet)
} else {
let rc = resolve_user_home().join(".bashrc");
let snippet = mlstack_global_dir().join("activate-global.sh");
(rc, snippet)
};
let marker = "# mlstack-global-python";
if rc_path.exists() {
let content = std::fs::read_to_string(&rc_path).unwrap_or_default();
if content.contains(marker) {
return Ok(false); }
}
let act = if yes_default {
true
} else {
let is_tty = std::io::stdin().is_terminal();
if !is_tty {
true } else {
eprint!(
"Make ~/.mlstack/global the default Python in your {} rc? [Y/n] ",
if shell.contains("fish") {
"fish"
} else if shell.contains("zsh") {
"zsh"
} else {
"bash"
}
);
let mut input = String::new();
std::io::stdin().read_line(&mut input).ok();
let trimmed = input.trim().to_lowercase();
trimmed.is_empty() || trimmed == "y" || trimmed == "yes"
}
};
if !act {
return Ok(false);
}
if let Some(parent) = rc_path.parent() {
std::fs::create_dir_all(parent)?;
}
if !rc_path.exists() {
std::fs::write(&rc_path, "")?;
}
let source_line = if shell.contains("fish") {
format!(
"source \"$HOME/.mlstack/global/activate-global.fish\" {}",
marker
)
} else {
format!(
"source \"$HOME/.mlstack/global/activate-global.sh\" {}",
marker
)
};
let mut content = std::fs::read_to_string(&rc_path).unwrap_or_default();
if !content.ends_with('\n') {
content.push('\n');
}
content.push_str(&source_line);
content.push('\n');
std::fs::write(&rc_path, content)?;
Ok(true)
}
pub fn offer_mlstack_env_source() -> anyhow::Result<()> {
const MARKER: &str = "# mlstack-rocm-env";
let source_line = format!(
"[ -f \"$HOME/.mlstack_env\" ] && source \"$HOME/.mlstack_env\" {}",
MARKER
);
for name in ["bashrc", "zshrc"] {
let rc_path = resolve_user_home().join(format!(".{}", name));
let mut content = std::fs::read_to_string(&rc_path).unwrap_or_default();
if content.contains(".mlstack_env") {
continue; }
if !content.is_empty() && !content.ends_with('\n') {
content.push('\n');
}
content.push_str("# mlstack ROCm env — auto-sourced on Global install\n");
content.push_str(&source_line);
content.push('\n');
if let Some(parent) = rc_path.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(&rc_path, content)?;
}
Ok(())
}
pub fn normalize_mlstack_env(
user_home: &Path,
python_bin: &str,
rocm_path: &Path,
) -> anyhow::Result<EnvUpdateResult> {
let env_path = user_home.join(".mlstack_env");
let rocm_str = rocm_path.to_string_lossy();
let rocm_lib_path = rocm_path.join("lib");
let rocm_lib = rocm_lib_path.to_string_lossy();
let user_home_str = user_home.to_string_lossy().to_string();
if env_path.exists() {
let contents = fs::read_to_string(&env_path)?;
let (normalized, changed) =
normalize_env_contents(&contents, python_bin, &rocm_str, &rocm_lib, &user_home_str);
if changed {
fs::write(&env_path, normalized)?;
return Ok(EnvUpdateResult::Updated);
}
return Ok(EnvUpdateResult::Unchanged);
}
if let Some(parent) = env_path.parent() {
fs::create_dir_all(parent)?;
}
let content = generate_env_file(python_bin, &rocm_str, &rocm_lib, &user_home_str);
fs::write(&env_path, content)?;
Ok(EnvUpdateResult::Created)
}
pub fn normalize_env_contents(
contents: &str,
python_bin: &str,
rocm_home: &str,
rocm_lib: &str,
user_home: &str,
) -> (String, bool) {
let mut changed = false;
let mut lines: Vec<String> = Vec::new();
let mut seen_keys: HashSet<String> = HashSet::new();
for line in contents.lines() {
let trimmed = line.trim();
if trimmed.is_empty() || trimmed.starts_with('#') {
lines.push(line.to_string());
continue;
}
if let Some(key) = extract_export_key(trimmed) {
seen_keys.insert(key.clone());
let (normalized_line, line_changed) =
normalize_env_line(line, &key, python_bin, rocm_home, rocm_lib, user_home);
changed |= line_changed;
lines.push(normalized_line);
} else {
lines.push(line.to_string());
}
}
let required = [
(
"MLSTACK_PYTHON_BIN",
format!("export MLSTACK_PYTHON_BIN={}", python_bin),
),
("ROCM_PATH", format!("export ROCM_PATH={}", rocm_home)),
(
"ORT_MIGRAPHX_FP16_ENABLE",
"export ORT_MIGRAPHX_FP16_ENABLE=0".to_string(),
),
];
for (key, line) in &required {
if !seen_keys.contains(*key) {
lines.push(line.clone());
changed = true;
}
}
(lines.join("\n"), changed)
}
pub fn generate_env_file(
python_bin: &str,
rocm_home: &str,
rocm_lib: &str,
user_home: &str,
) -> String {
format!(
"# ML Stack Environment File (generated by Rusty-Stack)\n\
export MLSTACK_PYTHON_BIN={python_bin}\n\
export UV_PYTHON={python_bin}\n\
export ROCM_HOME={rocm_home}\n\
export ROCM_PATH={rocm_home}\n\
export HIP_PATH={rocm_home}\n\
export PYTHONPATH={rocm_lib}:$PYTHONPATH\n\
export PATH=\"{rocm_home}/bin:{rocm_home}/hip/bin:$PATH\"\n\
export LD_LIBRARY_PATH=\"{rocm_home}/lib:{rocm_home}/hip/lib:{rocm_home}/opencl/lib:{user_home}/.mlstack/lib-compat:$LD_LIBRARY_PATH\"\n\
export ORT_MIGRAPHX_FP16_ENABLE=0\n"
)
}
fn extract_export_key(line: &str) -> Option<String> {
let line = line.trim();
let after_export = line.strip_prefix("export ")?;
let (key, _) = after_export.split_once('=')?;
Some(key.trim().to_string())
}
fn normalize_env_line(
line: &str,
key: &str,
python_bin: &str,
rocm_home: &str,
rocm_lib: &str,
user_home: &str,
) -> (String, bool) {
match key {
"MLSTACK_PYTHON_BIN" | "UV_PYTHON" => {
let desired = format!("export {}={}", key, python_bin);
let trimmed = line.trim();
if trimmed != desired {
(desired, true)
} else {
(line.to_string(), false)
}
}
"ROCM_PATH" | "ROCM_HOME" | "HIP_PATH" => {
let desired = format!("export {}={}", key, rocm_home);
let trimmed = line.trim();
if trimmed != desired {
(desired, true)
} else {
(line.to_string(), false)
}
}
"PYTHONPATH" => {
let desired = format!("export PYTHONPATH={}:$PYTHONPATH", rocm_lib);
let trimmed = line.trim();
if trimmed != desired {
(desired, true)
} else {
(line.to_string(), false)
}
}
"PATH" => {
let desired = format!(
"export PATH=\"{}/bin:{}/hip/bin:$PATH\"",
rocm_home, rocm_home
);
let trimmed = line.trim();
if trimmed != desired {
(desired, true)
} else {
(line.to_string(), false)
}
}
"LD_LIBRARY_PATH" => {
let desired = format!(
"export LD_LIBRARY_PATH=\"{}/lib:{}/hip/lib:{}/opencl/lib:{}/.mlstack/lib-compat:$LD_LIBRARY_PATH\"",
rocm_home, rocm_home, rocm_home, user_home
);
let trimmed = line.trim();
if trimmed != desired {
(desired, true)
} else {
(line.to_string(), false)
}
}
_ => (line.to_string(), false),
}
}
pub fn dedup_path_var(path_var: &str) -> String {
let mut seen = HashSet::new();
let mut result = Vec::new();
for entry in path_var.split(':') {
if !entry.is_empty() && seen.insert(entry.to_string()) {
result.push(entry.to_string());
}
}
result.join(":")
}
pub fn command_on_path(cmd: &str) -> bool {
let path = match env::var_os("PATH") {
Some(p) => p,
None => return false,
};
for dir in env::split_paths(&path) {
let candidate = dir.join(cmd);
if candidate.is_file() {
return true;
}
}
false
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_resolve_user_home_returns_valid_path() {
let home = resolve_user_home();
assert!(
!home.as_os_str().is_empty(),
"Home directory should not be empty"
);
}
#[test]
fn test_resolve_user_home_env_override() {
let _env = crate::test_support::lock_env();
let saved = std::env::var("MLSTACK_USER_HOME").ok();
let dir = tempfile::tempdir().unwrap();
let dir_path = dir.path().to_path_buf();
std::env::set_var("MLSTACK_USER_HOME", dir_path.to_string_lossy().to_string());
let result = resolve_user_home();
assert_eq!(result, dir_path);
match saved {
Some(v) => std::env::set_var("MLSTACK_USER_HOME", v),
None => std::env::remove_var("MLSTACK_USER_HOME"),
}
}
#[test]
fn test_resolve_user_home_from_passwd() {
if let Ok(user) = env::var("USER") {
if let Some(home) = lookup_home_from_passwd(&user) {
assert!(
home.is_dir(),
"Home directory from passwd should exist: {:?}",
home
);
assert!(
!home.as_os_str().is_empty(),
"Home directory should not be empty"
);
}
}
}
#[test]
fn test_lookup_home_from_passwd_root() {
let home = lookup_home_from_passwd("root");
assert!(home.is_some(), "root should be in /etc/passwd");
assert_eq!(home.unwrap(), PathBuf::from("/root"));
}
#[test]
fn test_lookup_home_from_passwd_nonexistent() {
let home = lookup_home_from_passwd("nonexistent_user_xyz_12345");
assert!(home.is_none());
}
#[test]
fn test_python_interpreters_returns_existing_paths() {
let interpreters = python_interpreters();
let valid: Vec<&PathBuf> = interpreters.iter().filter(|p| p.exists()).collect();
assert!(
!valid.is_empty(),
"Should find at least one existing Python interpreter"
);
}
#[test]
fn test_python_interpreters_deduplication() {
let interpreters = python_interpreters();
let unique: HashSet<PathBuf> = interpreters.clone().into_iter().collect();
assert_eq!(
interpreters.len(),
unique.len(),
"Python interpreters should be deduplicated"
);
}
#[test]
fn test_python_interpreters_env_var_override() {
let _env = crate::test_support::lock_env();
let saved_py = std::env::var("MLSTACK_PYTHON_BIN").ok();
let saved_home = std::env::var("MLSTACK_USER_HOME").ok();
let dir = tempfile::tempdir().unwrap();
let fake_python = dir.path().join("bin/python3");
fs::create_dir_all(dir.path().join("bin")).unwrap();
fs::write(&fake_python, "#!/bin/sh").unwrap();
std::env::set_var(
"MLSTACK_USER_HOME",
dir.path().to_string_lossy().to_string(),
);
std::env::set_var(
"MLSTACK_PYTHON_BIN",
fake_python.to_string_lossy().to_string(),
);
let interpreters = python_interpreters();
assert!(
interpreters[0] == fake_python,
"First interpreter should be the env var override: {:?} vs {:?}",
interpreters[0],
fake_python
);
match saved_py {
Some(v) => std::env::set_var("MLSTACK_PYTHON_BIN", v),
None => std::env::remove_var("MLSTACK_PYTHON_BIN"),
}
match saved_home {
Some(v) => std::env::set_var("MLSTACK_USER_HOME", v),
None => std::env::remove_var("MLSTACK_USER_HOME"),
}
}
#[test]
fn test_python_interpreters_for_home_with_component_venvs() {
let dir = tempfile::tempdir().unwrap();
let home = dir.path();
for venv_name in &["rocm_venv", "pytorch"] {
let bin = home.join(venv_name).join("bin");
fs::create_dir_all(&bin).unwrap();
fs::write(bin.join("python3"), "#!/bin/sh").unwrap();
}
let interpreters = python_interpreters_for_home(home);
assert!(
interpreters.len() >= 2,
"Should find at least 2 component venv interpreters, found {}",
interpreters.len()
);
let rocm_venv_pos = interpreters
.iter()
.position(|p| p.to_string_lossy().contains("rocm_venv"));
let system_pos = interpreters
.iter()
.position(|p| p.to_string_lossy().starts_with("/usr/bin"));
if let (Some(rv), Some(sp)) = (rocm_venv_pos, system_pos) {
assert!(rv < sp, "Component venvs should come before system Pythons");
}
}
#[test]
fn test_extract_python_version() {
assert_eq!(
extract_python_version(Path::new("/usr/bin/python3.13")),
(3, 13)
);
assert_eq!(
extract_python_version(Path::new("/usr/bin/python3.12")),
(3, 12)
);
assert_eq!(
extract_python_version(Path::new("/usr/bin/python3")),
(3, 0)
);
assert_eq!(extract_python_version(Path::new("/usr/bin/python")), (0, 0));
}
#[test]
fn test_normalize_mlstack_env_creates_new_file() {
let dir = tempfile::tempdir().unwrap();
let python_bin = "/usr/bin/python3";
let rocm_path = Path::new("/opt/rocm");
let result = normalize_mlstack_env(dir.path(), python_bin, rocm_path).unwrap();
assert_eq!(result, EnvUpdateResult::Created);
let env_path = dir.path().join(".mlstack_env");
assert!(env_path.exists());
let contents = fs::read_to_string(&env_path).unwrap();
assert!(contents.contains("export MLSTACK_PYTHON_BIN=/usr/bin/python3"));
assert!(contents.contains("export ROCM_PATH=/opt/rocm"));
assert!(contents.contains("export PYTHONPATH=/opt/rocm/lib:$PYTHONPATH"));
}
#[test]
fn test_normalize_mlstack_env_idempotent() {
let dir = tempfile::tempdir().unwrap();
let python_bin = "/usr/bin/python3";
let rocm_path = Path::new("/opt/rocm");
normalize_mlstack_env(dir.path(), python_bin, rocm_path).unwrap();
let first_contents = fs::read_to_string(dir.path().join(".mlstack_env")).unwrap();
let result = normalize_mlstack_env(dir.path(), python_bin, rocm_path).unwrap();
assert_eq!(result, EnvUpdateResult::Unchanged);
let second_contents = fs::read_to_string(dir.path().join(".mlstack_env")).unwrap();
assert_eq!(
first_contents, second_contents,
"Idempotent: second run should produce identical file"
);
}
#[test]
fn test_normalize_mlstack_env_updates_stale_values() {
let dir = tempfile::tempdir().unwrap();
let env_path = dir.path().join(".mlstack_env");
fs::write(
&env_path,
"# ML Stack Environment File\n\
export MLSTACK_PYTHON_BIN=/usr/bin/python3.11\n\
export ROCM_PATH=/opt/rocm-6.0\n\
export PYTHONPATH=/opt/rocm-6.0/lib:$PYTHONPATH\n\
export PATH=\"/opt/rocm-6.0/bin:/opt/rocm-6.0/hip/bin:$PATH\"\n\
export LD_LIBRARY_PATH=\"/opt/rocm-6.0/lib:/opt/rocm-6.0/hip/lib:/opt/rocm-6.0/opencl/lib:$LD_LIBRARY_PATH\"\n",
)
.unwrap();
let result =
normalize_mlstack_env(dir.path(), "/usr/bin/python3.13", Path::new("/opt/rocm"))
.unwrap();
assert_eq!(result, EnvUpdateResult::Updated);
let contents = fs::read_to_string(&env_path).unwrap();
assert!(
contents.contains("export MLSTACK_PYTHON_BIN=/usr/bin/python3.13"),
"Should update python bin"
);
assert!(
contents.contains("export ROCM_PATH=/opt/rocm\n")
|| contents.contains("export ROCM_PATH=/opt/rocm\r"),
"Should update ROCM_PATH"
);
}
#[test]
fn test_normalize_mlstack_env_adds_missing_keys() {
let dir = tempfile::tempdir().unwrap();
let env_path = dir.path().join(".mlstack_env");
fs::write(&env_path, "# Minimal env\nexport PATH=/usr/bin:$PATH\n").unwrap();
let result =
normalize_mlstack_env(dir.path(), "/usr/bin/python3", Path::new("/opt/rocm")).unwrap();
assert_eq!(result, EnvUpdateResult::Updated);
let contents = fs::read_to_string(&env_path).unwrap();
assert!(contents.contains("export MLSTACK_PYTHON_BIN=/usr/bin/python3"));
assert!(contents.contains("export ROCM_PATH=/opt/rocm"));
}
#[test]
fn test_extract_export_key() {
assert_eq!(
extract_export_key("export MLSTACK_PYTHON_BIN=/usr/bin/python3"),
Some("MLSTACK_PYTHON_BIN".to_string())
);
assert_eq!(
extract_export_key("export ROCM_PATH=/opt/rocm"),
Some("ROCM_PATH".to_string())
);
assert_eq!(extract_export_key("# comment"), None);
assert_eq!(extract_export_key(""), None);
assert_eq!(extract_export_key("not an export"), None);
}
#[test]
fn test_dedup_path_var() {
assert_eq!(
dedup_path_var("/usr/bin:/usr/local/bin:/usr/bin"),
"/usr/bin:/usr/local/bin"
);
assert_eq!(dedup_path_var("/usr/bin"), "/usr/bin");
assert_eq!(dedup_path_var(""), "");
assert_eq!(
dedup_path_var("/usr/bin::/usr/local/bin"),
"/usr/bin:/usr/local/bin"
);
}
#[test]
fn test_generate_env_file() {
let content = generate_env_file(
"/usr/bin/python3",
"/opt/rocm",
"/opt/rocm/lib",
"/home/test",
);
assert!(content.contains("export MLSTACK_PYTHON_BIN=/usr/bin/python3"));
assert!(content.contains("export UV_PYTHON=/usr/bin/python3"));
assert!(content.contains("export ROCM_HOME=/opt/rocm"));
assert!(content.contains("export ROCM_PATH=/opt/rocm"));
assert!(content.contains("export HIP_PATH=/opt/rocm"));
assert!(content.contains("export PYTHONPATH=/opt/rocm/lib:$PYTHONPATH"));
assert!(content.contains("/opt/rocm/bin"));
assert!(content.contains(".mlstack/lib-compat"));
assert!(
content.contains("/home/test/.mlstack/lib-compat"),
"LD_LIBRARY_PATH must use the passed user_home: {content}"
);
assert!(content.contains("export ORT_MIGRAPHX_FP16_ENABLE=0"));
}
#[test]
fn test_normalize_env_contents_no_change() {
let contents =
"# ML Stack Environment File\nexport MLSTACK_PYTHON_BIN=/usr/bin/python3\nexport ROCM_PATH=/opt/rocm\nexport ORT_MIGRAPHX_FP16_ENABLE=0\n";
let (result, changed) = normalize_env_contents(
contents,
"/usr/bin/python3",
"/opt/rocm",
"/opt/rocm/lib",
"/home/test",
);
assert!(
!changed,
"Should not report changes when content is already normalized"
);
assert!(result.contains("export MLSTACK_PYTHON_BIN=/usr/bin/python3"));
assert!(result.contains("export ROCM_PATH=/opt/rocm"));
assert!(result.contains("export ORT_MIGRAPHX_FP16_ENABLE=0"));
}
#[test]
fn test_normalize_env_contents_with_change() {
let contents =
"# ML Stack Environment File\nexport MLSTACK_PYTHON_BIN=/usr/bin/python3.11\n"
.to_string();
let (result, changed) = normalize_env_contents(
&contents,
"/usr/bin/python3.13",
"/opt/rocm",
"/opt/rocm/lib",
"/home/test",
);
assert!(changed);
assert!(result.contains("export MLSTACK_PYTHON_BIN=/usr/bin/python3.13"));
assert!(result.contains("export ROCM_PATH=/opt/rocm"));
}
#[test]
fn test_parse_python_version_cap_3_12() {
let (major, minor) = parse_python_version_cap("Python 3.12.12").unwrap();
assert_eq!(major, 3);
assert_eq!(minor, 12);
}
#[test]
fn test_parse_python_version_cap_3_14_rejected() {
let parsed = parse_python_version_cap("Python 3.14.5");
assert!(
parsed.is_some(),
"3.14.5 parses to (3, 14) but cap check rejects later"
);
let (major, minor) = parsed.unwrap();
assert_eq!(major, 3);
assert_eq!(minor, 14);
}
#[test]
fn test_parse_python_version_cap_invalid() {
assert!(parse_python_version_cap("Not a version string").is_none());
}
#[test]
fn test_check_python_version_cap_rejects_3_14() {
let (major, minor) = parse_python_version_cap("Python 3.14.0").unwrap();
assert!(
major > 3 || (major == 3 && minor > 13),
"3.14.0 exceeds cap (≤3.13)"
);
}
#[test]
fn test_check_python_version_cap_accepts_3_12() {
let (major, minor) = parse_python_version_cap("Python 3.12.12").unwrap();
assert!(major <= 3 && minor <= 13, "3.12.12 is within cap (≤3.13)");
}
#[test]
fn test_write_activate_snippets_creates_files() {
let dir = tempfile::tempdir().unwrap();
let _env = crate::test_support::lock_env();
let saved_home = std::env::var("HOME").ok();
std::env::set_var("HOME", dir.path().to_string_lossy().to_string());
let (sh_path, fish_path) = write_activate_snippets().unwrap();
assert!(sh_path.exists());
assert!(fish_path.exists());
let sh_content = std::fs::read_to_string(&sh_path).unwrap();
assert!(sh_content.contains("export PATH=\"$HOME/.mlstack/global/bin:$PATH\""));
assert!(sh_content.contains("# mlstack-global-python"));
let fish_content = std::fs::read_to_string(&fish_path).unwrap();
assert!(fish_content.contains("fish_add_path $HOME/.mlstack/global/bin"));
assert!(fish_content.contains("# mlstack-global-python"));
match saved_home {
Some(v) => std::env::set_var("HOME", v),
None => std::env::remove_var("HOME"),
}
}
#[test]
fn test_offer_global_prepend_idempotent_marker() {
let dir = tempfile::tempdir().unwrap();
let _env = crate::test_support::lock_env();
let saved_home = std::env::var("HOME").ok();
let saved_shell = std::env::var("SHELL").ok();
std::env::set_var("HOME", dir.path().to_string_lossy().to_string());
std::env::set_var("SHELL", "/bin/bash");
let rc_path = dir.path().join(".bashrc");
std::fs::write(&rc_path, "# Existing rc\n").unwrap();
let act1 = offer_global_prepend(true).unwrap();
assert!(act1, "First call should append marker line");
let content1 = std::fs::read_to_string(&rc_path).unwrap();
assert!(
content1.contains("# mlstack-global-python"),
"Marker line should be present"
);
let act2 = offer_global_prepend(true).unwrap();
assert!(!act2, "Second call should detect existing marker and skip");
let content2 = std::fs::read_to_string(&rc_path).unwrap();
let marker_count = content2.matches("# mlstack-global-python").count();
assert_eq!(
marker_count, 1,
"Marker should appear exactly once (idempotent)"
);
match saved_home {
Some(v) => std::env::set_var("HOME", v),
None => std::env::remove_var("HOME"),
}
match saved_shell {
Some(v) => std::env::set_var("SHELL", v),
None => std::env::remove_var("SHELL"),
}
}
#[test]
fn test_offer_global_prepend_detects_shell() {
let dir = tempfile::tempdir().unwrap();
let _env = crate::test_support::lock_env();
let saved_home = std::env::var("HOME").ok();
let saved_shell = std::env::var("SHELL").ok();
std::env::set_var("HOME", dir.path().to_string_lossy().to_string());
std::env::set_var("SHELL", "/usr/bin/fish");
let fish_config = dir.path().join(".config/fish/config.fish");
std::fs::create_dir_all(fish_config.parent().unwrap()).unwrap();
std::fs::write(&fish_config, "# Fish rc\n").unwrap();
let act_fish = offer_global_prepend(true).unwrap();
assert!(act_fish);
let fish_content = std::fs::read_to_string(&fish_config).unwrap();
assert!(fish_content.contains("source \"$HOME/.mlstack/global/activate-global.fish\""));
assert!(fish_content.contains("# mlstack-global-python"));
match saved_home {
Some(v) => std::env::set_var("HOME", v),
None => std::env::remove_var("HOME"),
}
match saved_shell {
Some(v) => std::env::set_var("SHELL", v),
None => std::env::remove_var("SHELL"),
}
}
}