MLflow, Model Serving and Governance for Generative AI

Table of Contents

1. Professional MLflow Workflows for LLMs, RAG, and Agents
   • 1.1 MLflow for Generative AI
   • 1.2 Logging RAG Chains with MLflow
   • 1.3 GenAI Evaluation with LLM-as-judge
   • 1.4 MLflow Tracing: Identifying and Fixing Failing Requests
2. Zero-downtime Model Serving for Generative AI
   • 2.1 Mosaic AI Model Serving Architecture
   • 2.2 Deploying and Monitoring a RAG Endpoint in Production
3. Mosaic AI Gateway: Enterprise LLM Routing and Governance
   • 3.1 AI Gateway Overview and Unified Access
   • 3.2 Cost Control, Rate Limiting, and Fallbacks
   • 3.3 Configuring Routes, Limits, and Fallbacks
4. End-to-end Governance and Lineage with Unity Catalog
   • 4.1 Registering GenAI Assets in Unity Catalog
   • 4.2 Access Control and Permissions for AI Assets
   • 4.3 Complete Lineage and Audit Trail

1. Professional MLflow Workflows for LLMs, RAG, and Agents

1.1 MLflow for Generative AI

MLflow has evolved significantly to meet the unique requirements of generative AI applications. The platform builds on traditional machine learning tracking while adding specialized features for modern workflows.

MLflow Architecture for GenAI

graph TB
    subgraph Traditional["🔵 Traditional ML Tracking"]
        P[Parameters & Metrics Logging]
        A[Model Artifacts Storage]
        E[Basic Experiment Comparison]
    end

    subgraph GenAI["🟣 GenAI-optimized Features"]
        PV[Prompt Versioning & Templates]
        LJ[LLM-as-judge Evaluation]
        TR[Real-time Tracing for RAG / Agents]
    end

    subgraph Core["MLflow Core"]
        Traditional --> GenAI
    end

    subgraph Integrations["Integrations"]
        UC[Unity Catalog Integration]
        LC[LangChain Flavor]
        TF[Transformers Flavor]
        PF[PyFunc Flavor]
    end

    GenAI --> Integrations

MLflow Flavors for Generative AI

The three available flavors cover all GenAI use cases:

graph LR
    subgraph Flavors["MLflow Flavors"]
        TF["🤗 Transformers\n─────────────\nHuggingFace models\nAutomatic tokenizer\nhandling"]
        LC["🦜 LangChain\n─────────────\nFull chain serialization\nIncluding retrievers\nand prompts"]
        PF["⚙️ PyFunc\n─────────────\nCustom inference\nlogic for complex\nRAG pipelines"]
    end

Logged Models and Prompt Engineering

graph LR
    subgraph LM["Logged Models"]
        B1[Bundle code, metrics & params]
        B2[Automatic signature inference]
        B3[One-click deployment ready]
    end

    subgraph PT["Prompt Templates"]
        P1[Version-controlled prompts]
        P2[A/B testing capabilities]
        P3[Template variables for\ndynamic content]
    end

    LM -->|deployment| PROD[Production Endpoint]
    PT -->|versioning| LM

LLM-as-judge Evaluation and Unity Catalog Integration

graph TB
    subgraph LLMJudge["LLM-as-judge Evaluation"]
        QS[Automated quality scoring]
        RM[Relevance metrics]
        CT[Cost & token tracking]
    end

    subgraph UC["Unity Catalog Integration"]
        CMR[Centralized model registry]
        GAC[Governed access control]
        FLT[Full lineage tracking]
    end

    LLMJudge --> UC
    UC --> COMPLIANCE[✅ Compliance & Governance]

1.2 Logging RAG Chains with MLflow

Logging RAG chains is a foundational skill for building production-ready retrieval-augmented generation systems.

Best Practices for MLflow Experiments

graph LR
    subgraph ALWAYS["✅ DO"]
        A1[Use explicit run contexts]
        A2[Log input_example for signatures]
        A3[Version prompt templates]
    end

    subgraph AVOID["❌ AVOID"]
        B1[Hardcoded model parameters]
        B2[Ignoring signature inference]
        B3[Experiments without tags]
    end

RAG Chain Architecture

graph TD
    DOC[📄 Source Documents] --> VS[Vector Store\nChroma / FAISS / Databricks VS]
    VS --> RET[Retriever\nk top documents]
    QUERY[❓ User question] --> RET
    RET --> CHAIN[RAG Chain\nLangChain]
    PROMPT[📝 Prompt Template\nversion-controlled] --> CHAIN
    LLM[🤖 LLM\nGPT-4o / DBRX] --> CHAIN
    CHAIN --> ANS[✅ Answer]

    MLFLOW[🔵 MLflow Logging] -.->|log_model| CHAIN
    MLFLOW -.->|log_param| RET
    MLFLOW -.->|log_param| LLM
    MLFLOW -.->|log_param| PROMPT

Code Pattern: Logging a RAG Chain

import mlflow
import mlflow.langchain
from langchain.chains import RetrievalQA
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings, ChatOpenAI

# Configure MLflow with Unity Catalog
mlflow.set_registry_uri("databricks-uc")
mlflow.set_experiment("/Users/your-user/rag-chain-logging")

# Build the RAG chain
embedding_model = OpenAIEmbeddings(model="text-embedding-ada-002")
vectorstore = Chroma.from_documents(documents, embedding_model)
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
llm = ChatOpenAI(model="gpt-4o", temperature=0.1)

rag_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=retriever
)

# Input example for signature
sample_query = {"query": "What are MLflow's tracking capabilities?"}

# Logging with explicit run context
with mlflow.start_run(run_name="rag-chain-demo-v1"):
    # Log important parameters
    mlflow.log_param("embedding_model", "text-embedding-ada-002")
    mlflow.log_param("llm_endpoint", "gpt-4o")
    mlflow.log_param("prompt_version", "v1.0")
    mlflow.log_param("retriever_k", 3)
    mlflow.log_param("temperature", 0.1)

    # Log the RAG chain model
    mlflow.langchain.log_model(
        rag_chain,
        artifact_path="rag_chain",
        input_example=sample_query
    )

