MLflow 3: Rebuilt for the GenAI Era

MLflow has been the default ML lifecycle tool for years. Model tracking, experiment logging, artifact management — if you’ve done ML on Databricks (or anywhere else), you’ve probably used it.

But here’s the thing: MLflow 2 was built for a world where “model” meant a scikit-learn classifier or a PyTorch checkpoint. That world still exists, but it now shares the stage with LLM chains, RAG pipelines, and autonomous agents that call tools, make decisions, and occasionally hallucinate.

MLflow 3 is Databricks’ answer to this shift. Not a patch. Not a feature flag. A ground-up rethinking of what an ML platform needs to be when your “model” is a multi-step agent with a prompt, a retriever, and a mind of its own.

What Actually Changed: MLflow 2 vs. MLflow 3

#

Let’s cut through the marketing and look at the structural differences.

Models Are Now First-Class Citizens

#

In MLflow 2, models were artifacts attached to runs. You’d log a model inside mlflow.start_run(), and the model lived under that run’s artifact path. Want to find your model? Navigate through experiments → runs → artifacts. It worked, but it was like finding your car keys by retracing your entire day.

MLflow 3 introduces LoggedModel — a first-class entity that exists independently of runs:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

import mlflow

# MLflow 2: models live inside runs
with mlflow.start_run():
    mlflow.sklearn.log_model(model, "my_model")

# MLflow 3: models are standalone entities
model_info = mlflow.pyfunc.log_model(
    name="fraud_detector",
    python_model=my_model
)
# No start_run() needed. The model IS the entity.

This isn’t just syntactic sugar. LoggedModel carries its own metadata, parameters, metrics, and tags. You can search, compare, and organize models directly — without digging through run hierarchies.

Tracing Replaces Logging for GenAI

#

Traditional ML logging captures inputs, outputs, hyperparameters, and metrics. That’s fine when your model does one thing: takes input, produces output.

But a GenAI agent might:

1. Parse the user query
2. Retrieve context from a vector store
3. Build a prompt from a template
4. Call an LLM
5. Parse the response
6. Decide to call a tool
7. Call another LLM with tool results
8. Return a final answer

Logging step 1 and step 8 tells you almost nothing. You need tracing — hierarchical, step-by-step visibility into the entire execution flow.

MLflow 3 makes tracing a first-class concept with auto-instrumentation for 50+ frameworks:

1
2
3
4
5
6
7
8
9

import mlflow

# One line. That's it. Every OpenAI call is now traced.
mlflow.openai.autolog()

# Or for LangChain:
mlflow.langchain.autolog()

# Or for Anthropic, LlamaIndex, PydanticAI, smolagents...

Each trace captures:

• Latency at every step (where is your agent spending time?)
• Token usage and cost (which LLM calls are burning your budget?)
• Input/output pairs at each node (what did the retriever actually return?)
• Error propagation (where in the chain did things go wrong?)

flowchart TD
    subgraph Trace["MLflow Trace: Customer Support Agent"]
        A["User Query"] --> B["Intent Classification\n12ms | 150 tokens"]
        B --> C["RAG Retrieval\n45ms | Vector DB"]
        C --> D["Prompt Assembly\n2ms | Template v3"]
        D --> E["LLM Call\n890ms | 1,200 tokens"]
        E --> F{Tool Call?}
        F -->|Yes| G["Tool Execution\n200ms | API Call"]
        G --> H["Follow-up LLM\n650ms | 800 tokens"]
        F -->|No| I["Response"]
        H --> I
    end

    style A fill:#4CAF50,color:#fff
    style E fill:#FF9800,color:#fff
    style H fill:#FF9800,color:#fff
    style I fill:#2196F3,color:#fff

OpenTelemetry Compatibility

#

Here’s the part that matters if you run agents outside Databricks: MLflow traces are OpenTelemetry-compatible. This means you can export traces to any OTel-compatible backend — Jaeger, Grafana Tempo, Datadog, whatever your ops team already uses.

Your agent doesn’t need to run on Databricks to be observable. It just needs MLflow’s tracing SDK.

