Skip to content

Model Registries

zenml.model_registries special

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.

base_model_registry

Base class for all ZenML model registries.

BaseModelRegistry (StackComponent, ABC)

Base class for all ZenML model registries.

Source code in 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: BaseModelRegistryConfig property readonly

Returns the config of the model registries.

Returns:

Type Description
BaseModelRegistryConfig

The config of the model registries.

delete_model(self, name)

Deletes a registered model from the model registry.

Parameters:

Name Type Description Default
name str

The name of the registered model.

required

Exceptions:

Type Description
KeyError

If the model does not exist.

RuntimeError

If deletion fails.

Source code in 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(self, name, version)

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

Exceptions:

Type Description
KeyError

If the model version does not exist.

RuntimeError

If deletion fails.

Source code in 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(self, 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[zenml.model_registries.base_model_registry.ModelVersionStage]

The stage of the model version.

None

Returns:

Type Description
Optional[zenml.model_registries.base_model_registry.RegistryModelVersion]

The latest model version.

Source code in 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(self, name)

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.

Exceptions:

Type Description
zenml.exceptions.EntityExistsError

If the model does not exist.

RuntimeError

If retrieval fails.

Source code in 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(self, model_version)

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 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(self, name, version)

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.

Exceptions:

Type Description
KeyError

If the model version does not exist.

RuntimeError

If retrieval fails.

Source code in 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(self, name=None, model_source_uri=None, metadata=None, stage=None, count=None, created_after=None, created_before=None, order_by_date=None, **kwargs)

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[zenml.model_registries.base_model_registry.ModelRegistryModelMetadata]

Metadata associated with this model version.

None
stage Optional[zenml.model_registries.base_model_registry.ModelVersionStage]

The stage of the model version.

None
count Optional[int]

The number of model versions to return.

None
created_after Optional[datetime.datetime]

The timestamp after which to list model versions.

None
created_before Optional[datetime.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[zenml.model_registries.base_model_registry.RegistryModelVersion]]

A list of model versions.

Source code in 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(self, name=None, metadata=None)

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[zenml.model_registries.base_model_registry.RegisteredModel]

A list of registered models.

Source code in 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(self, name, version, **kwargs)

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.

Exceptions:

Type Description
KeyError

If the model version does not exist.

RuntimeError

If loading fails.

Source code in 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(self, name, description=None, metadata=None)

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.

Exceptions:

Type Description
zenml.exceptions.EntityExistsError

If a model with the same name already exists.

RuntimeError

If registration fails.

Source code in 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(self, name, version=None, model_source_uri=None, description=None, metadata=None, **kwargs)

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[zenml.model_registries.base_model_registry.ModelRegistryModelMetadata]

The metadata associated with the model version.

None
**kwargs Any

Additional keyword arguments.

{}

Returns:

Type Description
RegistryModelVersion

The registered model version.

Exceptions:

Type Description
RuntimeError

If registration fails.

Source code in 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(self, name, description=None, metadata=None, remove_metadata=None)

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

Exceptions:

Type Description
KeyError

If the model does not exist.

RuntimeError

If update fails.

Source code in 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(self, name, version, description=None, metadata=None, remove_metadata=None, stage=None)

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[zenml.model_registries.base_model_registry.ModelRegistryModelMetadata]

Metadata associated with this model version.

None
remove_metadata Optional[List[str]]

The metadata to remove from the model version.

None
stage Optional[zenml.model_registries.base_model_registry.ModelVersionStage]

The stage of the model version.

None

Returns:

Type Description
RegistryModelVersion

The updated model version.

Exceptions:

Type Description
KeyError

If the model version does not exist.

RuntimeError

If update fails.

Source code in 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 (StackComponentConfig) pydantic-model

Base config for model registries.

Source code in zenml/model_registries/base_model_registry.py
class BaseModelRegistryConfig(StackComponentConfig):
    """Base config for model registries."""

BaseModelRegistryFlavor (Flavor)

Base class for all ZenML model registry flavors.

Source code in 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: Type[zenml.model_registries.base_model_registry.BaseModelRegistryConfig] property readonly

Config class for this flavor.

Returns:

Type Description
Type[zenml.model_registries.base_model_registry.BaseModelRegistryConfig]

The config class for this flavor.

implementation_class: Type[zenml.stack.stack_component.StackComponent] property readonly

Returns the implementation class for this flavor.

Returns:

Type Description
Type[zenml.stack.stack_component.StackComponent]

The implementation class for this flavor.

type: StackComponentType property readonly

Type of the flavor.

Returns:

Type Description
StackComponentType

The type of the flavor.

ModelRegistryModelMetadata (BaseModel) pydantic-model

Base class for all ZenML model registry model metadata.

The ModelRegistryModelMetadata class represents metadata associated with a registered model version, including information such as the associated pipeline name, pipeline run ID, step name, ZenML version, and custom attributes. It serves as a blueprint for creating concrete model metadata implementations in a registry, and provides a record of the history of a model and its development process.

Source code in zenml/model_registries/base_model_registry.py
class ModelRegistryModelMetadata(BaseModel):
    """Base class for all ZenML model registry model metadata.

    The `ModelRegistryModelMetadata` class represents metadata associated with
    a registered model version, including information such as the associated
    pipeline name, pipeline run ID, step name, ZenML version, and custom
    attributes. It serves as a blueprint for creating concrete model metadata
    implementations in a registry, and provides a record of the history of a
    model and its development process.
    """

    zenml_version: Optional[str] = None
    zenml_run_name: Optional[str] = None
    zenml_pipeline_name: Optional[str] = None
    zenml_pipeline_uuid: Optional[str] = None
    zenml_pipeline_run_uuid: Optional[str] = None
    zenml_step_name: Optional[str] = None
    zenml_workspace: Optional[str] = None

    @property
    def custom_attributes(self) -> Dict[str, str]:
        """Returns a dictionary of custom attributes.

        Returns:
            A dictionary of custom attributes.
        """
        # Return all attributes that are not explicitly defined as Pydantic
        # fields in this class
        return {
            k: str(v)
            for k, v in self.__dict__.items()
            if k not in self.__fields__.keys()
        }

    def dict(
        self,
        *,
        exclude_unset: bool = False,
        exclude_none: bool = True,
        **kwargs: Any,
    ) -> Dict[str, str]:
        """Returns a dictionary representation of the metadata.

        This method overrides the default Pydantic `dict` method to allow
        for the exclusion of fields with a value of None.

        Args:
            exclude_unset: Whether to exclude unset attributes.
            exclude_none: Whether to exclude None attributes.
            **kwargs: Additional keyword arguments.

        Returns:
            A dictionary representation of the metadata.
        """
        if exclude_none:
            return {
                k: v
                for k, v in super()
                .dict(exclude_unset=exclude_unset, **kwargs)
                .items()
                if v is not None
            }
        else:
            return super().dict(exclude_unset=exclude_unset, **kwargs)

    class Config:
        """Pydantic configuration class."""

        # Allow extra attributes to be set in the metadata
        extra = "allow"
custom_attributes: Dict[str, str] property readonly

Returns a dictionary of custom attributes.

Returns:

Type Description
Dict[str, str]

A dictionary of custom attributes.

Config

Pydantic configuration class.

Source code in zenml/model_registries/base_model_registry.py
class Config:
    """Pydantic configuration class."""

    # Allow extra attributes to be set in the metadata
    extra = "allow"
dict(self, *, exclude_unset=False, exclude_none=True, **kwargs)

Returns a dictionary representation of the metadata.

This method overrides the default Pydantic dict method to allow for the exclusion of fields with a value of None.

Parameters:

Name Type Description Default
exclude_unset bool

Whether to exclude unset attributes.

False
exclude_none bool

Whether to exclude None attributes.

True
**kwargs Any

Additional keyword arguments.

{}

Returns:

Type Description
Dict[str, str]

A dictionary representation of the metadata.

Source code in zenml/model_registries/base_model_registry.py
def dict(
    self,
    *,
    exclude_unset: bool = False,
    exclude_none: bool = True,
    **kwargs: Any,
) -> Dict[str, str]:
    """Returns a dictionary representation of the metadata.

    This method overrides the default Pydantic `dict` method to allow
    for the exclusion of fields with a value of None.

    Args:
        exclude_unset: Whether to exclude unset attributes.
        exclude_none: Whether to exclude None attributes.
        **kwargs: Additional keyword arguments.

    Returns:
        A dictionary representation of the metadata.
    """
    if exclude_none:
        return {
            k: v
            for k, v in super()
            .dict(exclude_unset=exclude_unset, **kwargs)
            .items()
            if v is not None
        }
    else:
        return super().dict(exclude_unset=exclude_unset, **kwargs)

ModelVersionStage (Enum)

Enum of the possible stages of a registered model.

Source code in zenml/model_registries/base_model_registry.py
class ModelVersionStage(Enum):
    """Enum of the possible stages of a registered model."""

    NONE = "None"
    STAGING = "Staging"
    PRODUCTION = "Production"
    ARCHIVED = "Archived"

RegisteredModel (BaseModel) pydantic-model

Base class for all ZenML registered models.

Model Registration are the top-level entities in the model registry. They serve as a container for all the versions of a model.

Attributes:

Name Type Description
name str

Name of the registered model.

description Optional[str]

Description of the registered model.

metadata Optional[Dict[str, str]]

metadata associated with the registered model.

Source code in zenml/model_registries/base_model_registry.py
class RegisteredModel(BaseModel):
    """Base class for all ZenML registered models.

    Model Registration are the top-level entities in the model registry.
    They serve as a container for all the versions of a model.

    Attributes:
        name: Name of the registered model.
        description: Description of the registered model.
        metadata: metadata associated with the registered model.
    """

    name: str
    description: Optional[str] = None
    metadata: Optional[Dict[str, str]] = None

RegistryModelVersion (BaseModel) pydantic-model

Base class for all ZenML model versions.

The RegistryModelVersion class represents a version or snapshot of a registered model, including information such as the associated ModelBundle, version number, creation time, pipeline run information, and metadata. It serves as a blueprint for creating concrete model version implementations in a registry, and provides a record of the history of a model and its development process.

All model registries must extend this class with their own specific fields.

Attributes:

Name Type Description
registered_model RegisteredModel

The registered model associated with this model

model_source_uri str

The URI of the model bundle associated with this model, The model source can not be changed after the model version is created. If the model source is changed, a new model version must be created.

model_format str

The format of the model bundle associated with this model, The model format is set automatically by the model registry integration and can not be changed after the model version is created.

model_library Optional[str]

The library used to create the model bundle associated with this model, The model library refers to the library used to create the model source, e.g. TensorFlow, PyTorch, etc. For some model registries, the model library is set retrieved automatically by the model registry.

version str

The version number of this model version

description Optional[str]

The description of this model version

created_at Optional[datetime.datetime]

The creation time of this model version

last_updated_at Optional[datetime.datetime]

The last updated time of this model version

stage ModelVersionStage

The current stage of this model version

metadata Optional[zenml.model_registries.base_model_registry.ModelRegistryModelMetadata]

Metadata associated with this model version

Source code in zenml/model_registries/base_model_registry.py
class RegistryModelVersion(BaseModel):
    """Base class for all ZenML model versions.

    The `RegistryModelVersion` class represents a version or snapshot of a registered
    model, including information such as the associated `ModelBundle`, version
    number, creation time, pipeline run information, and metadata. It serves as
    a blueprint for creating concrete model version implementations in a registry,
    and provides a record of the history of a model and its development process.

    All model registries must extend this class with their own specific fields.

    Attributes:
        registered_model: The registered model associated with this model
        model_source_uri: The URI of the model bundle associated with this model,
            The model source can not be changed after the model version is created.
            If the model source is changed, a new model version must be created.
        model_format: The format of the model bundle associated with this model,
            The model format is set automatically by the model registry integration
            and can not be changed after the model version is created.
        model_library: The library used to create the model bundle associated with
            this model, The model library refers to the library used to create the
            model source, e.g. TensorFlow, PyTorch, etc. For some model registries,
            the model library is set retrieved automatically by the model registry.
        version: The version number of this model version
        description: The description of this model version
        created_at: The creation time of this model version
        last_updated_at: The last updated time of this model version
        stage: The current stage of this model version
        metadata: Metadata associated with this model version
    """

    version: str
    model_source_uri: str
    model_format: str
    model_library: Optional[str] = None
    registered_model: RegisteredModel
    description: Optional[str] = None
    created_at: Optional[datetime] = None
    last_updated_at: Optional[datetime] = None
    stage: ModelVersionStage = ModelVersionStage.NONE
    metadata: Optional[ModelRegistryModelMetadata] = None