zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/materializers/base_materializer.py

@@ -13,9 +13,13 @@
 #  permissions and limitations under the License.
 """Metaclass implementation for registering ZenML BaseMaterializer subclasses."""
 
+import contextlib
 import inspect
-from typing import Any, ClassVar, Dict, 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
 from zenml.exceptions import MaterializerInterfaceError
 from zenml.io import fileio
@@ -51,51 +55,55 @@ class BaseMaterializerMeta(type):
         cls = cast(
             Type["BaseMaterializer"], super().__new__(mcs, name, bases, dct)
         )
-
-        # Skip the following validation and registration for base classes.
-        if cls.SKIP_REGISTRATION:
-            # Reset the flag so subclasses don't have it set automatically.
-            cls.SKIP_REGISTRATION = False
-            return cls
-
-        # Validate that the class is properly defined.
-        if not cls.ASSOCIATED_TYPES:
-            raise MaterializerInterfaceError(
-                f"Invalid materializer class '{name}'. When creating a "
-                f"custom materializer, make sure to specify at least one "
-                f"type in its ASSOCIATED_TYPES class variable.",
-                url="https://docs.zenml.io/user-guide/advanced-guide/artifact-management/handle-custom-data-types",
-            )
-
-        # Validate associated artifact type.
-        if cls.ASSOCIATED_ARTIFACT_TYPE:
-            try:
-                cls.ASSOCIATED_ARTIFACT_TYPE = ArtifactType(
-                    cls.ASSOCIATED_ARTIFACT_TYPE
-                )
-            except ValueError:
+        if not cls._DOCS_BUILDING_MODE:
+            # Skip the following validation and registration for base classes.
+            if cls.SKIP_REGISTRATION:
+                # Reset the flag so subclasses don't have it set automatically.
+                cls.SKIP_REGISTRATION = False
+                return cls
+
+            # Validate that the class is properly defined.
+            if not cls.ASSOCIATED_TYPES:
                 raise MaterializerInterfaceError(
                     f"Invalid materializer class '{name}'. When creating a "
-                    f"custom materializer, make sure to specify a valid "
-                    f"artifact type in its ASSOCIATED_ARTIFACT_TYPE class "
-                    f"variable.",
-                    url="https://docs.zenml.io/user-guide/advanced-guide/artifact-management/handle-custom-data-types",
+                    f"custom materializer, make sure to specify at least one "
+                    f"type in its ASSOCIATED_TYPES class variable.",
+                    url="https://docs.zenml.io/how-to/handle-data-artifacts/handle-custom-data-types",
                 )
 
