Model Registries

Initialization of the MLflow Service.

Model registries are centralized repositories that facilitate the collaboration and management of machine learning models. They provide functionalities such as version control, metadata tracking, and storage of model artifacts, enabling data scientists to efficiently share and keep track of their models within a team or organization.

BaseModelRegistry

Bases: StackComponent, ABC

Base class for all ZenML model registries.

Source code in src/zenml/model_registries/base_model_registry.py

class BaseModelRegistry(StackComponent, ABC):
    """Base class for all ZenML model registries."""

    @property
    def config(self) -> BaseModelRegistryConfig:
        """Returns the config of the model registries.

        Returns:
            The config of the model registries.
        """
        return cast(BaseModelRegistryConfig, self._config)

    # ---------
    # Model Registration Methods
    # ---------

    @abstractmethod
    def register_model(
        self,
        name: str,
        description: Optional[str] = None,
        metadata: Optional[Dict[str, str]] = None,
    ) -> RegisteredModel:
        """Registers a model in the model registry.

        Args:
            name: The name of the registered model.
            description: The description of the registered model.
            metadata: The metadata associated with the registered model.

        Returns:
            The registered model.

        Raises:
            zenml.exceptions.EntityExistsError: If a model with the same name already exists.
            RuntimeError: If registration fails.
        """

    @abstractmethod
    def delete_model(
        self,
        name: str,
    ) -> None:
        """Deletes a registered model from the model registry.

        Args:
            name: The name of the registered model.

        Raises:
            KeyError: If the model does not exist.
            RuntimeError: If deletion fails.
        """

    @abstractmethod
    def update_model(
        self,
        name: str,
        description: Optional[str] = None,
        metadata: Optional[Dict[str, str]] = None,
        remove_metadata: Optional[List[str]] = None,
    ) -> RegisteredModel:
        """Updates a registered model in the model registry.

        Args:
            name: The name of the registered model.
            description: The description of the registered model.
            metadata: The metadata associated with the registered model.
            remove_metadata: The metadata to remove from the registered model.

        Raises:
            KeyError: If the model does not exist.
            RuntimeError: If update fails.
        """

    @abstractmethod
    def get_model(self, name: str) -> RegisteredModel:
        """Gets a registered model from the model registry.

        Args:
            name: The name of the registered model.

        Returns:
            The registered model.

        Raises:
            zenml.exceptions.EntityExistsError: If the model does not exist.
            RuntimeError: If retrieval fails.
        """

    @abstractmethod
    def list_models(
        self,
        name: Optional[str] = None,
        metadata: Optional[Dict[str, str]] = None,
    ) -> List[RegisteredModel]:
        """Lists all registered models in the model registry.

        Args:
            name: The name of the registered model.
            metadata: The metadata associated with the registered model.

        Returns:
            A list of registered models.
        """

    # ---------
    # Model Version Methods
    # ---------

    @abstractmethod
    def register_model_version(
        self,
        name: str,
        version: Optional[str] = None,
        model_source_uri: Optional[str] = None,
        description: Optional[str] = None,
        metadata: Optional[ModelRegistryModelMetadata] = None,
        **kwargs: Any,
    ) -> RegistryModelVersion:
        """Registers a model version in the model registry.

        Args:
            name: The name of the registered model.
            model_source_uri: The source URI of the model.
            version: The version of the model version.
            description: The description of the model version.
            metadata: The metadata associated with the model
                version.
            **kwargs: Additional keyword arguments.

        Returns:
            The registered model version.

        Raises:
            RuntimeError: If registration fails.
        """

    @abstractmethod
    def delete_model_version(
        self,
        name: str,
        version: str,
    ) -> None:
        """Deletes a model version from the model registry.

        Args:
            name: The name of the registered model.
            version: The version of the model version to delete.

        Raises:
            KeyError: If the model version does not exist.
            RuntimeError: If deletion fails.
        """

    @abstractmethod
    def update_model_version(
        self,
        name: str,
        version: str,
        description: Optional[str] = None,
        metadata: Optional[ModelRegistryModelMetadata] = None,
        remove_metadata: Optional[List[str]] = None,
        stage: Optional[ModelVersionStage] = None,
    ) -> RegistryModelVersion:
        """Updates a model version in the model registry.

        Args:
            name: The name of the registered model.
            version: The version of the model version to update.
            description: The description of the model version.
            metadata: Metadata associated with this model version.
            remove_metadata: The metadata to remove from the model version.
            stage: The stage of the model version.

        Returns:
            The updated model version.

        Raises:
            KeyError: If the model version does not exist.
            RuntimeError: If update fails.
        """

    @abstractmethod
    def list_model_versions(
        self,
        name: Optional[str] = None,
        model_source_uri: Optional[str] = None,
        metadata: Optional[ModelRegistryModelMetadata] = None,
        stage: Optional[ModelVersionStage] = None,
        count: Optional[int] = None,
        created_after: Optional[datetime] = None,
        created_before: Optional[datetime] = None,
        order_by_date: Optional[str] = None,
        **kwargs: Any,
    ) -> Optional[List[RegistryModelVersion]]:
        """Lists all model versions for a registered model.

        Args:
            name: The name of the registered model.
            model_source_uri: The model source URI of the registered model.
            metadata: Metadata associated with this model version.
            stage: The stage of the model version.
            count: The number of model versions to return.
            created_after: The timestamp after which to list model versions.
            created_before: The timestamp before which to list model versions.
            order_by_date: Whether to sort by creation time, this can
                be "asc" or "desc".
            kwargs: Additional keyword arguments.

        Returns:
            A list of model versions.
        """

    def get_latest_model_version(
        self,
        name: str,
        stage: Optional[ModelVersionStage] = None,
    ) -> Optional[RegistryModelVersion]:
        """Gets the latest model version for a registered model.

        This method is used to get the latest model version for a registered
        model. If no stage is provided, the latest model version across all
        stages is returned. If a stage is provided, the latest model version
        for that stage is returned.

        Args:
            name: The name of the registered model.
            stage: The stage of the model version.

        Returns:
            The latest model version.
        """
        model_versions = self.list_model_versions(
            name=name, stage=stage, order_by_date="desc", count=1
        )
        if model_versions:
            return model_versions[0]
        return None

    @abstractmethod
    def get_model_version(
        self, name: str, version: str
    ) -> RegistryModelVersion:
        """Gets a model version for a registered model.

        Args:
            name: The name of the registered model.
            version: The version of the model version to get.

        Returns:
            The model version.

        Raises:
            KeyError: If the model version does not exist.
            RuntimeError: If retrieval fails.
        """

    @abstractmethod
    def load_model_version(
        self,
        name: str,
        version: str,
        **kwargs: Any,
    ) -> Any:
        """Loads a model version from the model registry.

        Args:
            name: The name of the registered model.
            version: The version of the model version to load.
            **kwargs: Additional keyword arguments.

        Returns:
            The loaded model version.

        Raises:
            KeyError: If the model version does not exist.
            RuntimeError: If loading fails.
        """

    @abstractmethod
    def get_model_uri_artifact_store(
        self,
        model_version: RegistryModelVersion,
    ) -> str:
        """Gets the URI artifact store for a model version.

        This method retrieves the URI of the artifact store for a specific model
        version. Its purpose is to ensure that the URI is in the correct format
        for the specific artifact store being used. This is essential for the
        model serving component, which relies on the URI to serve the model
        version. In some cases, the URI may be stored in a different format by
        certain model registry integrations. This method allows us to obtain the
        URI in the correct format, regardless of the integration being used.

        Note: In some cases the URI artifact store may not be available to the
        user, the method should save the target model in one of the other
        artifact stores supported by ZenML and return the URI of that artifact
        store.

        Args:
            model_version: The model version for which to get the URI artifact
                store.

        Returns:
            The URI artifact store for the model version.
        """

