MLflow for LLMOps

What MLflow is (for LLM teams)

MLflow is an open MLOps platform: experiment tracking, artifact storage, model registry, and (increasingly) GenAI/LLM evaluation hooks. Databricks and many lakehouse shops standardize on it; the habits transfer even if you use Weights & Biases for day-to-day compare UX.

LLMOps is MLOps with messier artifacts. You version more than .pt / .safetensors weights:

flowchart LR
  Exp[Experiment runs] --> Reg[Model / prompt registry]
  Reg --> Stage[Staging]
  Stage --> Gate[Eval gate]
  Gate --> Prod[Production endpoint]
  Prod --> Mon[Monitoring / traces]
  Mon --> Exp

Interview cue: List five artifacts you’d version for a RAG bot and how you’d gate a production alias flip.

The engineering problem

Without a registry + tracking loop:

• “What beat baseline last Tuesday?” lives in Slack
• Prod prompt ≠ the prompt that passed eval
• LoRA adapters ship without pointer to base revision
• Offline and online metrics disagree with no shared IDs
• Rollback means “find the old file on someone’s laptop”

MLflow (or equivalent) makes lineage + promotion first-class — the release control plane for GenAI.

Architecture: pieces you will touch

flowchart TB
  Code[Git SHA] --> Run[mlflow run]
  Data[Golden set version] --> Run
  Run --> Arts[Artifacts: prompt / adapter / preds]
  Run --> Reg[Registry version]
  Reg -->|alias staging| Gate[Eval gate]
  Gate -->|alias production| Serve[Serving endpoint]

Mental model: a run is one experiment; a registered version is a promote-able pointer; an alias (staging / production) is what serving should resolve — not a floating “latest file on disk.”

How it fits LLM apps

flowchart LR
  Dev[Dev iterate] --> Track[MLflow / W&B runs]
  Track --> Reg[Registry]
  Reg --> Deploy[Deploy alias]
  Deploy --> OTel[Prod OTel traces]
  OTel --> Sample[Online eval sample]
  Sample --> Track

Prod traces (OpenTelemetry for LLMs) and experiment trackers are complements: traces debug a request; MLflow/W&B compare versions.

How to use (shape)

# Shape only — APIs evolve; pin mlflow version
import mlflow

mlflow.set_experiment("rag-refund-bot")
with mlflow.start_run(tags={"git": "abc123", "suite": "refund-v3"}):
    mlflow.log_params({
        "embed_model": "text-embedding-3-small",
        "chunk_size": 512,
        "top_k": 5,
        "prompt_hash": "p_7f3a",
        "generator": "gpt-4.1-mini",
    })
    # ... run harness ...
    mlflow.log_metrics({"recall_at_5": 0.81, "judge_pass": 0.74})
    mlflow.log_text(prompt_template, "prompt.txt")
    mlflow.log_dict(preds, "preds.json")
    # mlflow.log_model(...) / register_model(...) when packaging a chain

Signature habit: declare expected inputs (query, tenant_id) and outputs (answer, citation_ids) so serving and eval disagree loudly instead of silently.

Promotion sketch

1. Register version from the winning run
2. Assign staging alias
3. Run gated eval suite (evals fundamentals; guided Evals, guardrails, safety)
4. Flip production alias only on pass
5. Keep previous alias for instant rollback

sequenceDiagram
  participant Dev as Engineer
  participant MF as MLflow registry
  participant Gate as Eval gate
  participant Svc as Serving
  Dev->>MF: Register v12 from run
  Dev->>MF: alias staging = v12
  Dev->>Gate: Run suite refund-v3
  Gate-->>Dev: Pass
  Dev->>MF: alias production = v12
  Svc->>MF: Resolve production
  MF-->>Svc: v12

What “good lineage” looks like

Serving should resolve alias → immutable version, then stamp that version on every OTel span.

Decision tree: tracker vs registry vs traces

Failure modes

• Logging PII in preds tables — redact or hash
• Unversioned prompts — registry points at model weights only
• Floating latest in prod — always resolve alias → version number
• Eval suite drift — bump suite version when labels/rubrics change
• Offline win, online loss — tie prod prompt_version / model_version into OTel
• Adapter without base pin — cannot reproduce LoRA
• Treating tracker as source of truth for code — code stays in git; tracker stores pointers
• One giant experiment — cannot filter; use tags + naming conventions

Production checklist

1. Every promote goes through registry aliases — no ad-hoc file copies.
2. Prompts, chunkers, embedders, and eval suites versioned as first-class artifacts.
3. Signatures on logged chains.
4. Redaction policy for preds / traces artifacts.
5. Rollback drill: flip production to previous version in staging monthly.
6. Prod spans include model_version + prompt_version.
7. Dual-track with W&B for research UX is fine — one promote path to prod.

Alternatives and companions

Databricks GenAI paths often wire MLflow registry → serving. You can recreate the pattern elsewhere: immutable versions + aliases + eval gates.

Micro-project

1. Log a RAG chain run: params + metrics + prompt.txt artifact.
2. Register two versions; assign staging / production aliases.
3. Fail a deliberate eval gate; show rollback to previous alias.
4. Pair with Weights & Biases compare habits and Deploy, cost, latency, observability runbooks.

Related

Key Tech: Weights & Biases. Core: evals fundamentals. Guided Evals, guardrails, safety and Deploy, cost, latency, observability. Advanced: fine-tuning LoRA/QLoRA.