Note: Configuring the experiment with Unity Catalog (databricks-uc) ensures centralized and governed tracking.

Artifact Structure in the MLflow UI

mlflow-run/
├── artifacts/
│   └── rag_chain/
│       ├── MLmodel          ← signature, flavors, metadata
│       ├── model/           ← chain serialization
│       ├── requirements.txt
│       └── python_env.yaml
├── params/
│   ├── embedding_model
│   ├── llm_endpoint
│   ├── prompt_version
│   ├── retriever_k
│   └── temperature
└── tags/

1.3 GenAI Evaluation with LLM-as-judge

The LLM-as-judge concept uses a powerful language model to automatically evaluate generated outputs, enabling quality assurance at scale.

Evaluation Process Overview

graph TD
    EVAL_DATA[📊 Evaluation dataset\nQuestions + expected answers]
    MODEL[🤖 RAG model to evaluate]
    JUDGE[⚖️ LLM Judge\nGPT-4o / Claude]

    EVAL_DATA --> MLF_EVAL[mlflow.evaluate]
    MODEL --> MLF_EVAL
    JUDGE --> MLF_EVAL

    MLF_EVAL --> M1[📈 Relevance Score]
    MLF_EVAL --> M2[🎯 Groundedness Score]
    MLF_EVAL --> M3[🚫 Toxicity Score]
    MLF_EVAL --> M4[💰 Cost & Token Tracking]
    MLF_EVAL --> M5[🔍 Retrieval Quality Score]

    M1 & M2 & M3 & M4 & M5 --> UI[MLflow UI\nComparative dashboard]

Key LLM-as-judge Metrics

Code Pattern: mlflow.evaluate with LLM-as-judge

import mlflow
import pandas as pd

# Evaluation dataset
eval_data = pd.DataFrame({
    "inputs": [
        "What are MLflow's tracking capabilities?",
        "How does model versioning work?",
        "What is an MLflow model signature?",
    ],
    "ground_truth": [
        "MLflow tracks parameters, metrics and artifacts...",
        "Model versioning allows you to...",
        "The signature defines the input/output schema...",
    ]
})

# Reference to the logged model
logged_model = "runs:/<run_id>/rag_chain"

# Launch evaluation with LLM-as-judge
with mlflow.start_run(run_name="evaluation-v1"):
    results = mlflow.evaluate(
        model=logged_model,
        data=eval_data,
        targets="ground_truth",
        model_type="question-answering",
        evaluators="default",
        extra_metrics=[
            mlflow.metrics.genai.relevance(),
            mlflow.metrics.genai.groundedness(),
            mlflow.metrics.genai.correctness(),
        ]
    )

print(f"Average Relevance: {results.metrics['relevance/v1/mean']:.2f}")
print(f"Average Groundedness: {results.metrics['groundedness/v1/mean']:.2f}")
print(results.tables["eval_results_table"])

Comparing Two Runs in the MLflow UI

graph LR
    subgraph Run1["Run v1 - Baseline"]
        R1_REL["Relevance: 3.2"]
        R1_GRD["Groundedness: 2.8"]
        R1_COR["Correctness: 3.0"]
        R1_K["retriever_k = 1"]
    end

    subgraph Run2["Run v2 - Improved"]
        R2_REL["Relevance: 4.1 ✅"]
        R2_GRD["Groundedness: 3.9 ✅"]
        R2_COR["Correctness: 4.2 ✅"]
        R2_K["retriever_k = 3"]
    end

    Run1 -->|MLflow UI Comparison| Run2

1.4 MLflow Tracing: Identifying and Fixing Failing Requests

MLflow Tracing provides complete visibility into every step of the GenAI pipeline, enabling precise identification and correction of retrieval and generation issues.

What MLflow Tracing Captures

graph TB
    QUERY[❓ User request] --> TRACE_ROOT[Span: RetrievalQA Root]

    TRACE_ROOT --> SPAN1[Span: VectorStoreRetriever\n📊 Timing: 245ms\n🔍 k=1 document retrieved]
    TRACE_ROOT --> SPAN2[Span: StuffDocumentsChain\n📊 Timing: 15ms\n📄 Context formatting]
    TRACE_ROOT --> SPAN3[Span: LLMChain\n📊 Timing: 1.2s]
    SPAN3 --> SPAN4[Span: ChatDatabricks\n📊 Tokens: 847 input / 124 output\n💰 Cost: $0.003]

    SPAN4 --> RESPONSE[✅ Generated response]

    style SPAN1 fill:#ff9999
    style SPAN4 fill:#99ff99

Autologging vs. Manual Tracing

graph LR
    subgraph Auto["🤖 Autologging\n(Zero-code setup)"]
        A1[mlflow.langchain.autolog]
        A2[Automatic span creation]
        A3[Compatible with LangChain / OpenAI]
    end

    subgraph Manual["⚙️ Manual Tracing\n(Fine-grained control)"]
        M1[mlflow.start_span]
        M2[Custom span attributes]
        M3[Complex custom logic]
    end

Code Pattern: Autologging with Tracing

import mlflow

# Single line to enable full tracing
mlflow.langchain.autolog(log_traces=True)

# Configure the experiment
mlflow.set_experiment("/Users/your-user/rag-tracing")

# Requests are automatically traced
response1 = rag_chain.invoke(
    {"query": "What are MLflow's tracking capabilities?"}
)
print(f"Trace ID: {mlflow.get_last_active_trace().info.trace_id}")

response2 = rag_chain.invoke(
    {"query": "How does model versioning work?"}
)

# Potentially failing request
response3 = rag_chain.invoke(
    {"query": "What are the rate limiting features?"}
)

Code Pattern: Manual Tracing for Custom Logic

import mlflow
from mlflow.tracing import trace

