Experiment - Model registry

Problem:You have trained multiple versions of a machine learning model, but you lack a system to organize, track, and manage these models effectively.

Current Metrics:No model tracking or versioning system is in place, leading to confusion about which model is best or currently deployed.

Issue:Without a model registry, it is hard to manage model versions, track metadata, and deploy the best model reliably.

Your Task

Set up a simple model registry system to save, version, and load models with metadata, enabling easy tracking and deployment.

Use Python and a local file system for storing models and metadata.

Do not use any external model registry services or databases.

Keep the solution simple and beginner-friendly.

Hint 1

Hint 2

Hint 3

Solution

ML Python

import os
import json
import joblib
from datetime import datetime

class SimpleModelRegistry:
    def __init__(self, registry_dir='model_registry'):
        self.registry_dir = registry_dir
        os.makedirs(self.registry_dir, exist_ok=True)
        self.metadata_file = os.path.join(self.registry_dir, 'metadata.json')
        if not os.path.exists(self.metadata_file):
            with open(self.metadata_file, 'w') as f:
                json.dump([], f)

    def register_model(self, model, version, metrics):
        model_path = os.path.join(self.registry_dir, f'model_v{version}.joblib')
        joblib.dump(model, model_path)
        with open(self.metadata_file, 'r') as f:
            metadata = json.load(f)
        metadata.append({
            'version': version,
            'model_path': model_path,
            'metrics': metrics,
            'registered_at': datetime.now().isoformat()
        })
        with open(self.metadata_file, 'w') as f:
            json.dump(metadata, f, indent=2)
        print(f'Model version {version} registered successfully.')

    def load_model(self, version=None):
        with open(self.metadata_file, 'r') as f:
            metadata = json.load(f)
        if not metadata:
            print('No models registered yet.')
            return None
        if version is None:
            # Load latest version
            latest = max(metadata, key=lambda x: x['version'])
            version = latest['version']
        else:
            latest = next((m for m in metadata if m['version'] == version), None)
            if latest is None:
                print(f'Model version {version} not found.')
                return None
        model = joblib.load(latest['model_path'])
        print(f'Model version {version} loaded.')
        return model

# Example usage:
if __name__ == '__main__':
    from sklearn.datasets import load_iris
    from sklearn.tree import DecisionTreeClassifier
    from sklearn.model_selection import train_test_split
    from sklearn.metrics import accuracy_score

    data = load_iris()
    X_train, X_test, y_train, y_test = train_test_split(data.data, data.target, test_size=0.3, random_state=42)

    model = DecisionTreeClassifier(random_state=42)
    model.fit(X_train, y_train)
    preds = model.predict(X_test)
    acc = accuracy_score(y_test, preds)

    registry = SimpleModelRegistry()
    registry.register_model(model, version=1, metrics={'accuracy': acc})

    loaded_model = registry.load_model()  # Loads latest
    loaded_preds = loaded_model.predict(X_test)
    loaded_acc = accuracy_score(y_test, loaded_preds)
    print(f'Loaded model accuracy: {loaded_acc:.2f}')

Created a SimpleModelRegistry class to save models and metadata in a local folder.

Used joblib to save and load model files with version numbers.

Stored metadata including version, path, metrics, and registration time in a JSON file.

Added methods to register new models and load models by version or latest.

Provided example usage with a decision tree on iris dataset showing registration and loading.

Results Interpretation

Before: No model tracking, no version control, confusion about best model.

After: Models saved with versions and metrics, easy to load specific or latest model reliably.

A model registry helps organize and manage machine learning models by saving versions and metadata, making deployment and tracking easier and less error-prone.

Bonus Experiment

Extend the model registry to keep track of model deployment status (e.g., 'staging', 'production') and allow switching the production model.

💡 Hint

Add a 'status' field in metadata and create methods to update and retrieve models by status.