zenml-nightly 0.68.1.dev20241103__py3-none-any.whl → 0.70.0.dev20241115__py3-none-any.whl

zenml/materializers/base_materializer.py

@@ -13,8 +13,11 @@
 #  permissions and limitations under the License.
 """Metaclass implementation for registering ZenML BaseMaterializer subclasses."""
 
+import contextlib
 import inspect
-from typing import Any, ClassVar, Dict, Optional, Tuple, Type, cast
+import shutil
+import tempfile
+from typing import Any, ClassVar, Dict, Iterator, Optional, Tuple, Type, cast
 
 from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
 from zenml.enums import ArtifactType, VisualizationType
@@ -326,3 +329,67 @@ class BaseMaterializer(metaclass=BaseMaterializerMeta):
         if isinstance(storage_size, int):
             return {"storage_size": StorageSize(storage_size)}
         return {}
+
+    @contextlib.contextmanager
+    def get_temporary_directory(
+        self,
+        delete_at_exit: bool,
+        delete_after_step_execution: bool = True,
+    ) -> Iterator[str]:
+        """Context manager to get a temporary directory.
+
+        Args:
+            delete_at_exit: If set to True, the temporary directory will be
+                deleted after the context manager exits.
+            delete_after_step_execution: If `delete_at_exit` is set to False and
+                this is set to True, the temporary directory will be deleted
+                after the step finished executing. If a materializer is being
+                used outside of the context of a step execution, the temporary
+                directory will not be deleted and the user is responsible for
+                deleting it themselves.
+
+        Yields:
+            Path to the temporary directory.
+        """
+        temp_dir = tempfile.mkdtemp(prefix="zenml-")
+
+        if delete_after_step_execution and not delete_at_exit:
+            # We should not delete the directory when the context manager
+            # exits, but cleanup once the step has finished executing.
+            self._register_directory_for_deletion_after_step_execution(
+                temp_dir
+            )
+
+        try:
+            yield temp_dir
+        finally:
+            if delete_at_exit:
+                shutil.rmtree(temp_dir)
+
+    def _register_directory_for_deletion_after_step_execution(
+        self, directory: str
+    ) -> None:
+        """Register directory to be deleted after the current step finishes.
+
+        If no step is currently being executed, this method does nothing.
+
+        Args:
+            directory: The directory to register for deletion.
+        """
+        from zenml import get_step_context
+
+        try:
+            step_context = get_step_context()
+        except RuntimeError:
+            logger.debug(
+                "Materializer called outside of step execution, not cleaning "
+                "up directory %s",
+                directory,
+            )
+            return
+
+        def _callback() -> None:
+            shutil.rmtree(directory)
+            logger.debug("Cleaned up materializer directory %s", directory)
+
+        step_context._cleanup_registry.register_callback(_callback)

zenml/metadata/lazy_load.py

@@ -13,10 +13,25 @@
 #  permissions and limitations under the License.
 """Run Metadata Lazy Loader definition."""
 
-from typing import TYPE_CHECKING, Optional
+from typing import Optional
 
-if TYPE_CHECKING:
-    from zenml.models import RunMetadataResponse
+from pydantic import BaseModel
+
+from zenml.metadata.metadata_types import MetadataType
+
+
+class LazyRunMetadataResponse(BaseModel):
+    """Lazy run metadata response.
+
+    Used if the run metadata is accessed from the model in
+    a pipeline context available only during pipeline compilation.
+    """
+
+    lazy_load_artifact_name: Optional[str] = None
+    lazy_load_artifact_version: Optional[str] = None
+    lazy_load_metadata_name: Optional[str] = None
+    lazy_load_model_name: str
+    lazy_load_model_version: Optional[str] = None
 
 
 class RunMetadataLazyGetter:
@@ -47,7 +62,7 @@ class RunMetadataLazyGetter:
         self._lazy_load_artifact_name = _lazy_load_artifact_name
         self._lazy_load_artifact_version = _lazy_load_artifact_version
 
