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)`

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)`

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
        if self.model_extra:
            return {k: str(v) for k, v in self.model_extra.items()}
        return {}

    def model_dump(
        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 `model_dump` 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()
                .model_dump(exclude_unset=exclude_unset, **kwargs)
                .items()
                if v is not None
            }
        else:
            return super().model_dump(exclude_unset=exclude_unset, **kwargs)

    model_config = ConfigDict(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.

`model_dump(self, *, exclude_unset=False, exclude_none=True, **kwargs)`

Returns a dictionary representation of the metadata.

This method overrides the default Pydantic model_dump 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 model_dump(
    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 `model_dump` 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()
            .model_dump(exclude_unset=exclude_unset, **kwargs)
            .items()
            if v is not None
        }
    else:
        return super().model_dump(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)`

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)`

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

    # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
    #  fields defined under base models. If not handled, this raises a warning.
    #  It is possible to suppress this warning message with the following
    #  configuration, however the ultimate solution is to rename these fields.
    #  Even though they do not cause any problems right now, if we are not
    #  careful we might overwrite some fields protected by pydantic.
    model_config = ConfigDict(protected_namespaces=())

Model Registries

zenml.model_registries special

base_model_registry

BaseModelRegistry (StackComponent, ABC)

config: BaseModelRegistryConfig property readonly

delete_model(self, name)

delete_model_version(self, name, version)

get_latest_model_version(self, name, stage=None)

get_model(self, name)

get_model_uri_artifact_store(self, model_version)

get_model_version(self, name, version)

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)

list_models(self, name=None, metadata=None)

load_model_version(self, name, version, **kwargs)

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

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

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

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

BaseModelRegistryConfig (StackComponentConfig)

BaseModelRegistryFlavor (Flavor)

config_class: Type[zenml.model_registries.base_model_registry.BaseModelRegistryConfig] property readonly

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

type: StackComponentType property readonly

ModelRegistryModelMetadata (BaseModel)

custom_attributes: Dict[str, str] property readonly

model_dump(self, *, exclude_unset=False, exclude_none=True, **kwargs)

ModelVersionStage (Enum)

RegisteredModel (BaseModel)

RegistryModelVersion (BaseModel)

`zenml.model_registries` `special`

`base_model_registry`

`BaseModelRegistry (StackComponent, ABC)`

`config: BaseModelRegistryConfig` `property` `readonly`

`delete_model(self, name)`

`delete_model_version(self, name, version)`

`get_latest_model_version(self, name, stage=None)`

`get_model(self, name)`

`get_model_uri_artifact_store(self, model_version)`

`get_model_version(self, name, version)`

`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)`

`list_models(self, name=None, metadata=None)`

`load_model_version(self, name, version, **kwargs)`

`register_model(self, name, description=None, metadata=None)`

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

`update_model(self, name, description=None, metadata=None, remove_metadata=None)`

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

`BaseModelRegistryConfig (StackComponentConfig)`

`BaseModelRegistryFlavor (Flavor)`

`config_class: Type[zenml.model_registries.base_model_registry.BaseModelRegistryConfig]` `property` `readonly`

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

`type: StackComponentType` `property` `readonly`

`ModelRegistryModelMetadata (BaseModel)`

`custom_attributes: Dict[str, str]` `property` `readonly`

`model_dump(self, *, exclude_unset=False, exclude_none=True, **kwargs)`

`ModelVersionStage (Enum)`

`RegisteredModel (BaseModel)`

`RegistryModelVersion (BaseModel)`