Model Registries

zenml.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.

Attributes

__all__ = ['BaseModelRegistry', 'BaseModelRegistryConfig', 'BaseModelRegistryFlavor'] module-attribute

Classes

BaseModelRegistry(name: str, id: UUID, config: StackComponentConfig, flavor: str, type: StackComponentType, user: Optional[UUID], workspace: UUID, created: datetime, updated: datetime, labels: Optional[Dict[str, Any]] = None, connector_requirements: Optional[ServiceConnectorRequirements] = None, connector: Optional[UUID] = None, connector_resource_id: Optional[str] = None, *args: Any, **kwargs: Any)

Bases: StackComponent, ABC

Base class for all ZenML model registries.

Source code in src/zenml/stack/stack_component.py

def __init__(
    self,
    name: str,
    id: UUID,
    config: StackComponentConfig,
    flavor: str,
    type: StackComponentType,
    user: Optional[UUID],
    workspace: UUID,
    created: datetime,
    updated: datetime,
    labels: Optional[Dict[str, Any]] = None,
    connector_requirements: Optional[ServiceConnectorRequirements] = None,
    connector: Optional[UUID] = None,
    connector_resource_id: Optional[str] = None,
    *args: Any,
    **kwargs: Any,
):
    """Initializes a StackComponent.

    Args:
        name: The name of the component.
        id: The unique ID of the component.
        config: The config of the component.
        flavor: The flavor of the component.
        type: The type of the component.
        user: The ID of the user who created the component.
        workspace: The ID of the workspace the component belongs to.
        created: The creation time of the component.
        updated: The last update time of the component.
        labels: The labels of the component.
        connector_requirements: The requirements for the connector.
        connector: The ID of a connector linked to the component.
        connector_resource_id: The custom resource ID to access through
            the connector.
        *args: Additional positional arguments.
        **kwargs: Additional keyword arguments.

    Raises:
        ValueError: If a secret reference is passed as name.
    """
    if secret_utils.is_secret_reference(name):
        raise ValueError(
            "Passing the `name` attribute of a stack component as a "
            "secret reference is not allowed."
        )

    self.id = id
    self.name = name
    self._config = config
    self.flavor = flavor
    self.type = type
    self.user = user
    self.workspace = workspace
    self.created = created
    self.updated = updated
    self.labels = labels
    self.connector_requirements = connector_requirements
    self.connector = connector
    self.connector_resource_id = connector_resource_id
    self._connector_instance: Optional[ServiceConnector] = None

Attributes

config: BaseModelRegistryConfig property

Returns the config of the model registries.

Returns:

Type	Description
BaseModelRegistryConfig	The config of the model registries.

Functions

delete_model(name: str) -> None 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: str, version: str) -> None 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: 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.

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: str) -> RegisteredModel 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: RegistryModelVersion) -> str 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: str, version: str) -> RegistryModelVersion 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: 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]] 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: Optional[str] = None, metadata: Optional[Dict[str, str]] = None) -> List[RegisteredModel] 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: str, version: str, **kwargs: Any) -> Any 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: str, description: Optional[str] = None, metadata: Optional[Dict[str, str]] = None) -> RegisteredModel 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: str, version: Optional[str] = None, model_source_uri: Optional[str] = None, description: Optional[str] = None, metadata: Optional[ModelRegistryModelMetadata] = None, **kwargs: Any) -> RegistryModelVersion 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: str, description: Optional[str] = None, metadata: Optional[Dict[str, str]] = None, remove_metadata: Optional[List[str]] = None) -> RegisteredModel 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: str, version: str, description: Optional[str] = None, metadata: Optional[ModelRegistryModelMetadata] = None, remove_metadata: Optional[List[str]] = None, stage: Optional[ModelVersionStage] = None) -> RegistryModelVersion 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(warn_about_plain_text_secrets: bool = False, **kwargs: Any)

Bases: StackComponentConfig

Base config for model registries.

Source code in src/zenml/stack/stack_component.py

