Develop a Custom Model Registry

Learning how to develop a custom model registry.

Before diving into the specifics of this component type, it is beneficial to familiarize yourself with our general guide to writing custom component flavors in ZenML. This guide provides an essential understanding of ZenML's component flavor concepts.

Base abstraction in progress!

The Model registry stack component is relatively new in ZenML. While it is fully functional, it can be challenging to cover all the ways ML systems deal with model versioning. This means that the API might change in the future. We will keep this page up-to-date with the latest changes.

If you are writing a custom model registry flavor, and you found that the base abstraction is lacking or not flexible enough, please let us know by messaging us on Slack, or by opening an issue on GitHub

Base Abstraction

The BaseModelRegistry is the abstract base class that needs to be subclassed in order to create a custom component that can be used to register and retrieve models. As model registries can come in many shapes and forms, the base class exposes a deliberately basic and generic interface:

from abc import ABC, abstractmethod

from enum import Enum

from typing import Any, Dict, List, Optional, Type, cast

from pydantic import BaseModel, Field, root_validator

from zenml.enums import StackComponentType

from zenml.stack import Flavor, StackComponent

from zenml.stack.stack_component import StackComponentConfig

class BaseModelRegistryConfig(StackComponentConfig):

"""Base config for model registries."""

class BaseModelRegistry(StackComponent, ABC):

"""Base class for all ZenML model registries."""

@property

def config(self) -> BaseModelRegistryConfig:

"""Returns the config of the model registry."""

return cast(BaseModelRegistryConfig, self._config)

# ---------

# Model Registration Methods

# ---------

@abstractmethod

def register_model(

self,

name: str,

description: Optional[str] = None,