zenml-nightly 0.67.0.dev20241002__py3-none-any.whl → 0.68.0.dev20241027__py3-none-any.whl

zenml/model/model.py

@@ -21,6 +21,7 @@ from typing import (
     Dict,
     List,
     Optional,
+    Tuple,
     Union,
 )
 from uuid import UUID
@@ -31,6 +32,7 @@ from zenml.constants import MAX_RETRIES_FOR_VERSIONED_ENTITY_CREATION
 from zenml.enums import MetadataResourceTypes, ModelStages
 from zenml.exceptions import EntityExistsError
 from zenml.logger import get_logger
+from zenml.pipelines.pipeline_context import get_pipeline_context
 from zenml.utils.pydantic_utils import before_validator_handler
 from zenml.utils.string_utils import format_name_template
 
@@ -64,8 +66,6 @@ class Model(BaseModel):
         to a specific version/stage. If skipped new version will be created.
     save_models_to_registry: Whether to save all ModelArtifacts to Model Registry,
         if available in active stack.
-    model_version_id: The ID of a specific Model Version, if given - it will override
-        `name` and `version` settings. Used mostly internally.
     """
 
     name: str
@@ -81,13 +81,14 @@ class Model(BaseModel):
         default=None, union_mode="smart"
     )
     save_models_to_registry: bool = True
-    model_version_id: Optional[UUID] = None
 
+    # technical attributes
+    model_version_id: Optional[UUID] = None
     suppress_class_validation_warnings: bool = False
     was_created_in_this_run: bool = False
-
     _model_id: UUID = PrivateAttr(None)
-    _number: int = PrivateAttr(None)
+    _number: Optional[int] = PrivateAttr(None)
+    _created_model_version: bool = PrivateAttr(False)
 
     # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
     #  fields defined under base models. If not handled, this raises a warning.
@@ -145,15 +146,20 @@ class Model(BaseModel):
                 doesn't exist and can only be read given current
                 config (you used stage name or number as
                 a version name).
+
+        Raises:
+            KeyError: if model version doesn't exist and
+                cannot be fetched from the Model Control Plane.
         """
         if self._number is None:
             try:
-                self._get_or_create_model_version()
-            except RuntimeError:
-                logger.info(
+                mv = self._get_or_create_model_version()
+                self._number = mv.number
+            except RuntimeError as e:
+                raise KeyError(
                     f"Version `{self.version}` of `{self.name}` model doesn't "
                     "exist and cannot be fetched from the Model Control Plane."
-                )
+                ) from e
         return self._number
 
     @property
@@ -354,17 +360,13 @@ class Model(BaseModel):
             RuntimeError: If the model version run metadata cannot be fetched.
         """
         from zenml.metadata.lazy_load import RunMetadataLazyGetter
-        from zenml.new.pipelines.pipeline_context import (
-            get_pipeline_context,
-        )
 
         try:
             get_pipeline_context()
             # avoid exposing too much of internal details by keeping the return type
             return RunMetadataLazyGetter(  # type: ignore[return-value]
-                self,
-                None,
-                None,
+                self.name,
+                self._lazy_version,
             )
         except RuntimeError:
             pass
@@ -376,20 +378,6 @@ class Model(BaseModel):
             )
         return response.run_metadata
 
-    # TODO: deprecate me
-    @property
-    def metadata(self) -> Dict[str, "MetadataType"]:
-        """DEPRECATED, use `run_metadata` instead.
-
-        Returns:
-            The model version run metadata.
-        """
-        logger.warning(
-            "Model `metadata` property is deprecated. Please use "
-            "`run_metadata` instead."
-        )
-        return {k: v.value for k, v in self.run_metadata.items()}
-
     def delete_artifact(
         self,
         name: str,
@@ -461,7 +449,6 @@ class Model(BaseModel):
         name: str,
         version: Optional[str] = None,
     ) -> Optional["ArtifactVersionResponse"]:
-        from zenml import get_pipeline_context
         from zenml.models.v2.core.artifact_version import (
             LazyArtifactVersionResponse,
         )
@@ -471,9 +458,8 @@ class Model(BaseModel):
             return LazyArtifactVersionResponse(
                 lazy_load_name=name,
                 lazy_load_version=version,
-                lazy_load_model=Model(
-                    name=self.name, version=self.version or self.number
-                ),
+                lazy_load_model_name=self.name,
+                lazy_load_model_version=self._lazy_version,
             )
         except RuntimeError:
             pass
@@ -510,14 +496,21 @@ class Model(BaseModel):
 
         Returns:
             Dict of validated values.
+
+        Raises:
+            ValueError: If the model version id, but call is not internal.
         """
