Model Registry Governance skill

Model Registry Governance

Create a trustworthy system of record for model artifacts, prompts, adapters, and evaluation evidence.

When to Use This Skill

• Setting up a centralized model registry for your organization
• Defining metadata standards for model artifacts
• Building approval workflows for model promotion to production
• Implementing lifecycle policies for model retirement
• Preparing for compliance audits of AI systems

Prerequisites

• MLflow Tracking Server or Weights & Biases instance deployed
• Object storage for model artifacts (S3, GCS, or MinIO)
• CI/CD pipeline with access to the registry API
• OPA or similar policy engine for governance checks
• Git repository for policy definitions and promotion scripts

Core Principles

• Traceability: every production model maps to source code, data snapshot, and evaluation results.
• Reproducibility: builds are deterministic with pinned dependencies.
• Policy-driven promotion: no manual bypass for critical safety checks.
• Lifecycle hygiene: stale, vulnerable, or unowned models are retired automatically.

MLflow Registry Setup

# Install MLflow with required backends
pip install mlflow[extras] psycopg2-binary boto3

# Start MLflow tracking server with PostgreSQL backend and S3 artifact store
mlflow server \
  --backend-store-uri postgresql://mlflow:password@db:5432/mlflow \
  --default-artifact-root s3://mlflow-artifacts/models \
  --host 0.0.0.0 \
  --port 5000 \
  --serve-artifacts

# docker-compose.yaml for MLflow
services:
  mlflow:
    image: ghcr.io/mlflow/mlflow:2.12.0
    command: >
      mlflow server
      --backend-store-uri postgresql://mlflow:${DB_PASSWORD}@db:5432/mlflow
      --default-artifact-root s3://mlflow-artifacts/models
      --host 0.0.0.0
      --port 5000
      --serve-artifacts
    ports:
      - "5000:5000"
    environment:
      AWS_ACCESS_KEY_ID: ${AWS_ACCESS_KEY_ID}
      AWS_SECRET_ACCESS_KEY: ${AWS_SECRET_ACCESS_KEY}
    depends_on:
      - db

  db:
    image: postgres:16-alpine
    environment:
      POSTGRES_DB: mlflow
      POSTGRES_USER: mlflow
      POSTGRES_PASSWORD: ${DB_PASSWORD}
    volumes:
      - pgdata:/var/lib/postgresql/data

volumes:
  pgdata:

Required Metadata Schema

# model_metadata_schema.py
from pydantic import BaseModel, Field
from typing import List, Optional
from datetime import datetime
from enum import Enum

class LifecycleState(str, Enum):
    DRAFT = "draft"
    CANDIDATE = "candidate"
    APPROVED = "approved"
    DEPRECATED = "deprecated"
    RETIRED = "retired"

class RiskRating(str, Enum):
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    CRITICAL = "critical"

class ModelMetadata(BaseModel):
    """Required metadata for every registered model."""
    # Identity
    name: str = Field(description="Model name matching registry key")
    version: str = Field(description="Semantic version")
    checksum: str = Field(description="SHA-256 of model artifact")
    storage_uri: str = Field(description="Artifact store path")

    # Lineage
    base_model: str = Field(description="Parent model identifier")
    fine_tune_method: Optional[str] = Field(default=None)
    training_dataset: Optional[str] = Field(default=None)
    training_date: Optional[datetime] = Field(default=None)
    source_commit: str = Field(description="Git SHA of training code")

    # Evaluation
    eval_datasets: List[str] = Field(description="Evaluation dataset IDs")
    eval_report_uri: str = Field(description="Path to evaluation results")
    quality_score: float = Field(ge=0, le=1)
    safety_score: float = Field(ge=0, le=1)

    # Governance
    license: str = Field(description="SPDX license identifier")
    allowed_use_cases: List[str]
    prohibited_use_cases: List[str]
    risk_rating: RiskRating
    security_controls: List[str]

    # Ownership
    owner: str = Field(description="Primary owner email")
    backup_owner: str = Field(description="Backup owner email")
    escalation_contact: str
    team: str

    # Lifecycle
    state: LifecycleState = LifecycleState.DRAFT
    created_at: datetime = Field(default_factory=datetime.utcnow)
    approved_at: Optional[datetime] = None
    approved_by: Optional[str] = None
    expires_at: Optional[datetime] = None

Model Registration Script