The Prompt Registry: Version Control for Your Most Important Code

#

If you’re building GenAI applications, your prompts are production code. They determine output quality more than any hyperparameter ever did. Yet most teams manage prompts in… strings. In notebooks. In Slack messages that say “hey try this version.”

MLflow 3’s Prompt Registry treats prompts as versioned, governed artifacts with a proper lifecycle.

Creating and Versioning Prompts

#

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

import mlflow

# Register a new prompt
prompt = mlflow.register_prompt(
    name="customer_support_responder",
    template=(
        "You are a helpful support agent for {{company_name}}.\n"
        "Answer the customer's question using ONLY the provided context.\n\n"
        "Context: {{retrieved_context}}\n\n"
        "Customer question: {{question}}\n\n"
        "If you cannot answer from the context, say so honestly."
    ),
    commit_message="Initial version: conservative, context-only responses"
)

# Later, create a new version with improvements
prompt_v2 = mlflow.register_prompt(
    name="customer_support_responder",
    template=(
        "You are a helpful support agent for {{company_name}}.\n"
        "Answer the customer's question using the provided context.\n"
        "You may use general knowledge to supplement, but clearly "
        "distinguish between verified (from context) and general information.\n\n"
        "Context: {{retrieved_context}}\n\n"
        "Customer question: {{question}}\n\n"
        "Format: Start with the direct answer, then provide details."
    ),
    commit_message="v2: Allow general knowledge with clear attribution"
)

Every version is immutable once created. You get a full history with commit messages — Git for prompts, essentially.

Aliases: Safe Deployment Without Code Changes

#

The real power is in aliases. Instead of hardcoding prompt version numbers, you reference aliases:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

# In your application code (never changes):
prompt = mlflow.load_prompt("customer_support_responder@production")

# Deployment workflow:
# 1. Test v2 in staging
mlflow.set_prompt_alias("customer_support_responder", "staging", version=2)

# 2. Run evals, compare quality scores
# 3. Promote to production
mlflow.set_prompt_alias("customer_support_responder", "production", version=2)

# 4. If something goes wrong, instant rollback
mlflow.set_prompt_alias("customer_support_responder", "production", version=1)

This means your product manager can update a prompt through the MLflow UI, test it in staging, and promote it to production — without a single code deployment. No PR, no CI/CD pipeline, no “can you rebuild the Docker image.”

Prompt Evaluation Before Promotion

#

Before promoting a prompt to production, evaluate it:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

import mlflow

# Load both versions
prompt_v1 = mlflow.load_prompt("customer_support_responder", version=1)
prompt_v2 = mlflow.load_prompt("customer_support_responder", version=2)

# Evaluate against a test dataset
results_v1 = mlflow.evaluate(
    model=my_agent_with_prompt(prompt_v1),
    data=eval_dataset,
    evaluators="default",
)

results_v2 = mlflow.evaluate(
    model=my_agent_with_prompt(prompt_v2),
    data=eval_dataset,
    evaluators="default",
)

# Compare quality scores side-by-side before promoting

MLflow 3 includes built-in LLM judges for automated evaluation — relevance, faithfulness, toxicity, and custom rubrics. You can set up a prompt evaluation pipeline that runs before any alias promotion, catching regressions before they reach users.

Agent Observability: Beyond Databricks

#

One of MLflow 3’s strongest moves is making observability platform-agnostic. Your agent can run on:

• Databricks Model Serving
• AWS SageMaker
• A Kubernetes pod in your own data center
• A laptop during development

…and MLflow traces it all the same way.

Auto-Tracing Integrations

#

MLflow 3 ships with auto-tracing for the frameworks that matter:

One line of code. No manual span creation. No decorator spaghetti. The auto-tracing captures the full execution hierarchy, token counts, latencies, and costs.

Self-Hosted Observability

#

Here’s what makes this compelling: MLflow is open source. Your trace data stays on your infrastructure. No SaaS vendor lock-in. No per-trace pricing surprises.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

import mlflow

# Point to your own tracking server
mlflow.set_tracking_uri("http://your-mlflow-server:5000")

