Model Registry Management
Managing ML model versions, staging, approval workflows, and deployment through a model registry — from tracking model artifacts through governance and production monitoring.
When to Use
- Tracking multiple model versions across environments
- Managing model promotion (staging → production)
- Implementing model governance and audit trails
- Automating model deployment pipelines
- Monitoring model performance in production
Registry Architecture
from typing import Dict, List, Optional, Any
from datetime import datetime
from enum import Enum
import json
import os
class ModelStage(str, Enum):
NONE = 'none'
STAGING = 'staging'
PRODUCTION = 'production'
ARCHIVED = 'archived'
class ModelStatus(str, Enum):
PENDING = 'pending_review'
APPROVED = 'approved'
REJECTED = 'rejected'
ROLLED_BACK = 'rolled_back'
class ModelRegistry:
"""Track and manage ML model versions."""
def __init__(self, registry_path: str = './model_registry'):
self.path = registry_path
os.makedirs(registry_path, exist_ok=True)
def register_model(self, name: str, version: str,
model_path: str, metrics: Dict,
params: Dict = None, framework: str = 'pytorch',
description: str = '') -> Dict:
"""Register a new model version."""
metadata = {
'name': name, 'version': version,
'model_path': model_path, 'framework': framework,
'metrics': metrics, 'params': params or {},
'description': description,
'stage': ModelStage.NONE.value,
'status': ModelStatus.PENDING.value,
'created_at': datetime.now().isoformat(),
'updated_at': datetime.now().isoformat(),
'registered_by': None,
'approvals': [],
'lineage': {'source_run': None, 'dataset_version': None},
}
# Save to registry
reg_path = f"{self.path}/{name}/{version}"
os.makedirs(reg_path, exist_ok=True)
with open(f"{reg_path}/metadata.json", 'w') as f:
json.dump(metadata, f, indent=2)
return metadata
def promote_to_staging(self, name: str, version: str,
approved_by: str = None) -> bool:
"""Promote model to staging."""
meta = self._load_metadata(name, version)
if not meta: return False
meta['stage'] = ModelStage.STAGING.value
meta['updated_at'] = datetime.now().isoformat()
meta['approvals'].append({
'action': 'promote_staging', 'by': approved_by,
'time': datetime.now().isoformat(),
})
self._save_metadata(name, version, meta)
return True
def promote_to_production(self, name: str, version: str,
approved_by: str = None) -> bool:
"""Promote model to production. Demotes current production."""
meta = self._load_metadata(name, version)
if not meta: return False
# Archive current production model
current_prod = self.get_production_model(name)
if current_prod:
self._load_metadata(name, current_prod['version'])['stage'] = ModelStage.ARCHIVED.value
meta['stage'] = ModelStage.PRODUCTION.value
meta['status'] = ModelStatus.APPROVED.value
meta['updated_at'] = datetime.now().isoformat()
meta['approvals'].append({
'action': 'promote_production', 'by': approved_by,
'time': datetime.now().isoformat(),
})
self._save_metadata(name, version, meta)
return True
def get_production_model(self, name: str) -> Optional[Dict]:
"""Get the current production model for a name."""
versions = self.list_versions(name)
for v in reversed(versions):
meta = self._load_metadata(name, v)
if meta and meta['stage'] == ModelStage.PRODUCTION.value:
return meta
return None
def list_versions(self, name: str) -> List[str]:
model_dir = f"{self.path}/{name}"
if not os.path.exists(model_dir): return []
return sorted(os.listdir(model_dir))
def compare_versions(self, name: str, versions: List[str]) -> str:
report = f"📊 Model: {name} — Version Comparison\n" + "=" * 50 + "\n"
for v in versions:
meta = self._load_metadata(name, v)
if meta:
report += f"\nv{v}: {meta['stage']}"
for metric, value in meta.get('metrics', {}).items():
report += f"\n {metric}: {value}"
report += f"\n Params: {meta.get('params', {})}"
return report
Governance and Approvals
class ModelGovernance:
"""Model governance and approval workflows."""
def __init__(self, required_approvals: int = 2):
self.required = required_approvals
def approve(self, registry: ModelRegistry, name: str,
version: str, reviewer: str, notes: str = '') -> Dict:
meta = registry._load_metadata(name, version)
if not meta: return {'error': 'Model not found'}
meta['approvals'].append({
'action': 'approve', 'by': reviewer, 'notes': notes,
'time': datetime.now().isoformat(),
})
# Check if enough approvals
approvals = [a for a in meta['approvals'] if a['action'] == 'approve']
if len(approvals) >= self.required:
meta['status'] = ModelStatus.APPROVED.value
registry._save_metadata(name, version, meta)
return meta
def audit_trail(self, registry: ModelRegistry, name: str,
version: str) -> str:
meta = registry._load_metadata(name, version)
if not meta: return "Model not found"
trail = f"📋 Audit Trail: {name} v{version}\n" + "=" * 50 + "\n"
trail += f"Created: {meta['created_at']}\n"
trail += f"Metrics: {meta['metrics']}\n"
trail += f"Stage: {meta['stage']}\n"
trail += "\nApprovals:\n"
for a in meta.get('approvals', []):
trail += f" {a['time']} — {a['action']} by {a.get('by', 'system')}\n"
return trail
Common Pitfalls
- No versioning — overwriting model files loses history; always version
- Manual promotion — human error in moving models to production; automate the pipeline
- No staging validation — promoting to production without staging validation causes incidents
- Missing lineage — can't trace which training run or data produced a model
- No rollback plan — if production model fails, need to quickly revert
Verification Checklist
- Model registry tracks all model versions
- Staging → production promotion workflow defined
- Approval gates for production deployment
- Rollback procedure documented and tested
- Model lineage tracked (training run, dataset version)
- Model performance monitored in production
- Audit trail available for compliance
See Also
- ml-experiment-tracking — experiment tracking feeding into registry
- ml-deployment-serving — deploying registered models
- ml-pipeline-design — CI/CD for model registry
- model-evaluation-metrics — evaluating models before registration