# register_model.py
import mlflow
from mlflow.tracking import MlflowClient
import json
import hashlib

def register_model(
    model_path: str,
    model_name: str,
    metadata: dict,
    mlflow_uri: str = "http://mlflow:5000"
):
    """Register a model with full metadata and governance tags."""
    mlflow.set_tracking_uri(mlflow_uri)
    client = MlflowClient()

    # Compute artifact checksum
    with open(model_path, "rb") as f:
        checksum = hashlib.sha256(f.read()).hexdigest()
    metadata["checksum"] = checksum

    # Log model with metadata
    with mlflow.start_run(run_name=f"register-{model_name}-{metadata['version']}") as run:
        # Log all metadata as params
        mlflow.log_params({
            "model_name": model_name,
            "version": metadata["version"],
            "base_model": metadata["base_model"],
            "risk_rating": metadata["risk_rating"],
            "owner": metadata["owner"],
            "license": metadata["license"],
        })

        # Log quality metrics
        mlflow.log_metrics({
            "quality_score": metadata["quality_score"],
            "safety_score": metadata["safety_score"],
        })

        # Log full metadata as artifact
        with open("metadata.json", "w") as f:
            json.dump(metadata, f, indent=2, default=str)
        mlflow.log_artifact("metadata.json")

        # Log model artifact
        mlflow.log_artifact(model_path)

        # Register in model registry
        model_uri = f"runs:/{run.info.run_id}/model"
        result = mlflow.register_model(model_uri, model_name)

        # Set lifecycle tags
        client.set_model_version_tag(
            model_name, result.version, "state", "draft"
        )
        client.set_model_version_tag(
            model_name, result.version, "risk_rating", metadata["risk_rating"]
        )
        client.set_model_version_tag(
            model_name, result.version, "checksum", checksum
        )

    return result

Approval Workflow

1. Registration request created from CI.
2. Security checks (artifact scan, dependency scan, provenance).
3. Evaluation package uploaded (quality, toxicity, jailbreak, bias, latency, cost).
4. Required approvals: platform + product + security (as policy dictates).
5. Promotion to stage/prod based on signed decision record.

Promotion Script

# promote_model.py
import mlflow
from mlflow.tracking import MlflowClient
from datetime import datetime
import sys

def promote_model(
    model_name: str,
    version: str,
    target_stage: str,
    approver: str,
    mlflow_uri: str = "http://mlflow:5000"
):
    """Promote a model version after governance checks pass."""
    mlflow.set_tracking_uri(mlflow_uri)
    client = MlflowClient()

    # Verify current state allows promotion
    mv = client.get_model_version(model_name, version)
    current_state = mv.tags.get("state", "draft")

    valid_transitions = {
        "draft": ["candidate"],
        "candidate": ["approved", "draft"],
        "approved": ["deprecated"],
        "deprecated": ["retired"],
    }

    if target_stage not in valid_transitions.get(current_state, []):
        raise ValueError(
            f"Invalid transition: {current_state} -> {target_stage}. "
            f"Allowed: {valid_transitions.get(current_state, [])}"
        )

    # Verify required eval scores for production promotion
    if target_stage == "approved":
        run = client.get_run(mv.run_id)
        quality = float(run.data.metrics.get("quality_score", 0))
        safety = float(run.data.metrics.get("safety_score", 0))

        if quality < 0.85:
            raise ValueError(f"Quality score {quality} below threshold 0.85")
        if safety < 0.95:
            raise ValueError(f"Safety score {safety} below threshold 0.95")

    # Record promotion
    now = datetime.utcnow().isoformat()
    client.set_model_version_tag(model_name, version, "state", target_stage)
    client.set_model_version_tag(model_name, version, f"promoted_to_{target_stage}_at", now)
    client.set_model_version_tag(model_name, version, f"promoted_to_{target_stage}_by", approver)

    # Transition MLflow stage alias
    stage_map = {
        "candidate": "Staging",
        "approved": "Production",
        "deprecated": "Archived",
    }
    if target_stage in stage_map:
        client.transition_model_version_stage(
            model_name, version, stage_map[target_stage]
        )

    print(f"Model {model_name} v{version}: {current_state} -> {target_stage}")
    print(f"Approved by: {approver} at {now}")

if __name__ == "__main__":
    promote_model(
        model_name=sys.argv[1],
        version=sys.argv[2],
        target_stage=sys.argv[3],
        approver=sys.argv[4],
    )