-    def __getitem__(self, key: str) -> "RunMetadataResponse":
+    def __getitem__(self, key: str) -> MetadataType:
         """Get the metadata for the given key.
 
         Args:
@@ -56,9 +71,7 @@ class RunMetadataLazyGetter:
         Returns:
             The metadata lazy loader wrapper for the given key.
         """
-        from zenml.models.v2.core.run_metadata import LazyRunMetadataResponse
-
-        return LazyRunMetadataResponse(
+        return LazyRunMetadataResponse(  # type: ignore[return-value]
             lazy_load_model_name=self._lazy_load_model_name,
             lazy_load_model_version=self._lazy_load_model_version,
             lazy_load_artifact_name=self._lazy_load_artifact_name,

zenml/model/model.py

@@ -14,7 +14,6 @@
 """Model user facing interface to pass into pipeline or step."""
 
 import datetime
-import time
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -28,7 +27,6 @@ from uuid import UUID
 
 from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_validator
 
-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
@@ -43,7 +41,6 @@ if TYPE_CHECKING:
         ModelResponse,
         ModelVersionResponse,
         PipelineRunResponse,
-        RunMetadataResponse,
         StepRunResponse,
     )
 
@@ -349,7 +346,7 @@ class Model(BaseModel):
         )
 
     @property
-    def run_metadata(self) -> Dict[str, "RunMetadataResponse"]:
+    def run_metadata(self) -> Dict[str, "MetadataType"]:
         """Get model version run metadata.
 
         Returns:
@@ -528,14 +525,6 @@ class Model(BaseModel):
         data["suppress_class_validation_warnings"] = True
         return data
 
-    def _validate_config_in_runtime(self) -> "ModelVersionResponse":
-        """Validate that config doesn't conflict with runtime environment.
-
-        Returns:
-            The model version based on configuration.
-        """
-        return self._get_or_create_model_version()
-
     def _get_or_create_model(self) -> "ModelResponse":
         """This method should get or create a model from Model Control Plane.
 
@@ -679,16 +668,6 @@ class Model(BaseModel):
         if isinstance(self.version, str):
             self.version = format_name_template(self.version)
 
-        zenml_client = Client()
-        model_version_request = ModelVersionRequest(
-            user=zenml_client.active_user.id,
-            workspace=zenml_client.active_workspace.id,
-            name=str(self.version) if self.version else None,
-            description=self.description,
-            model=model.id,
-            tags=self.tags,
-        )
-        mv_request = ModelVersionRequest.model_validate(model_version_request)
         try:
             if self.version or self.model_version_id:
                 model_version = self._get_model_version()
@@ -718,60 +697,34 @@ class Model(BaseModel):
                     " as an example. You can explore model versions using "
                     f"`zenml model version list -n {self.name}` CLI command."
                 )
-            retries_made = 0
-            for i in range(MAX_RETRIES_FOR_VERSIONED_ENTITY_CREATION):
-                try:
-                    model_version = (
-                        zenml_client.zen_store.create_model_version(
-                            model_version=mv_request
-                        )
-                    )
-                    break
-                except EntityExistsError as e:
-                    if i == MAX_RETRIES_FOR_VERSIONED_ENTITY_CREATION - 1:
-                        raise RuntimeError(
-                            f"Failed to create model version "
-                            f"`{self.version if self.version else 'new'}` "
-                            f"in model `{self.name}`. Retried {retries_made} times. "
-                            "This could be driven by exceptionally high concurrency of "
-                            "pipeline runs. Please, reach out to us on ZenML Slack for support."
-                        ) from e
-                    # smoothed exponential back-off, it will go as 0.2, 0.3,
-                    # 0.45, 0.68, 1.01, 1.52, 2.28, 3.42, 5.13, 7.69, ...
-                    sleep = 0.2 * 1.5**i
-                    logger.debug(
-                        f"Failed to create new model version for "
-                        f"model `{self.name}`. Retrying in {sleep}..."
-                    )
-                    time.sleep(sleep)
-                    retries_made += 1
-            self.version = model_version.name
+
+            client = Client()
+            model_version_request = ModelVersionRequest(
+                user=client.active_user.id,
+                workspace=client.active_workspace.id,
+                name=str(self.version) if self.version else None,
+                description=self.description,
+                model=model.id,
+                tags=self.tags,
+            )
+            model_version = client.zen_store.create_model_version(
+                model_version=model_version_request
+            )
+
             self._created_model_version = True
 
             logger.info(
                 "Created new model version `%s` for model `%s`.",
-                self.version,
+                model_version.name,
                 self.name,
             )
 
