Model

zenml.model

Concepts related to the Model Control Plane feature.

Modules

lazy_load

Model Version Data Lazy Loader definition.

Classes

ModelVersionDataLazyLoader

Bases: BaseModel

Model Version Data Lazy Loader helper class.

It helps the inner codes to fetch proper artifact, model version metadata or artifact metadata from the model version during runtime time of the step.

Functions

model

Model user facing interface to pass into pipeline or step.

Classes

Model

Bases: BaseModel

Model class to pass into pipeline or step to set it into a model context.

name: The name of the model. license: The license under which the model is created. description: The description of the model. audience: The target audience of the model. use_cases: The use cases of the model. limitations: The known limitations of the model. trade_offs: The tradeoffs of the model. ethics: The ethical implications of the model. tags: Tags associated with the model. version: The version name, version number or stage is optional and points model context to a specific version/stage. If skipped new version will be created. version also supports placeholders: standard {date} and {time} and any custom placeholders that are passed as substitutions in the pipeline or step decorators. save_models_to_registry: Whether to save all ModelArtifacts to Model Registry, if available in active stack.

Attributes

id: UUID property

Get version id from the Model Control Plane.

Returns:

Type	Description
UUID	ID of the model version or None, if model version doesn't exist and can only be read given current config (you used stage name or number as a version name).

Raises:

Type	Description
RuntimeError	if model version doesn't exist and cannot be fetched from the Model Control Plane.

model_id: UUID property

Get model id from the Model Control Plane.

Returns:

Type	Description
UUID	The UUID of the model containing this model version.

number: int property

Get version number from the Model Control Plane.

Returns:

Type	Description
int	Number of the model version or None, if model version doesn't exist and can only be read given current config (you used stage name or number as a version name).

Raises:

Type	Description
KeyError	if model version doesn't exist and cannot be fetched from the Model Control Plane.

run_metadata: Dict[str, MetadataType] property

Get model version run metadata.

Returns:

Type	Description
Dict[str, MetadataType]	The model version run metadata.

Raises:

Type	Description
RuntimeError	If the model version run metadata cannot be fetched.

stage: Optional[ModelStages] property

Get version stage from the Model Control Plane.

Returns:

Type	Description
Optional[ModelStages]	Stage of the model version or None, if model version doesn't exist and can only be read given current config (you used stage name or number as a version name).

Functions

delete_all_artifacts(only_link: bool = True, delete_from_artifact_store: bool = False) -> None

Delete all artifacts linked to this model version.

Parameters:

Name	Type	Description	Default
only_link	bool	Whether to only delete the link to the artifact.	True
delete_from_artifact_store	bool	Whether to delete the artifact from the artifact store.	False

Source code in src/zenml/model/model.py

def delete_all_artifacts(
    self,
    only_link: bool = True,
    delete_from_artifact_store: bool = False,
) -> None:
    """Delete all artifacts linked to this model version.

    Args:
        only_link: Whether to only delete the link to the artifact.
        delete_from_artifact_store: Whether to delete the artifact from
            the artifact store.
    """
    from zenml.client import Client

    client = Client()

    if not only_link and delete_from_artifact_store:
        mv = self._get_model_version()
        artifact_responses = mv.data_artifacts
        artifact_responses.update(mv.model_artifacts)
        artifact_responses.update(mv.deployment_artifacts)

        for artifact_ in artifact_responses.values():
            for artifact_response_ in artifact_.values():
                client._delete_artifact_from_artifact_store(
                    artifact_version=artifact_response_
                )

    client.delete_all_model_version_artifact_links(self.id, only_link)

delete_artifact(name: str, version: Optional[str] = None, only_link: bool = True, delete_metadata: bool = True, delete_from_artifact_store: bool = False) -> None

Delete the artifact linked to this model version.

Parameters:

Name	Type	Description	Default
name	str	The name of the artifact to delete.	required
version	Optional[str]	The version of the artifact to delete (None for latest/non-versioned)	None
only_link	bool	Whether to only delete the link to the artifact.	True
delete_metadata	bool	Whether to delete the metadata of the artifact.	True
delete_from_artifact_store	bool	Whether to delete the artifact from the artifact store.	False

Source code in src/zenml/model/model.py

def delete_artifact(
    self,
    name: str,
    version: Optional[str] = None,
    only_link: bool = True,
    delete_metadata: bool = True,
    delete_from_artifact_store: bool = False,
) -> None:
    """Delete the artifact linked to this model version.

    Args:
        name: The name of the artifact to delete.
        version: The version of the artifact to delete (None for
            latest/non-versioned)
        only_link: Whether to only delete the link to the artifact.
        delete_metadata: Whether to delete the metadata of the artifact.
        delete_from_artifact_store: Whether to delete the artifact from the
            artifact store.
    """
    from zenml.client import Client
    from zenml.models import ArtifactVersionResponse

    artifact_version = self.get_artifact(name, version)
    if isinstance(artifact_version, ArtifactVersionResponse):
        client = Client()
        client.delete_model_version_artifact_link(
            model_version_id=self.id,
            artifact_version_id=artifact_version.id,
        )
        if not only_link:
            client.delete_artifact_version(
                name_id_or_prefix=artifact_version.id,
                delete_metadata=delete_metadata,
                delete_from_artifact_store=delete_from_artifact_store,
            )