Lifecycle States

State	Description	Serving Allowed	New Usage
draft	Internal experimentation	Dev only	Dev only
candidate	Passed baseline tests	Staging	Staging
approved	Authorized for production	All environments	Yes
deprecated	Replacement announced	Existing only	Blocked
retired	Archived for audit	None	None

Lifecycle Automation

# lifecycle_policy.py
from mlflow.tracking import MlflowClient
from datetime import datetime, timedelta

def enforce_lifecycle_policies(mlflow_uri: str = "http://mlflow:5000"):
    """Run periodic lifecycle enforcement."""
    client = MlflowClient()

    for rm in client.search_registered_models():
        for mv in client.search_model_versions(f"name='{rm.name}'"):
            tags = mv.tags
            state = tags.get("state", "draft")

            # Auto-deprecate models with expired approvals (90 days)
            if state == "approved":
                approved_at = tags.get("promoted_to_approved_at")
                if approved_at:
                    approved_date = datetime.fromisoformat(approved_at)
                    if datetime.utcnow() - approved_date > timedelta(days=90):
                        print(f"Auto-deprecating {rm.name} v{mv.version}: approval expired")
                        client.set_model_version_tag(rm.name, mv.version, "state", "deprecated")
                        client.set_model_version_tag(
                            rm.name, mv.version, "auto_deprecated_reason", "approval_expired"
                        )

            # Auto-retire deprecated models after 30 days
            if state == "deprecated":
                deprecated_at = tags.get("promoted_to_deprecated_at")
                if deprecated_at:
                    deprecated_date = datetime.fromisoformat(deprecated_at)
                    if datetime.utcnow() - deprecated_date > timedelta(days=30):
                        print(f"Auto-retiring {rm.name} v{mv.version}")
                        client.set_model_version_tag(rm.name, mv.version, "state", "retired")
                        client.transition_model_version_stage(
                            rm.name, mv.version, "Archived"
                        )

            # Flag drafts with no activity for 14 days
            if state == "draft":
                created = datetime.fromisoformat(mv.creation_timestamp / 1000)
                if datetime.utcnow() - created > timedelta(days=14):
                    print(f"Stale draft: {rm.name} v{mv.version}")

Governance Policies (OPA/Rego)

# policy/model_governance.rego
package model.governance

# Reject artifacts without SBOM
deny[msg] {
    not input.metadata.sbom_uri
    msg := "Model must include SBOM artifact URI"
}

# Block promotion if critical CVEs remain
deny[msg] {
    input.target_state == "approved"
    input.security_scan.critical_cves > 0
    msg := sprintf("Cannot promote: %d critical CVEs unresolved", [input.security_scan.critical_cves])
}

# Require refreshed evals after prompt changes
deny[msg] {
    input.target_state == "approved"
    input.prompt_changed
    not input.eval_refreshed_after_prompt_change
    msg := "Evaluation must be re-run after prompt template changes"
}

# Require minimum eval scores for production
deny[msg] {
    input.target_state == "approved"
    input.metadata.quality_score < 0.85
    msg := sprintf("Quality score %.2f below threshold 0.85", [input.metadata.quality_score])
}

# Require dual approval for high-risk models
deny[msg] {
    input.target_state == "approved"
    input.metadata.risk_rating == "high"
    count(input.approvals) < 2
    msg := "High-risk models require at least 2 approvals"
}

Audit Readiness

Maintain immutable records of:

• Who approved and when
• Which policy checks executed
• Which exceptions were granted
• What model/version served each customer request window

Troubleshooting

Issue	Diagnosis	Resolution
Model registration fails	Check MLflow server connectivity and artifact store permissions	Verify S3/GCS credentials and bucket policy
Promotion blocked by policy	Review OPA deny messages in CI output	Fix metadata gaps or request policy exception
Stale models not auto-retiring	Lifecycle cron job not running	Check CronJob status in Kubernetes
Duplicate model versions	Race condition in CI pipeline	Add locking via registry API or database
Missing eval evidence	Eval pipeline skipped or failed	Re-run eval suite and re-register

Related Skills

• sbom-supply-chain - Provenance and signing
• policy-as-code - Enforce governance with policy engines
• llm-fine-tuning - Version adapters and training outputs
• llmops-platform-engineering - Platform CI/CD and promotion workflows
• ai-sre-incident-response - Incident response for model issues