def __init__(
    self, warn_about_plain_text_secrets: bool = False, **kwargs: Any
) -> None:
    """Ensures that secret references don't clash with pydantic validation.

    StackComponents allow the specification of all their string attributes
    using secret references of the form `{{secret_name.key}}`. This however
    is only possible when the stack component does not perform any explicit
    validation of this attribute using pydantic validators. If this were
    the case, the validation would run on the secret reference and would
    fail or in the worst case, modify the secret reference and lead to
    unexpected behavior. This method ensures that no attributes that require
    custom pydantic validation are set as secret references.

    Args:
        warn_about_plain_text_secrets: If true, then warns about using
            plain-text secrets.
        **kwargs: Arguments to initialize this stack component.

    Raises:
        ValueError: If an attribute that requires custom pydantic validation
            is passed as a secret reference, or if the `name` attribute
            was passed as a secret reference.
    """
    for key, value in kwargs.items():
        try:
            field = self.__class__.model_fields[key]
        except KeyError:
            # Value for a private attribute or non-existing field, this
            # will fail during the upcoming pydantic validation
            continue

        if value is None:
            continue

        if not secret_utils.is_secret_reference(value):
            if (
                secret_utils.is_secret_field(field)
                and warn_about_plain_text_secrets
            ):
                logger.warning(
                    "You specified a plain-text value for the sensitive "
                    f"attribute `{key}` for a `{self.__class__.__name__}` "
                    "stack component. This is currently only a warning, "
                    "but future versions of ZenML will require you to pass "
                    "in sensitive information as secrets. Check out the "
                    "documentation on how to configure your stack "
                    "components with secrets here: "
                    "https://docs.zenml.io/getting-started/deploying-zenml/secret-management"
                )
            continue

        if pydantic_utils.has_validators(
            pydantic_class=self.__class__, field_name=key
        ):
            raise ValueError(
                f"Passing the stack component attribute `{key}` as a "
                "secret reference is not allowed as additional validation "
                "is required for this attribute."
            )

    super().__init__(**kwargs)

BaseModelRegistryFlavor

Bases: Flavor

Base class for all ZenML model registry flavors.

Attributes

config_class: Type[BaseModelRegistryConfig] property

Config class for this flavor.

Returns:

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

implementation_class: Type[StackComponent] abstractmethod property

Returns the implementation class for this flavor.

Returns:

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

type: StackComponentType property

Type of the flavor.

Returns:

Name	Type	Description
StackComponentType	StackComponentType	The type of the flavor.

Modules

base_model_registry

Base class for all ZenML model registries.

Classes

BaseModelRegistry(name: str, id: UUID, config: StackComponentConfig, flavor: str, type: StackComponentType, user: Optional[UUID], workspace: UUID, created: datetime, updated: datetime, labels: Optional[Dict[str, Any]] = None, connector_requirements: Optional[ServiceConnectorRequirements] = None, connector: Optional[UUID] = None, connector_resource_id: Optional[str] = None, *args: Any, **kwargs: Any)

Bases: StackComponent, ABC

Base class for all ZenML model registries.

Source code in src/zenml/stack/stack_component.py

def __init__(
    self,
    name: str,
    id: UUID,
    config: StackComponentConfig,
    flavor: str,
    type: StackComponentType,
    user: Optional[UUID],
    workspace: UUID,
    created: datetime,
    updated: datetime,
    labels: Optional[Dict[str, Any]] = None,
    connector_requirements: Optional[ServiceConnectorRequirements] = None,
    connector: Optional[UUID] = None,
    connector_resource_id: Optional[str] = None,
    *args: Any,
    **kwargs: Any,
):
    """Initializes a StackComponent.

    Args:
        name: The name of the component.
        id: The unique ID of the component.
        config: The config of the component.
        flavor: The flavor of the component.
        type: The type of the component.
        user: The ID of the user who created the component.
        workspace: The ID of the workspace the component belongs to.
        created: The creation time of the component.
        updated: The last update time of the component.
        labels: The labels of the component.
        connector_requirements: The requirements for the connector.
        connector: The ID of a connector linked to the component.
        connector_resource_id: The custom resource ID to access through
            the connector.
        *args: Additional positional arguments.
        **kwargs: Additional keyword arguments.

    Raises:
        ValueError: If a secret reference is passed as name.
    """
    if secret_utils.is_secret_reference(name):
        raise ValueError(
            "Passing the `name` attribute of a stack component as a "
            "secret reference is not allowed."
        )

    self.id = id
    self.name = name
    self._config = config
    self.flavor = flavor
    self.type = type
    self.user = user
    self.workspace = workspace
    self.created = created
    self.updated = updated
    self.labels = labels
    self.connector_requirements = connector_requirements
    self.connector = connector
    self.connector_resource_id = connector_resource_id
    self._connector_instance: Optional[ServiceConnector] = None

Attributes

config: BaseModelRegistryConfig property

Returns the config of the model registries.

Returns:

Type	Description
BaseModelRegistryConfig	The config of the model registries.

