zenml-nightly 0.68.1.dev20241105__py3-none-any.whl → 0.68.1.dev20241107__py3-none-any.whl

zenml/integrations/lightgbm/materializers/lightgbm_dataset_materializer.py

@@ -14,7 +14,6 @@
 """Implementation of the LightGBM materializer."""
 
 import os
-import tempfile
 from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type
 
 import lightgbm as lgb
@@ -46,16 +45,14 @@ class LightGBMDatasetMaterializer(BaseMaterializer):
         """
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
 
-        # Create a temporary folder
-        temp_dir = tempfile.mkdtemp(prefix="zenml-temp-")
-        temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
+        with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
+            temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
 
-        # Copy from artifact store to temporary file
-        fileio.copy(filepath, temp_file)
-        matrix = lgb.Dataset(temp_file, free_raw_data=False)
+            # Copy from artifact store to temporary file
+            fileio.copy(filepath, temp_file)
+            matrix = lgb.Dataset(temp_file, free_raw_data=False)
 
-        # No clean up this time because matrix is lazy loaded
-        return matrix
+            return matrix
 
     def save(self, matrix: lgb.Dataset) -> None:
         """Creates a binary serialization for a lightgbm.Dataset object.
@@ -65,14 +62,12 @@ class LightGBMDatasetMaterializer(BaseMaterializer):
         """
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
 
-        # Make a temporary phantom artifact
-        temp_dir = tempfile.mkdtemp(prefix="zenml-temp-")
-        temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
-        matrix.save_binary(temp_file)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
+            matrix.save_binary(temp_file)
 
-        # Copy it into artifact store
-        fileio.copy(temp_file, filepath)
-        fileio.rmtree(temp_dir)
+            # Copy it into artifact store
+            fileio.copy(temp_file, filepath)
 
     def extract_metadata(
         self, matrix: lgb.Dataset

zenml/integrations/pillow/materializers/pillow_image_materializer.py

@@ -14,7 +14,6 @@
 """Materializer for Pillow Image objects."""
 
 import os
-import tempfile
 from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type
 
 from PIL import Image
@@ -57,16 +56,15 @@ class PillowImageMaterializer(BaseMaterializer):
         files = io_utils.find_files(self.uri, f"{DEFAULT_IMAGE_FILENAME}.*")
         filepath = [file for file in files if not fileio.isdir(file)][0]
 
-        # create a temporary folder
-        temp_dir = tempfile.TemporaryDirectory(prefix="zenml-temp-")
-        temp_file = os.path.join(
-            temp_dir.name,
-            f"{DEFAULT_IMAGE_FILENAME}{os.path.splitext(filepath)[1]}",
-        )
+        with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
+            temp_file = os.path.join(
+                temp_dir,
+                f"{DEFAULT_IMAGE_FILENAME}{os.path.splitext(filepath)[1]}",
+            )
 
-        # copy from artifact store to temporary file
-        fileio.copy(filepath, temp_file)
-        return Image.open(temp_file)
+            # copy from artifact store to temporary file
+            fileio.copy(filepath, temp_file)
+            return Image.open(temp_file)
 
     def save(self, image: Image.Image) -> None:
         """Write to artifact store.
@@ -74,18 +72,17 @@ class PillowImageMaterializer(BaseMaterializer):
         Args:
             image: An Image.Image object.
         """
-        temp_dir = tempfile.TemporaryDirectory(prefix="zenml-temp-")
-        file_extension = image.format or DEFAULT_IMAGE_EXTENSION
-        full_filename = f"{DEFAULT_IMAGE_FILENAME}.{file_extension}"
-        temp_image_path = os.path.join(temp_dir.name, full_filename)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            file_extension = image.format or DEFAULT_IMAGE_EXTENSION
+            full_filename = f"{DEFAULT_IMAGE_FILENAME}.{file_extension}"
+            temp_image_path = os.path.join(temp_dir, full_filename)
 
-        # save the image in a temporary directory
-        image.save(temp_image_path)
+            # save the image in a temporary directory
+            image.save(temp_image_path)
 
-        # copy the saved image to the artifact store
-        artifact_store_path = os.path.join(self.uri, full_filename)
-        io_utils.copy(temp_image_path, artifact_store_path, overwrite=True)  # type: ignore[attr-defined]
-        temp_dir.cleanup()
+            # copy the saved image to the artifact store
+            artifact_store_path = os.path.join(self.uri, full_filename)
+            io_utils.copy(temp_image_path, artifact_store_path, overwrite=True)  # type: ignore[attr-defined]
 
     def save_visualizations(
         self, image: Image.Image

zenml/integrations/polars/materializers/dataframe_materializer.py

@@ -14,7 +14,6 @@
 """Polars materializer."""
 
 import os
-import tempfile
 from typing import Any, ClassVar, Tuple, Type, Union
 
 import polars as pl
@@ -22,7 +21,6 @@ import pyarrow as pa  # type: ignore
 import pyarrow.parquet as pq  # type: ignore
 
 from zenml.enums import ArtifactType
-from zenml.io import fileio
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.utils import io_utils
 
@@ -45,35 +43,29 @@ class PolarsMaterializer(BaseMaterializer):
         Returns:
             A Polars data frame or series.
         """
