05-logged-model-and-uc-registration

name: 05-logged-model-and-uc-registration description: > Use when you need to log an agent to MLflow, register it in Unity Catalog, or promote a new version as champion. Covers the full model lifecycle — even if you just want "log my agent and make it the production version." Also use when setting up MLFLOW_ACTIVE_MODEL_ID for production trace linking, metric-based promotion gating, or rollback. SDLC Step 5. license: Apache-2.0 compatibility: "Requires Databricks workspace with MLflow 3.10+ and Unity Catalog. Scripts use uv." clients: [ide_cli, genie_code] bundle_resource: none deploy_verb: none deploy_note: "Agent logged to MLflow and registered in UC (with champion alias) via the MLflow SDK; the UC model lands under the per-user prefixed catalog/schema. No bundle resource. Identical on both clients; on Genie Code run any CLI step through runDatabricksCli. See skills/genie-code-environment." coverage: full metadata: last_verified: "2026-06-05" volatility: high upstream_sources: [] author: "prashanth-subrahmanyam" version: "4.0.0" domain: "genai-agents" pipeline_position: "S5" consumes: "evaluation_results, thresholds_met, mlflow_run_id" produces: "logged_model_id, uc_model_version, champion_alias" grounded_in: "https://docs.databricks.com/en/mlflow/models-in-uc.html, https://docs.databricks.com/aws/en/mlflow3/genai/prompt-version-mgmt/version-tracking/link-production-traces-to-app-versions, https://docs.databricks.com/aws/en/mlflow3/genai/eval-monitor/concepts/eval-harness" upstream_sources: - name: "ai-dev-kit" repo: "databricks-solutions/ai-dev-kit" paths: - "databricks-skills/databricks-mlflow-evaluation/SKILL.md" relationship: "reference" last_synced: "2026-04-27" sync_commit: "281d9acd92d936bd5294f78bd7ec68fb12d4a696"

Logged Model and Unity Catalog Registration

Upstream Lineage

This skill references AI-Dev-Kit's databricks-mlflow-evaluation skill for metric-based promotion gates, evaluation result interpretation, and production trace/model lineage. If registration or champion promotion depends on scorer outputs or eval harness behavior, consult the upstream skill first, then apply this skill's LoggedModel and UC alias conventions.

When to Use

Use this skill when you need to:

• Persist a GenAI agent (or agent-like artifact) to MLflow with correct resources (serving endpoints, vector search, and so on) for lineage and deployment.
• Create LoggedModel entities that tie iterations, configs, and evaluation metrics together.
• Register a model version in Unity Catalog under a stable naming convention.
• Promote a champion only when evaluation metrics meet thresholds (gating), or use aliases for rollback and discovery.
• Link production traces to a specific logged model via MLFLOW_ACTIVE_MODEL_ID.

This step assumes upstream evaluation results, thresholds, and an MLflow run (SDLC Step 4). Downstream artifacts include logged model id, UC model version, and champion / alias state.

Track Applicability

• Track A (Custom Agent): Log the agent with mlflow.pyfunc.log_model(), register in UC, apply champion gating.
• Track B (Supervisor API): Model registration is optional — you may still register a wrapper for versioning and lineage.
• Track C (Model Serving): Model registration is typically mandatory for governed deployment.

Part A: Agent Logging

Official guidance: log agents with explicit dependencies via mlflow.models.resources, validate with mlflow.models.predict before production rollout. See Log an agent to MLflow.

Recommended default: Use mlflow.pyfunc.log_model() with explicit resources for most agents. Use mlflow.langchain.log_model() only if your agent is a LangChain chain/graph.

Code-based logging

import mlflow
from mlflow.models.resources import DatabricksServingEndpoint, DatabricksVectorSearchIndex

mlflow.models.set_model(MyAgent())

with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        artifact_path="agent",
        python_model="/path/to/agent.py",
        input_example={"input": [{"role": "user", "content": "Hello"}]},
        resources=[
            DatabricksServingEndpoint(endpoint_name="my-llm"),
            DatabricksVectorSearchIndex(index_name="prod.docs.index"),
        ],
        pip_requirements=["mlflow[databricks]>=3.10.1", "databricks-sdk>=0.40.0"],
    )

Declare every external dependency the agent needs at inference time. Load references/agent-logging.md if you need resource types, input examples, or error-handling details.

LangChain logging

mlflow.langchain.log_model(
    lc_model="/path/to/agent.py",
    artifact_path="agent",
    resources=[...],
)

DO / DON'T: Agent Logging

Part B: LoggedModel Lifecycle

Create a tracked model: mlflow.create_logged_model(name=...) creates a LoggedModel entity associated with the active run (see Models in Unity Catalog for how logged models relate to registration and lineage).

Config snapshot (generic): Log structured artifacts alongside the model — for example agent config, prompt version IDs, tool definitions, and feature flags as JSON under the run — so rollback and audits have full context without overloading run params.

Aliases on the LoggedModel:

mlflow.set_logged_model_alias(model_id=model_id, alias="champion")