config property

Returns the config of the model registries.

Returns:

Type	Description
BaseModelRegistryConfig	The config of the model registries.

delete_model(name) abstractmethod

Deletes a registered model from the model registry.

Parameters:

Name	Type	Description	Default
name	str	The name of the registered model.	required

Raises:

Type	Description
KeyError	If the model does not exist.
RuntimeError	If deletion fails.

Source code in src/zenml/model_registries/base_model_registry.py

@abstractmethod
def delete_model(
    self,
    name: str,
) -> None:
    """Deletes a registered model from the model registry.

    Args:
        name: The name of the registered model.

    Raises:
        KeyError: If the model does not exist.
        RuntimeError: If deletion fails.
    """

delete_model_version(name, version) abstractmethod

Deletes a model version from the model registry.

Parameters:

Name	Type	Description	Default
name	str	The name of the registered model.	required
version	str	The version of the model version to delete.	required

Raises:

Type	Description
KeyError	If the model version does not exist.
RuntimeError	If deletion fails.

Source code in src/zenml/model_registries/base_model_registry.py

@abstractmethod
def delete_model_version(
    self,
    name: str,
    version: str,
) -> None:
    """Deletes a model version from the model registry.

    Args:
        name: The name of the registered model.
        version: The version of the model version to delete.

    Raises:
        KeyError: If the model version does not exist.
        RuntimeError: If deletion fails.
    """

