io.github.us-all/mlflow

MLflow MCP Server

The widest-coverage MLflow MCP — including MLflow 3 traces, prompt-optimization, webhooks, and Databricks trace attachments that no other MCP exposes.

82 tools across experiments, runs, registry, logged models, traces, assessments, webhooks, prompt-optimization. Aggregation tools (summarize-experiment, summarize-run) fold 3–5 round-trips into one structured response with already-fetched metric stats.

What it does that others don't

• Full coverage — only third-party MLflow MCP shipping prompt-optimization-jobs (5 tools), webhooks (6), MLflow 3 LoggedModel (8), and Databricks trace attachments (list-trace-attachments, get-trace-attachment — Databricks MLflow only; OSS returns 404).
• Aggregation tools — summarize-experiment returns experiment + topN runs + metric stats (min/max/mean) in one call from already-fetched data, zero extra round-trips. summarize-run dedups metricHistory.history.*.key (~100KB savings on 4k-point series).
• MCP Prompts (4) — debug-failed-traces, promote-best-run, compare-top-runs, annotate-trace-quality. Workflow templates the model invokes directly.
• MCP Resources (6) — mlflow://run/{runId}, mlflow://experiment/{expId}, mlflow://run/{runId}/artifacts, mlflow://experiment/{expId}/runs, mlflow://registered-model/{name}/versions, mlflow://trace/{traceId}.
• Token-efficient by design — extractFields projection on get-run / search-runs / search-traces / get-trace / fat reads, MLFLOW_TOOLS / MLFLOW_DISABLE 8 categories, search-tools meta-tool.
• Apps SDK card — compare-runs renders as a side-by-side card on ChatGPT clients (run summary + metric/param tables with diff highlight) via _meta["openai/outputTemplate"]. Claude clients receive the same JSON content.
• stdio + Streamable HTTP — defaults to stdio. Set MCP_TRANSPORT=http for ChatGPT Apps SDK or remote clients (Bearer auth via MCP_HTTP_TOKEN).

Try this — 5 prompts

Connect the server to Claude Desktop or Claude Code, then paste any of these:

1. Best run — "In the customer-churn-v3 experiment, find the run with the highest val_accuracy. Show its hyperparameters and metric history."
2. Failure mode clustering — "Find traces with status=ERROR from the last 24h in experiment 12. Group the failures by exception type and surface the 3 most common."
3. Run comparison — "Compare the top 5 runs of experiment 12 by validation_loss. Show differing hyperparameters in a table."
4. Model promotion — "Get the latest version of recommendation_v2 registered model with the champion alias. Show its training metrics + lineage to the source run."
5. Trace deep-dive — "Pull trace tr-abc123. Highlight slow spans and any failed feedback annotations." (Add list-trace-attachments on Databricks workspaces.)

When to use this vs alternatives

The official mlflow[mcp] is bundled inside MLflow itself and intentionally trace-narrow. Use it for quick managed-MLflow trace inspection. Use this server for end-to-end coverage, especially MLflow 3 entities, prompt-optimization workflows, and aggregation-driven AI debugging.

Install

Claude Desktop

{
  "mcpServers": {
    "mlflow": {
      "command": "npx",
      "args": ["-y", "@us-all/mlflow-mcp"],
      "env": {
        "MLFLOW_TRACKING_URI": "http://localhost:5000"
      }
    }
  }
}

Claude Code

claude mcp add mlflow -s user \
  -e MLFLOW_TRACKING_URI=http://localhost:5000 \
  -- npx -y @us-all/mlflow-mcp

Docker

docker run --rm -i \
  -e MLFLOW_TRACKING_URI=http://your-host:5000 \
  ghcr.io/us-all/mlflow-mcp-server

Build from source

git clone https://github.com/us-all/mlflow-mcp-server.git
cd mlflow-mcp-server && pnpm install && pnpm build
node dist/index.js

Configuration

Categories (8): experiments, runs, registry, logged-models, traces, assessments, webhooks, prompts.

When MCP_TRANSPORT=http: POST /mcp (Bearer-auth JSON-RPC) + GET /health (public liveness).

Databricks managed MLflow

For Databricks-hosted MLflow:

MLFLOW_TRACKING_URI=https://<workspace>.cloud.databricks.com
MLFLOW_TRACKING_TOKEN=dapi...   # PAT or service-principal token

The MLflow REST API path (/api/2.0/mlflow/...) is identical between OSS and Databricks. Bearer auth handles both PAT and service-principal flows.

Token efficiency

Plus extractFields on get-run / search-runs / search-traces / get-trace / summarize-experiment — caller can scope response fields per call.

Read-only mode

By default, all writes are blocked. The following require MLFLOW_ALLOW_WRITE=true:

create-experiment, update-experiment, delete-experiment, restore-experiment, set-experiment-tag, delete-experiment-tag, create-run, update-run, delete-run, restore-run, log-metric, log-param, log-batch, log-inputs, set-run-tag, delete-run-tag, create-registered-model, rename-registered-model, update-registered-model, delete-registered-model, plus all model-version, logged-model, trace, assessment, webhook, and prompt-optimization writes.