Guard with getattr(mlflow, "set_logged_model_alias", None) if you must support older MLflow builds.

Link evaluation results: After mlflow.genai.evaluate(), metrics are automatically linked to the active LoggedModel for that evaluation context when using the GenAI eval harness (eval harness concepts). You can also log metrics explicitly:

with mlflow.start_run():
    logged = mlflow.create_logged_model(name="my-agent-iteration-7")
    model_id = logged.model_id
    mlflow.log_metrics({"accuracy": 0.92, "latency_p95_ms": 800}, model_id=model_id)

Rollback (generic): Point the champion (or production) alias at a prior model_id, or re-resolve runs:/... / UC version in your deployment config. Document which alias means “live” vs “candidate” in your team’s runbook.

Load references/logged-model-patterns.md if you need initialization/finalization or artifact layout patterns. Load references/model-promotion.md if you need promotion workflow details.

DO / DON'T: LoggedModel

MLFLOW_ACTIVE_MODEL_ID (production trace linking)

Set MLFLOW_ACTIVE_MODEL_ID=<logged_model_id> in the environment of your deployed app (app config, Docker ENV, Databricks App app.yml, or job cluster env) so production traces attach to that LoggedModel. The Databricks UI can then show production traces for that model/version.

• Pattern: export or declare MLFLOW_ACTIVE_MODEL_ID to the id returned when you created the LoggedModel (or the id of the version you promoted).
• Docs: Link production traces to app versions.

This is separate from UC registered model aliases: LoggedModel id ties tracing; UC version ties catalog consumers and serving.

Part C: UC Model Registration

Use the Unity Catalog registry URI for all UC operations (Models in Unity Catalog).

Naming convention

Use a stable, three-level UC name: {catalog}.{schema}.{model_name}. Choose model_name to match your product or agent family (e.g. support_agent or agent_{tenant}). Document the convention in your repo; avoid ad-hoc string concatenation across teams.

Standard registration pattern

import mlflow
from mlflow import MlflowClient

client = MlflowClient(registry_uri="databricks-uc")
catalog, schema = "main", "ml_models"
uc_model_name = f"{catalog}.{schema}.my_agent"

try:
    client.create_registered_model(uc_model_name)
except Exception:
    pass  # already exists

run_id = "..."  # run where pyfunc/langchain model was logged
result = client.create_model_version(
    uc_model_name,
    source=f"runs:/{run_id}/model",
)

You can also use mlflow.register_model(model_uri, uc_model_name) with registry_uri set to databricks-uc — equivalent intent; keep one style per pipeline for clarity.

Custom aliases and metric-based gating

# After comparing new eval metrics to thresholds or to the current @champion:
client.set_registered_model_alias(uc_model_name, "champion", result.version)

Gating pattern (conceptual):

def promote_if_better(client, uc_model_name, new_version, new_metrics, thresholds, judges):
    champion = None
    try:
        champion = client.get_model_version_by_alias(uc_model_name, "champion")
    except Exception:
        pass
    if champion is None:
        client.set_registered_model_alias(uc_model_name, "champion", new_version)
        return
    old_metrics = load_metrics_for_version(client, champion)  # your helper
    if all(new_metrics[k] >= thresholds[k] for k in thresholds) and beats(new_metrics, old_metrics, judges):
        client.set_registered_model_alias(uc_model_name, "champion", new_version)

Align metric names and thresholds with your eval harness (eval harness).

DO / DON'T: UC Registration

Load references/uc-gating.md if you need extended UC registration, score comparison, or gating walkthroughs.

Deployment job linkage

Pattern: Store deployment metadata on the model version (tags or version description): serving endpoint name, Databricks App URL, bundle job name, or Git SHA. When promotion sets @champion, a downstream job (or orchestration webhook) reads that metadata and deploys or updates the endpoint/app to that artifact.

• On promotion: trigger CI/CD or a Databricks Job that resolves models:/catalog.schema.model@champion and updates inference infrastructure.
• Keep rollback as re-pointing the alias plus the same job with the previous version id.

Common Mistakes

Validation Checklist

• Agent logged with pip_requirements including mlflow[databricks] (and SDK as needed).
• All inference dependencies in resources.
• mlflow.models.predict succeeds on a smoke input before deploy.
• LoggedModel created in the run that owns iteration artifacts (if applicable).
• Eval metrics linked via GenAI eval harness and/or log_metrics(..., model_id=...).
• UC naming convention documented; client uses databricks-uc.
• Champion promotion policy defined: LoggedModel alias vs UC @champion vs both.
• Production app sets MLFLOW_ACTIVE_MODEL_ID when trace-to-version linkage is required.

References

Related skills

• Foundation Step 1 — MLflow GenAI Foundation: ../../foundation/01-mlflow-genai-foundation/SKILL.md
• Foundation Step 2 — Experiment Tracing: ../../foundation/02-experiment-tracing-and-uc-storage/SKILL.md
• SDLC Step 4 — Evaluation Runs: Produces evaluation_results and thresholds_met for this step

Official documentation (grounding)

Local reference files (this skill)

Version History