@trace
def custom_rag_pipeline(query: str) -> str:
    """RAG pipeline with detailed manual tracing."""

    # Retrieval span
    with mlflow.start_span("retrieval") as retrieval_span:
        retrieval_span.set_attributes({
            "query": query,
            "k": 3,
            "vector_store": "databricks-vector-search"
        })
        docs = retriever.get_relevant_documents(query)
        retrieval_span.set_attributes({
            "num_docs_retrieved": len(docs),
            "doc_sources": [d.metadata.get("source", "") for d in docs]
        })

    # Generation span
    with mlflow.start_span("generation") as gen_span:
        context = "\n\n".join([doc.page_content for doc in docs])
        prompt = f"Context:\n{context}\n\nQuestion: {query}\n\nAnswer:"
        response = llm.invoke(prompt)

        gen_span.set_attributes({
            "tokens_input": response.usage_metadata.get("input_tokens", 0),
            "tokens_output": response.usage_metadata.get("output_tokens", 0),
            "model": "gpt-4o"
        })

    return response.content

# Usage
result = custom_rag_pipeline("What is the fallback mechanism?")

Diagnosing and Fixing a Retrieval Issue

graph TD
    PROBLEM[🔴 Request with incorrect response] --> TRACE_UI[Open MLflow Tracing UI]
    TRACE_UI --> INSPECT[Inspect the span\nVectorStoreRetriever]
    INSPECT --> DOCS_FOUND{Relevant\ndocuments\nretrieved?}
    DOCS_FOUND -->|No - k=1 insufficient| FIX[Increase k\nfrom 1 to 3]
    DOCS_FOUND -->|Yes| CHECK_PROMPT[Check the prompt]
    FIX --> NEW_CHAIN[Recreate the chain\nwith k=3]
    NEW_CHAIN --> VERIFY[Verify improvement\nnew trace]
    VERIFY --> COMPARE[✅ Compare the traces]

# Fix: increase k to retrieve more documents
improved_retriever = vectorstore.as_retriever(
    search_kwargs={"k": 3}  # Increased from 1 to 3
)

improved_rag_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=improved_retriever
)

# Verify improvement on the failing request
improved_response = improved_rag_chain.invoke(
    {"query": "What are the rate limiting features?"}
)

Privacy note: On Databricks, telemetry is disabled by default, ensuring data security.

2. Zero-downtime Model Serving for Generative AI

2.1 Mosaic AI Model Serving Architecture

Mosaic AI Model Serving provides enterprise-grade infrastructure for deploying GenAI models in production with no service interruption.

Architecture Overview

graph TB
    subgraph CLIENT["Client Applications"]
        APP1[Web Application]
        APP2[Backend Service]
        APP3[Batch Job]
    end

    subgraph GATEWAY["Load Balancer / Gateway"]
        LB[Intelligent Router\nToken-based Authentication]
    end

    subgraph SERVING["Mosaic AI Model Serving"]
        subgraph RT["Real-time Inference"]
            EP1[Endpoint 1\nServerless GPU]
            EP2[Endpoint 2\nServerless GPU]
        end
        subgraph BATCH["Batch Inference"]
            BJ[Batch Job\nHigh-throughput]
        end
        subgraph CACHE["Optimizations"]
            SC[Semantic Cache]
            WP[Warm Pools]
        end
    end

    subgraph UC_INT["Unity Catalog"]
        MR[Model Registry]
        LIN[Lineage Tracking]
    end

    APP1 & APP2 --> LB
    APP3 --> BJ
    LB --> RT
    RT --> SC
    SC --> WP
    RT -.->|governance| UC_INT

Serverless GPU Infrastructure

graph LR
    subgraph GPU["🖥️ Serverless GPU Infrastructure"]
        N1[No cluster provisioning\nrequired]
        N2[Automatic GPU allocation\non demand]
        N3[Pay-per-use\nonly]
    end

    subgraph SCALE["📈 Autoscaling"]
        S1[Scale-to-zero\nwhen idle]
        S2[Warm Pools\nfor fast startup]
        S3[Configurable triggers]
    end

    subgraph BATCH_CONT["⚡ Continuous Batching"]
        B1[Dynamic request\nbatching]
        B2[Maximize GPU\nutilization]
        B3[Reduce latency\nper request]
    end

Pricing Model and Latency Management

Key Metrics for GenAI Endpoints

graph TD
    subgraph METRICS["📊 Monitoring Metrics"]
        M1["🔢 Tokens/second\n(throughput)"]
        M2["📋 Request queue depth\n(bottlenecks)"]
        M3["🖥️ GPU utilization %\n(resource optimization)"]
        M4["⏱️ Latency percentiles\nP50 / P95 / P99"]
    end

    ENDPOINT[Production Endpoint] --> METRICS
    METRICS --> ALERTING[🔔 Alerting &\nScaling Decisions]

Real-time vs. Batch Inference

graph LR
    subgraph RT["⚡ Real-time Inference"]
        RT1[Sub-second responses]
        RT2[Interactive applications]
        RT3[Streaming support]
    end

    subgraph BT["📦 Batch Inference"]
        BT1[High-volume processing]
        BT2[Offline workloads]
        BT3[Cost-optimized]
    end

    USE_CASE{Use case} -->|Chatbot / API| RT
    USE_CASE -->|Nightly processing| BT

2.2 Deploying and Monitoring a RAG Endpoint in Production

5-Step Deployment Workflow

graph TD
    STEP1["1️⃣ Register the model\nin Unity Catalog"]
    STEP2["2️⃣ Create a serving endpoint\n(UI or API)"]
    STEP3["3️⃣ Configure compute\nand scaling"]
    STEP4["4️⃣ Test via SDK\nand REST API"]
    STEP5["5️⃣ Monitor the\nmetrics dashboard"]

    STEP1 --> STEP2 --> STEP3 --> STEP4 --> STEP5
    STEP5 -->|Continuous optimization| STEP3

Complete Code: Registration and Deployment

import mlflow
from mlflow.deployments import get_deploy_client

# ─── Step 1: Register in Unity Catalog ──────────────────────────────────────
mlflow.set_registry_uri("databricks-uc")