# Enable auto-tracing
mlflow.openai.autolog()

# Every traced call now goes to YOUR server
response = client.chat.completions.create(
    model="gpt-4",
    messages=[{"role": "user", "content": "Explain quantum computing"}]
)
# Trace automatically captured: latency, tokens, cost, full I/O

For teams already running Databricks, MLflow 3 traces integrate natively with Unity Catalog for governance — who accessed what trace data, data lineage from prompt to production response, and audit trails that compliance teams actually accept.

Mosaic AI Agent Framework Integration

#

On Databricks, MLflow 3 is the backbone of the Mosaic AI Agent Framework. The integration is deeper than “they work together” — MLflow provides the tracking, evaluation, and governance layer that the Agent Framework builds on.

The workflow looks like this:

flowchart LR
    subgraph Dev["Development"]
        A["Author Agent"] --> B["Log with MLflow"]
        B --> C["Evaluate"]
    end

    subgraph Govern["Governance"]
        C --> D["Unity Catalog\nModel Registry"]
        D --> E["Prompt Registry"]
    end

    subgraph Deploy["Deployment"]
        D --> F["Model Serving"]
        F --> G["MLflow Tracing\nObservability"]
    end

    style A fill:#4CAF50,color:#fff
    style D fill:#FF9800,color:#fff
    style G fill:#2196F3,color:#fff

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

import mlflow

# Log your agent with MLflow — standard API
mlflow.pyfunc.log_model(
    name="support_agent",
    python_model=my_agent,
    registered_model_name="production_support_agent"
)

# Evaluate with built-in LLM judges
eval_results = mlflow.evaluate(
    model="models:/production_support_agent/latest",
    data=test_dataset,
    evaluators="default",
)

# Deploy to Model Serving — traces flow automatically
# No additional instrumentation needed

The key insight: once your agent is logged with MLflow, tracing is automatic in Databricks Model Serving. Every inference request generates a trace. Every trace is governed by Unity Catalog. You get observability without writing observability code.

Practical Migration: What to Do Monday Morning

#

If you’re on MLflow 2.x, here’s the pragmatic upgrade path:

1. Install MLflow 3 Alongside Existing Workflows

#

1

pip install --upgrade mlflow>=3.0

MLflow 3 reads MLflow 2 resources — your existing experiments, runs, and registered models continue to work.

2. Start with Tracing (Highest ROI)

#

Don’t try to migrate everything at once. Add auto-tracing to your GenAI applications first — it’s the feature with the most immediate value and zero disruption to existing workflows:

1
2

import mlflow
mlflow.openai.autolog()  # Done. You now have observability.

3. Move Prompts to the Registry

#

Identify prompts that are hardcoded in your application. Move them to the Prompt Registry one at a time. Start with the prompts that change most frequently — those benefit most from versioning and aliases.

4. Adopt LoggedModel for New Work

#

For new models and agents, use the LoggedModel pattern (no start_run() wrapper). Keep existing pipelines on the run-based approach until you have a natural reason to refactor.

Key Takeaways

#

MLflow 3 isn’t just an upgrade — it’s a repositioning. The platform that tracked your XGBoost experiments now wants to be the control plane for your entire LLM operation: prompts, agents, traces, and evaluations.

The three features that matter most:

1. Prompt Registry — Version control, aliases, and evaluation for prompts. Treats prompt engineering as the production discipline it actually is.

2. Native Tracing — Hierarchical, OpenTelemetry-compatible observability for GenAI applications. Works everywhere, not just on Databricks.

3. LoggedModel — Models as first-class entities, decoupled from runs. A cleaner abstraction for the world where “model” means many different things.

Whether you’re running agents on Databricks or deploying them independently, MLflow 3 provides the lifecycle management that GenAI applications desperately need. The gap between “cool demo” and “production-grade AI system” is mostly observability, evaluation, and governance — and that’s exactly where MLflow 3 aims.

MLflow 3 is Generally Available. The Prompt Registry, tracing SDK, and LoggedModel entity are all production-ready. For the full migration guide, see the MLflow 3 documentation.