Limitations & gotchas

• search-traces.maxResults is clamped to 500. MLflow 3.12+ rejects per-page max_results > 500 with INVALID_PARAMETER_VALUE. For larger result sets, loop on nextPageToken — total trace count is unbounded.
• Trace attachments are Databricks-only. list-trace-attachments / get-trace-attachment call routes that OSS MLflow (verified through 3.12.0) returns 404 for. Tool descriptions surface this; calls against OSS return a structured MlflowError.
• search-traces.maxResults cap applies per page, not per call — pagination still gets you the full set.
• Bearer / Basic auth code paths are not yet validated against live Databricks (see open roadmap item). Works against OSS MLflow 3.12 (Bearer optional).

MCP Prompts (4)

Workflow templates available via MCP prompts/list:

• debug-failed-traces — find failed traces, group failure modes
• promote-best-run — find best run, register, set champion alias
• compare-top-runs — top-N comparison by metric
• annotate-trace-quality — guided feedback annotation loop

MCP Resources

URI-based read-only access:

mlflow://run/{runId}, mlflow://experiment/{expId}, mlflow://experiment-by-name/{name}, mlflow://registered-model/{name}, mlflow://model-version/{name}/{version}, mlflow://trace/{traceId}, mlflow://run/{runId}/artifacts, mlflow://experiment/{expId}/runs, mlflow://registered-model/{name}/versions.

Tools (82)

8 categories. Use search-tools to discover at runtime; full list collapsed below.

get-run, search-runs, search-traces, get-trace, and summarize-experiment accept extractFields for response slicing.

Full tool list

Experiments (9)

create-experiment, search-experiments, get-experiment, get-experiment-by-name, update-experiment, delete-experiment, restore-experiment, set-experiment-tag, delete-experiment-tag

Runs (18)

create-run, get-run, search-runs, update-run, delete-run, restore-run, log-metric, log-param, log-batch, log-inputs, get-metric-history, set-run-tag, delete-run-tag, list-artifacts, get-best-run, compare-runs, search-runs-by-tags, summarize-run (aggregation)

Registered Models (12)

create-registered-model, get-registered-model, search-registered-models, rename-registered-model, update-registered-model, delete-registered-model, get-latest-model-versions, set-registered-model-tag, delete-registered-model-tag, set-registered-model-alias, delete-registered-model-alias, get-model-version-by-alias

Model Versions (9)

create-model-version, get-model-version, search-model-versions, update-model-version, delete-model-version, transition-model-version-stage, get-model-version-download-uri, set-model-version-tag, delete-model-version-tag

Logged Models — MLflow 3 (8)

create-logged-model, search-logged-models, get-logged-model, finalize-logged-model, delete-logged-model, set-logged-model-tags, delete-logged-model-tag, log-logged-model-params

Traces (8)

search-traces, get-trace, get-trace-info, delete-traces, set-trace-tag, delete-trace-tag, list-trace-attachments, get-trace-attachment

Assessments (5)

log-feedback, log-expectation, get-assessment, update-assessment, delete-assessment

Webhooks (6)

create-webhook, list-webhooks, get-webhook, update-webhook, delete-webhook, test-webhook

Prompt Optimization (5)

create-prompt-optimization-job, get-prompt-optimization-job, search-prompt-optimization-jobs, cancel-prompt-optimization-job, delete-prompt-optimization-job

Aggregations

summarize-experiment, summarize-run — fold 3–5 round-trips into one structured response with caveats array.

Meta

search-tools — query other tools by keyword; always enabled.

Local validation with docker compose

# 1. start MLflow (UI at http://localhost:5050)
docker compose up -d mlflow

# 2. seed demo experiment, runs, registered model, traces
docker compose run --rm seed

# 3a. probe the MCP server locally against the compose'd MLflow
MLFLOW_TRACKING_URI=http://localhost:5050 \
  MLFLOW_EXPERIMENT_ID=1 \
  MLFLOW_ALLOW_WRITE=true \
  node dist/index.js

# 3b. or run inside compose (stdio)
docker compose run --rm mcp

# tear down
docker compose down -v

./dev/seed.py is idempotent — skips if demo experiment already has runs.

Architecture

Claude → MCP stdio → src/index.ts → src/tools/*.ts → MlflowClient (fetch) → MLflow REST API

Built on @us-all/mcp-toolkit:

• extractFields — token-efficient response projections
• aggregate(fetchers, caveats) — fan-out helper for summarize-experiment
• createWrapToolHandler — Bearer/basic credential redaction + MlflowError extraction
• search-tools meta-tool

Targets MLflow 3.5.1+ (uses v3 traces/assessments REST). Dev compose pinned to MLflow 3.12.0 (multimodal trace attachments + paginated trace search).

Tech stack

Node.js 22+ • TypeScript strict ESM • pnpm • @modelcontextprotocol/sdk • zod • dotenv • vitest.

References

• MLflow MCP overview: https://mlflow.org/docs/latest/genai/mcp/
• MLflow REST API: https://mlflow.org/docs/latest/api_reference/rest-api.html

License

MIT