Skip to content

Make MLflow port configurable and prevent killing non-MLflow processes #14

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
deacon-mp wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Make MLflow port configurable and prevent killing non-MLflow processes #14

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 16 additions & 11 deletions hook.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@

from app.utility.base_world import BaseWorld

MLFLOW_PORT = int(os.environ.get("MCP_MLFLOW_PORT", "5000"))

name = 'mcp'
description = 'Attachment for Model Context Protocol'
address = '/plugin/mcp/gui'
Expand All @@ -31,20 +33,23 @@
]:
logging.getLogger(mod).setLevel(logging.WARNING)

def kill_existing_mlflow(port=5000):
"""Find and kill any process listening on the specified port (usually MLflow)."""
for proc in psutil.process_iter(attrs=["pid", "name"]):
def kill_existing_mlflow(port=MLFLOW_PORT):
"""Find and kill any MLflow process listening on the specified port."""
for proc in psutil.process_iter(attrs=["pid", "name", "cmdline"]):
try:
cmdline = " ".join(proc.info.get("cmdline") or [])
if "mlflow" not in cmdline.lower():
continue
connections = proc.connections(kind='inet')
for conn in connections:
if conn.status == psutil.CONN_LISTEN and conn.laddr.port == port:
print(f"[MCP] Killing existing MLflow process (PID {proc.pid}) on port {port}")
log.info(f"[MCP] Killing existing MLflow process (PID {proc.pid}) on port {port}")
proc.kill()
break
except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
continue
except Exception as e:
print(f"[MCP] Unexpected error while checking connections: {e}")
log.warning(f"[MCP] Unexpected error while checking connections: {e}")
continue


Expand All @@ -54,9 +59,9 @@ def is_port_open(port):
return sock.connect_ex(('127.0.0.1', port)) == 0

# 🔁 Start MLflow server if it's not already running
if not is_port_open(5000):
if not is_port_open(MLFLOW_PORT):
# 🧼 Kill old MLflow server if it exists
kill_existing_mlflow(5000)
kill_existing_mlflow(MLFLOW_PORT)
try:
plugin_dir = os.path.dirname(os.path.abspath(__file__))
mlruns_path = os.path.join(plugin_dir, 'mlruns')
Expand All @@ -66,18 +71,18 @@ def is_port_open(port):
"--backend-store-uri", f"sqlite:///{db_path}",
"--default-artifact-root", mlruns_path,
"--host", "127.0.0.1",
"--port", "5000"
"--port", str(MLFLOW_PORT)
], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
log.debug("[MCP] Starting MLflow server at http://localhost:5000")
log.debug(f"[MCP] Starting MLflow server at http://localhost:{MLFLOW_PORT}")
except Exception as e:
log.error(f"[MCP] Failed to start MLflow server: {e}")
traceback.print_exc()
else:
log.info("[MCP] MLflow server already running on port 5000")
log.info(f"[MCP] MLflow server already running on port {MLFLOW_PORT}")

# 💤 Optional: Wait until server is reachable
for i in range(10):
if is_port_open(5000):
if is_port_open(MLFLOW_PORT):
log.debug("[MCP] MLflow is ready.")
break
time.sleep(1)
Expand Down
Loading