# Find the latest run
client_mlflow = mlflow.tracking.MlflowClient()
experiment = client_mlflow.get_experiment_by_name(
    "/Users/your-user/rag-chain-logging"
)
runs = client_mlflow.search_runs(
    experiment_ids=[experiment.experiment_id],
    order_by=["start_time DESC"],
    max_results=1
)
run_id = runs[0].info.run_id
print(f"Run found: {run_id} — {runs[0].data.tags.get('mlflow.runName')}")

# Register in Unity Catalog
model_uri = f"runs:/{run_id}/rag_chain"
registered = mlflow.register_model(
    model_uri=model_uri,
    name="workspace.genai_models.rag_chain_model"
)
print(f"Model registered at version: {registered.version}")

# Add description and alias
client_mlflow.update_registered_model(
    name="workspace.genai_models.rag_chain_model",
    description="RAG chain for Q&A on internal documents — production"
)
client_mlflow.set_registered_model_alias(
    name="workspace.genai_models.rag_chain_model",
    alias="production",
    version=str(registered.version)
)

# ─── Step 2: Create / update the endpoint ────────────────────────────────────
deploy_client = get_deploy_client("databricks")

endpoint_config = {
    "served_entities": [{
        "entity_name": "workspace.genai_models.rag_chain_model",
        "entity_version": str(registered.version),
        "workload_size": "Small",
        "scale_to_zero_enabled": True,
        "environment_vars": {
            "DATABRICKS_TOKEN": "{{secrets/rag-scope/databricks-token}}"
        }
    }],
    "traffic_config": {
        "routes": [{
            "served_model_name": f"rag_chain_model-{registered.version}",
            "traffic_percentage": 100
        }]
    }
}

try:
    endpoint = deploy_client.create_endpoint(
        name="rag-chain-endpoint-demo",
        config=endpoint_config
    )
    print("Endpoint created successfully")
except Exception:
    endpoint = deploy_client.update_endpoint(
        endpoint="rag-chain-endpoint-demo",
        config=endpoint_config
    )
    print("Endpoint updated successfully")

Code Pattern: Testing via SDK and REST API

import requests
import json

ENDPOINT_NAME = "rag-chain-endpoint-demo"
WORKSPACE_URL = "https://<your-workspace>.azuredatabricks.net"

# ─── Test via MLflow SDK ─────────────────────────────────────────────────────
deploy_client = get_deploy_client("databricks")

response_sdk = deploy_client.predict(
    endpoint=ENDPOINT_NAME,
    inputs={"query": "What are MLflow's tracking capabilities?"}
)
print("SDK Response:", response_sdk)

# ─── Test via REST API ───────────────────────────────────────────────────────
TOKEN = dbutils.secrets.get(scope="rag-scope", key="databricks-token")

headers = {
    "Authorization": f"Bearer {TOKEN}",
    "Content-Type": "application/json"
}
payload = {
    "inputs": {"query": "How does model versioning work?"}
}

response_rest = requests.post(
    f"{WORKSPACE_URL}/serving-endpoints/{ENDPOINT_NAME}/invocations",
    headers=headers,
    data=json.dumps(payload)
)
print(f"Status: {response_rest.status_code}")
print("REST Response:", response_rest.json())

# ─── Generate test traffic for metrics ───────────────────────────────────────
test_questions = [
    "What is MLflow?",
    "How to register a model in Unity Catalog?",
    "What inference types are supported?",
    "How to configure autoscaling?",
    "What metrics to monitor in production?"
]

for i, question in enumerate(test_questions, 1):
    r = deploy_client.predict(
        endpoint=ENDPOINT_NAME,
        inputs={"query": question}
    )
    print(f"Request {i}/{len(test_questions)} completed")

Retrieve Endpoint Configuration Programmatically

# Inspect the current configuration
endpoint_info = deploy_client.get_endpoint(ENDPOINT_NAME)

print(f"Name: {endpoint_info['name']}")
print(f"Status: {endpoint_info['state']['ready']}")
print(f"Created by: {endpoint_info.get('creator', 'N/A')}")

for entity in endpoint_info['config']['served_entities']:
    print(f"\nServed model: {entity['entity_name']}")
    print(f"  Version: {entity['entity_version']}")
    print(f"  Workload size: {entity['workload_size']}")
    print(f"  Scale-to-zero: {entity['scale_to_zero_enabled']}")

3. Mosaic AI Gateway: Enterprise LLM Routing and Governance

3.1 AI Gateway Overview and Unified Access

The AI Gateway provides a unified interface to manage access to foundation models from different providers, centralizing security, costs, and monitoring.

AI Gateway Architecture

graph TB
    subgraph APPS["Applications"]
        A1[Python App]
        A2[Databricks Notebook]
        A3[External Service]
    end

    subgraph GATEWAY["🛡️ Mosaic AI Gateway\nSingle entry point"]
        AUTH[Authentication\nToken-based]
        PII[PII Filtering\nDetection & Redaction]
        CACHE[Semantic Cache\nReducing API calls]
        RL[Rate Limiting\nPer user / team / endpoint]
        ROUTE[Route Engine\nPremium / Cost / Fallback]
    end

    subgraph PROVIDERS["LLM Providers"]
        OAI[OpenAI\nGPT-4o / GPT-3.5]
        ANT[Anthropic\nClaude 3.x]
        DBX[Databricks\nDBRX / LLaMA]
        AZ[Azure OpenAI]
    end

    subgraph TRACKING["Tracking & Governance"]
        UC[Unity Catalog\nUsage logs]
        DASH[Cost Dashboard]
    end

    APPS --> GATEWAY
    GATEWAY --> PROVIDERS
    GATEWAY --> TRACKING

    style GATEWAY fill:#e8f4f8,stroke:#2196F3

Key Gateway Capabilities

3.2 Cost Control, Rate Limiting, and Fallbacks

