Model Registry

The Model Versioning Problem

After a few months of ML work, you'll find yourself asking:

• "Which model version is in production?"

• "What data did I use to train model v1.3?"

• "What were the hyperparameters for that model with 95% accuracy?"

• "Who deployed this model and when?"

Without a model registry, this information lives in notebooks, Slack messages, or someone's memory.

Kubeflow Model Registry solves this by providing a central repository for models with full lineage tracking.

What is Model Registry?

A model registry is a central hub that stores:

• Model artifacts: The trained model files

• Metadata: Framework, version, hyperparameters

• Lineage: Training data, code version, pipeline run

• Performance metrics: Accuracy, precision, recall

• Deployment history: Where and when deployed

• Lifecycle stage: Development, staging, production, archived

Setting Up Model Registry

Model Registry is part of Kubeflow's core components. If you installed Kubeflow following the earlier guide, it's already available.

Verify Installation

# Check if Model Registry is running
kubectl get pods -n kubeflow | grep model-registry

# Expected output:
# model-registry-xxx Running

Access via Python SDK

# Install ML Metadata (Model Registry backend)
pip install ml-metadata==1.14.0 kfp==2.6.0

# Connect to registry
from kfp.registry import RegistryClient

client = RegistryClient(host='http://localhost:8080')

Registering Models

Method 1: From Pipeline

Best practice—register models automatically when training:

from kfp import dsl
from kfp.dsl import Output, Model, Metrics

@dsl.component(
    base_image='python:3.12-slim',
    packages_to_install=['scikit-learn==1.3.0', 'joblib==1.3.2']
)
def train_and_register(
    dataset_path: str,
    model_output: Output[Model],
    metrics: Output[Metrics],
    model_name: str = 'iris-classifier',
    model_version: str = 'v1.0.0'
):
    """Train model and register with metadata."""
    import joblib
    import json
    from sklearn.datasets import load_iris
    from sklearn.ensemble import RandomForestClassifier
    from sklearn.model_selection import train_test_split
    from sklearn.metrics import accuracy_score
    
    # Load data
    X, y = load_iris(return_X_y=True)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
    
    # Train
    model = RandomForestClassifier(n_estimators=100, random_state=42)
    model.fit(X_train, y_train)
    
    # Evaluate
    y_pred = model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    
    # Save model
    joblib.dump(model, model_output.path)
    
    # Log metrics
    metrics.log_metric('accuracy', accuracy)
    metrics.log_metric('n_samples', len(X_train))
    
    # Add metadata
    model_output.metadata['framework'] = 'sklearn'
    model_output.metadata['model_type'] = 'RandomForestClassifier'
    model_output.metadata['n_estimators'] = 100
    model_output.metadata['version'] = model_version
    model_output.metadata['accuracy'] = accuracy
    
    print(f"Model registered: {model_name} {model_version}")
    print(f"Accuracy: {accuracy:.3f}")

Method 2: Manual Registration

For models trained outside pipelines:

from kfp.registry import RegistryClient
import datetime

client = RegistryClient(host='http://localhost:8080')

# Register model
model_id = client.upload_model(
    model_name='iris-classifier',
    model_path='model.pkl',
    framework='sklearn',
    framework_version='1.3.0',
    model_type='RandomForestClassifier',
    tags=['production', 'iris', 'classification'],
    description='RandomForest model for Iris classification',
    metadata={
        'n_estimators': 100,
        'max_depth': 10,
        'training_date': datetime.datetime.now().isoformat(),
        'data_version': 'v2.1',
        'accuracy': 0.95
    }
)

print(f"Model registered with ID: {model_id}")

Querying Models

List All Models

# Get all models
models = client.list_models()

for model in models:
    print(f"{model.name} - {model.version}")
    print(f"  Framework: {model.framework}")
    print(f"  Created: {model.create_time}")
    print(f"  Stage: {model.stage}")

Search by Criteria

# Find production models
prod_models = client.list_models(
    filter_query="stage='production'"
)

# Find by tag
sklearn_models = client.list_models(
    filter_query="tags LIKE '%sklearn%'"
)

# Find recent models
from datetime import datetime, timedelta

week_ago = datetime.now() - timedelta(days=7)
recent_models = client.list_models(
    filter_query=f"create_time > '{week_ago.isoformat()}'"
)

Get Model Details

# Get specific model version
model = client.get_model('iris-classifier', version='v1.0.0')

print(f"Model: {model.name}")
print(f"Version: {model.version}")
print(f"Accuracy: {model.metadata['accuracy']}")
print(f"Stage: {model.stage}")

# Download model
model_path = client.download_model(model.id, destination='./model.pkl')

Model Lifecycle Management

Transition Model Stages

Models progress through stages:

1. Development: Being developed and tested

2. Staging: Deployed to staging environment

3. Production: Serving production traffic

4. Archived: No longer in use

# Promote to staging
client.transition_model_stage(
    model_name='iris-classifier',
    version='v1.0.0',
    stage='staging'
)

# After testing in staging, promote to production
client.transition_model_stage(
    model_name='iris-classifier',
    version='v1.0.0',
    stage='production'
)

# Archive old version
client.transition_model_stage(
    model_name='iris-classifier',
    version='v0.9.0',
    stage='archived'
)

Version Management

# Create new version
new_version = client.create_model_version(
    model_name='iris-classifier',
    version='v1.1.0',
    model_path='new_model.pkl',
    metadata={
        'accuracy': 0.97,  # Improved!
        'changes': 'Added feature engineering',
        'parent_version': 'v1.0.0'
    }
)