-        # Create a temporary directory to store the model
-        temp_dir = tempfile.TemporaryDirectory()
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            io_utils.copy_dir(self.uri, temp_dir)
 
-        # Copy from artifact store to temporary directory
-        io_utils.copy_dir(self.uri, temp_dir.name)
-
-        # Load the data from the temporary directory
-        table = pq.read_table(
-            os.path.join(temp_dir.name, "dataframe.parquet").replace("\\", "/")
-        )
-
-        # If the data is of type pl.Series, convert it back to a pyarrow array
-        # instead of a table.
-        if (
-            table.schema.metadata
-            and b"zenml_is_pl_series" in table.schema.metadata
-        ):
-            isinstance_bytes = table.schema.metadata[b"zenml_is_pl_series"]
-            isinstance_series = bool.from_bytes(isinstance_bytes, "big")
-            if isinstance_series:
-                table = table.column(0)
+            # Load the data from the temporary directory
+            table = pq.read_table(
+                os.path.join(temp_dir, "dataframe.parquet").replace("\\", "/")
+            )
 
-        # Convert the table to a Polars data frame or series
-        data = pl.from_arrow(table)
+            # If the data is of type pl.Series, convert it back to a pyarrow array
+            # instead of a table.
+            if (
+                table.schema.metadata
+                and b"zenml_is_pl_series" in table.schema.metadata
+            ):
+                isinstance_bytes = table.schema.metadata[b"zenml_is_pl_series"]
+                isinstance_series = bool.from_bytes(isinstance_bytes, "big")
+                if isinstance_series:
+                    table = table.column(0)
 
-        # Cleanup and return
-        fileio.rmtree(temp_dir.name)
+            # Convert the table to a Polars data frame or series
+            data = pl.from_arrow(table)
 
-        return data
+            return data
 
     def save(self, data: Union[pl.DataFrame, pl.Series]) -> None:
         """Writes Polars data to the artifact store.
@@ -107,15 +99,10 @@ class PolarsMaterializer(BaseMaterializer):
             {b"zenml_is_pl_series": isinstance_bytes}
         )
 
-        # Create a temporary directory to store the model
-        temp_dir = tempfile.TemporaryDirectory()
-
-        # Write the table to a Parquet file
-        path = os.path.join(temp_dir.name, "dataframe.parquet").replace(
-            "\\", "/"
-        )
-        pq.write_table(table, path)  # Uses lz4 compression by default
-        io_utils.copy_dir(temp_dir.name, self.uri)
-
-        # Remove the temporary directory
-        fileio.rmtree(temp_dir.name)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            # Write the table to a Parquet file
+            path = os.path.join(temp_dir, "dataframe.parquet").replace(
+                "\\", "/"
+            )
+            pq.write_table(table, path)  # Uses lz4 compression by default
+            io_utils.copy_dir(temp_dir, self.uri)

zenml/integrations/pycaret/materializers/model_materializer.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """PyCaret materializer."""
 
-import tempfile
 from typing import (
     Any,
     Type,
@@ -65,7 +64,6 @@ from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
 from xgboost import XGBClassifier, XGBRegressor
 
 from zenml.enums import ArtifactType
-from zenml.io import fileio
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.utils import io_utils
 
@@ -133,19 +131,10 @@ class PyCaretMaterializer(BaseMaterializer):
         Returns:
             A PyCaret model.
         """