graph TB
    subgraph RATE["🚦 Rate Limiting"]
        RL1[Per-user limits]
        RL2[Per-team quotas]
        RL3[Per-endpoint controls]
    end

    subgraph SPEND["💰 Spend Caps"]
        SC1[Budget enforcement]
        SC2[Preventing excessive costs]
        SC3[Threshold alerts]
    end

    subgraph FALLBACK["🔄 Automatic Fallback"]
        FB1[Provider failover\nfor high availability]
        FB2[Graceful degradation\nif provider unavailable]
        FB3[Zero-downtime routing]
    end

    subgraph COST_MGMT["📊 Cost Management"]
        CM1[Unified billing visibility]
        CM2[Cross-provider comparison]
        CM3[Tracking in Unity Catalog]
    end

3.3 Configuring Routes, Limits, and Fallbacks

Route Types

graph LR
    subgraph ROUTES["AI Gateway Routes"]
        PREM["🏆 Premium Route\nGPT-4o / Claude 3 Opus\nQuality-critical applications\n80% traffic"]
        COST["💡 Cost-optimized Route\nGPT-3.5 / LLaMA 3 8B\nHigh volumes\n15% traffic"]
        FALL["🔄 Fallback Route\nGPT-4o Mini\nHigh availability\n5% traffic"]
    end

    REQUEST[Incoming request] --> ROUTER{Route Engine}
    ROUTER -->|High quality required| PREM
    ROUTER -->|Volume / cost| COST
    ROUTER -->|Primary provider down| FALL

Code Pattern: YAML Route Configuration

# AI Gateway route configuration
routes:
  - name: premium
    model:
      name: gpt-4o
      provider: openai
    rate_limit:
      calls: 1000
      renewal_period: day
    traffic_percentage: 80

  - name: cost-optimized
    model:
      name: llama-3-8b-instruct
      provider: databricks
    rate_limit:
      calls: 10000
      renewal_period: day
    traffic_percentage: 15

  - name: fallback
    model:
      name: gpt-4o-mini
      provider: openai
    priority: 2
    traffic_percentage: 5

Code Pattern: Creating Routes via Python SDK

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import (
    EndpointCoreConfigInput,
    ServedEntityInput,
    ExternalModel,
    ExternalModelProvider,
    OpenAiConfig,
    AiGatewayConfig,
    AiGatewayRateLimit,
    AiGatewayRateLimitRenewalPeriod,
    AiGatewayUsageTrackingConfig,
    AiGatewayGuardrails,
    AiGatewayGuardrailParameters,
    AiGatewayGuardrailPiiBehavior,
    AiGatewayGuardrailPiiBehaviorBehavior,
)

w = WorkspaceClient()

# ─── Premium Route ───────────────────────────────────────────────────────────
premium_config = {
    "name": "gateway-premium-route",
    "config": {
        "served_entities": [{
            "external_model": {
                "name": "gpt-4o",
                "provider": "openai",
                "task": "llm/v1/chat",
                "openai_config": {
                    "openai_api_key": "{{secrets/openai/api_key}}"
                }
            }
        }],
        "ai_gateway": {
            "rate_limits": [
                {"calls": 1000, "renewal_period": "day", "key": "endpoint"},
                {"calls": 100,  "renewal_period": "minute", "key": "user"}
            ],
            "usage_tracking_config": {"enabled": True}
        }
    }
}

# ─── Cost-optimized Route ────────────────────────────────────────────────────
cost_config = {
    "name": "gateway-cost-route",
    "config": {
        "served_entities": [{
            "entity_name": "databricks-meta-llama-3-1-8b-instruct",
            "workload_size": "Small",
            "scale_to_zero_enabled": True
        }],
        "ai_gateway": {
            "rate_limits": [
                {"calls": 10000, "renewal_period": "day", "key": "endpoint"}
            ],
            "usage_tracking_config": {"enabled": True}
        }
    }
}

# ─── Fallback Route ──────────────────────────────────────────────────────────
fallback_config = {
    "name": "gateway-fallback-route",
    "config": {
        "served_entities": [
            {
                "external_model": {
                    "name": "gpt-4o",
                    "provider": "openai",
                    "task": "llm/v1/chat",
                    "openai_config": {"openai_api_key": "{{secrets/openai/api_key}}"}
                }
            },
            {
                "external_model": {
                    "name": "gpt-4o-mini",
                    "provider": "openai",
                    "task": "llm/v1/chat",
                    "openai_config": {"openai_api_key": "{{secrets/openai/api_key}}"}
                }
            }
        ],
        "ai_gateway": {
            "usage_tracking_config": {"enabled": True}
        }
    }
}

# Create endpoints
for config in [premium_config, cost_config, fallback_config]:
    w.serving_endpoints.create(**config)
    print(f"Route '{config['name']}' created")

Configuring Guardrails and PII Filtering

# Configuration with complete PII guardrails
guardrails_config = {
    "name": "protected-llm-endpoint",
    "config": {
        "served_entities": [{
            "external_model": {
                "name": "gpt-4o",
                "provider": "openai",
                "task": "llm/v1/chat",
                "openai_config": {
                    "openai_api_key": "{{secrets/openai/api_key}}"
                }
            }
        }],
        "ai_gateway": {
            "guardrails": {
                "input": {
                    "pii": {
                        "behavior": "BLOCK",          # Block if PII detected in input
                        "pii_types": [
                            "EMAIL_ADDRESS",
                            "PHONE_NUMBER",
                            "SSN",
                            "CREDIT_CARD",
                            "IP_ADDRESS",
                            "PERSON_NAME"
                        ]
                    },
                    "safety": "NONE"
                },
                "output": {
                    "pii": {
                        "behavior": "REDACT"          # Redact PII in output
                    }
                }
            },
            "rate_limits": [
                {"calls": 100, "renewal_period": "minute", "key": "user"},
                {"calls": 5000, "renewal_period": "day",   "key": "endpoint"}
            ],
            "usage_tracking_config": {"enabled": True}
        }
    }
}

Using the AI Gateway for Requests

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import ChatMessage, ChatMessageRole

w = WorkspaceClient()

# Send a message via the gateway
response = w.serving_endpoints.query(
    name="databricks-meta-llama",
    messages=[
        ChatMessage(
            role=ChatMessageRole.SYSTEM,
            content="You are an expert assistant in MLflow and Databricks."
        ),
        ChatMessage(
            role=ChatMessageRole.USER,
            content="Explain the concept of model serving with Mosaic AI."
        )
    ],
    max_tokens=512,
    temperature=0.7
)