get_artifact(name: str, version: Optional[str] = None) -> Optional[ArtifactVersionResponse]

Get the artifact linked to this model version.

Parameters:

Name	Type	Description	Default
name	str	The name of the artifact to retrieve.	required
version	Optional[str]	The version of the artifact to retrieve (None for latest/non-versioned)	None

Returns:

Type	Description
Optional[ArtifactVersionResponse]	Specific version of the artifact or placeholder in the design time of the pipeline.

Source code in src/zenml/model/model.py

def get_artifact(
    self,
    name: str,
    version: Optional[str] = None,
) -> Optional["ArtifactVersionResponse"]:
    """Get the artifact linked to this model version.

    Args:
        name: The name of the artifact to retrieve.
        version: The version of the artifact to retrieve (None for
            latest/non-versioned)

    Returns:
        Specific version of the artifact or placeholder in the design time
            of the pipeline.
    """
    if lazy := self._lazy_artifact_get(name, version):
        return lazy

    return self._get_or_create_model_version().get_artifact(
        name=name,
        version=version,
    )

get_data_artifact(name: str, version: Optional[str] = None) -> Optional[ArtifactVersionResponse]

Get the data artifact linked to this model version.

Parameters:

Name	Type	Description	Default
name	str	The name of the data artifact to retrieve.	required
version	Optional[str]	The version of the data artifact to retrieve (None for latest/non-versioned)	None

Returns:

Type	Description
Optional[ArtifactVersionResponse]	Specific version of the data artifact or placeholder in the design
Optional[ArtifactVersionResponse]	time of the pipeline.

Source code in src/zenml/model/model.py

def get_data_artifact(
    self,
    name: str,
    version: Optional[str] = None,
) -> Optional["ArtifactVersionResponse"]:
    """Get the data artifact linked to this model version.

    Args:
        name: The name of the data artifact to retrieve.
        version: The version of the data artifact to retrieve (None for
            latest/non-versioned)

    Returns:
        Specific version of the data artifact or placeholder in the design
        time of the pipeline.
    """
    if lazy := self._lazy_artifact_get(name, version):
        return lazy

    return self._get_or_create_model_version().get_data_artifact(
        name=name,
        version=version,
    )

get_deployment_artifact(name: str, version: Optional[str] = None) -> Optional[ArtifactVersionResponse]

Get the deployment artifact linked to this model version.

Parameters:

Name	Type	Description	Default
name	str	The name of the deployment artifact to retrieve.	required
version	Optional[str]	The version of the deployment artifact to retrieve (None for latest/non-versioned)	None

Returns:

Type	Description
Optional[ArtifactVersionResponse]	Specific version of the deployment artifact or placeholder in the
Optional[ArtifactVersionResponse]	design time of the pipeline.

Source code in src/zenml/model/model.py

def get_deployment_artifact(
    self,
    name: str,
    version: Optional[str] = None,
) -> Optional["ArtifactVersionResponse"]:
    """Get the deployment artifact linked to this model version.

    Args:
        name: The name of the deployment artifact to retrieve.
        version: The version of the deployment artifact to retrieve (None
            for latest/non-versioned)

    Returns:
        Specific version of the deployment artifact or placeholder in the
        design time of the pipeline.
    """
    if lazy := self._lazy_artifact_get(name, version):
        return lazy

    return self._get_or_create_model_version().get_deployment_artifact(
        name=name,
        version=version,
    )

get_model_artifact(name: str, version: Optional[str] = None) -> Optional[ArtifactVersionResponse]

Get the model artifact linked to this model version.

Parameters:

Name	Type	Description	Default
name	str	The name of the model artifact to retrieve.	required
version	Optional[str]	The version of the model artifact to retrieve (None for latest/non-versioned)	None

Returns:

Type	Description
Optional[ArtifactVersionResponse]	Specific version of the model artifact or placeholder in the design time of the pipeline.

Source code in src/zenml/model/model.py

def get_model_artifact(
    self,
    name: str,
    version: Optional[str] = None,
) -> Optional["ArtifactVersionResponse"]:
    """Get the model artifact linked to this model version.

    Args:
        name: The name of the model artifact to retrieve.
        version: The version of the model artifact to retrieve (None for
            latest/non-versioned)

    Returns:
        Specific version of the model artifact or placeholder in the design
            time of the pipeline.
    """
    if lazy := self._lazy_artifact_get(name, version):
        return lazy

    return self._get_or_create_model_version().get_model_artifact(
        name=name,
        version=version,
    )