-        # Validate associated data types.
-        for associated_type in cls.ASSOCIATED_TYPES:
-            if not inspect.isclass(associated_type):
-                raise MaterializerInterfaceError(
-                    f"Associated type {associated_type} for materializer "
-                    f"{name} is not a class.",
-                    url="https://docs.zenml.io/user-guide/advanced-guide/artifact-management/handle-custom-data-types",
+            # Validate associated artifact type.
+            if cls.ASSOCIATED_ARTIFACT_TYPE:
+                try:
+                    cls.ASSOCIATED_ARTIFACT_TYPE = ArtifactType(
+                        cls.ASSOCIATED_ARTIFACT_TYPE
+                    )
+                except ValueError:
+                    raise MaterializerInterfaceError(
+                        f"Invalid materializer class '{name}'. When creating a "
+                        f"custom materializer, make sure to specify a valid "
+                        f"artifact type in its ASSOCIATED_ARTIFACT_TYPE class "
+                        f"variable.",
+                        url="https://docs.zenml.io/how-to/handle-data-artifacts/handle-custom-data-types",
+                    )
+
+            # Validate associated data types.
+            for associated_type in cls.ASSOCIATED_TYPES:
+                if not inspect.isclass(associated_type):
+                    raise MaterializerInterfaceError(
+                        f"Associated type {associated_type} for materializer "
+                        f"{name} is not a class.",
+                        url="https://docs.zenml.io/how-to/handle-data-artifacts/handle-custom-data-types",
+                    )
+
+            # Register the materializer.
+            for associated_type in cls.ASSOCIATED_TYPES:
+                materializer_registry.register_materializer_type(
+                    associated_type, cls
                 )
 
-        # Register the materializer.
-        for associated_type in cls.ASSOCIATED_TYPES:
-            materializer_registry.register_materializer_type(
-                associated_type, cls
-            )
+            from zenml.utils import notebook_utils
+
+            notebook_utils.try_to_save_notebook_cell_code(cls)
 
         return cls
 
@@ -112,13 +120,36 @@ class BaseMaterializer(metaclass=BaseMaterializerMeta):
     # it themselves.
     SKIP_REGISTRATION: ClassVar[bool] = True
 
-    def __init__(self, uri: str):
+    _DOCS_BUILDING_MODE: ClassVar[bool] = False
+
+    def __init__(
+        self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
+    ):
         """Initializes a materializer with the given URI.
 
         Args:
             uri: The URI where the artifact data will be stored.
+            artifact_store: The artifact store used to store this artifact.
         """
         self.uri = uri
+        self._artifact_store = artifact_store
+
+    @property
+    def artifact_store(self) -> BaseArtifactStore:
+        """Returns the artifact store used to store this artifact.
+
+        It either comes from the configuration of the materializer or from the
+        active stack.
+
+        Returns:
+            The artifact store used to store this artifact.
+        """
+        if self._artifact_store:
+            return self._artifact_store
+        else:
+            from zenml.client import Client
+
+            return Client().active_stack.artifact_store
 
     # ================
     # Public Interface
@@ -155,7 +186,7 @@ class BaseMaterializer(metaclass=BaseMaterializerMeta):
         Example:
         ```
         visualization_uri = os.path.join(self.uri, "visualization.html")
-        with fileio.open(visualization_uri, "w") as f:
+        with self.artifact_store.open(visualization_uri, "w") as f:
             f.write("<html><body>data</body></html>")
 
         visualization_uri_2 = os.path.join(self.uri, "visualization.png")
@@ -201,38 +232,72 @@ class BaseMaterializer(metaclass=BaseMaterializerMeta):
     # ================
     # Internal Methods
     # ================
+    def validate_save_type_compatibility(self, data_type: Type[Any]) -> None:
+        """Checks whether the materializer can save the given type.
+
+        Args:
+            data_type: The type to check.
+
+        Raises:
+            TypeError: If the materializer cannot save the given type.
+        """
+        if not self.can_save_type(data_type):
+            raise TypeError(
+                f"Unable to save type {data_type}. {self.__class__.__name__} "
+                f"can only save artifacts of the following types: "
+                f"{self.ASSOCIATED_TYPES}."
+            )
 
-    def validate_type_compatibility(self, data_type: Type[Any]) -> None:
-        """Checks whether the materializer can read/write the given type.
+    def validate_load_type_compatibility(self, data_type: Type[Any]) -> None:
+        """Checks whether the materializer can load the given type.
 
         Args:
             data_type: The type to check.
 
         Raises:
-            TypeError: If the materializer cannot read/write the given type.
+            TypeError: If the materializer cannot load the given type.
         """
-        if not self.can_handle_type(data_type):
+        if not self.can_load_type(data_type):
             raise TypeError(
-                f"Unable to handle type {data_type}. {self.__class__.__name__} "
-                f"can only read/write artifacts of the following types: "
+                f"Unable to load type {data_type}. {self.__class__.__name__} "
+                f"can only load artifacts of the following types: "
                 f"{self.ASSOCIATED_TYPES}."
             )
 
     @classmethod
-    def can_handle_type(cls, data_type: Type[Any]) -> bool:
-        """Whether the materializer can read/write a certain type.
+    def can_save_type(cls, data_type: Type[Any]) -> bool:
+        """Whether the materializer can save a certain type.
 
         Args:
             data_type: The type to check.
 
         Returns:
-            Whether the materializer can read/write the given type.
+            Whether the materializer can save the given type.
         """
         return any(
             issubclass(data_type, associated_type)
             for associated_type in cls.ASSOCIATED_TYPES
         )
 
+    @classmethod
+    def can_load_type(cls, data_type: Type[Any]) -> bool:
+        """Whether the materializer can load an artifact as the given type.
+
+        Args:
+            data_type: The type to check.
+
+        Returns:
+            Whether the materializer can load an artifact as the given type.
+        """
+        return any(
+            issubclass(associated_type, data_type)
+            # This next condition is not always correct, but better to have a
+            # false positive here instead of failing for cases where it would
+            # have worked.
+            or issubclass(data_type, associated_type)
+            for associated_type in cls.ASSOCIATED_TYPES
+        )
+
     def extract_full_metadata(self, data: Any) -> Dict[str, "MetadataType"]:
         """Extract both base and custom metadata from the given data.
 
@@ -264,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/materializers/built_in_materializer.py

@@ -21,13 +21,14 @@ from typing import (
     Dict,
     Iterable,
     List,
+    Optional,
     Tuple,
     Type,
     Union,
 )
 
-from zenml.enums import ArtifactType
-from zenml.io import fileio
+from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
+from zenml.enums import ArtifactType, VisualizationType
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.materializers.materializer_registry import materializer_registry
@@ -55,13 +56,16 @@ class BuiltInMaterializer(BaseMaterializer):
     ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.DATA
     ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = BASIC_TYPES
 
-    def __init__(self, uri: str):
+    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)
+        super().__init__(uri, artifact_store)
         self.data_path = os.path.join(self.uri, DEFAULT_FILENAME)
 
     def load(
@@ -76,7 +80,7 @@ class BuiltInMaterializer(BaseMaterializer):
             The data read.
         """
         contents = yaml_utils.read_json(self.data_path)
-        if type(contents) != data_type:
+        if type(contents) is not data_type:
             # TODO [ENG-142]: Raise error or try to coerce
             logger.debug(
                 f"Contents {contents} was type {type(contents)} but expected "
@@ -117,13 +121,16 @@ class BytesMaterializer(BaseMaterializer):
     ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.DATA
     ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = (bytes,)
 
-    def __init__(self, uri: str):
+    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)
+        super().__init__(uri, artifact_store)
         self.data_path = os.path.join(self.uri, DEFAULT_BYTES_FILENAME)
 
     def load(self, data_type: Type[Any]) -> Any:
@@ -135,7 +142,7 @@ class BytesMaterializer(BaseMaterializer):
         Returns:
             The data read.
         """
-        with fileio.open(self.data_path, "rb") as file_:
+        with self.artifact_store.open(self.data_path, "rb") as file_:
             return file_.read()
 
     def save(self, data: Any) -> None:
@@ -144,7 +151,7 @@ class BytesMaterializer(BaseMaterializer):
         Args:
             data: The data to store.
         """
-        with fileio.open(self.data_path, "wb") as file_:
+        with self.artifact_store.open(self.data_path, "wb") as file_:
             file_.write(data)
 
 
@@ -252,13 +259,16 @@ class BuiltInContainerMaterializer(BaseMaterializer):
         tuple,
     )
 
-    def __init__(self, uri: str):
+    def __init__(
+        self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
+    ):
         """Define `self.data_path` and `self.metadata_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)