Functions

delete_model(name: str) -> None 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: str, version: str) -> None 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: 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.

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: str) -> RegisteredModel 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: RegistryModelVersion) -> str 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: str, version: str) -> RegistryModelVersion 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: 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]] 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: Optional[str] = None, metadata: Optional[Dict[str, str]] = None) -> List[RegisteredModel] 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: str, version: str, **kwargs: Any) -> Any 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: str, description: Optional[str] = None, metadata: Optional[Dict[str, str]] = None) -> RegisteredModel 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: str, version: Optional[str] = None, model_source_uri: Optional[str] = None, description: Optional[str] = None, metadata: Optional[ModelRegistryModelMetadata] = None, **kwargs: Any) -> RegistryModelVersion 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: str, description: Optional[str] = None, metadata: Optional[Dict[str, str]] = None, remove_metadata: Optional[List[str]] = None) -> RegisteredModel 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: str, version: str, description: Optional[str] = None, metadata: Optional[ModelRegistryModelMetadata] = None, remove_metadata: Optional[List[str]] = None, stage: Optional[ModelVersionStage] = None) -> RegistryModelVersion 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(warn_about_plain_text_secrets: bool = False, **kwargs: Any)

Bases: StackComponentConfig

Base config for model registries.

Source code in src/zenml/stack/stack_component.py

def __init__(
    self, warn_about_plain_text_secrets: bool = False, **kwargs: Any
) -> None:
    """Ensures that secret references don't clash with pydantic validation.

    StackComponents allow the specification of all their string attributes
    using secret references of the form `{{secret_name.key}}`. This however
    is only possible when the stack component does not perform any explicit
    validation of this attribute using pydantic validators. If this were
    the case, the validation would run on the secret reference and would
    fail or in the worst case, modify the secret reference and lead to
    unexpected behavior. This method ensures that no attributes that require
    custom pydantic validation are set as secret references.

    Args:
        warn_about_plain_text_secrets: If true, then warns about using
            plain-text secrets.
        **kwargs: Arguments to initialize this stack component.

    Raises:
        ValueError: If an attribute that requires custom pydantic validation
            is passed as a secret reference, or if the `name` attribute
            was passed as a secret reference.
    """
    for key, value in kwargs.items():
        try:
            field = self.__class__.model_fields[key]
        except KeyError:
            # Value for a private attribute or non-existing field, this
            # will fail during the upcoming pydantic validation
            continue

        if value is None:
            continue

        if not secret_utils.is_secret_reference(value):
            if (
                secret_utils.is_secret_field(field)
                and warn_about_plain_text_secrets
            ):
                logger.warning(
                    "You specified a plain-text value for the sensitive "
                    f"attribute `{key}` for a `{self.__class__.__name__}` "
                    "stack component. This is currently only a warning, "
                    "but future versions of ZenML will require you to pass "
                    "in sensitive information as secrets. Check out the "
                    "documentation on how to configure your stack "
                    "components with secrets here: "
                    "https://docs.zenml.io/getting-started/deploying-zenml/secret-management"
                )
            continue

        if pydantic_utils.has_validators(
            pydantic_class=self.__class__, field_name=key
        ):
            raise ValueError(
                f"Passing the stack component attribute `{key}` as a "
                "secret reference is not allowed as additional validation "
                "is required for this attribute."
            )

    super().__init__(**kwargs)

BaseModelRegistryFlavor

Bases: Flavor

Base class for all ZenML model registry flavors.

Attributes

config_class: Type[BaseModelRegistryConfig] property

Config class for this flavor.

Returns:

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

implementation_class: Type[StackComponent] abstractmethod property

Returns the implementation class for this flavor.

Returns:

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

type: StackComponentType property

Type of the flavor.

Returns:

Name	Type	Description
StackComponentType	StackComponentType	The type of the flavor.

ModelRegistryModelMetadata

Bases: 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.

Attributes

custom_attributes: Dict[str, str] property

Returns a dictionary of custom attributes.

Returns:

Type	Description
Dict[str, str]	A dictionary of custom attributes.

Functions

model_dump(*, 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.

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 src/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

Bases: Enum

Enum of the possible stages of a registered model.

RegisteredModel

Bases: 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.

RegistryModelVersion

Bases: 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]	The creation time of this model version
last_updated_at	Optional[datetime]	The last updated time of this model version
stage	ModelVersionStage	The current stage of this model version
metadata	Optional[ModelRegistryModelMetadata]	Metadata associated with this model version