diff --git a/app/mcp_svc.py b/app/mcp_svc.py index 30ae224..64f25a7 100644 --- a/app/mcp_svc.py +++ b/app/mcp_svc.py @@ -9,6 +9,8 @@ from app.utility.base_service import BaseService import asyncio +from mlflow.tracking import get_tracking_uri + from plugins.mcp.app.config import resolve_llm_config from plugins.mcp.app.mlflow_run import ( RunTracker, @@ -479,9 +481,19 @@ async def execute(self, focus: str = None, prompt: str = "", model_config: dict # never share mlflow's thread-local active-run stack. Minting into # the workflow's own experiment also keeps History grouped the way # each workflow declares it. - run_id = RunTracker.start( - workflow.mlflow_experiment, f"MCP {workflow.display_name}" - ).run_id + try: + run_id = RunTracker.start( + workflow.mlflow_experiment, f"MCP {workflow.display_name}" + ).run_id + except Exception as e: + # Every other MLflow write is best-effort, so a tracking server + # that is down surfaces only here, as a raw provider exception + # that names neither MLflow nor the port it came from. + raise RuntimeError( + f"MLflow tracking is unavailable at {get_tracking_uri()}, so the " + f"run could not be started: {e}. Check the caldera log for an " + f"MLflow port conflict." + ) from e # Resolve the session this run belongs to. The first turn auto-starts # a session keyed by its own run_id; follow-up turns echo the same diff --git a/hook.py b/hook.py index 2d35efb..6018992 100644 --- a/hook.py +++ b/hook.py @@ -1,10 +1,16 @@ import atexit +import json import os import subprocess import socket +import sys import time import traceback import logging +import urllib.error +import urllib.request + +import psutil from dotenv import load_dotenv from app.utility.base_world import BaseWorld @@ -53,6 +59,11 @@ def is_port_open(port, host='127.0.0.1'): _MLFLOW_URI = _mlflow['tracking_uri'] os.environ.setdefault('MLFLOW_TRACKING_URI', _MLFLOW_URI) +# This tree's tracking store. The artifact root doubles as the fingerprint +# that tells our own server apart from another tree's on the same port. +_MLFLOW_BACKEND_URI = f"sqlite:///{os.path.join(_PLUGIN_DIR, 'mlruns.db')}" +_MLFLOW_ARTIFACT_ROOT = os.path.join(_PLUGIN_DIR, 'mlruns') + # Only set when this process started MLflow, so shutdown never kills a server # the operator is running themselves. _mlflow_proc = None @@ -69,19 +80,101 @@ def _stop_mlflow_server(): _mlflow_proc.kill() -# Started from enable(), not at import: this spawns a process and blocks for -# up to 10s, and importing the module should do neither. -def _ensure_mlflow_server(): +def _probe_listener(timeout=2): + """The artifact root of the server on our port, or None if not MLflow. + + A TCP connect only proves something answers. Adopting a listener that + is not MLflow -- macOS binds :5000 to AirPlay Receiver -- makes every + create_run fail with a 403 the operator only sees as a dead run. + Returns "" when it answers the tracking API but names no root. + """ + url = f"{_MLFLOW_URI}/api/2.0/mlflow/experiments/get?experiment_id=0" + try: + with urllib.request.urlopen(url, timeout=timeout) as response: + body = response.read() + except urllib.error.HTTPError as e: + # A deleted default experiment still answers with an MLflow error body. + body = e.read() + except Exception: + return None + try: + payload = json.loads(body.decode()) + except Exception: + return None + if "experiment" not in payload and "error_code" not in payload: + return None + return str((payload.get("experiment") or {}).get("artifact_location") or "") + + +def _serves_our_store(artifact_root): + """Whether that artifact root is this tree's, so the server is ours.""" + if not artifact_root: + return False + path = artifact_root.split("://", 1)[-1] + try: + root = os.path.realpath(_MLFLOW_ARTIFACT_ROOT) + return os.path.realpath(path).startswith(root + os.sep) + except Exception: + return False + + +def _mlflow_server_port(cmdline): + """The port an ``mlflow server``/``mlflow ui`` argv serves, else None.""" + if not any("mlflow" in part for part in cmdline): + return None + if "server" not in cmdline and "ui" not in cmdline: + return None + if "--port" in cmdline: + try: + return int(cmdline[cmdline.index("--port") + 1]) + except (IndexError, ValueError): + return None + return 5000 # mlflow's own default when the flag is absent + + +def _reclaim_port(): + """Stop the foreign MLflow on our port, returning whether it freed. + + Only an mlflow server bound to our port is a candidate. macOS denies + a port-to-pid lookup without root, so the owner is found by cmdline. + """ + victims = [] + for proc in psutil.process_iter(["cmdline"]): + try: + if _mlflow_server_port(proc.info.get("cmdline") or []) == _MLFLOW_PORT: + victims.append(proc) + except psutil.Error: + continue + for proc in victims: + # The CLI parent forks uvicorn workers that hold the socket, so + # terminating it alone leaves the port bound. + try: + targets = proc.children(recursive=True) + [proc] + for target in targets: + target.terminate() + _, alive = psutil.wait_procs(targets, timeout=5) + for target in alive: + target.kill() + except psutil.Error as e: + log.warning(f"[MCP] Could not stop MLflow pid {proc.pid}: {e}") + + for _ in range(10): + if not is_port_open(_MLFLOW_PORT, _MLFLOW_HOST): + return True + time.sleep(0.5) + return False + + +def _start_mlflow_server(): global _mlflow_proc - if is_port_open(_MLFLOW_PORT, _MLFLOW_HOST): - log.info(f"[MCP] MLflow already running on {_MLFLOW_HOST}:{_MLFLOW_PORT}") - return try: - plugin_dir = os.path.dirname(os.path.abspath(__file__)) + # The interpreter running caldera, not a bare "mlflow": restarting a + # reclaimed server is now a normal path, and PATH may not carry the + # venv that holds the mlflow client this process imports. _mlflow_proc = subprocess.Popen([ - "mlflow", "server", - "--backend-store-uri", f"sqlite:///{os.path.join(plugin_dir, 'mlruns.db')}", - "--default-artifact-root", os.path.join(plugin_dir, 'mlruns'), + sys.executable, "-m", "mlflow", "server", + "--backend-store-uri", _MLFLOW_BACKEND_URI, + "--default-artifact-root", _MLFLOW_ARTIFACT_ROOT, "--host", _MLFLOW_HOST, "--port", str(_MLFLOW_PORT), ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) @@ -102,6 +195,38 @@ def _ensure_mlflow_server(): log.error(f"[MCP] MLflow did not start within 10s on {_MLFLOW_HOST}:{_MLFLOW_PORT}.") +# Started from enable(), not at import: this spawns a process and blocks for +# up to 10s, and importing the module should do neither. +def _ensure_mlflow_server(): + if is_port_open(_MLFLOW_PORT, _MLFLOW_HOST): + artifact_root = _probe_listener() + if artifact_root is None: + log.error( + f"[MCP] {_MLFLOW_HOST}:{_MLFLOW_PORT} is held by something that is " + f"not MLflow, so every run will fail to start. Free the port, or " + f"set mlflow.port in plugins/mcp/conf/default.yml." + ) + return + if _serves_our_store(artifact_root): + log.info(f"[MCP] MLflow already running on {_MLFLOW_HOST}:{_MLFLOW_PORT}") + return + # Sharing another tree's store crosses both trees' history, and the + # boot reconcile then marks that tree's live runs KILLED. + log.warning( + f"[MCP] The MLflow server on {_MLFLOW_HOST}:{_MLFLOW_PORT} tracks " + f"{artifact_root or 'an unknown store'}, not {_MLFLOW_ARTIFACT_ROOT}. " + f"Restarting it against this tree's store." + ) + if not _reclaim_port(): + log.error( + f"[MCP] Could not free {_MLFLOW_HOST}:{_MLFLOW_PORT}. Runs would land " + f"in another tree's store, so that server is left alone and runs will " + f"fail. Stop it, or set mlflow.port in plugins/mcp/conf/default.yml." + ) + return + _start_mlflow_server() + + # ✅ Now import modules that depend on MLflow try: from plugins.mcp.app.mcp_svc import MCPService diff --git a/tests/test_mlflow_lifecycle.py b/tests/test_mlflow_lifecycle.py index 2c70836..ff9f7cb 100644 --- a/tests/test_mlflow_lifecycle.py +++ b/tests/test_mlflow_lifecycle.py @@ -3,7 +3,14 @@ Caldera has no plugin teardown hook, so a Popen with no cleanup is reparented to init and holds the port across every restart. It must also never stop a server the operator is running themselves. + +Adoption is the other half: a bare TCP connect proved only that something +answered on the port, so a non-MLflow listener made every run fail at +create_run, and another tree's server took our runs into its store and +killed that tree's live runs at the next boot reconcile. """ +import logging +import os import subprocess import pytest @@ -64,9 +71,77 @@ def _hang(timeout=None): assert proc.killed -def test_an_already_running_server_is_left_alone(monkeypatch): +def _no_spawn(monkeypatch, why): + monkeypatch.setattr(hook, "_start_mlflow_server", + lambda: pytest.fail(why)) + + +def test_our_own_running_server_is_adopted(monkeypatch): monkeypatch.setattr(hook, "is_port_open", lambda *a, **k: True) - monkeypatch.setattr(subprocess, "Popen", lambda *a, **k: pytest.fail("spawned over a running server")) - hook._mlflow_proc = None + monkeypatch.setattr(hook, "_probe_listener", + lambda *a, **k: os.path.join(hook._MLFLOW_ARTIFACT_ROOT, "0")) + _no_spawn(monkeypatch, "spawned over our own running server") + monkeypatch.setattr(hook, "_reclaim_port", lambda: pytest.fail("killed our own server")) + hook._ensure_mlflow_server() + + +def test_a_listener_that_is_not_mlflow_is_neither_adopted_nor_killed(monkeypatch, caplog): + # macOS binds :5000 to AirPlay Receiver, and adopting it made every + # create_run fail with a 403. + monkeypatch.setattr(hook, "is_port_open", lambda *a, **k: True) + monkeypatch.setattr(hook, "_probe_listener", lambda *a, **k: None) + _no_spawn(monkeypatch, "spawned onto an occupied port") + monkeypatch.setattr(hook, "_reclaim_port", lambda: pytest.fail("killed an unidentified process")) + with caplog.at_level(logging.ERROR, logger=hook.log.name): + hook._ensure_mlflow_server() + # Silent adoption was the defect: not spawning is not enough, the + # operator has to be told why their runs are about to fail. + assert "not MLflow" in caplog.text + + +def test_another_trees_server_is_reclaimed_and_restarted(monkeypatch): + monkeypatch.setattr(hook, "is_port_open", lambda *a, **k: True) + monkeypatch.setattr(hook, "_probe_listener", lambda *a, **k: "/somewhere/else/mlruns/0") + started = [] + monkeypatch.setattr(hook, "_reclaim_port", lambda: True) + monkeypatch.setattr(hook, "_start_mlflow_server", lambda: started.append(True)) hook._ensure_mlflow_server() - assert hook._mlflow_proc is None + assert started + + +def test_a_foreign_server_that_will_not_die_is_not_written_to(monkeypatch, caplog): + monkeypatch.setattr(hook, "is_port_open", lambda *a, **k: True) + monkeypatch.setattr(hook, "_probe_listener", lambda *a, **k: "/somewhere/else/mlruns/0") + monkeypatch.setattr(hook, "_reclaim_port", lambda: False) + _no_spawn(monkeypatch, "spawned onto a port still held by another server") + with caplog.at_level(logging.ERROR, logger=hook.log.name): + hook._ensure_mlflow_server() + assert "Could not free" in caplog.text + + +@pytest.mark.parametrize("artifact_root, ours", [ + ("", False), + ("/somewhere/else/mlruns/0", False), + ("mlflow-artifacts:/0", False), +]) +def test_a_foreign_artifact_root_is_not_ours(artifact_root, ours): + assert hook._serves_our_store(artifact_root) is ours + + +@pytest.mark.parametrize("prefix", ["", "file://"]) +def test_our_artifact_root_is_ours(prefix): + assert hook._serves_our_store(prefix + os.path.join(hook._MLFLOW_ARTIFACT_ROOT, "0")) + + +@pytest.mark.parametrize("cmdline, port", [ + (["mlflow", "server", "--port", "5050"], 5050), + (["/venv/bin/mlflow", "ui", "--port", "5050"], 5050), + # mlflow's own default, which is what an operator typing `mlflow ui` gets. + (["mlflow", "server"], 5000), + # The uvicorn workers carry the port too, but the CLI parent owns them. + (["python", "-m", "uvicorn", "--port", "5050", "mlflow.server.fastapi_app:app"], None), + (["python", "server.py", "--port", "5050"], None), + (["mlflow", "server", "--port", "not-a-port"], None), +]) +def test_only_an_mlflow_cli_on_our_port_is_a_reclaim_candidate(cmdline, port): + assert hook._mlflow_server_port(cmdline) == port diff --git a/tests/test_run_execution_lifecycle.py b/tests/test_run_execution_lifecycle.py index 372004f..700b4ab 100644 --- a/tests/test_run_execution_lifecycle.py +++ b/tests/test_run_execution_lifecycle.py @@ -136,3 +136,26 @@ def test_reconcile_spares_runs_the_live_cache_still_owns(svc): assert asyncio.run(svc.reconcile_orphaned_runs()) == [] assert MlflowClient().get_run(in_flight).info.status == "RUNNING" + + +def test_an_unreachable_tracking_server_names_mlflow(svc, monkeypatch): + """The defect: /execute 500'd with a bare provider exception. + + A port adopted from a non-MLflow listener fails at create_run, and + that was the only MLflow write not wrapped, so the chat UI showed a + 403 with nothing tying it to the tracking server. + """ + async def never_runs(prompt, lm_obj, run_id=None, **kwargs): + pytest.fail("started a workflow with no run to track it") + + svc.workflow_registry = {"slow": _workflow("slow", never_runs)} + + def _refuse(*a, **k): + raise Exception( + "API request to endpoint /api/2.0/mlflow/experiments/get-by-name " + "failed with error code 403 != 200" + ) + monkeypatch.setattr(RunTracker, "start", _refuse) + + with pytest.raises(RuntimeError, match="MLflow tracking is unavailable"): + asyncio.run(svc.execute(workflow_id="slow", prompt="p"))