print(response.choices[0].message.content)

4. End-to-end Governance and Lineage with Unity Catalog

4.1 Registering GenAI Assets in Unity Catalog

Unity Catalog is the centralized registry for all GenAI assets, ensuring discoverability, versioning, and access control.

Supported Asset Types

graph TB
    subgraph UC["🏛️ Unity Catalog"]
        subgraph CATALOG["Catalog: genai_catalog"]
            subgraph MODELS_SCHEMA["Schema: genai_models"]
                M1[rag_chain_v1\nFine-tuned RAG model]
                M2[summarizer_v2\nSummarization model]
            end
            subgraph PROMPTS_SCHEMA["Schema: genai_prompts"]
                P1[rag_qa_prompt\nQ&A template]
                P2[summarize_prompt\nSummary template]
            end
            subgraph VECTORS_SCHEMA["Schema: genai_vectors"]
                V1[docs_vector_index\nVector search index]
                V2[faq_vector_index]
            end
            subgraph DATA_SCHEMA["Schema: genai_data"]
                D1[documents_table\nSource data]
                D2[eval_dataset]
            end
        end
    end

Organization Best Practices

Complete Code: Registering Assets

import mlflow
from pyspark.sql import SparkSession

spark = SparkSession.builder.getOrCreate()

# ─── Step 1: Prepare the Unity Catalog environment ──────────────────────────
mlflow.set_registry_uri("databricks-uc")

# Create catalog and schema via SQL
spark.sql("CREATE CATALOG IF NOT EXISTS genai_catalog")
spark.sql("CREATE SCHEMA IF NOT EXISTS genai_catalog.genai_models")
spark.sql("CREATE SCHEMA IF NOT EXISTS genai_catalog.genai_prompts")
spark.sql("CREATE SCHEMA IF NOT EXISTS genai_catalog.genai_vectors")

# ─── Step 2: Register the model ─────────────────────────────────────────────
client = mlflow.tracking.MlflowClient()

# Find the latest run for the RAG experiment
experiment = client.get_experiment_by_name(
    "/Users/your-user/rag-chain-logging"
)
runs = client.search_runs(
    experiment_ids=[experiment.experiment_id],
    filter_string="tags.mlflow.runName = 'rag-chain-demo-v1'",
    order_by=["start_time DESC"],
    max_results=1
)
run_id = runs[0].info.run_id
print(f"Run found: {runs[0].data.tags['mlflow.runName']}")

# Register in Unity Catalog
registered = mlflow.register_model(
    model_uri=f"runs:/{run_id}/rag_chain",
    name="genai_catalog.genai_models.rag_chain_v1"
)
print(f"Model created: version {registered.version}")

# ─── Step 3: Metadata and aliases ───────────────────────────────────────────
client.update_registered_model(
    name="genai_catalog.genai_models.rag_chain_v1",
    description="RAG chain for Q&A on internal documentation. "
                "Uses DBRX as LLM and Databricks Vector Search."
)

# Aliases for different environments
client.set_registered_model_alias(
    name="genai_catalog.genai_models.rag_chain_v1",
    alias="production",
    version=str(registered.version)
)

# ─── Step 4: Register the prompt as a Unity Catalog function ─────────────────
spark.sql("""
CREATE OR REPLACE FUNCTION genai_catalog.genai_prompts.rag_qa_prompt(
    context STRING,
    question STRING
)
RETURNS STRING
LANGUAGE PYTHON
AS $$
    template = f'''You are an expert assistant. Answer the question based solely on the provided context.

Context:
{context}

Question: {question}

Instructions:
- Use only the information from the provided context
- If the answer is not in the context, respond "I don't know"
- Be precise and concise

Answer:'''
    return template
$$
""")

# Test the function immediately
result = spark.sql("""
    SELECT genai_catalog.genai_prompts.rag_qa_prompt(
        'MLflow is a platform for managing the ML model lifecycle.',
        'What is MLflow?'
    ) AS formatted_prompt
""")
result.show(truncate=False)

# ─── Step 5: View all registered assets ─────────────────────────────────────
print("\n=== Registered models ===")
for model in client.search_registered_models(
    filter_string="name LIKE 'genai_catalog.genai_models.%'"
):
    print(f"  {model.name} — {len(model.latest_versions)} version(s)")

print("\n=== Registered functions (prompts) ===")
functions = spark.sql(
    "SHOW FUNCTIONS IN genai_catalog.genai_prompts"
)
functions.show()

Recommended Unity Catalog Structure

genai_catalog/
├── genai_models/                 ← ML/LLM models
│   ├── rag_chain_v1              (alias: production → v4)
│   └── summarizer_v1             (alias: staging → v2)
├── genai_prompts/                ← Prompt templates
│   ├── rag_qa_prompt()
│   └── summarize_prompt()
├── genai_vectors/                ← Vector search indexes
│   └── docs_vector_index
└── genai_data/                   ← Data tables
    ├── documents
    └── eval_dataset

4.2 Access Control and Permissions for AI Assets

RBAC Hierarchy in Unity Catalog

graph TD
    META[🏢 Metastore\nWorkspace level]
    META --> CAT["📚 Catalog\ngenai_catalog\nUSE CATALOG / BROWSE"]
    CAT --> SCH1["📂 Schema: genai_models\nUSE SCHEMA / CREATE MODEL"]
    CAT --> SCH2["📂 Schema: genai_prompts\nUSE SCHEMA"]
    CAT --> SCH3["📂 Schema: genai_vectors\nUSE SCHEMA"]
    SCH1 --> OBJ1["🤖 Model: rag_chain_v1\nEXECUTE"]
    SCH1 --> OBJ2["🤖 Model: summarizer_v1\nEXECUTE"]
    SCH2 --> OBJ3["⚙️ Function: rag_qa_prompt\nEXECUTE"]
    SCH3 --> OBJ4["🔍 Index: docs_vector_index\nSELECT"]

    style META fill:#fff3cd
    style CAT fill:#d4edda
    style SCH1 fill:#cce5ff
    style SCH2 fill:#cce5ff
    style SCH3 fill:#cce5ff

