Model Registry

A model registry is the single source of truth for trained models — tracking versions, metadata, evaluation results, and promotion status. Without it, production runs a model nobody can identify, and rollback means “which checkpoint was it again?”

---

1. What the Registry Stores

class ModelVersion:
    """Schema for a registered model version."""
    model_name: str           # "llama-3-70b-sft"
    version: str              # "v2.3"
    
    # Artifact location
    weights_uri: str          # "s3://models/llama-3-70b-sft/v2.3/model.safetensors"
    format: str               # "safetensors", "gguf", "tensorrt"
    size_gb: float            # 35.2
    
    # Provenance
    training_run_id: str      # Link to experiment tracker
    base_model: str           # "meta-llama/Llama-3.1-70B"
    training_data_hash: str   # Reproducibility
    git_sha: str              # Code version
    
    # Quality
    eval_scores: dict         # {"mmlu": 0.82, "humaneval": 0.67, "safety": 0.99}
    eval_date: datetime
    
    # Deployment
    hardware_requirements: dict  # {"tp_size": 1, "gpu_type": "H100", "quantization": "int4"}
    serving_config: dict         # {"max_batch_size": 256, "max_seq_len": 8192}
    
    # Lifecycle
    status: str               # "experimental" → "staging" → "canary" → "production" → "deprecated"
    promoted_by: str          # "ml-team@company.com"
    promoted_at: datetime
    deprecation_date: datetime | None

---

2. Promotion Workflow

Experimental → Staging → Canary → Production → Deprecated
     │             │         │         │            │
     │         Eval suite   5% traffic  100%    Replaced by
     │         passes       24h monitor traffic  newer version
     │             │         │         │
     └─ reject ←───┘    rollback ←────┘

Gate Conditions

PROMOTION_GATES = {
    "experimental_to_staging": {
        "mmlu": (">=", 0.80),
        "humaneval": (">=", 0.65),
        "safety_refusal_rate": (">=", 0.95),
        "latency_p99_ms": ("<=", 2000),
    },
    "staging_to_canary": {
        "all_staging_gates": True,
        "regression_vs_production": ("delta_mmlu", ">=", -0.01),  # Max 1% regression
        "format_converted": True,  # TensorRT/quantized version ready
    },
    "canary_to_production": {
        "canary_error_rate": ("<=", 0.001),
        "canary_latency_regression": ("<=", 1.10),  # Max 10% slower
        "canary_quality_score": (">=", "production_baseline"),
        "canary_duration_hours": (">=", 24),
    },
}

async def promote(model_version: str, target_stage: str):
    """Promote model through gates."""
    version = registry.get(model_version)
    gates = PROMOTION_GATES[f"{version.status}_to_{target_stage}"]
    
    for metric, condition in gates.items():
        if not check_condition(version, metric, condition):
            raise PromotionBlocked(f"Gate failed: {metric} {condition}")
    
    version.status = target_stage
    version.promoted_at = datetime.utcnow()
    registry.update(version)
    
    if target_stage == "canary":
        await deploy_canary(version, traffic_pct=5)
    elif target_stage == "production":
        await deploy_production(version)

---

3. Registry Implementation Options

Tool	Type	Strengths	Best for
MLflow Model Registry	Open-source	Integrated with MLflow tracking, stage transitions	MLflow/Databricks ecosystem
HuggingFace Hub	Hosted	Model cards, community sharing, Git-based	Open model distribution
W&B Model Registry	Hosted	Links to experiments, artifact lineage	W&B users
Custom (S3 + Postgres)	DIY	Full control, custom metadata	Specific compliance needs

Custom Registry with S3 + Postgres

class ModelRegistry:
    def __init__(self, db, storage):
        self.db = db          # Postgres for metadata
        self.storage = storage # S3 for weights
    
    async def register(self, version: ModelVersion):
        # Upload weights to S3
        await self.storage.upload(
            version.local_path,
            f"models/{version.model_name}/{version.version}/"
        )
        version.weights_uri = f"s3://models/{version.model_name}/{version.version}/"
        
        # Store metadata in Postgres
        await self.db.insert("model_versions", version.dict())
    
    async def get_production(self, model_name: str) -> ModelVersion:
        """Get the current production version of a model."""
        return await self.db.query_one(
            "SELECT * FROM model_versions WHERE model_name = %s AND status = 'production'",
            model_name,
        )
    
    async def rollback(self, model_name: str):
        """Rollback to previous production version."""
        current = await self.get_production(model_name)
        previous = await self.db.query_one(
            "SELECT * FROM model_versions WHERE model_name = %s AND status = 'deprecated' ORDER BY promoted_at DESC LIMIT 1",
            model_name,
        )
        
        current.status = "deprecated"
        previous.status = "production"
        await self.db.update("model_versions", current.dict())
        await self.db.update("model_versions", previous.dict())
        await deploy_production(previous)

---

4. Multi-Format Management

Production models often need multiple formats:

Base weights (SafeTensors, BF16)
  ├── INT4 AWQ (vLLM serving)
  ├── FP8 (TensorRT-LLM engine)
  ├── GGUF Q4_K_M (llama.cpp / edge)
  └── ONNX (cross-platform)

Each format is a derived artifact linked to the same model version

class ModelArtifacts:
    model_version: str
    formats: dict = {
        "safetensors_bf16": "s3://models/llama-70b/v2.3/bf16/",
        "awq_int4": "s3://models/llama-70b/v2.3/awq-int4/",
        "tensorrt_fp8": "s3://models/llama-70b/v2.3/trt-fp8/",
        "gguf_q4km": "s3://models/llama-70b/v2.3/gguf-q4km/",
    }

---

5. Monitoring Registered Models

registry_metrics = {
    "models_in_production": count_by_status("production"),
    "models_in_canary": count_by_status("canary"),
    "avg_time_to_production_days": avg_promotion_time(),
    "rollbacks_last_30d": count_rollbacks(days=30),
    "deprecated_models_with_traffic": find_deprecated_with_traffic(),
}

# Alert if deprecated model still receives traffic
if deprecated_with_traffic:
    alert("Deprecated models still serving traffic — update routing")

For automated evaluation that feeds registry gates, see Offline Evaluation. For CI/CD automation, see CI/CD for Models.

---

Self-Check

• Design a model registry schema with provenance, quality, and deployment metadata
• Implement a 4-stage promotion workflow with gate conditions
• Build rollback capability to revert to the previous production model
• Manage multiple model formats (SafeTensors, AWQ, TensorRT, GGUF) as linked artifacts
• Set up monitoring for deprecated models that still receive traffic

---

References

• MLflow Model Registry
• HuggingFace Hub Documentation
• For experiment tracking, see Experiment Tracking