-        # Create a temporary directory to store the model
-        temp_dir = tempfile.TemporaryDirectory()
-
-        # Copy from artifact store to temporary directory
-        io_utils.copy_dir(self.uri, temp_dir.name)
-
-        # Load the model from the temporary directory
-        model = load_model(temp_dir.name)
-
-        # Cleanup and return
-        fileio.rmtree(temp_dir.name)
-
-        return model
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            io_utils.copy_dir(self.uri, temp_dir)
+            model = load_model(temp_dir)
+            return model
 
     def save(self, model: Any) -> None:
         """Writes a PyCaret model to the artifact store.
@@ -153,10 +142,6 @@ class PyCaretMaterializer(BaseMaterializer):
         Args:
             model: Any of the supported models.
         """
-        # Create a temporary directory to store the model
-        temp_dir = tempfile.TemporaryDirectory()
-        save_model(model, temp_dir.name)
-        io_utils.copy_dir(temp_dir.name, self.uri)
-
-        # Remove the temporary directory
-        fileio.rmtree(temp_dir.name)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            save_model(model, temp_dir)
+            io_utils.copy_dir(temp_dir, self.uri)

zenml/integrations/tensorflow/materializers/keras_materializer.py

@@ -14,7 +14,6 @@
 """Implementation of the TensorFlow Keras materializer."""
 
 import os
-import tempfile
 from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type
 
 import tensorflow as tf
@@ -22,7 +21,6 @@ from tensorflow.python import keras as tf_keras
 from tensorflow.python.keras.utils.layer_utils import count_params
 
 from zenml.enums import ArtifactType
-from zenml.io import fileio
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.utils import io_utils
 
@@ -49,20 +47,15 @@ class KerasMaterializer(BaseMaterializer):
         Returns:
             A keras.Model model.
         """
-        # Create a temporary directory to store the model
-        temp_dir = tempfile.TemporaryDirectory()
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            # Copy from artifact store to temporary directory
+            temp_model_file = os.path.join(temp_dir, self.MODEL_FILE_NAME)
+            io_utils.copy_dir(self.uri, temp_dir)
 
-        # Copy from artifact store to temporary directory
-        temp_model_file = os.path.join(temp_dir.name, self.MODEL_FILE_NAME)
-        io_utils.copy_dir(self.uri, temp_dir.name)
+            # Load the model from the temporary directory
+            model = tf.keras.models.load_model(temp_model_file)
 
-        # Load the model from the temporary directory
-        model = tf.keras.models.load_model(temp_model_file)
-
-        # Cleanup and return
-        fileio.rmtree(temp_dir.name)
-
-        return model
+            return model
 
     def save(self, model: tf_keras.Model) -> None:
         """Writes a keras model to the artifact store.
@@ -70,14 +63,10 @@ class KerasMaterializer(BaseMaterializer):
         Args:
             model: A keras.Model model.
         """
-        # Create a temporary directory to store the model
-        temp_dir = tempfile.TemporaryDirectory()
-        temp_model_file = os.path.join(temp_dir.name, self.MODEL_FILE_NAME)
-        model.save(temp_model_file)
-        io_utils.copy_dir(temp_dir.name, self.uri)
-
-        # Remove the temporary directory
-        fileio.rmtree(temp_dir.name)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            temp_model_file = os.path.join(temp_dir, self.MODEL_FILE_NAME)
+            model.save(temp_model_file)
+            io_utils.copy_dir(temp_dir, self.uri)
 
     def extract_metadata(
         self, model: tf_keras.Model

zenml/integrations/tensorflow/materializers/tf_dataset_materializer.py

@@ -14,13 +14,11 @@
 """Implementation of the TensorFlow dataset materializer."""
 
 import os
-import tempfile
 from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type
 
 import tensorflow as tf
 
 from zenml.enums import ArtifactType
-from zenml.io import fileio
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.utils import io_utils
 
@@ -45,13 +43,11 @@ class TensorflowDatasetMaterializer(BaseMaterializer):
         Returns:
             A tf.data.Dataset object.
         """
