zenml-nightly 0.66.0.dev20240927__py3-none-any.whl → 0.68.1.dev20241102__py3-none-any.whl

zenml/materializers/built_in_materializer.py

@@ -403,7 +403,7 @@ class BuiltInContainerMaterializer(BaseMaterializer):
             yaml_utils.write_json(self.metadata_path, metadata)
             # Materialize each element.
             for element, materializer in zip(data, materializers):
-                materializer.validate_type_compatibility(type(element))
+                materializer.validate_save_type_compatibility(type(element))
                 materializer.save(element)
         # If an error occurs, delete all created files.
         except Exception as e:

zenml/materializers/uuid_materializer.py

@@ -0,0 +1,79 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Implementation of ZenML's UUID materializer."""
+
+import os
+import uuid
+from typing import Any, ClassVar, Dict, Optional, Tuple, Type
+
+from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
+from zenml.enums import ArtifactType
+from zenml.materializers.base_materializer import BaseMaterializer
+from zenml.metadata.metadata_types import MetadataType
+
+DEFAULT_FILENAME = "uuid.txt"
+
+
+class UUIDMaterializer(BaseMaterializer):
+    """Materializer to handle UUID objects."""
+
+    ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = (uuid.UUID,)
+    ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.DATA
+
+    def __init__(
+        self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
+    ):
+        """Define `self.data_path`.
+
+        Args:
+            uri: The URI where the artifact data is stored.
+            artifact_store: The artifact store where the artifact data is stored.
+        """
+        super().__init__(uri, artifact_store)
+        self.data_path = os.path.join(self.uri, DEFAULT_FILENAME)
+
+    def load(self, _: Type[uuid.UUID]) -> uuid.UUID:
+        """Read UUID from artifact store.
+
+        Args:
+            _: The type of the data to be loaded.
+
+        Returns:
+            The loaded UUID.
+        """
+        with self.artifact_store.open(self.data_path, "r") as f:
+            uuid_str = f.read().strip()
+        return uuid.UUID(uuid_str)
+
+    def save(self, data: uuid.UUID) -> None:
+        """Write UUID to artifact store.
+
+        Args:
+            data: The UUID to be saved.
+        """
+        with self.artifact_store.open(self.data_path, "w") as f:
+            f.write(str(data))
+
+    def extract_metadata(self, data: uuid.UUID) -> Dict[str, MetadataType]:
+        """Extract metadata from the UUID.
+
+        Args:
+            data: The UUID to extract metadata from.
+
+        Returns:
+            A dictionary of metadata extracted from the UUID.
+        """
+        return {
+            "string_representation": str(data),
+        }

zenml/metadata/lazy_load.py

@@ -16,7 +16,6 @@
 from typing import TYPE_CHECKING, Optional
 
 if TYPE_CHECKING:
-    from zenml.model.model import Model
     from zenml.models import RunMetadataResponse
 
 
@@ -30,18 +29,21 @@ class RunMetadataLazyGetter:
 
     def __init__(
         self,
-        _lazy_load_model: "Model",
-        _lazy_load_artifact_name: Optional[str],
-        _lazy_load_artifact_version: Optional[str],
+        _lazy_load_model_name: str,
+        _lazy_load_model_version: Optional[str],
+        _lazy_load_artifact_name: Optional[str] = None,
+        _lazy_load_artifact_version: Optional[str] = None,
     ):
         """Initialize a RunMetadataLazyGetter.
 
         Args:
-            _lazy_load_model: The model version.
+            _lazy_load_model_name: The model name.
+            _lazy_load_model_version: The model version.
             _lazy_load_artifact_name: The artifact name.
             _lazy_load_artifact_version: The artifact version.
         """
-        self._lazy_load_model = _lazy_load_model
+        self._lazy_load_model_name = _lazy_load_model_name
+        self._lazy_load_model_version = _lazy_load_model_version
         self._lazy_load_artifact_name = _lazy_load_artifact_name
         self._lazy_load_artifact_version = _lazy_load_artifact_version
 
