model-registry-governance

What this skill does

# 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

```bash
# 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
```

```yaml
# 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

```python
# 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

```python
# 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

```python
# 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 ValueEr