From 43f2bd4bcf182a790a1ee5fcef65b337b3fa8795 Mon Sep 17 00:00:00 2001 From: deacon Date: Mon, 16 Mar 2026 17:24:11 -0400 Subject: [PATCH] Make MLflow port configurable and prevent killing non-MLflow processes The hook.py kill_existing_mlflow() function previously killed any process on port 5000 without verifying it was MLflow. This change: - Checks process cmdline contains 'mlflow' before killing - Makes the port configurable via MCP_MLFLOW_PORT env var - Replaces hardcoded port references with the configurable value Co-Authored-By: Claude Opus 4.6 (1M context) --- hook.py | 27 ++++++++++++++++----------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/hook.py b/hook.py index 822b11e..b4739c6 100644 --- a/hook.py +++ b/hook.py @@ -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' @@ -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 @@ -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') @@ -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)