-        temp_dir = tempfile.mkdtemp()
-        io_utils.copy_dir(self.uri, temp_dir)
-        path = os.path.join(temp_dir, DEFAULT_FILENAME)
-        dataset = tf.data.Dataset.load(path)
-        # Don't delete the temporary directory here as the dataset is lazily
-        # loaded and needs to read it when the object gets used
-        return dataset
+        with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
+            io_utils.copy_dir(self.uri, temp_dir)
+            path = os.path.join(temp_dir, DEFAULT_FILENAME)
+            dataset = tf.data.Dataset.load(path)
+            return dataset
 
     def save(self, dataset: tf.data.Dataset) -> None:
         """Persists a tf.data.Dataset object.
@@ -59,15 +55,12 @@ class TensorflowDatasetMaterializer(BaseMaterializer):
         Args:
             dataset: The dataset to persist.
         """
-        temp_dir = tempfile.TemporaryDirectory()
-        path = os.path.join(temp_dir.name, DEFAULT_FILENAME)
-        try:
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            path = os.path.join(temp_dir, DEFAULT_FILENAME)
             tf.data.Dataset.save(
                 dataset, path, compression=None, shard_func=None
             )
-            io_utils.copy_dir(temp_dir.name, self.uri)
-        finally:
-            fileio.rmtree(temp_dir.name)
+            io_utils.copy_dir(temp_dir, self.uri)
 
     def extract_metadata(
         self, dataset: tf.data.Dataset

zenml/integrations/vllm/services/vllm_deployment.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """Implementation of the vLLM Inference Server Service."""
 
+import argparse
 import os
 from typing import Any, List, Optional, Union
 
@@ -137,15 +138,23 @@ class VLLMDeploymentService(LocalDaemonService, BaseDeploymentService):
         self.endpoint.prepare_for_start()
 
         import uvloop
-        from vllm.entrypoints.openai.api_server import run_server
-        from vllm.entrypoints.openai.cli_args import make_arg_parser
-        from vllm.utils import FlexibleArgumentParser
+        from vllm.entrypoints.openai.api_server import (
+            run_server,
+        )
+        from vllm.entrypoints.openai.cli_args import (
+            make_arg_parser,
+        )
+        from vllm.utils import (
+            FlexibleArgumentParser,
+        )
 
         try:
-            parser = make_arg_parser(FlexibleArgumentParser())
-            args = parser.parse_args()
+            parser: argparse.ArgumentParser = make_arg_parser(
+                FlexibleArgumentParser()
+            )
+            args: argparse.Namespace = parser.parse_args()
             # Override port with the available port
-            self.config.port = self.endpoint.status.port
+            self.config.port = self.endpoint.status.port or self.config.port
             # Update the arguments in place
             args.__dict__.update(self.config.model_dump())
             uvloop.run(run_server(args=args))
@@ -161,7 +170,7 @@ class VLLMDeploymentService(LocalDaemonService, BaseDeploymentService):
         """
         if not self.is_running:
             return None
-        return self.endpoint.prediction_url_path
+        return self.endpoint.prediction_url
 
     def predict(self, data: "Any") -> "Any":
         """Make a prediction using the service.

zenml/integrations/whylogs/materializers/whylogs_materializer.py

@@ -14,7 +14,6 @@
 """Implementation of the whylogs materializer."""
 
 import os
-import tempfile
 from typing import Any, ClassVar, Dict, Tuple, Type, cast
 
 from whylogs.core import DatasetProfileView  # type: ignore
@@ -51,18 +50,14 @@ class WhylogsMaterializer(BaseMaterializer):
         """
         filepath = os.path.join(self.uri, PROFILE_FILENAME)
 
-        # Create a temporary folder
-        temp_dir = tempfile.mkdtemp(prefix="zenml-temp-")
-        temp_file = os.path.join(str(temp_dir), PROFILE_FILENAME)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            temp_file = os.path.join(str(temp_dir), PROFILE_FILENAME)
 
-        # Copy from artifact store to temporary file
-        fileio.copy(filepath, temp_file)
-        profile_view = DatasetProfileView.read(temp_file)
+            # Copy from artifact store to temporary file
+            fileio.copy(filepath, temp_file)
+            profile_view = DatasetProfileView.read(temp_file)
 
-        # Cleanup and return
-        fileio.rmtree(temp_dir)
-
-        return profile_view
+            return profile_view
 
     def save(self, profile_view: DatasetProfileView) -> None:
         """Writes a whylogs dataset profile view.
@@ -72,15 +67,13 @@ class WhylogsMaterializer(BaseMaterializer):
         """
         filepath = os.path.join(self.uri, PROFILE_FILENAME)
 
-        # Create a temporary folder
-        temp_dir = tempfile.mkdtemp(prefix="zenml-temp-")
-        temp_file = os.path.join(str(temp_dir), PROFILE_FILENAME)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            temp_file = os.path.join(str(temp_dir), PROFILE_FILENAME)
 
-        profile_view.write(temp_file)
+            profile_view.write(temp_file)
 
-        # Copy it into artifact store
-        fileio.copy(temp_file, filepath)
-        fileio.rmtree(temp_dir)
+            # Copy it into artifact store
+            fileio.copy(temp_file, filepath)
 
         try:
             self._upload_to_whylabs(profile_view)

zenml/integrations/xgboost/materializers/xgboost_booster_materializer.py

@@ -14,7 +14,6 @@
 """Implementation of an XGBoost booster materializer."""
 
 import os
-import tempfile
 from typing import Any, ClassVar, Tuple, Type
 
 import xgboost as xgb
@@ -43,18 +42,15 @@ class XgboostBoosterMaterializer(BaseMaterializer):
         """
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
 
-        # Create a temporary folder
-        temp_dir = tempfile.mkdtemp(prefix="zenml-temp-")
-        temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
 
-        # Copy from artifact store to temporary file
-        fileio.copy(filepath, temp_file)
-        booster = xgb.Booster()
-        booster.load_model(temp_file)
+            # Copy from artifact store to temporary file
+            fileio.copy(filepath, temp_file)
+            booster = xgb.Booster()
+            booster.load_model(temp_file)
 
-        # Cleanup and return
-        fileio.rmtree(temp_dir)
-        return booster
+            return booster
 
     def save(self, booster: xgb.Booster) -> None:
         """Creates a JSON serialization for a xgboost Booster model.
@@ -64,14 +60,7 @@ class XgboostBoosterMaterializer(BaseMaterializer):
         """
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
 
-        # Make a temporary phantom artifact
-        with tempfile.NamedTemporaryFile(
-            mode="w", suffix=".json", delete=False
-        ) as f:
-            booster.save_model(f.name)
-            # Copy it into artifact store
-            fileio.copy(f.name, filepath)
-
-        # Close and remove the temporary file
-        f.close()
-        fileio.remove(f.name)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
+            booster.save_model(temp_file)
+            fileio.copy(temp_file, filepath)

zenml/integrations/xgboost/materializers/xgboost_dmatrix_materializer.py

@@ -14,7 +14,6 @@
 """Implementation of the XGBoost dmatrix materializer."""
 
 import os
-import tempfile
 from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type
 
 import xgboost as xgb
@@ -46,17 +45,14 @@ class XgboostDMatrixMaterializer(BaseMaterializer):
         """
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
 
-        # Create a temporary folder
-        temp_dir = tempfile.mkdtemp(prefix="zenml-temp-")
-        temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
 
-        # Copy from artifact store to temporary file
-        fileio.copy(filepath, temp_file)
-        matrix = xgb.DMatrix(temp_file)
+            # Copy from artifact store to temporary file
+            fileio.copy(filepath, temp_file)
+            matrix = xgb.DMatrix(temp_file)
 
-        # Cleanup and return
-        fileio.rmtree(temp_dir)
-        return matrix
+            return matrix
 
     def save(self, matrix: xgb.DMatrix) -> None:
         """Creates a binary serialization for a xgboost.DMatrix object.
@@ -66,15 +62,10 @@ class XgboostDMatrixMaterializer(BaseMaterializer):
         """
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
 
-        # Make a temporary phantom artifact
-        with tempfile.NamedTemporaryFile(mode="wb", delete=False) as f:
-            matrix.save_binary(f.name)
-            # Copy it into artifact store
-            fileio.copy(f.name, filepath)
-
-        # Close and remove the temporary file
-        f.close()
-        fileio.remove(f.name)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
+            matrix.save_binary(temp_file)
+            fileio.copy(temp_file, filepath)
 
     def extract_metadata(
         self, dataset: xgb.DMatrix

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)