# Compare versions
v1 = client.get_model('iris-classifier', version='v1.0.0')
v2 = client.get_model('iris-classifier', version='v1.1.0')

print(f"v1.0.0 accuracy: {v1.metadata['accuracy']}")
print(f"v1.1.0 accuracy: {v2.metadata['accuracy']}")
print(f"Improvement: {(v2.metadata['accuracy'] - v1.metadata['accuracy']) * 100:.2f}%")

Model Lineage

Track where models come from:

# Register with lineage information
model_id = client.upload_model(
    model_name='iris-classifier',
    model_path='model.pkl',
    metadata={
        # Data lineage
        'training_data': 's3://data/iris_train_2024-01.csv',
        'data_version': 'v2.1',
        'data_hash': 'sha256:abc123...',
        
        # Code lineage
        'git_commit': '7f3a9c2',
        'git_repo': 'https://github.com/user/ml-project',
        'pipeline_run_id': 'run-12345',
        
        # Training lineage
        'training_script': 'train_model.py',
        'hyperparameters': {
            'n_estimators': 100,
            'max_depth': 10
        },
        
        # Environment
        'python_version': '3.12',
        'sklearn_version': '1.3.0',
        'cuda_version': None
    }
)

Integration with KServe

Deploy models from registry:

def deploy_from_registry(model_name: str, version: str):
    """Deploy model from registry to KServe."""
    
    # Get model from registry
    model = client.get_model(model_name, version=version)
    model_uri = model.uri
    
    # Create InferenceService
    from kubernetes import client as k8s_client
    
    inference_service = {
        'apiVersion': 'serving.kserve.io/v1beta1',
        'kind': 'InferenceService',
        'metadata': {
            'name': f'{model_name}-{version}'.replace('.', '-'),
            'namespace': 'ml-workspace'
        },
        'spec': {
            'predictor': {
                'sklearn': {
                    'storageUri': model_uri,
                    'resources': {
                        'limits': {
                            'cpu': '1',
                            'memory': '2Gi'
                        }
                    }
                }
            }
        }
    }
    
    # Deploy
    custom_api = k8s_client.CustomObjectsApi()
    custom_api.create_namespaced_custom_object(
        group='serving.kserve.io',
        version='v1beta1',
        namespace='ml-workspace',
        plural='inferenceservices',
        body=inference_service
    )
    
    print(f"Deployed {model_name} {version} to KServe")

# Deploy production model
deploy_from_registry('iris-classifier', 'v1.0.0')

Model Comparison

Compare multiple models:

def compare_models(model_name: str, versions: list):
    """Compare model versions."""
    
    comparison = []
    for version in versions:
        model = client.get_model(model_name, version=version)
        comparison.append({
            'version': version,
            'accuracy': model.metadata.get('accuracy'),
            'created': model.create_time,
            'stage': model.stage
        })
    
    # Sort by accuracy
    comparison.sort(key=lambda x: x['accuracy'], reverse=True)
    
    print(f"Model Comparison for {model_name}:")
    for m in comparison:
        print(f"  {m['version']}: {m['accuracy']:.3f} ({m['stage']})")
    
    return comparison

# Compare versions
compare_models('iris-classifier', ['v1.0.0', 'v1.1.0', 'v1.2.0'])

Best Practices

1. Semantic Versioning

Use semantic versioning (MAJOR.MINOR.PATCH):

# v1.0.0 -> Initial version
# v1.1.0 -> New features, backward compatible
# v2.0.0 -> Breaking changes

version = 'v1.0.0'

2. Rich Metadata

Store comprehensive metadata:

metadata = {
    # Performance
    'accuracy': 0.95,
    'precision': 0.94,
    'recall': 0.96,
    'f1_score': 0.95,
    
    # Training
    'n_samples': 10000,
    'training_time_seconds': 120,
    'hyperparameters': {...},
    
    # Data
    'data_version': 'v2.1',
    'features': ['feat1', 'feat2'],
    
    # Lineage
    'git_commit': '7f3a9c2',
    'pipeline_run': 'run-123',
    
    # Deployment
    'deployment_date': '2024-01-03',
    'deployed_by': '[email protected]'
}

3. Automate Registration

Always register models from pipelines:

@dsl.pipeline(name='train-and-register')
def training_pipeline():
    train_task = train_model()
    register_task = register_model(model=train_task.outputs['model'])
    deploy_task = deploy_model(model=register_task.outputs['registered_model'])

4. Track Production Models

Tag production deployments:

# When deploying to production
client.transition_model_stage(
    model_name='iris-classifier',
    version='v1.2.0',
    stage='production'
)

client.add_tags(
    model_name='iris-classifier',
    version='v1.2.0',
    tags=['deployed', 'cluster-us-east-1', '2024-01-03']
)

5. Regular Audits

Periodically review registered models:

# Find old development models
old_dev_models = client.list_models(
    filter_query="stage='development' AND create_time < '2023-01-01'"
)

print(f"Found {len(old_dev_models)} old development models")
print("Consider archiving or deleting these")

Key Takeaways

1. Model Registry provides single source of truth for models

2. Track lineage: data, code, hyperparameters

3. Use lifecycle stages to manage deployments

4. Automate registration from pipelines

5. Rich metadata enables better decision-making

Next Steps

With models tracked and deployed, we need to ensure they perform well in production. In Monitoring & Observability, we'll learn how to track model performance and detect issues early.

---

Resources:

• Kubeflow Model Registry

• ML Metadata

• Model Management Best Practices