get_latest_model_version(name, stage=None)

Gets the latest model version for a registered model.

This method is used to get the latest model version for a registered model. If no stage is provided, the latest model version across all stages is returned. If a stage is provided, the latest model version for that stage is returned.

Parameters:

Name	Type	Description	Default
name	str	The name of the registered model.	required
stage	Optional[ModelVersionStage]	The stage of the model version.	None

Returns:

Type	Description
Optional[RegistryModelVersion]	The latest model version.

Source code in src/zenml/model_registries/base_model_registry.py

def get_latest_model_version(
    self,
    name: str,
    stage: Optional[ModelVersionStage] = None,
) -> Optional[RegistryModelVersion]:
    """Gets the latest model version for a registered model.

    This method is used to get the latest model version for a registered
    model. If no stage is provided, the latest model version across all
    stages is returned. If a stage is provided, the latest model version
    for that stage is returned.

    Args:
        name: The name of the registered model.
        stage: The stage of the model version.

    Returns:
        The latest model version.
    """
    model_versions = self.list_model_versions(
        name=name, stage=stage, order_by_date="desc", count=1
    )
    if model_versions:
        return model_versions[0]
    return None

get_model(name) abstractmethod

Gets a registered model from the model registry.

Parameters:

Name	Type	Description	Default
name	str	The name of the registered model.	required

Returns:

Type	Description
RegisteredModel	The registered model.

Raises:

Type	Description
EntityExistsError	If the model does not exist.
RuntimeError	If retrieval fails.

Source code in src/zenml/model_registries/base_model_registry.py

@abstractmethod
def get_model(self, name: str) -> RegisteredModel:
    """Gets a registered model from the model registry.

    Args:
        name: The name of the registered model.

    Returns:
        The registered model.

    Raises:
        zenml.exceptions.EntityExistsError: If the model does not exist.
        RuntimeError: If retrieval fails.
    """

get_model_uri_artifact_store(model_version) abstractmethod

Gets the URI artifact store for a model version.

This method retrieves the URI of the artifact store for a specific model version. Its purpose is to ensure that the URI is in the correct format for the specific artifact store being used. This is essential for the model serving component, which relies on the URI to serve the model version. In some cases, the URI may be stored in a different format by certain model registry integrations. This method allows us to obtain the URI in the correct format, regardless of the integration being used.

Note: In some cases the URI artifact store may not be available to the user, the method should save the target model in one of the other artifact stores supported by ZenML and return the URI of that artifact store.

Parameters:

Name	Type	Description	Default
model_version	RegistryModelVersion	The model version for which to get the URI artifact store.	required