+        super().__init__(uri, artifact_store)
         self.data_path = os.path.join(self.uri, DEFAULT_FILENAME)
         self.metadata_path = os.path.join(self.uri, DEFAULT_METADATA_FILENAME)
 
@@ -283,16 +293,16 @@ class BuiltInContainerMaterializer(BaseMaterializer):
             RuntimeError: If the data was not found.
         """
         # If the data was not serialized, there must be metadata present.
-        if not fileio.exists(self.data_path) and not fileio.exists(
-            self.metadata_path
-        ):
+        if not self.artifact_store.exists(
+            self.data_path
+        ) and not self.artifact_store.exists(self.metadata_path):
             raise RuntimeError(
                 f"Materialization of type {data_type} failed. Expected either"
                 f"{self.data_path} or {self.metadata_path} to exist."
             )
 
         # If the data was serialized as JSON, deserialize it.
-        if fileio.exists(self.data_path):
+        if self.artifact_store.exists(self.data_path):
             outputs = yaml_utils.read_json(self.data_path)
 
         # Otherwise, use the metadata to reconstruct the data as a list.
@@ -329,11 +339,11 @@ class BuiltInContainerMaterializer(BaseMaterializer):
         # Cast the data to the correct type.
         if issubclass(data_type, dict) and not isinstance(outputs, dict):
             keys, values = outputs
-            return dict(zip(keys, values))
+            return data_type(zip(keys, values))
         if issubclass(data_type, tuple) and not isinstance(outputs, tuple):
-            return tuple(outputs)
+            return data_type(outputs)
         if issubclass(data_type, set) and not isinstance(outputs, set):
-            return set(outputs)
+            return data_type(outputs)
         return outputs
 
     def save(self, data: Any) -> None:
@@ -375,7 +385,7 @@ class BuiltInContainerMaterializer(BaseMaterializer):
         try:
             for i, element in enumerate(data):
                 element_path = os.path.join(self.uri, str(i))
-                fileio.mkdir(element_path)
+                self.artifact_store.mkdir(element_path)
                 type_ = type(element)
                 materializer_class = materializer_registry[type_]
                 materializer = materializer_class(uri=element_path)
@@ -393,18 +403,35 @@ 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:
             # Delete metadata
-            if fileio.exists(self.metadata_path):
-                fileio.remove(self.metadata_path)
+            if self.artifact_store.exists(self.metadata_path):
+                self.artifact_store.remove(self.metadata_path)
             # Delete all elements that were already saved.
             for entry in metadata:
-                fileio.rmtree(entry["path"])
+                self.artifact_store.rmtree(entry["path"])
             raise e
 
+    # save dict type objects to JSON file with JSON visualization type
+    def save_visualizations(self, data: Any) -> Dict[str, "VisualizationType"]:
+        """Save visualizations for the given data.
+
+        Args:
+            data: The data to save visualizations for.
+
+        Returns:
+            A dictionary of visualization URIs and their types.
+        """
+        # dict/list type objects are always saved as JSON files
+        # doesn't work for non-serializable types as they
+        # are saved as list of lists in different files
+        if _is_serializable(data):
+            return {self.data_path.replace("\\", "/"): VisualizationType.JSON}
+        return {}
+
     def extract_metadata(self, data: Any) -> Dict[str, "MetadataType"]:
         """Extract metadata from the given built-in container object.
 