Permission Matrix by Role

Principle of least privilege: BROWSE allows metadata discovery without access to actual data.

Code Pattern: Applying RBAC

-- Step 1: Catalog access (top level)
GRANT USE CATALOG ON CATALOG genai_catalog TO `data-team`;
GRANT USE CATALOG ON CATALOG genai_catalog TO `app-team`;
GRANT BROWSE    ON CATALOG genai_catalog TO `all users`;

-- Step 2: Schema access
GRANT USE SCHEMA ON SCHEMA genai_catalog.genai_models  TO `data-team`;
GRANT USE SCHEMA ON SCHEMA genai_catalog.genai_models  TO `app-team`;
GRANT USE SCHEMA ON SCHEMA genai_catalog.genai_prompts TO `app-team`;

-- Step 3: Object-level permissions
-- data-team: read training data
GRANT SELECT ON TABLE genai_catalog.genai_data.documents TO `data-team`;
GRANT SELECT ON TABLE genai_catalog.genai_data.eval_dataset TO `data-team`;

-- app-team: execute models and functions
GRANT EXECUTE ON MODEL    genai_catalog.genai_models.rag_chain_v1   TO `app-team`;
GRANT EXECUTE ON FUNCTION genai_catalog.genai_prompts.rag_qa_prompt TO `app-team`;

-- ml-engineers: create and manage models
GRANT CREATE MODEL ON SCHEMA genai_catalog.genai_models TO `ml-engineers`;

# Verify current permissions
from pyspark.sql import SparkSession
spark = SparkSession.builder.getOrCreate()

print("=== Permissions on the catalog ===")
spark.sql("SHOW GRANTS ON CATALOG genai_catalog").show(truncate=False)

print("\n=== Permissions on schema genai_models ===")
spark.sql("SHOW GRANTS ON SCHEMA genai_catalog.genai_models").show(truncate=False)

print("\n=== Permissions on the model ===")
spark.sql(
    "SHOW GRANTS ON MODEL genai_catalog.genai_models.rag_chain_v1"
).show(truncate=False)

# Revoke permissions if necessary
spark.sql("""
    REVOKE EXECUTE ON MODEL genai_catalog.genai_models.rag_chain_v1
    FROM `former-team-member`
""")
print("Permissions revoked successfully")

4.3 Complete Lineage and Audit Trail

End-to-end Lineage Flow

graph LR
    subgraph DATA["📊 Source Data"]
        TBL[Table\ndocuments]
        RAW[Raw\nfiles]
    end

    subgraph EMBEDDINGS["🔢 Embeddings"]
        EMB[Embedding model\ntext-embedding-ada-002]
        VEC[Generated vectors]
    end

    subgraph VECTOR_IDX["🔍 Vector Index"]
        VI[docs_vector_index\nDatabricks Vector Search]
    end

    subgraph CHAIN["⛓️ RAG Chain"]
        RET[Retriever\nk=3]
        LLM[LLM\nDBRX / GPT-4o]
        PROMPT[Prompt Template\nv1.0]
        CHAIN_OBJ[rag_chain_v1\nv4 — alias: production]
    end

    subgraph ENDPOINT["🚀 Serving Endpoint"]
        EP[rag-chain-endpoint-demo\nversion 9 active]
    end

    RAW --> TBL
    TBL --> EMB
    EMB --> VEC
    VEC --> VI
    VI --> RET
    PROMPT --> CHAIN_OBJ
    RET --> CHAIN_OBJ
    LLM --> CHAIN_OBJ
    CHAIN_OBJ --> EP

    style EP fill:#d4edda,stroke:#28a745

MLflow Lineage Capabilities

graph TB
    subgraph LINEAGE["🔗 Lineage Tracking (Unity Catalog + MLflow)"]
        ALC[Automatic\nlineage capture]
        EPT[Experiment-to-production\ntracking]
        UAL[Unified audit\nlogging]
    end

    subgraph AUDIT["📋 Audit Capabilities"]
        WHO[Who accessed?]
        WHAT[What?]
        WHEN[When?]
        HOW[From where?\n(source IP)]
    end

    LINEAGE --> AUDIT
    AUDIT --> COMPLIANCE[✅ Compliance\nReporting]

Complete Code: Querying Lineage and Audit Logs

import mlflow
from mlflow.deployments import get_deploy_client

client = mlflow.tracking.MlflowClient()
deploy_client = get_deploy_client("databricks")
spark = SparkSession.builder.getOrCreate()

# ─── Step 1: Source → endpoint lineage ──────────────────────────────────────
print("=== Lineage: source tables ===")
lineage_query = spark.sql("""
    SELECT
        source_table_full_name,
        target_table_full_name,
        event_type
    FROM system.access.table_lineage
    WHERE target_table_full_name LIKE 'genai_catalog.%'
    ORDER BY event_time DESC
    LIMIT 20
""")
lineage_query.show(truncate=False)

# ─── Step 2: Audit logs ────────────────────────────────────────────────────
print("\n=== Audit logs for the last 30 days ===")
audit_logs = spark.sql("""
    SELECT
        event_time,
        user_name,
        action_name,
        request_params['asset_name']  AS asset_name,
        response.status_code          AS status_code,
        source_ip_address
    FROM system.access.audit
    WHERE action_name IN (
        'getModel',
        'predictServingEndpoint',
        'updatePermissions',
        'createServingEndpoint',
        'updateRegisteredModel'
    )
    AND event_time >= DATEADD(DAY, -30, CURRENT_TIMESTAMP())
    ORDER BY event_time DESC
    LIMIT 50
""")
audit_logs.show(truncate=False)

# ─── Step 3: Model version history ──────────────────────────────────────────
print("\n=== Model version history ===")
model_versions = client.search_model_versions(
    filter_string="name='genai_catalog.genai_models.rag_chain_v1'",
    order_by=["creation_timestamp DESC"],
    max_results=10
)
for v in model_versions:
    print(f"  v{v.version} | Status: {v.status} | "
          f"Created: {v.creation_timestamp} | Run: {v.run_id[:8]}...")