+        self.version = model_version.name
         self.model_version_id = model_version.id
         self._model_id = model_version.model.id
         self._number = model_version.number
         return model_version
 
-    def _merge(self, model: "Model") -> None:
-        self.license = self.license or model.license
-        self.description = self.description or model.description
-        self.audience = self.audience or model.audience
-        self.use_cases = self.use_cases or model.use_cases
-        self.limitations = self.limitations or model.limitations
-        self.trade_offs = self.trade_offs or model.trade_offs
-        self.ethics = self.ethics or model.ethics
-        if model.tags is not None:
-            self.tags = list(
-                {t for t in self.tags or []}.union(set(model.tags))
-            )
-
     def __hash__(self) -> int:
         """Get hash of the `Model`.
 

zenml/model/utils.py

@@ -56,6 +56,11 @@ def log_model_metadata(
         ValueError: If no model name/version is provided and the function is not
             called inside a step with configured `model` in decorator.
     """
+    logger.warning(
+        "The `log_model_metadata` function is deprecated and will soon be "
+        "removed. Please use `log_metadata` instead."
+    )
+
     if model_name and model_version:
         from zenml import Model
 

zenml/models/__init__.py

@@ -239,12 +239,7 @@ from zenml.models.v2.core.run_template import (
 )
 from zenml.models.v2.base.base_plugin_flavor import BasePluginFlavorResponse
 from zenml.models.v2.core.run_metadata import (
-    LazyRunMetadataResponse,
     RunMetadataRequest,
-    RunMetadataFilter,
-    RunMetadataResponse,
-    RunMetadataResponseBody,
-    RunMetadataResponseMetadata,
 )
 from zenml.models.v2.core.schedule import (
     ScheduleRequest,
@@ -418,7 +413,6 @@ EventSourceResponseResources.model_rebuild()
 FlavorResponseBody.model_rebuild()
 FlavorResponseMetadata.model_rebuild()
 LazyArtifactVersionResponse.model_rebuild()
-LazyRunMetadataResponse.model_rebuild()
 ModelResponseBody.model_rebuild()
 ModelResponseMetadata.model_rebuild()
 ModelVersionResponseBody.model_rebuild()
@@ -444,8 +438,6 @@ RunTemplateResponseBody.model_rebuild()
 RunTemplateResponseMetadata.model_rebuild()
 RunTemplateResponseResources.model_rebuild()
 RunTemplateResponseBody.model_rebuild()
-RunMetadataResponseBody.model_rebuild()
-RunMetadataResponseMetadata.model_rebuild()
 ScheduleResponseBody.model_rebuild()
 ScheduleResponseMetadata.model_rebuild()
 SecretResponseBody.model_rebuild()
@@ -637,10 +629,6 @@ __all__ = [
     "RunTemplateResponseResources",
     "RunTemplateFilter",
     "RunMetadataRequest",
-    "RunMetadataFilter",
-    "RunMetadataResponse",
-    "RunMetadataResponseBody",
-    "RunMetadataResponseMetadata",
     "ScheduleRequest",
     "ScheduleUpdate",
     "ScheduleFilter",

zenml/models/v2/base/filter.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """Base filter model definitions."""
 
+import json
 from abc import ABC, abstractmethod
 from datetime import datetime
 from typing import (
@@ -36,7 +37,7 @@ from pydantic import (
     field_validator,
     model_validator,
 )
-from sqlalchemy import asc, desc
+from sqlalchemy import Float, and_, asc, cast, desc
 from sqlmodel import SQLModel
 
 from zenml.constants import (
@@ -63,6 +64,11 @@ logger = get_logger(__name__)
 
 AnyQuery = TypeVar("AnyQuery", bound=Any)
 
+ONEOF_ERROR = (
+    "When you are using the 'oneof:' filtering make sure that the "
+    "provided value is a json formatted list."
+)
+
 
 class Filter(BaseModel, ABC):
     """Filter for all fields.
@@ -171,8 +177,28 @@ class StrFilter(Filter):
         GenericFilterOps.STARTSWITH,
         GenericFilterOps.CONTAINS,
         GenericFilterOps.ENDSWITH,
+        GenericFilterOps.ONEOF,
+        GenericFilterOps.GT,
+        GenericFilterOps.GTE,
+        GenericFilterOps.LT,
+        GenericFilterOps.LTE,
     ]
 
+    @model_validator(mode="after")
+    def check_value_if_operation_oneof(self) -> "StrFilter":
+        """Validator to check if value is a list if oneof operation is used.
+
+        Raises:
+            ValueError: If the value is not a list
+
+        Returns:
+            self
+        """
+        if self.operation == GenericFilterOps.ONEOF:
+            if not isinstance(self.value, list):
+                raise ValueError(ONEOF_ERROR)
+        return self
+
     def generate_query_conditions_from_column(self, column: Any) -> Any:
         """Generate query conditions for a string column.
 
@@ -181,6 +207,9 @@ class StrFilter(Filter):
 
         Returns:
             A list of query conditions.
+
+        Raises:
+            ValueError: the comparison of the column to a numeric value fails.
         """
         if self.operation == GenericFilterOps.CONTAINS:
             return column.like(f"%{self.value}%")
@@ -190,6 +219,40 @@ class StrFilter(Filter):
             return column.endswith(f"{self.value}")
         if self.operation == GenericFilterOps.NOT_EQUALS:
             return column != self.value
+        if self.operation == GenericFilterOps.ONEOF:
+            return column.in_(self.value)
+        if self.operation in {
+            GenericFilterOps.GT,
+            GenericFilterOps.LT,
+            GenericFilterOps.GTE,
+            GenericFilterOps.LTE,
+        }:
+            try:
+                numeric_column = cast(column, Float)
+
+                assert self.value is not None
+
+                if self.operation == GenericFilterOps.GT:
+                    return and_(
+                        numeric_column, numeric_column > float(self.value)
+                    )
+                if self.operation == GenericFilterOps.LT:
+                    return and_(
+                        numeric_column, numeric_column < float(self.value)
+                    )
+                if self.operation == GenericFilterOps.GTE:
+                    return and_(
+                        numeric_column, numeric_column >= float(self.value)
+                    )
+                if self.operation == GenericFilterOps.LTE:
+                    return and_(
+                        numeric_column, numeric_column <= float(self.value)
+                    )
+            except Exception as e:
+                raise ValueError(
+                    f"Failed to compare the column '{column}' to the "
+                    f"value '{self.value}' (must be numeric): {e}"
+                )
 
         return column == self.value
 
@@ -211,6 +274,9 @@ class UUIDFilter(StrFilter):
         if isinstance(value, str):
             return value.replace("-", "")
 
+        if isinstance(value, list):
+            return [str(v).replace("-", "") for v in value]
+
         return value
 
     def generate_query_conditions_from_column(self, column: Any) -> Any:
@@ -588,6 +654,10 @@ class BaseFilter(BaseModel):
 
         Returns:
             A tuple of the filter value and the operator.
+
+        Raises:
+            ValueError: when we try to use the `oneof` operator with the wrong
+                value.
         """
         operator = GenericFilterOps.EQUALS  # Default operator
         if isinstance(value, str):
@@ -598,6 +668,15 @@ class BaseFilter(BaseModel):
             ):
                 value = split_value[1]
                 operator = GenericFilterOps(split_value[0])
+
+            if operator == operator.ONEOF:
+                try:
+                    value = json.loads(value)
+                    if not isinstance(value, list):
+                        raise ValueError
+                except ValueError:
+                    raise ValueError(ONEOF_ERROR)
+
         return value, operator
 
     def generate_name_or_id_query_conditions(
@@ -648,8 +727,8 @@ class BaseFilter(BaseModel):
 
         return or_(*conditions)
 
+    @staticmethod
     def generate_custom_query_conditions_for_column(
-        self,
         value: Any,
         table: Type[SQLModel],
         column: str,
@@ -833,16 +912,17 @@ class FilterGenerator:
 
         # Create str filters
         if self.is_str_field(column):
-            return StrFilter(
-                operation=GenericFilterOps(operator),
+            return self._define_str_filter(
+                operator=GenericFilterOps(operator),
                 column=column,
                 value=value,
             )
 
         # Handle unsupported datatypes
         logger.warning(
-            f"The Datatype {self._model_class.model_fields[column].annotation} might "
-            "not be supported for filtering. Defaulting to a string filter."
+            f"The Datatype {self._model_class.model_fields[column].annotation} "
+            "might not be supported for filtering. Defaulting to a string "
+            "filter."
         )
         return StrFilter(
             operation=GenericFilterOps(operator),
@@ -1032,8 +1112,9 @@ class FilterGenerator:
                     "Invalid value passed as UUID query parameter."
                 ) from e
 
-        # Cast the value to string for further comparisons.
-        value = str(value)
+        # For equality checks, ensure that the value is a valid UUID.
+        if operator == GenericFilterOps.ONEOF and not isinstance(value, list):
+            raise ValueError(ONEOF_ERROR)
 
         # Generate the filter.
         uuid_filter = UUIDFilter(
@@ -1043,6 +1124,38 @@ class FilterGenerator:
         )
         return uuid_filter
 
+    @staticmethod
+    def _define_str_filter(
+        column: str, value: Any, operator: GenericFilterOps
+    ) -> StrFilter:
+        """Define a str filter for a given column.
+
+        Args:
+            column: The column to filter on.
+            value: The UUID value by which to filter.
+            operator: The operator to use for filtering.
+
+        Returns:
+            A Filter object.
+
+        Raises:
+            ValueError: If the value is not a proper value.
+        """
+        # For equality checks, ensure that the value is a valid UUID.
+        if operator == GenericFilterOps.ONEOF and not isinstance(value, list):
+            raise ValueError(
+                "If you are using `oneof:` as a filtering op, the value needs "
+                "to be a json formatted list string."
+            )
+
+        # Generate the filter.
+        str_filter = StrFilter(
+            operation=GenericFilterOps(operator),
+            column=column,
+            value=value,
+        )
+        return str_filter
+
     @staticmethod
     def _define_bool_filter(
         column: str, value: Any, operator: GenericFilterOps

zenml/models/v2/core/artifact_version.py

@@ -34,7 +34,7 @@ from pydantic import (
 
 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.enums import ArtifactSaveType, ArtifactType, GenericFilterOps
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
 from zenml.models.v2.base.filter import StrFilter
@@ -57,9 +57,6 @@ if TYPE_CHECKING:
         ArtifactVisualizationResponse,
     )
     from zenml.models.v2.core.pipeline_run import PipelineRunResponse
-    from zenml.models.v2.core.run_metadata import (
-        RunMetadataResponse,
-    )
     from zenml.models.v2.core.step_run import StepRunResponse
 
 logger = get_logger(__name__)
@@ -107,6 +104,9 @@ class ArtifactVersionRequest(WorkspaceScopedRequest):
     visualizations: Optional[List["ArtifactVisualizationRequest"]] = Field(
         default=None, title="Visualizations of the artifact."
     )
+    save_type: ArtifactSaveType = Field(
+        title="The save type of the artifact version.",
+    )
     metadata: Optional[Dict[str, MetadataType]] = Field(
         default=None, title="Metadata of the artifact version."
     )
@@ -193,6 +193,9 @@ class ArtifactVersionResponseBody(WorkspaceScopedResponseBody):
         title="The ID of the pipeline run that generated this artifact version.",
         default=None,
     )
+    save_type: ArtifactSaveType = Field(
+        title="The save type of the artifact version.",
+    )
     artifact_store_id: Optional[UUID] = Field(
         title="ID of the artifact store in which this artifact is stored.",
         default=None,
@@ -230,7 +233,7 @@ class ArtifactVersionResponseMetadata(WorkspaceScopedResponseMetadata):
     visualizations: Optional[List["ArtifactVisualizationResponse"]] = Field(
         default=None, title="Visualizations of the artifact."
     )
-    run_metadata: Dict[str, "RunMetadataResponse"] = Field(
+    run_metadata: Dict[str, MetadataType] = Field(
         default={}, title="Metadata of the artifact."
     )
 
@@ -313,6 +316,15 @@ class ArtifactVersionResponse(
         """
         return self.get_body().producer_pipeline_run_id
 
+    @property
+    def save_type(self) -> ArtifactSaveType:
+        """The `save_type` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().save_type
+
     @property
     def artifact_store_id(self) -> Optional[UUID]:
         """The `artifact_store_id` property.
@@ -343,7 +355,7 @@ class ArtifactVersionResponse(
         return self.get_metadata().visualizations
 
     @property
-    def run_metadata(self) -> Dict[str, "RunMetadataResponse"]:
+    def run_metadata(self) -> Dict[str, MetadataType]:
         """The `metadata` property.
 
         Returns:
@@ -462,6 +474,7 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
         "user",
         "model",
         "pipeline_run",
+        "run_metadata",
     ]
     artifact_id: Optional[Union[UUID, str]] = Field(
         default=None,
@@ -533,6 +546,10 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
         description="Name/ID of a pipeline run that is associated with this "
         "artifact version.",
     )
+    run_metadata: Optional[Dict[str, str]] = Field(
+        default=None,
+        description="The run_metadata to filter the artifact versions by.",
+    )
 
     model_config = ConfigDict(protected_namespaces=())
 
@@ -552,6 +569,7 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
             ModelSchema,
             ModelVersionArtifactSchema,
             PipelineRunSchema,
+            RunMetadataSchema,
             StepRunInputArtifactSchema,
             StepRunOutputArtifactSchema,
             StepRunSchema,
@@ -633,6 +651,23 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
             )
             custom_filters.append(pipeline_run_filter)
 
+        if self.run_metadata is not None:
+            from zenml.enums import MetadataResourceTypes
+
+            for key, value in self.run_metadata.items():
+                additional_filter = and_(
+                    RunMetadataSchema.resource_id == ArtifactVersionSchema.id,
+                    RunMetadataSchema.resource_type
+                    == MetadataResourceTypes.ARTIFACT_VERSION,
+                    RunMetadataSchema.key == key,
+                    self.generate_custom_query_conditions_for_column(
+                        value=value,
+                        table=RunMetadataSchema,
+                        column="value",
+                    ),
+                )
+                custom_filters.append(additional_filter)
+
         return custom_filters
 
 
@@ -671,7 +706,7 @@ class LazyArtifactVersionResponse(ArtifactVersionResponse):
         )
 
     @property
-    def run_metadata(self) -> Dict[str, "RunMetadataResponse"]:
+    def run_metadata(self) -> Dict[str, MetadataType]:
         """The `metadata` property in lazy loading mode.
 
         Returns: