use std::collections::HashSet;
use std::path::PathBuf;
use std::process::Command;
use crate::core::registry::InstalledComponentRegistry;
use crate::platform::environment::{
command_on_path, mlstack_global_dir, resolve_canonical_python_bin, resolve_user_home,
};
const ROCM_SYSTEM_PACKAGES: &[&str] = &[
"rocm-hip-sdk",
"rocm-opencl-sdk",
"rocm-core",
"rocm-dev",
"rocm-opencl-runtime",
"amdgpu",
"amdgpu-dkms",
];
const ML_PIP_PACKAGES: &[&str] = &[
"torch",
"torchvision",
"torchaudio",
"triton",
"pytorch-triton-rocm",
"pytorch-triton",
"vllm",
"deepspeed",
"megatron-core",
"bitsandbytes",
"flash-attn",
"flash_attn",
"aiter",
"onnxruntime",
"onnxruntime-rocm",
"migraphx",
"mpi4py",
"wandb",
"xformers",
"accelerate",
"transformers",
];
#[derive(Debug, Clone, Default)]
pub struct UninstallOptions {
pub keep_rocm: bool,
pub yes: bool,
pub purge_mlstack_dir: bool,
pub sudo_password: Option<String>,
}
#[derive(Debug, Clone, Default)]
pub struct UninstallReport {
pub pip_uninstall_attempted: Vec<String>,
pub system_packages_purged: bool,
pub opt_rocm_removed: bool,
pub env_files_removed: Vec<String>,
pub sourcing_lines_stripped: Vec<String>,
pub registry_cleared: bool,
pub warnings: Vec<String>,
}
impl UninstallReport {
fn note(&mut self, msg: impl Into<String>) {
self.warnings.push(msg.into());
}
}
pub fn uninstall_stack(opts: &UninstallOptions) -> anyhow::Result<UninstallReport> {
let mut report = UninstallReport::default();
let home = resolve_user_home();
if !opts.yes {
eprintln!(
"[uninstall] Proceeding (use --yes to skip this notice). keep_rocm={}, purge_dir={}",
opts.keep_rocm, opts.purge_mlstack_dir
);
}
let python = resolve_canonical_python_bin();
clean_corrupted_pip_distributions(&python, opts.sudo_password.as_deref(), &mut report);
println!("[uninstall] Removing ML pip packages via {python} …");
let mut pkgs: Vec<String> = ML_PIP_PACKAGES.iter().map(|s| s.to_string()).collect();
pkgs.extend(crate::core::registry::InstalledComponentRegistry::load().all_pip_packages());
pkgs.sort();
pkgs.dedup();
report.pip_uninstall_attempted = pkgs.clone();
let mut pip_args = vec![
"-m".to_string(),
"pip".to_string(),
"uninstall".to_string(),
"-y".to_string(),
"--break-system-packages".to_string(),
];
pip_args.extend(pkgs);
let pip_status = Command::new(&python).args(&pip_args).status();
match pip_status {
Ok(s) if s.success() => {}
Ok(s) => report.note(format!(
"pip uninstall exited non-zero ({s}) — some packages may not have been installed"
)),
Err(e) => report.note(format!("could not run {python} -m pip ({e})")),
}
if !opts.keep_rocm {
let installed = filter_installed_system_packages(ROCM_SYSTEM_PACKAGES);
if installed.is_empty() {
report.note(
"no ROCm/amdgpu system packages are installed — skipping system-package purge",
);
} else if let Some((pm, purge_args)) = build_system_purge_cmd(&installed) {
println!("[uninstall] Purging ROCm/amdgpu system packages via {pm} …");
let ran = run_privileged(&pm, &purge_args, opts.sudo_password.as_deref(), &mut report);
report.system_packages_purged = ran;
} else {
report.note("no supported system package manager detected (apt/dnf/pacman/zypper) — skipping ROCm system-package purge");
}
if PathBuf::from("/opt/rocm").exists() {
println!("[uninstall] Removing /opt/rocm …");
report.opt_rocm_removed = run_privileged(
"rm",
&["-rf".to_string(), "/opt/rocm".to_string()],
opts.sudo_password.as_deref(),
&mut report,
);
}
} else {
println!("[uninstall] --keep-rocm: leaving ROCm system packages and /opt/rocm in place");
}
remove_env_files(&home, &mut report);
strip_shell_sourcing(&home, &mut report);
let mut registry = InstalledComponentRegistry::load();
if !registry.is_empty() {
registry.clear();
if let Err(e) = registry.save() {
report.note(format!("failed to clear registry: {e}"));
} else {
report.registry_cleared = true;
}
}
if opts.purge_mlstack_dir {
let root = crate::platform::environment::mlstack_root();
if root.exists() {
if std::fs::remove_dir_all(&root).is_ok() {
report
.env_files_removed
.push(format!("(dir) {}", root.display()));
} else {
let removed = run_privileged(
"rm",
&["-rf".to_string(), root.to_string_lossy().into_owned()],
opts.sudo_password.as_deref(),
&mut report,
);
if removed {
report
.env_files_removed
.push(format!("(dir, sudo) {}", root.display()));
}
}
}
}
Ok(report)
}
fn build_system_purge_cmd(packages: &[&str]) -> Option<(String, Vec<String>)> {
let (pm, verb) = if command_on_path("apt") {
(
"apt",
vec!["purge".to_string(), "-y".to_string(), "-qq".to_string()],
)
} else if command_on_path("pacman") {
let full_set = enumerate_rocm_explicit_purge_set(packages);
if full_set.is_empty() {
return None;
}
let mut args = vec!["-Rn".to_string(), "--noconfirm".to_string()];
args.extend(full_set.into_iter().map(|s| s.to_string()));
return Some(("pacman".to_string(), args));
} else if command_on_path("dnf") {
("dnf", vec!["remove".to_string(), "-y".to_string()])
} else if command_on_path("zypper") {
("zypper", vec!["remove".to_string(), "-y".to_string()])
} else if command_on_path("yum") {
("yum", vec!["remove".to_string(), "-y".to_string()])
} else {
return None;
};
let mut args = verb;
args.extend(packages.iter().map(|s| s.to_string()));
Some((pm.to_string(), args))
}
fn enumerate_rocm_explicit_purge_set(targets: &[&str]) -> Vec<String> {
let allowlist = rocm_allowlist_patterns();
let mut to_remove: HashSet<String> = targets.iter().map(|s| s.to_string()).collect();
let mut skipped_non_rocm: Vec<String> = Vec::new();
for target in targets {
let dependents = reverse_tree_pacman(target);
for dep in dependents {
if matches_rocm_allowlist(&dep, &allowlist) {
to_remove.insert(dep);
} else {
skipped_non_rocm.push(dep);
}
}
}
if !skipped_non_rocm.is_empty() {
eprintln!(
"[uninstall] Skipping {} non-ROCm dependent(s) that would be caught by cascade: {:?}. Only ROCm packages are removed.",
skipped_non_rocm.len(),
skipped_non_rocm
);
}
let mut sorted: Vec<_> = to_remove.into_iter().collect();
sorted.sort();
sorted.dedup();
sorted
}
fn rocm_allowlist_patterns() -> Vec<&'static str> {
vec![
"rocm-",
"hip",
"roc",
"hsa",
"comgr",
"migraphx",
"miopen",
"rccl",
"amdgpu",
"libdrm_amd",
"rocwmma",
"amd-comgr",
"hsakmt",
"hipfft",
"hipsparse",
"hipcub",
"hiprand",
"rocsolver",
"rocprim",
"rocrand",
"rocfft",
"llvm-libs",
]
}
fn matches_rocm_allowlist(pkg: &str, allowlist: &[&str]) -> bool {
allowlist
.iter()
.any(|pat| pkg.starts_with(pat) || pkg == pat.trim_end_matches('-'))
}
fn reverse_tree_pacman(package: &str) -> Vec<String> {
if !command_on_path("pactree") {
eprintln!("[uninstall] pactree not found — cannot enumerate ROCm dependents for {package}. Skipping dependent enumeration.");
return Vec::new();
}
Command::new("pactree")
.args(["-r", "-u", package])
.output()
.ok()
.and_then(|o| {
if o.status.success() {
String::from_utf8(o.stdout).ok()
} else {
None
}
})
.map(|s| {
s.lines()
.filter_map(|line| {
let trimmed = line.trim();
if trimmed.is_empty() || trimmed.starts_with('#') {
None
} else {
Some(trimmed.to_string())
}
})
.collect()
})
.unwrap_or_default()
}
fn filter_installed_system_packages<'a>(candidates: &[&'a str]) -> Vec<&'a str> {
candidates
.iter()
.copied()
.filter(|pkg| is_system_package_installed(pkg))
.collect()
}
fn is_system_package_installed(pkg: &str) -> bool {
if command_on_path("dpkg-query") {
return Command::new("dpkg-query")
.args(["-W", "-f=${Status}", pkg])
.output()
.map(|o| {
o.status.success()
&& String::from_utf8_lossy(&o.stdout).contains("install ok installed")
})
.unwrap_or(false);
}
if command_on_path("pacman") {
return Command::new("pacman")
.args(["-Q", pkg])
.output()
.map(|o| o.status.success())
.unwrap_or(false);
}
if command_on_path("rpm") {
return Command::new("rpm")
.args(["-q", pkg])
.output()
.map(|o| o.status.success())
.unwrap_or(false);
}
true
}
fn run_privileged(
program: &str,
args: &[String],
sudo_password: Option<&str>,
report: &mut UninstallReport,
) -> bool {
if is_root() {
return run_command_status(program, args, None, report);
}
if !command_on_path("sudo") {
report.note(format!(
"skipping `{program} {args:?}` — not root and sudo unavailable",
));
return false;
}
if let Some(pw) = sudo_password.filter(|p| !p.is_empty()) {
match crate::installers::common::askpass::Askpass::new(pw) {
Ok(helper) => {
let mut sudo_args = vec!["-A".to_string(), program.to_string()];
sudo_args.extend_from_slice(args);
return run_command_status("sudo", &sudo_args, Some(helper.path()), report);
}
Err(e) => {
report.note(format!(
"askpass setup failed ({e}); falling back to sudo -n"
));
}
}
}
let mut a = vec!["-n".to_string(), program.to_string()];
a.extend_from_slice(args);
run_command_status("sudo", &a, None, report)
}
fn run_command_status(
program: &str,
args: &[String],
askpass_path: Option<&str>,
report: &mut UninstallReport,
) -> bool {
let mut cmd = Command::new(program);
cmd.args(args);
if let Some(p) = askpass_path {
cmd.env("SUDO_ASKPASS", p);
}
match cmd.status() {
Ok(s) if s.success() => true,
Ok(s) => {
report.note(format!("`{program}` exited non-zero ({s})"));
false
}
Err(e) => {
report.note(format!(
"could not run `{program}` ({e}) — sudo may require a password (no TTY)"
));
false
}
}
}
#[cfg(unix)]
fn is_root() -> bool {
unsafe { libc::geteuid() == 0 }
}
#[cfg(not(unix))]
fn is_root() -> bool {
false
}
fn clean_corrupted_pip_distributions(
python: &str,
sudo_password: Option<&str>,
report: &mut UninstallReport,
) {
let sites_out = Command::new(python)
.args([
"-c",
"import site,sys; sys.stdout.write('\\n'.join(site.getsitepackages()+[site.getusersitepackages()]))",
])
.output();
let sites: Vec<String> = match sites_out {
Ok(o) if o.status.success() => String::from_utf8_lossy(&o.stdout)
.lines()
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.collect(),
_ => return,
};
let mut needs_sudo: Vec<String> = Vec::new();
for site_dir in &sites {
let Ok(entries) = std::fs::read_dir(site_dir) else {
continue;
};
for ent in entries.flatten() {
let file_name = ent.file_name();
let Some(name) = file_name.to_str() else {
continue;
};
if !name.starts_with('~') {
continue;
}
let path = ent.path();
let user_removed = match std::fs::metadata(&path) {
Ok(m) if m.is_dir() => std::fs::remove_dir_all(&path).is_ok(),
Ok(_) => std::fs::remove_file(&path).is_ok(),
Err(_) => false,
};
if user_removed {
report
.env_files_removed
.push(format!("(corrupted pip dist) {site_dir}/{name}"));
} else if let Some(p) = path.to_str() {
needs_sudo.push(p.to_string());
}
}
}
if needs_sudo.is_empty() {
return;
}
let count = needs_sudo.len();
let mut args = vec!["-rf".to_string()];
args.extend(needs_sudo.iter().cloned());
if run_privileged("rm", &args, sudo_password, report) {
for p in &needs_sudo {
report
.env_files_removed
.push(format!("(corrupted pip dist, sudo) {p}"));
}
} else {
report.note(format!(
"failed to remove {count} root-owned corrupted pip dist dir(s) — sudo required"
));
}
}
fn remove_env_files(home: &std::path::Path, report: &mut UninstallReport) {
let targets = [
home.join(".mlstack_env"),
home.join(".config/fish/conf.d/mlstack_env.fish"),
home.join(".rocm_env"),
];
for t in targets {
if t.exists() {
match std::fs::remove_file(&t) {
Ok(()) => report
.env_files_removed
.push(t.to_string_lossy().into_owned()),
Err(e) => report.note(format!("could not remove {}: {e}", t.display())),
}
}
}
}
fn strip_shell_sourcing(home: &std::path::Path, report: &mut UninstallReport) {
let rcs = [
home.join(".bashrc"),
home.join(".zshrc"),
home.join(".config/fish/config.fish"),
];
let had_trailing_newline = |content: &str| content.ends_with('\n');
for rc in rcs {
let Ok(content) = std::fs::read_to_string(&rc) else {
continue;
};
let mut filtered = String::with_capacity(content.len());
let mut changed = false;
for line in content.lines() {
let t = line.trim();
let is_source_line = t.contains(".mlstack_env")
&& (t.contains("source") || t.contains("bass") || t.starts_with(". "));
if is_source_line {
changed = true;
} else {
filtered.push_str(line);
filtered.push('\n');
}
}
if changed {
if !had_trailing_newline(&content) && filtered.ends_with('\n') {
filtered.pop();
}
let _ = std::fs::write(&rc, filtered);
report
.sourcing_lines_stripped
.push(rc.to_string_lossy().into_owned());
}
}
}
pub fn global_env_dir() -> PathBuf {
mlstack_global_dir()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn build_system_purge_cmd_apt_when_present() {
let cmd = build_system_purge_cmd(&["rocm-hip-sdk"]);
if command_on_path("apt") {
let (pm, args) = cmd.expect("apt present");
assert_eq!(pm, "apt");
assert!(args.contains(&"purge".to_string()));
assert!(args.contains(&"rocm-hip-sdk".to_string()));
}
}
#[test]
fn strip_sourcing_removes_mlstack_lines() {
let dir = tempfile::tempdir().unwrap();
let rc = dir.path().join(".bashrc");
std::fs::write(
&rc,
"# my rc\nsource \"$HOME/.mlstack_env\"\nexport FOO=bar\n. ~/.mlstack_env\n",
)
.unwrap();
let mut report = UninstallReport::default();
strip_shell_sourcing(dir.path(), &mut report);
let after = std::fs::read_to_string(&rc).unwrap();
assert!(
!after.contains(".mlstack_env"),
"sourcing lines removed: {after}"
);
assert!(
after.contains("export FOO=bar"),
"other lines preserved: {after}"
);
assert_eq!(report.sourcing_lines_stripped.len(), 1);
assert!(
after.ends_with('\n'),
"trailing newline preserved: {after:?}"
);
}
#[test]
fn strip_sourcing_preserves_non_source_mentions() {
let dir = tempfile::tempdir().unwrap();
let rc = dir.path().join(".bashrc");
let original = "# see ~/.mlstack_env for env\necho loaded .mlstack_env\nMLSTACK_ENV_FILE=.mlstack_env\n";
std::fs::write(&rc, original).unwrap();
let mut report = UninstallReport::default();
strip_shell_sourcing(dir.path(), &mut report);
let after = std::fs::read_to_string(&rc).unwrap();
assert_eq!(
report.sourcing_lines_stripped.len(),
0,
"no source lines, no rewrite: {after:?}"
);
assert!(
after.contains("# see ~/.mlstack_env for env"),
"comment preserved: {after}"
);
assert!(
after.contains("MLSTACK_ENV_FILE=.mlstack_env"),
"assignment preserved: {after}"
);
}
#[test]
fn remove_env_files_deletes_targets() {
let dir = tempfile::tempdir().unwrap();
std::fs::write(dir.path().join(".mlstack_env"), "x").unwrap();
let mut report = UninstallReport::default();
remove_env_files(dir.path(), &mut report);
assert!(!dir.path().join(".mlstack_env").exists());
assert!(!report.env_files_removed.is_empty());
}
}