get_pipeline_run(name: str) -> PipelineRunResponse

Get pipeline run linked to this version.

Parameters:

Name	Type	Description	Default
name	str	The name of the pipeline run to retrieve.	required

Returns:

Type	Description
PipelineRunResponse	PipelineRun as PipelineRunResponse

Source code in src/zenml/model/model.py

def get_pipeline_run(self, name: str) -> "PipelineRunResponse":
    """Get pipeline run linked to this version.

    Args:
        name: The name of the pipeline run to retrieve.

    Returns:
        PipelineRun as PipelineRunResponse
    """
    return self._get_or_create_model_version().get_pipeline_run(name=name)

load_artifact(name: str, version: Optional[str] = None) -> Any

Load artifact from the Model Control Plane.

Parameters:

Name	Type	Description	Default
name	str	Name of the artifact to load.	required
version	Optional[str]	Version of the artifact to load.	None

Returns:

Type	Description
Any	The loaded artifact.

Raises:

Type	Description
ValueError	if the model version is not linked to any artifact with the given name and version.

Source code in src/zenml/model/model.py

def load_artifact(self, name: str, version: Optional[str] = None) -> Any:
    """Load artifact from the Model Control Plane.

    Args:
        name: Name of the artifact to load.
        version: Version of the artifact to load.

    Returns:
        The loaded artifact.

    Raises:
        ValueError: if the model version is not linked to any artifact with
            the given name and version.
    """
    from zenml.artifacts.utils import load_artifact
    from zenml.models import ArtifactVersionResponse

    artifact = self.get_artifact(name=name, version=version)

    if not isinstance(artifact, ArtifactVersionResponse):
        raise ValueError(
            f"Version {self.version} of model {self.name} does not have "
            f"an artifact with name {name} and version {version}."
        )

    return load_artifact(artifact.id, str(artifact.version))

log_metadata(metadata: Dict[str, MetadataType]) -> None

Log model version metadata.

This function can be used to log metadata for current model version.

Parameters:

Name	Type	Description	Default
metadata	Dict[str, MetadataType]	The metadata to log.	required

Source code in src/zenml/model/model.py

def log_metadata(
    self,
    metadata: Dict[str, "MetadataType"],
) -> None:
    """Log model version metadata.

    This function can be used to log metadata for current model version.

    Args:
        metadata: The metadata to log.
    """
    from zenml.client import Client
    from zenml.models import RunMetadataResource

    response = self._get_or_create_model_version()
    Client().create_run_metadata(
        metadata=metadata,
        resources=[
            RunMetadataResource(
                id=response.id, type=MetadataResourceTypes.MODEL_VERSION
            )
        ],
    )

set_stage(stage: Union[str, ModelStages], force: bool = False) -> None

Sets this Model to a desired stage.

Parameters:

Name	Type	Description	Default
stage	Union[str, ModelStages]	the target stage for model version.	required
force	bool	whether to force archiving of current model version in target stage or raise.	False

Source code in src/zenml/model/model.py

def set_stage(
    self, stage: Union[str, ModelStages], force: bool = False
) -> None:
    """Sets this Model to a desired stage.

    Args:
        stage: the target stage for model version.
        force: whether to force archiving of current model version in
            target stage or raise.
    """
    self._get_or_create_model_version().set_stage(stage=stage, force=force)

Functions

utils

Utility functions for linking step outputs to model versions.

Classes

Functions

link_artifact_to_model(artifact_version: ArtifactVersionResponse, model: Optional[Model] = None) -> None

Link the artifact to the model.

Parameters:

Name	Type	Description	Default
artifact_version	ArtifactVersionResponse	The artifact version to link.	required
model	Optional[Model]	The model to link to.	None

Raises:

Type	Description
RuntimeError	If called outside a step.

Source code in src/zenml/model/utils.py

def link_artifact_to_model(
    artifact_version: ArtifactVersionResponse,
    model: Optional["Model"] = None,
) -> None:
    """Link the artifact to the model.

    Args:
        artifact_version: The artifact version to link.
        model: The model to link to.

    Raises:
        RuntimeError: If called outside a step.
    """
    if not model:
        is_issue = False
        try:
            step_context = get_step_context()
            model = step_context.model
        except StepContextError:
            is_issue = True

        if model is None or is_issue:
            raise RuntimeError(
                "`link_artifact_to_model` called without `model` parameter "
                "and configured model context cannot be identified. Consider "
                "passing the `model` explicitly or configuring it in "
                "@step or @pipeline decorator."
            )

    model_version = model._get_or_create_model_version()
    link_artifact_version_to_model_version(
        artifact_version=artifact_version,
        model_version=model_version,
    )