# ─── Step 4: MLflow experiment lineage ───────────────────────────────────────
print("\n=== Artifacts and parameters from the 5 latest runs ===")
import os
experiment = client.get_experiment_by_name(
    f"/Users/{os.environ.get('USER', 'your-user')}/rag-chain-logging"
)
recent_runs = client.search_runs(
    experiment_ids=[experiment.experiment_id],
    order_by=["start_time DESC"],
    max_results=5
)
for run in recent_runs:
    print(f"\n  Run: {run.data.tags.get('mlflow.runName', run.info.run_id[:8])}")
    artifacts = client.list_artifacts(run.info.run_id)
    print(f"  Artifacts: {[a.path for a in artifacts]}")
    print(f"  Key params: "
          f"llm={run.data.params.get('llm_endpoint', 'N/A')}, "
          f"prompt_v={run.data.params.get('prompt_version', 'N/A')}, "
          f"k={run.data.params.get('retriever_k', 'N/A')}")

# ─── Step 5: Endpoint → Unity Catalog model mapping ─────────────────────────
print("\n=== Active endpoint → UC model mapping ===")
endpoints = deploy_client.list_endpoints()
for ep in endpoints:
    ep_detail = deploy_client.get_endpoint(ep["name"])
    state = ep_detail.get("state", {}).get("ready", "UNKNOWN")
    for entity in ep_detail.get("config", {}).get("served_entities", []):
        model_name = entity.get("entity_name", "foundation-model")
        model_ver  = entity.get("entity_version", "N/A")
        print(f"  Endpoint: {ep['name']}")
        print(f"    Model: {model_name} (version {model_ver})")
        print(f"    Status: {state}\n")

Generate a Complete Lineage Report

def generate_lineage_report(catalog: str, models_schema: str, endpoint_name: str):
    """Generates a complete governance report for GenAI assets."""

    client = mlflow.tracking.MlflowClient()
    deploy_client = get_deploy_client("databricks")
    spark = SparkSession.builder.getOrCreate()

    print("=" * 60)
    print("GOVERNANCE REPORT — GENAI ASSETS")
    print("=" * 60)

    # 1. Data sources
    print("\n📊 DATA SOURCES")
    sources = spark.sql(f"""
        SELECT table_name, table_type, created
        FROM system.information_schema.tables
        WHERE table_catalog = '{catalog}'
        AND table_schema = 'genai_data'
    """)
    sources.show()

    # 2. Registered models
    print("\n🤖 REGISTERED MODELS")
    for model in client.search_registered_models(
        filter_string=f"name LIKE '{catalog}.{models_schema}.%'"
    ):
        print(f"  • {model.name}")
        for alias, ver in model.aliases.items():
            print(f"    Alias '{alias}' → version {ver}")

    # 3. Endpoint status
    print(f"\n🚀 ENDPOINT: {endpoint_name}")
    ep = deploy_client.get_endpoint(endpoint_name)
    print(f"  Status: {ep['state']['ready']}")

    # 4. Access controls
    print("\n🔒 ACCESS CONTROLS")
    grants = spark.sql(f"SHOW GRANTS ON CATALOG {catalog}")
    grants.show()

    # 5. Recent activity (audit)
    print("\n📋 RECENT ACCESS (7 days)")
    recent = spark.sql(f"""
        SELECT user_name, action_name, COUNT(*) AS access_count
        FROM system.access.audit
        WHERE action_name IN ('predictServingEndpoint', 'getModel')
        AND event_time >= DATEADD(DAY, -7, CURRENT_TIMESTAMP())
        GROUP BY user_name, action_name
        ORDER BY access_count DESC
    """)
    recent.show()

    print("\n✅ Compliance checklist:")
    print("  ✅ Data sources tracked")
    print("  ✅ Access patterns monitored")
    print("  ✅ Model changes audited")
    print("  ✅ Endpoint usage verified")

# Usage
generate_lineage_report(
    catalog="genai_catalog",
    models_schema="genai_models",
    endpoint_name="rag-chain-endpoint-demo"
)

Compliance Checklist

graph TD
    subgraph COMPLIANCE["✅ Compliance Checklist"]
        C1["✅ Data sources tracked\n(lineage source → endpoint)"]
        C2["✅ Access patterns monitored\n(who accesses what?)"]
        C3["✅ Model changes audited\n(version history)"]
        C4["✅ Endpoint usage verified\n(tokens, latency, cost)"]
        C5["✅ Permissions reviewed regularly\n(quarterly audit)"]
    end

General Summary

graph TB
    subgraph M1["Module 1 — MLflow Workflows"]
        LRC[Logging RAG Chains]
        EVAL[LLM-as-judge Evaluation]
        TRACE[MLflow Tracing]
    end

    subgraph M2["Module 2 — Model Serving"]
        MOSAIC[Mosaic AI\nServerless GPU]
        DEPLOY[Deploy RAG Endpoint]
        MONITOR[Metrics Monitoring]
    end

    subgraph M3["Module 3 — AI Gateway"]
        GATEWAY[Unified LLM Access]
        RATE[Rate Limiting & Fallback]
        GUARD[PII Filtering & Guardrails]
    end

    subgraph M4["Module 4 — Unity Catalog"]
        REGISTER[Register GenAI Assets]
        RBAC[RBAC & Permissions]
        LINEAGE[Lineage & Audit Trail]
    end

    M1 -->|versioned model| M2
    M2 -->|endpoint| M3
    M3 -->|governed access| M4
    M4 -->|compliance| M1

    style M1 fill:#e3f2fd
    style M2 fill:#f3e5f5
    style M3 fill:#e8f5e9
    style M4 fill:#fff3e0

Search Terms

mlflow · model · serving · governance · generative · ai · llmops · artificial · intelligence · pattern · lineage · architecture · catalog · gateway · rag · tracing · unity · genai · llm-as-judge · assets · evaluation · access · audit · autologging