Returns:

Type	Description
str	The URI artifact store for the model version.

Source code in src/zenml/model_registries/base_model_registry.py

@abstractmethod
def get_model_uri_artifact_store(
    self,
    model_version: RegistryModelVersion,
) -> str:
    """Gets the URI artifact store for a model version.

    This method retrieves the URI of the artifact store for a specific model
    version. Its purpose is to ensure that the URI is in the correct format
    for the specific artifact store being used. This is essential for the
    model serving component, which relies on the URI to serve the model
    version. In some cases, the URI may be stored in a different format by
    certain model registry integrations. This method allows us to obtain the
    URI in the correct format, regardless of the integration being used.

    Note: In some cases the URI artifact store may not be available to the
    user, the method should save the target model in one of the other
    artifact stores supported by ZenML and return the URI of that artifact
    store.

    Args:
        model_version: The model version for which to get the URI artifact
            store.

    Returns:
        The URI artifact store for the model version.
    """

get_model_version(name, version) abstractmethod

Gets a model version for a registered model.

Parameters:

Name	Type	Description	Default
name	str	The name of the registered model.	required
version	str	The version of the model version to get.	required

Returns:

Type	Description
RegistryModelVersion	The model version.

Raises:

Type	Description
KeyError	If the model version does not exist.
RuntimeError	If retrieval fails.

Source code in src/zenml/model_registries/base_model_registry.py

@abstractmethod
def get_model_version(
    self, name: str, version: str
) -> RegistryModelVersion:
    """Gets a model version for a registered model.

    Args:
        name: The name of the registered model.
        version: The version of the model version to get.

    Returns:
        The model version.

    Raises:
        KeyError: If the model version does not exist.
        RuntimeError: If retrieval fails.
    """

list_model_versions(name=None, model_source_uri=None, metadata=None, stage=None, count=None, created_after=None, created_before=None, order_by_date=None, **kwargs) abstractmethod

Lists all model versions for a registered model.

Parameters:

Name	Type	Description	Default
name	Optional[str]	The name of the registered model.	None
model_source_uri	Optional[str]	The model source URI of the registered model.	None
metadata	Optional[ModelRegistryModelMetadata]	Metadata associated with this model version.	None
stage	Optional[ModelVersionStage]	The stage of the model version.	None
count	Optional[int]	The number of model versions to return.	None
created_after	Optional[datetime]	The timestamp after which to list model versions.	None
created_before	Optional[datetime]	The timestamp before which to list model versions.	None
order_by_date	Optional[str]	Whether to sort by creation time, this can be "asc" or "desc".	None
kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
Optional[List[RegistryModelVersion]]	A list of model versions.

Source code in src/zenml/model_registries/base_model_registry.py

@abstractmethod
def list_model_versions(
    self,
    name: Optional[str] = None,
    model_source_uri: Optional[str] = None,
    metadata: Optional[ModelRegistryModelMetadata] = None,
    stage: Optional[ModelVersionStage] = None,
    count: Optional[int] = None,
    created_after: Optional[datetime] = None,
    created_before: Optional[datetime] = None,
    order_by_date: Optional[str] = None,
    **kwargs: Any,
) -> Optional[List[RegistryModelVersion]]:
    """Lists all model versions for a registered model.

    Args:
        name: The name of the registered model.
        model_source_uri: The model source URI of the registered model.
        metadata: Metadata associated with this model version.
        stage: The stage of the model version.
        count: The number of model versions to return.
        created_after: The timestamp after which to list model versions.
        created_before: The timestamp before which to list model versions.
        order_by_date: Whether to sort by creation time, this can
            be "asc" or "desc".
        kwargs: Additional keyword arguments.

    Returns:
        A list of model versions.
    """

list_models(name=None, metadata=None) abstractmethod

Lists all registered models in the model registry.

Parameters:

Name	Type	Description	Default
name	Optional[str]	The name of the registered model.	None
metadata	Optional[Dict[str, str]]	The metadata associated with the registered model.	None