zenml/materializers/cloudpickle_materializer.py

@@ -20,7 +20,6 @@ import cloudpickle
 
 from zenml.enums import ArtifactType
 from zenml.environment import Environment
-from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.utils.io_utils import (
@@ -72,7 +71,7 @@ class CloudpickleMaterializer(BaseMaterializer):
 
         # load data
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
-        with fileio.open(filepath, "rb") as fid:
+        with self.artifact_store.open(filepath, "rb") as fid:
             data = cloudpickle.load(fid)
         return data
 
@@ -95,7 +94,7 @@ class CloudpickleMaterializer(BaseMaterializer):
         """
         # Log a warning if this materializer was not explicitly specified for
         # the given data type.
-        if type(self) == CloudpickleMaterializer:
+        if type(self) is CloudpickleMaterializer:
             logger.warning(
                 f"No materializer is registered for type `{type(data)}`, so "
                 "the default Pickle materializer was used. Pickle is not "
@@ -103,7 +102,7 @@ class CloudpickleMaterializer(BaseMaterializer):
                 "the artifacts cannot be loaded when running with a different "
                 "Python version. Please consider implementing a custom "
                 f"materializer for type `{type(data)}` according to the "
-                "instructions at https://docs.zenml.io/user-guide/advanced-guide/artifact-management/handle-custom-data-types"
+                "instructions at https://docs.zenml.io/how-to/handle-data-artifacts/handle-custom-data-types"
             )
 
         # save python version for validation on loading
@@ -111,7 +110,7 @@ class CloudpickleMaterializer(BaseMaterializer):
 
         # save data
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
-        with fileio.open(filepath, "wb") as fid:
+        with self.artifact_store.open(filepath, "wb") as fid:
             cloudpickle.dump(data, fid)
 
     def _save_python_version(self) -> None: