kill_mlflow.py

import subprocess

try:
    import psutil

    # Find all MLflow server processes
    mlflow_processes = []
    for proc in psutil.process_iter(["pid", "name", "cmdline"]):
        try:
            cmdline = proc.info.get("cmdline", [])
            if cmdline and "mlflow.server" in " ".join(cmdline):
                mlflow_processes.append(proc)
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            continue

    if mlflow_processes:
        print(f"🔍 Found {len(mlflow_processes)} MLflow processes to clean up")
        for proc in mlflow_processes:
            try:
                print(f"   Terminating process {proc.pid}")
                proc.terminate()
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                pass

        # Wait for processes to terminate
        psutil.wait_procs(mlflow_processes, timeout=3)

        # Force kill any that didn't terminate
        remaining = [p for p in mlflow_processes if p.is_running()]
        if remaining:
            print(f"🔥 Force killing {len(remaining)} stubborn processes")
            for proc in remaining:
                try:
                    proc.kill()
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    pass
    else:
        print("✅ No MLflow processes found to clean up")

except ImportError:
    # Fallback if psutil is not available - use pkill
    print("📦 psutil not available, using pkill for cleanup")
    try:
        result = subprocess.run(
            ["pkill", "-f", "mlflow.server"],
            capture_output=True,
            text=True,
        )
        if result.returncode == 0:
            print("✅ Cleaned up processes with pkill")
        else:
            print("ℹ️ No processes found to kill")
    except FileNotFoundError:
        print("⚠️ pkill not available, manual cleanup may be needed")