@@ -57,7 +59,8 @@ class RunMetadataLazyGetter:
         from zenml.models.v2.core.run_metadata import LazyRunMetadataResponse
 
         return LazyRunMetadataResponse(
-            lazy_load_model=self._lazy_load_model,
+            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,
             lazy_load_artifact_version=self._lazy_load_artifact_version,
             lazy_load_metadata_name=key,

zenml/metadata/metadata_types.py

@@ -13,13 +13,18 @@
 #  permissions and limitations under the License.
 """Custom types that can be used as metadata of ZenML artifacts."""
 
+import json
 from typing import Any, Dict, List, Set, Tuple, Union
 
 from pydantic import GetCoreSchemaHandler
 from pydantic_core import CoreSchema, core_schema
 
+from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
+from zenml.logger import get_logger
 from zenml.utils.enum_utils import StrEnum
 
+logger = get_logger(__name__)
+
 
 class Uri(str):
     """Special string class to indicate a URI."""
@@ -203,3 +208,47 @@ def cast_to_metadata_type(
     metadata_type = metadata_enum_to_type_mapping[type_]
     typed_value = metadata_type(value)
     return typed_value  # type: ignore[no-any-return]
+
+
+def validate_metadata(
+    metadata: Dict[str, MetadataType],
+) -> Dict[str, MetadataType]:
+    """Validate metadata.
+
+    This function excludes and warns about metadata values that are too long
+    or of an unsupported type.
+
+    Args:
+        metadata: The metadata to validate.
+
+    Returns:
+        The validated metadata.
+    """
+    validated_metadata = {}
+
+    for key, value in metadata.items():
+        if len(key) > STR_FIELD_MAX_LENGTH:
+            logger.warning(
+                f"Metadata key '{key}' is too large to be "
+                "stored in the database. Skipping."
+            )
+            continue
+
+        if len(json.dumps(value)) > TEXT_FIELD_MAX_LENGTH:
+            logger.warning(
+                f"Metadata value for key '{key}' is too large to be "
+                "stored in the database. Skipping."
+            )
+            continue
+
+        try:
+            get_metadata_type(value)
+        except ValueError as e:
+            logger.warning(
+                f"Metadata value for key '{key}' is not of a supported "
+                f"type. Skipping. Full error: {e}"
+            )
+
+        validated_metadata[key] = value
+
+    return validated_metadata

zenml/model/lazy_load.py

@@ -13,11 +13,15 @@
 #  permissions and limitations under the License.
 """Model Version Data Lazy Loader definition."""
 
-from typing import Optional
+from typing import TYPE_CHECKING, Any, Dict, Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict, model_validator
 
-from zenml.model.model import Model
+from zenml.pipelines.pipeline_context import get_pipeline_context
+from zenml.utils.pydantic_utils import before_validator_handler
+
+if TYPE_CHECKING:
+    from zenml.models import ModelVersionResponse, PipelineRunResponse
 
 
 class ModelVersionDataLazyLoader(BaseModel):
@@ -28,7 +32,84 @@ class ModelVersionDataLazyLoader(BaseModel):
     model version during runtime time of the step.
     """
 
-    model: Model
+    model_name: str
+    model_version: Optional[str] = None
     artifact_name: Optional[str] = None
     artifact_version: Optional[str] = None
     metadata_name: Optional[str] = None
+
+    # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
+    #  fields defined under base models. If not handled, this raises a warning.
+    #  It is possible to suppress this warning message with the following
+    #  configuration, however the ultimate solution is to rename these fields.
+    #  Even though they do not cause any problems right now, if we are not
+    #  careful we might overwrite some fields protected by pydantic.
+    model_config = ConfigDict(protected_namespaces=())
+
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def _root_validator(cls, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Validate all in one.
+
+        Args:
+            data: Dict of values.
+
+        Returns:
+            Dict of validated values.
+
+        Raises:
+            ValueError: If the model version id, but call is not internal.
+        """
+        if data.get("model_version", None) is None:
+            try:
+                context = get_pipeline_context()
+                if (
+                    not context.model
+                    or context.model.name != data["model_name"]
+                ):
+                    raise ValueError(
+                        "`version` must be set if you use the `Model` class "
+                        "directly in the pipeline body, otherwise, you can use "
+                        "`get_pipeline_context().model` to lazy load the current "
+                        "Model Version from the pipeline context."
+                    )
+            except RuntimeError:
+                pass
+        data["suppress_class_validation_warnings"] = True
+        return data
+
+    def _get_model_response(
+        self, pipeline_run: "PipelineRunResponse"
+    ) -> "ModelVersionResponse":
+        # if the version/number is None -> return the model in context
+        if self.model_version is None:
+            if mv := pipeline_run.model_version:
+                if mv.model.name != self.model_name:
+                    raise RuntimeError(
+                        "Lazy loading of the model failed, since given name "
+                        f"`{self.model_name}` does not match the model name "
+                        f"in the pipeline context: `{mv.model.name}`."
+                    )
+                return mv
+            else:
+                raise RuntimeError(
+                    "Lazy loading of the model failed, since the model version "
+                    "is not set in the pipeline context."
+                )
+
+        # else return the model version by version
+        else:
+            from zenml.client import Client
+
+            try:
+                return Client().get_model_version(
+                    model_name_or_id=self.model_name,
+                    model_version_name_or_number_or_id=self.model_version,
+                )
+            except KeyError as e:
+                raise RuntimeError(
+                    "Lazy loading of the model version failed: "
+                    f"no model `{self.model_name}` with version "
+                    f"`{self.model_version}` could be found."
+                ) from e

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,13 @@ 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 +145,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 +359,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 +377,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 +448,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 +457,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 +495,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 +568,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,
@@ -699,42 +690,6 @@ class Model(BaseModel):
         )
         mv_request = ModelVersionRequest.model_validate(model_version_request)
         try:
-            if not self.version:
-                try:
-                    from zenml import get_step_context
-
-                    context = get_step_context()
-                except RuntimeError:
-                    pass
-                else:
-                    # if inside a step context we loop over all
-                    # model version configuration to find, if the
-                    # model version for current model was already
-                    # created in the current run, not to create
-                    # new model versions
-                    pipeline_mv = context.pipeline_run.config.model
-                    if (
-                        pipeline_mv
-                        and pipeline_mv.was_created_in_this_run
-                        and pipeline_mv.name == self.name
-                        and pipeline_mv.version is not None
-                    ):
-                        self.version = pipeline_mv.version
-                        self.model_version_id = pipeline_mv.model_version_id
-                    else:
-                        for step in context.pipeline_run.steps.values():
-                            step_mv = step.config.model
-                            if (
-                                step_mv
-                                and step_mv.was_created_in_this_run
-                                and step_mv.name == self.name
-                                and step_mv.version is not None
-                            ):
-                                self.version = step_mv.version
-                                self.model_version_id = (
-                                    step_mv.model_version_id
-                                )
-                                break
             if self.version or self.model_version_id:
                 model_version = self._get_model_version()
             else:
@@ -791,9 +746,13 @@ class Model(BaseModel):
                     time.sleep(sleep)
                     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 +795,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 +837,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 +867,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