link_artifact_version_to_model_version(artifact_version: ArtifactVersionResponse, model_version: ModelVersionResponse) -> None

Link an artifact version to a model version.

Parameters:

Name	Type	Description	Default
artifact_version	ArtifactVersionResponse	The artifact version to link.	required
model_version	ModelVersionResponse	The model version to link.	required

Source code in src/zenml/model/utils.py

def link_artifact_version_to_model_version(
    artifact_version: ArtifactVersionResponse,
    model_version: ModelVersionResponse,
) -> None:
    """Link an artifact version to a model version.

    Args:
        artifact_version: The artifact version to link.
        model_version: The model version to link.
    """
    client = Client()
    client.zen_store.create_model_version_artifact_link(
        ModelVersionArtifactRequest(
            artifact_version=artifact_version.id,
            model_version=model_version.id,
        )
    )

link_service_to_model(service_id: UUID, model: Optional[Model] = None, model_version_id: Optional[UUID] = None) -> None

Links a service to a model.

Parameters:

Name	Type	Description	Default
service_id	UUID	The ID of the service to link to the model.	required
model	Optional[Model]	The model to link the service to.	None
model_version_id	Optional[UUID]	The ID of the model version to link the service to.	None

Raises:

Type	Description
RuntimeError	If no model is provided and the model context cannot be identified.

Source code in src/zenml/model/utils.py

def link_service_to_model(
    service_id: UUID,
    model: Optional["Model"] = None,
    model_version_id: Optional[UUID] = None,
) -> None:
    """Links a service to a model.

    Args:
        service_id: The ID of the service to link to the model.
        model: The model to link the service to.
        model_version_id: The ID of the model version to link the service to.

    Raises:
        RuntimeError: If no model is provided and the model context cannot be
            identified.
    """
    client = Client()

    # If no model is provided, try to get it from the context
    if not model and not model_version_id:
        is_issue = False
        try:
            step_context = get_step_context()
            model = step_context.model
        except StepContextError:
            is_issue = True

        if model is None or is_issue:
            raise RuntimeError(
                "`link_service_to_model` called without `model` parameter "
                "and configured model context cannot be identified. Consider "
                "passing the `model` explicitly or configuring it in "
                "@step or @pipeline decorator."
            )

    model_version_id = (
        model_version_id or model._get_or_create_model_version().id
        if model
        else None
    )
    update_service = ServiceUpdate(model_version_id=model_version_id)
    client.zen_store.update_service(
        service_id=service_id, update=update_service
    )

log_model_metadata(metadata: Dict[str, MetadataType], model_name: Optional[str] = None, model_version: Optional[Union[ModelStages, int, str]] = None) -> None

Log model version metadata.

This function can be used to log metadata for existing model versions.

Parameters:

Name	Type	Description	Default
metadata	Dict[str, MetadataType]	The metadata to log.	required
model_name	Optional[str]	The name of the model to log metadata for. Can be omitted when being called inside a step with configured model in decorator.	None
model_version	Optional[Union[ModelStages, int, str]]	The version of the model to log metadata for. Can be omitted when being called inside a step with configured model in decorator.	None

Raises:

Type	Description
ValueError	If the function is not called with proper input.

Source code in src/zenml/model/utils.py

def log_model_metadata(
    metadata: Dict[str, "MetadataType"],
    model_name: Optional[str] = None,
    model_version: Optional[Union[ModelStages, int, str]] = None,
) -> None:
    """Log model version metadata.

    This function can be used to log metadata for existing model versions.

    Args:
        metadata: The metadata to log.
        model_name: The name of the model to log metadata for. Can
            be omitted when being called inside a step with configured
            `model` in decorator.
        model_version: The version of the model to log metadata for. Can
            be omitted when being called inside a step with configured
            `model` in decorator.

    Raises:
        ValueError: If the function is not called with proper input.
    """
    logger.warning(
        "The `log_model_metadata` function is deprecated and will soon be "
        "removed. Instead, you can consider using: "
        "`log_metadata(metadata={...}, infer_model=True)` instead. For more "
        "info: https://docs.zenml.io/how-to/model-management-metrics/track-metrics-metadata/attach-metadata-to-a-model"
    )

    from zenml import log_metadata

    if model_name and model_version:
        log_metadata(
            metadata=metadata,
            model_version=model_version,
            model_name=model_name,
        )
    elif model_name is None and model_version is None:
        log_metadata(
            metadata=metadata,
            infer_model=True,
        )
    else:
        raise ValueError(
            "You can call `log_model_metadata` by either providing both "
            "`model_name` and `model_version` or keeping both of them None."
        )