-        suppress_class_validation_warnings = (
-            data.get(
-                "suppress_class_validation_warnings",
-                False,
-            )
-            or data.get("model_version_id", None) is not None
+        suppress_class_validation_warnings = data.get(
+            "suppress_class_validation_warnings",
+            False,
         )
+        if not suppress_class_validation_warnings and data.get(
+            "model_version_id", None
+        ):
+            raise ValueError(
+                "`model_version_id` field is for internal use only"
+            )
+
         version = data.get("version", None)
 
         if (
@@ -576,7 +569,6 @@ class Model(BaseModel):
                     limitations=self.limitations,
                     trade_offs=self.trade_offs,
                     ethics=self.ethics,
-                    tags=self.tags,
                     user=zenml_client.active_user.id,
                     workspace=zenml_client.active_workspace.id,
                     save_models_to_registry=self.save_models_to_registry,
@@ -792,8 +784,13 @@ class Model(BaseModel):
                     retries_made += 1
             self.version = model_version.name
             self.was_created_in_this_run = True
+            self._created_model_version = True
 
-            logger.info(f"New model version `{self.version}` was created.")
+            logger.info(
+                "Created new model version `%s` for model `%s`.",
+                self.version,
+                self.name,
+            )
 
         self.model_version_id = model_version.id
         self._model_id = model_version.model.id
@@ -836,7 +833,7 @@ class Model(BaseModel):
         pipeline_run: "PipelineRunResponse",
         step_run: Optional["StepRunResponse"],
         return_logs: bool,
-    ) -> str:
+    ) -> Tuple[str, "PipelineRunResponse", Optional["StepRunResponse"]]:
         """Prepares model version inside pipeline run.
 
         Args:
@@ -878,16 +875,17 @@ class Model(BaseModel):
         if self_copy.model_version_id is None:
             model_version_response = self_copy._get_or_create_model_version()
 
+            client = Client()
             # update the configured model version id in runs accordingly
             if step_run:
-                Client().zen_store.update_run_step(
+                step_run = client.zen_store.update_run_step(
                     step_run_id=step_run.id,
                     step_run_update=StepRunUpdate(
                         model_version_id=model_version_response.id
                     ),
                 )
             else:
-                Client().zen_store.update_run(
+                pipeline_run = client.zen_store.update_run(
                     run_id=pipeline_run.id,
                     run_update=PipelineRunUpdate(
                         model_version_id=model_version_response.id
@@ -907,4 +905,22 @@ class Model(BaseModel):
                         "for a free trial at https://www.zenml.io/pro/"
                     )
         self.model_version_id = self_copy.model_version_id
-        return logs
+        return logs, pipeline_run, step_run
+
+    @property
+    def _lazy_version(self) -> Optional[str]:
+        """Get version name for lazy loader.
+
+        This getter ensures that new model version
+        creation is never triggered here.
+
+        Returns:
+            Version name or None if it was not set
+        """
+        if self._number is not None:
+            return str(self._number)
+        elif self.version is not None:
+            if isinstance(self.version, ModelStages):
+                return self.version.value
+            return str(self.version)
+        return None

zenml/model/utils.py

@@ -24,122 +24,16 @@ from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
 from zenml.model.model import Model
 from zenml.models import (
+    ArtifactVersionResponse,
     ModelVersionArtifactRequest,
+    ModelVersionResponse,
     ServiceUpdate,
 )
-from zenml.new.steps.step_context import get_step_context
+from zenml.steps.step_context import get_step_context
 
 logger = get_logger(__name__)
 
 
-def link_step_artifacts_to_model(
-    artifact_version_ids: Dict[str, UUID],
-) -> None:
-    """Links the output artifacts of a step to the model.
-
-    Args:
-        artifact_version_ids: The IDs of the published output artifacts.
-
-    Raises:
-        RuntimeError: If called outside of a step.
-    """
-    try:
-        step_context = get_step_context()
-    except StepContextError:
-        raise RuntimeError(
-            "`link_step_artifacts_to_model` can only be called from within a "
-            "step."
-        )
-    try:
-        model = step_context.model
-    except StepContextError:
-        model = None
-        logger.debug("No model context found, unable to auto-link artifacts.")
-
-    for artifact_name, artifact_version_id in artifact_version_ids.items():
-        artifact_config = step_context._get_output(
-            artifact_name
-        ).artifact_config
-
-        if artifact_config is None and model is not None:
-            artifact_config = ArtifactConfig(name=artifact_name)
-
-        if artifact_config:
-            link_artifact_config_to_model(
-                artifact_config=artifact_config,
-                artifact_version_id=artifact_version_id,
-                model=model,
-            )
-
-
-def link_artifact_config_to_model(
-    artifact_config: ArtifactConfig,
-    artifact_version_id: UUID,
-    model: Optional["Model"] = None,
-) -> None:
-    """Link an artifact config to its model version.
-
-    Args:
-        artifact_config: The artifact config to link.
-        artifact_version_id: The ID of the artifact to link.
-        model: The model version from the step or pipeline context.
-    """
-    client = Client()
-
-    # If the artifact config specifies a model itself then always use that
-    if artifact_config.model_name is not None:
-        from zenml.model.model import Model
-
-        model = Model(
-            name=artifact_config.model_name,
-            version=artifact_config.model_version,
-        )
-
-    if model:
-        logger.debug(
-            f"Linking artifact `{artifact_config.name}` to model "
-            f"`{model.name}` version `{model.version}` using config "
-            f"`{artifact_config}`."
-        )
-        request = ModelVersionArtifactRequest(
-            user=client.active_user.id,
-            workspace=client.active_workspace.id,
-            artifact_version=artifact_version_id,
-            model=model.model_id,
-            model_version=model.id,
-            is_model_artifact=artifact_config.is_model_artifact,
-            is_deployment_artifact=artifact_config.is_deployment_artifact,
-        )
-        client.zen_store.create_model_version_artifact_link(request)
-
-
-def log_model_version_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.
-    """
-    logger.warning(
-        "`log_model_version_metadata` is deprecated. Please use "
-        "`log_model_metadata` instead."
-    )
-    log_model_metadata(
-        metadata=metadata, model_name=model_name, model_version=model_version
-    )
-
-
 def log_model_metadata(
     metadata: Dict[str, "MetadataType"],
     model_name: Optional[str] = None,
@@ -179,8 +73,41 @@ def log_model_metadata(
     mv.log_metadata(metadata)
 
 
+def link_artifact_version_to_model_version(
+    artifact_version: ArtifactVersionResponse,
+    model_version: ModelVersionResponse,
+    artifact_config: Optional[ArtifactConfig] = None,
+) -> None:
+    """Link an artifact version to a model version.
+
+    Args:
+        artifact_version: The artifact version to link.
+        model_version: The model version to link.
+        artifact_config: Output artifact configuration.
+    """
+    if artifact_config:
+        is_model_artifact = artifact_config.is_model_artifact
+        is_deployment_artifact = artifact_config.is_deployment_artifact
+    else:
+        is_model_artifact = False
+        is_deployment_artifact = False
+
+    client = Client()
+    client.zen_store.create_model_version_artifact_link(
+        ModelVersionArtifactRequest(
+            user=client.active_user.id,
+            workspace=client.active_workspace.id,
+            artifact_version=artifact_version.id,
+            model=model_version.model.id,
+            model_version=model_version.id,
+            is_model_artifact=is_model_artifact,
+            is_deployment_artifact=is_deployment_artifact,
+        )
+    )
+
+
 def link_artifact_to_model(
-    artifact_version_id: UUID,
+    artifact_version: ArtifactVersionResponse,
     model: Optional["Model"] = None,
     is_model_artifact: bool = False,
     is_deployment_artifact: bool = False,
@@ -188,7 +115,7 @@ def link_artifact_to_model(
     """Link the artifact to the model.
 
     Args:
-        artifact_version_id: The ID of the artifact version.
+        artifact_version: The artifact version to link.
         model: The model to link to.
         is_model_artifact: Whether the artifact is a model artifact.
         is_deployment_artifact: Whether the artifact is a deployment artifact.
@@ -212,13 +139,15 @@ def link_artifact_to_model(
                 "@step or @pipeline decorator."
             )
 
-    link_artifact_config_to_model(
-        artifact_config=ArtifactConfig(
-            is_model_artifact=is_model_artifact,
-            is_deployment_artifact=is_deployment_artifact,
-        ),
-        artifact_version_id=artifact_version_id,
-        model=model,
+    model_version = model._get_or_create_model_version()
+    artifact_config = ArtifactConfig(
+        is_model_artifact=is_model_artifact,
+        is_deployment_artifact=is_deployment_artifact,
+    )
+    link_artifact_version_to_model_version(
+        artifact_version=artifact_version,
+        model_version=model_version,
+        artifact_config=artifact_config,
     )
 
 

zenml/models/__init__.py

@@ -292,6 +292,7 @@ from zenml.models.v2.core.step_run import (
     StepRunResponse,
     StepRunResponseBody,
     StepRunResponseMetadata,
+    StepRunResponseResources
 )
 from zenml.models.v2.core.tag import (
     TagFilter,
@@ -459,6 +460,7 @@ StackResponse.model_rebuild()
 StepRunRequest.model_rebuild()
 StepRunResponseBody.model_rebuild()
 StepRunResponseMetadata.model_rebuild()
+StepRunResponseResources.model_rebuild()
 TriggerExecutionResponseResources.model_rebuild()
 TriggerResponseBody.model_rebuild()
 TriggerResponseMetadata.model_rebuild()
@@ -685,6 +687,7 @@ __all__ = [
     "StepRunResponse",
     "StepRunResponseBody",
     "StepRunResponseMetadata",
+    "StepRunResponseResources",
     "TagFilter",
     "TagResourceResponse",
     "TagResourceResponseBody",

zenml/models/v2/core/artifact_version.py

@@ -30,7 +30,6 @@ from zenml.config.source import Source, SourceWithValidator
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
 from zenml.enums import ArtifactType, GenericFilterOps
 from zenml.logger import get_logger
-from zenml.model.model import Model
 from zenml.models.v2.base.filter import StrFilter
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedRequest,
@@ -401,18 +400,6 @@ class ArtifactVersionResponse(
             overwrite=overwrite,
         )
 
-    def read(self) -> Any:
-        """(Deprecated) Materializes (loads) the data stored in this artifact.
-
-        Returns:
-            The materialized data.
-        """
-        logger.warning(
-            "`artifact.read()` is deprecated and will be removed in a future "
-            "release. Please use `artifact.load()` instead."
-        )
-        return self.load()
-
     def visualize(self, title: Optional[str] = None) -> None:
         """Visualize the artifact in notebook environments.
 
@@ -623,7 +610,8 @@ class LazyArtifactVersionResponse(ArtifactVersionResponse):
     id: Optional[UUID] = None  # type: ignore[assignment]
     lazy_load_name: Optional[str] = None
     lazy_load_version: Optional[str] = None
-    lazy_load_model: Model
+    lazy_load_model_name: str
+    lazy_load_model_version: Optional[str] = None
 
     def get_body(self) -> None:  # type: ignore[override]
         """Protects from misuse of the lazy loader.
@@ -653,7 +641,8 @@ class LazyArtifactVersionResponse(ArtifactVersionResponse):
         from zenml.metadata.lazy_load import RunMetadataLazyGetter
 
         return RunMetadataLazyGetter(  # type: ignore[return-value]
-            self.lazy_load_model,
+            self.lazy_load_model_name,
+            self.lazy_load_model_version,
             self.lazy_load_name,
             self.lazy_load_version,
         )

zenml/models/v2/core/model.py

@@ -82,6 +82,7 @@ class ModelRequest(WorkspaceScopedRequest):
     )
     tags: Optional[List[str]] = Field(
         title="Tags associated with the model",
+        default=None,
     )
     save_models_to_registry: bool = Field(
         title="Whether to save all ModelArtifacts to Model Registry",

zenml/models/v2/core/model_version.py

@@ -326,7 +326,7 @@ class ModelVersionResponse(
     def to_model_class(
         self,
         was_created_in_this_run: bool = False,
-        suppress_class_validation_warnings: bool = False,
+        suppress_class_validation_warnings: bool = True,
     ) -> "Model":
         """Convert response model to Model object.
 

zenml/models/v2/core/pipeline_run.py

@@ -31,7 +31,6 @@ from zenml.config.pipeline_configurations import PipelineConfiguration
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.enums import ExecutionStatus
 from zenml.models.v2.base.scoped import (
-    WorkspaceScopedFilter,
     WorkspaceScopedRequest,
     WorkspaceScopedResponse,
     WorkspaceScopedResponseBody,
@@ -575,7 +574,7 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
     """Model to enable advanced filtering of all Workspaces."""
 
     FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
-        *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
+        *WorkspaceScopedTaggableFilter.FILTER_EXCLUDE_FIELDS,
         "unlisted",
         "code_repository_id",
         "build_id",

zenml/models/v2/core/run_metadata.py

@@ -21,7 +21,6 @@ from pydantic import Field, field_validator
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
 from zenml.enums import MetadataResourceTypes
 from zenml.metadata.metadata_types import MetadataType, MetadataTypeEnum
-from zenml.model.model import Model
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedFilter,
     WorkspaceScopedRequest,
@@ -239,7 +238,8 @@ class LazyRunMetadataResponse(RunMetadataResponse):
     lazy_load_artifact_name: Optional[str] = None
     lazy_load_artifact_version: Optional[str] = None
     lazy_load_metadata_name: Optional[str] = None
-    lazy_load_model: Model
+    lazy_load_model_name: str
+    lazy_load_model_version: Optional[str] = None
 
     def get_body(self) -> None:  # type: ignore[override]
         """Protects from misuse of the lazy loader.

zenml/models/v2/core/step_run.py

@@ -156,6 +156,14 @@ class StepRunResponseBody(WorkspaceScopedResponseBody):
     """Response body for step runs."""
 
     status: ExecutionStatus = Field(title="The status of the step.")
+    start_time: Optional[datetime] = Field(
+        title="The start time of the step run.",
+        default=None,
+    )
+    end_time: Optional[datetime] = Field(
+        title="The end time of the step run.",
+        default=None,
+    )
     inputs: Dict[str, "ArtifactVersionResponse"] = Field(
         title="The input artifact versions of the step run.",
         default={},
@@ -201,16 +209,6 @@ class StepRunResponseMetadata(WorkspaceScopedResponseMetadata):
         max_length=TEXT_FIELD_MAX_LENGTH,
     )
 
-    # Timestamps
-    start_time: Optional[datetime] = Field(
-        title="The start time of the step run.",
-        default=None,
-    )
-    end_time: Optional[datetime] = Field(
-        title="The end time of the step run.",
-        default=None,
-    )
-
     # References
     logs: Optional["LogsResponse"] = Field(
         title="Logs associated with this step run.",
@@ -409,7 +407,7 @@ class StepRunResponse(
         Returns:
             the value of the property.
         """
-        return self.get_metadata().start_time
+        return self.get_body().start_time
 
     @property
     def end_time(self) -> Optional[datetime]:
@@ -418,7 +416,7 @@ class StepRunResponse(
         Returns:
             the value of the property.
         """
-        return self.get_metadata().end_time
+        return self.get_body().end_time
 
     @property
     def logs(self) -> Optional["LogsResponse"]:
@@ -521,6 +519,11 @@ class StepRunFilter(WorkspaceScopedFilter):
         description="Pipeline run of this step run",
         union_mode="left_to_right",
     )
+    deployment_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Deployment of this step run",
+        union_mode="left_to_right",
+    )
     original_step_run_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="Original id for this step run",

zenml/models/v2/misc/server_models.py

@@ -99,10 +99,6 @@ class ServerModel(BaseModel):
         {},
         title="The metadata associated with the server.",
     )
-    use_legacy_dashboard: bool = Field(
-        False,
-        title="Flag to indicate whether the server is using the legacy dashboard.",
-    )
 
     last_user_activity: Optional[datetime] = Field(
         None,

zenml/orchestrators/base_orchestrator.py

@@ -18,6 +18,10 @@ from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional, Type, cast
 
 from pydantic import model_validator
 
+from zenml.constants import (
+    ENV_ZENML_PREVENT_CLIENT_SIDE_CACHING,
+    handle_bool_env_var,
+)
 from zenml.enums import ExecutionStatus, StackComponentType
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
@@ -185,6 +189,35 @@ class BaseOrchestrator(StackComponent, ABC):
 
         environment = get_config_environment_vars(deployment=deployment)
 
+        prevent_client_side_caching = handle_bool_env_var(
+            ENV_ZENML_PREVENT_CLIENT_SIDE_CACHING, default=False
+        )
+
+        if (
+            placeholder_run
+            and not deployment.schedule
+            and not prevent_client_side_caching
+        ):
+            from zenml.orchestrators import step_run_utils
+
+            cached_invocations = step_run_utils.create_cached_step_runs(
+                deployment=deployment,
+                pipeline_run=placeholder_run,
+                stack=stack,
+            )
+
+            for invocation_id in cached_invocations:
+                # Remove the cached step invocations from the deployment so
+                # the orchestrator does not try to run them
+                deployment.step_configurations.pop(invocation_id)
+
+            if len(deployment.step_configurations) == 0:
+                # All steps were cached, we update the pipeline run status and
+                # don't actually use the orchestrator to run the pipeline
+                self._cleanup_run()
+                logger.info("All steps of the pipeline run were cached.")
+                return
+
         try:
             if metadata_iterator := self.prepare_or_run_pipeline(
                 deployment=deployment,