Returns:

Type	Description
List[RegisteredModel]	A list of registered models.

Source code in src/zenml/model_registries/base_model_registry.py

@abstractmethod
def list_models(
    self,
    name: Optional[str] = None,
    metadata: Optional[Dict[str, str]] = None,
) -> List[RegisteredModel]:
    """Lists all registered models in the model registry.

    Args:
        name: The name of the registered model.
        metadata: The metadata associated with the registered model.

    Returns:
        A list of registered models.
    """

load_model_version(name, version, **kwargs) abstractmethod

Loads a model version from the model registry.

Parameters:

Name	Type	Description	Default
name	str	The name of the registered model.	required
version	str	The version of the model version to load.	required
**kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
Any	The loaded model version.

Raises:

Type	Description
KeyError	If the model version does not exist.
RuntimeError	If loading fails.

Source code in src/zenml/model_registries/base_model_registry.py

@abstractmethod
def load_model_version(
    self,
    name: str,
    version: str,
    **kwargs: Any,
) -> Any:
    """Loads a model version from the model registry.

    Args:
        name: The name of the registered model.
        version: The version of the model version to load.
        **kwargs: Additional keyword arguments.

    Returns:
        The loaded model version.

    Raises:
        KeyError: If the model version does not exist.
        RuntimeError: If loading fails.
    """

register_model(name, description=None, metadata=None) abstractmethod

Registers a model in the model registry.

Parameters:

Name	Type	Description	Default
name	str	The name of the registered model.	required
description	Optional[str]	The description of the registered model.	None
metadata	Optional[Dict[str, str]]	The metadata associated with the registered model.	None

Returns:

Type	Description
RegisteredModel	The registered model.

Raises:

Type	Description
EntityExistsError	If a model with the same name already exists.
RuntimeError	If registration fails.

Source code in src/zenml/model_registries/base_model_registry.py

@abstractmethod
def register_model(
    self,
    name: str,
    description: Optional[str] = None,
    metadata: Optional[Dict[str, str]] = None,
) -> RegisteredModel:
    """Registers a model in the model registry.

    Args:
        name: The name of the registered model.
        description: The description of the registered model.
        metadata: The metadata associated with the registered model.

    Returns:
        The registered model.

    Raises:
        zenml.exceptions.EntityExistsError: If a model with the same name already exists.
        RuntimeError: If registration fails.
    """

register_model_version(name, version=None, model_source_uri=None, description=None, metadata=None, **kwargs) abstractmethod

Registers a model version in the model registry.

Parameters:

Name	Type	Description	Default
name	str	The name of the registered model.	required
model_source_uri	Optional[str]	The source URI of the model.	None
version	Optional[str]	The version of the model version.	None
description	Optional[str]	The description of the model version.	None
metadata	Optional[ModelRegistryModelMetadata]	The metadata associated with the model version.	None
**kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
RegistryModelVersion	The registered model version.

Raises:

Type	Description
RuntimeError	If registration fails.

Source code in src/zenml/model_registries/base_model_registry.py

@abstractmethod
def register_model_version(
    self,
    name: str,
    version: Optional[str] = None,
    model_source_uri: Optional[str] = None,
    description: Optional[str] = None,
    metadata: Optional[ModelRegistryModelMetadata] = None,
    **kwargs: Any,
) -> RegistryModelVersion:
    """Registers a model version in the model registry.

    Args:
        name: The name of the registered model.
        model_source_uri: The source URI of the model.
        version: The version of the model version.
        description: The description of the model version.
        metadata: The metadata associated with the model
            version.
        **kwargs: Additional keyword arguments.

    Returns:
        The registered model version.

    Raises:
        RuntimeError: If registration fails.
    """

update_model(name, description=None, metadata=None, remove_metadata=None) abstractmethod

Updates a registered model in the model registry.

Parameters:

Name	Type	Description	Default
name	str	The name of the registered model.	required
description	Optional[str]	The description of the registered model.	None
metadata	Optional[Dict[str, str]]	The metadata associated with the registered model.	None
remove_metadata	Optional[List[str]]	The metadata to remove from the registered model.	None

Raises:

Type	Description
KeyError	If the model does not exist.
RuntimeError	If update fails.

Source code in src/zenml/model_registries/base_model_registry.py

@abstractmethod
def update_model(
    self,
    name: str,
    description: Optional[str] = None,
    metadata: Optional[Dict[str, str]] = None,
    remove_metadata: Optional[List[str]] = None,
) -> RegisteredModel:
    """Updates a registered model in the model registry.

    Args:
        name: The name of the registered model.
        description: The description of the registered model.
        metadata: The metadata associated with the registered model.
        remove_metadata: The metadata to remove from the registered model.

    Raises:
        KeyError: If the model does not exist.
        RuntimeError: If update fails.
    """

update_model_version(name, version, description=None, metadata=None, remove_metadata=None, stage=None) abstractmethod

Updates a model version in the model registry.

Parameters:

Name	Type	Description	Default
name	str	The name of the registered model.	required
version	str	The version of the model version to update.	required
description	Optional[str]	The description of the model version.	None
metadata	Optional[ModelRegistryModelMetadata]	Metadata associated with this model version.	None
remove_metadata	Optional[List[str]]	The metadata to remove from the model version.	None
stage	Optional[ModelVersionStage]	The stage of the model version.	None

Returns:

Type	Description
RegistryModelVersion	The updated model version.

Raises:

Type	Description
KeyError	If the model version does not exist.
RuntimeError	If update fails.

Source code in src/zenml/model_registries/base_model_registry.py

@abstractmethod
def update_model_version(
    self,
    name: str,
    version: str,
    description: Optional[str] = None,
    metadata: Optional[ModelRegistryModelMetadata] = None,
    remove_metadata: Optional[List[str]] = None,
    stage: Optional[ModelVersionStage] = None,
) -> RegistryModelVersion:
    """Updates a model version in the model registry.

    Args:
        name: The name of the registered model.
        version: The version of the model version to update.
        description: The description of the model version.
        metadata: Metadata associated with this model version.
        remove_metadata: The metadata to remove from the model version.
        stage: The stage of the model version.

    Returns:
        The updated model version.

    Raises:
        KeyError: If the model version does not exist.
        RuntimeError: If update fails.
    """

BaseModelRegistryConfig

Bases: StackComponentConfig

Base config for model registries.

Source code in src/zenml/model_registries/base_model_registry.py

class BaseModelRegistryConfig(StackComponentConfig):
    """Base config for model registries."""

BaseModelRegistryFlavor

Bases: Flavor

Base class for all ZenML model registry flavors.

Source code in src/zenml/model_registries/base_model_registry.py

class BaseModelRegistryFlavor(Flavor):
    """Base class for all ZenML model registry flavors."""

    @property
    def type(self) -> StackComponentType:
        """Type of the flavor.

        Returns:
            StackComponentType: The type of the flavor.
        """
        return StackComponentType.MODEL_REGISTRY

    @property
    def config_class(self) -> Type[BaseModelRegistryConfig]:
        """Config class for this flavor.

        Returns:
            The config class for this flavor.
        """
        return BaseModelRegistryConfig

    @property
    @abstractmethod
    def implementation_class(self) -> Type[StackComponent]:
        """Returns the implementation class for this flavor.

        Returns:
            The implementation class for this flavor.
        """
        return BaseModelRegistry

config_class property

Config class for this flavor.

Returns:

Type	Description
Type[BaseModelRegistryConfig]	The config class for this flavor.

implementation_class abstractmethod property

Returns the implementation class for this flavor.

Returns:

Type	Description
Type[StackComponent]	The implementation class for this flavor.

type property

Type of the flavor.

Returns:

Name	Type	Description
StackComponentType	StackComponentType	The type of the flavor.