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

zenml/integrations/huggingface/materializers/huggingface_pt_model_materializer.py

@@ -15,7 +15,6 @@
 
 import importlib
 import os
-from tempfile import TemporaryDirectory
 from typing import Any, ClassVar, Dict, Tuple, Type
 
 from transformers import (
@@ -46,17 +45,17 @@ class HFPTModelMaterializer(BaseMaterializer):
         Returns:
             The model read from the specified dir.
         """
-        temp_dir = TemporaryDirectory()
-        io_utils.copy_dir(
-            os.path.join(self.uri, DEFAULT_PT_MODEL_DIR), temp_dir.name
-        )
-
-        config = AutoConfig.from_pretrained(temp_dir.name)
-        architecture = config.architectures[0]
-        model_cls = getattr(
-            importlib.import_module("transformers"), architecture
-        )
-        return model_cls.from_pretrained(temp_dir.name)
+        with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
+            io_utils.copy_dir(
+                os.path.join(self.uri, DEFAULT_PT_MODEL_DIR), temp_dir
+            )
+
+            config = AutoConfig.from_pretrained(temp_dir)
+            architecture = config.architectures[0]
+            model_cls = getattr(
+                importlib.import_module("transformers"), architecture
+            )
+            return model_cls.from_pretrained(temp_dir)
 
     def save(self, model: PreTrainedModel) -> None:
         """Writes a Model to the specified dir.
@@ -64,12 +63,12 @@ class HFPTModelMaterializer(BaseMaterializer):
         Args:
             model: The Torch Model to write.
         """
-        temp_dir = TemporaryDirectory()
-        model.save_pretrained(temp_dir.name)
-        io_utils.copy_dir(
-            temp_dir.name,
-            os.path.join(self.uri, DEFAULT_PT_MODEL_DIR),
-        )
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            model.save_pretrained(temp_dir)
+            io_utils.copy_dir(
+                temp_dir,
+                os.path.join(self.uri, DEFAULT_PT_MODEL_DIR),
+            )
 
     def extract_metadata(
         self, model: PreTrainedModel

zenml/integrations/huggingface/materializers/huggingface_t5_materializer.py

@@ -14,7 +14,6 @@
 """Implementation of the Huggingface t5 materializer."""
 
 import os
-import tempfile
 from typing import Any, ClassVar, Type, Union
 
 from transformers import (
@@ -52,8 +51,7 @@ class HFT5Materializer(BaseMaterializer):
             ValueError: Unsupported data type used
         """
         filepath = self.uri
-
-        with tempfile.TemporaryDirectory(prefix="zenml-temp-") as temp_dir:
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
             # Copy files from artifact store to temporary directory
             for file in fileio.listdir(filepath):
                 src = os.path.join(filepath, file)
@@ -86,8 +84,7 @@ class HFT5Materializer(BaseMaterializer):
         Args:
             obj: A T5ForConditionalGeneration model or T5Tokenizer.
         """
-        # Create a temporary directory
-        with tempfile.TemporaryDirectory(prefix="zenml-temp-") as temp_dir:
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
             # Save the model or tokenizer
             obj.save_pretrained(temp_dir)
 

zenml/integrations/huggingface/materializers/huggingface_tf_model_materializer.py

@@ -15,7 +15,6 @@
 
 import importlib
 import os
-from tempfile import TemporaryDirectory
 from typing import Any, ClassVar, Dict, Tuple, Type
 
 from transformers import (
@@ -46,17 +45,17 @@ class HFTFModelMaterializer(BaseMaterializer):
         Returns:
             The model read from the specified dir.
         """
-        temp_dir = TemporaryDirectory()
-        io_utils.copy_dir(
-            os.path.join(self.uri, DEFAULT_TF_MODEL_DIR), temp_dir.name
-        )
-
-        config = AutoConfig.from_pretrained(temp_dir.name)
-        architecture = "TF" + config.architectures[0]
-        model_cls = getattr(
-            importlib.import_module("transformers"), architecture
-        )
-        return model_cls.from_pretrained(temp_dir.name)
+        with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
+            io_utils.copy_dir(
+                os.path.join(self.uri, DEFAULT_TF_MODEL_DIR), temp_dir
+            )
+
+            config = AutoConfig.from_pretrained(temp_dir)
+            architecture = "TF" + config.architectures[0]
+            model_cls = getattr(
+                importlib.import_module("transformers"), architecture
+            )
+            return model_cls.from_pretrained(temp_dir)
 
     def save(self, model: TFPreTrainedModel) -> None:
         """Writes a Model to the specified dir.
@@ -64,12 +63,12 @@ class HFTFModelMaterializer(BaseMaterializer):
         Args:
             model: The TF Model to write.
         """
-        temp_dir = TemporaryDirectory()
-        model.save_pretrained(temp_dir.name)
-        io_utils.copy_dir(
-            temp_dir.name,
-            os.path.join(self.uri, DEFAULT_TF_MODEL_DIR),
-        )
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            model.save_pretrained(temp_dir)
+            io_utils.copy_dir(
+                temp_dir,
+                os.path.join(self.uri, DEFAULT_TF_MODEL_DIR),
+            )
 
     def extract_metadata(
         self, model: TFPreTrainedModel

zenml/integrations/huggingface/materializers/huggingface_tokenizer_materializer.py

@@ -14,7 +14,6 @@
 """Implementation of the Huggingface tokenizer materializer."""
 
 import os
-from tempfile import TemporaryDirectory
 from typing import Any, ClassVar, Tuple, Type
 
 from transformers import AutoTokenizer
@@ -46,7 +45,7 @@ class HFTokenizerMaterializer(BaseMaterializer):
         Returns:
             The tokenizer read from the specified dir.
         """
-        with TemporaryDirectory() as temp_dir:
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
             io_utils.copy_dir(
                 os.path.join(self.uri, DEFAULT_TOKENIZER_DIR), temp_dir
             )
@@ -58,7 +57,7 @@ class HFTokenizerMaterializer(BaseMaterializer):
         Args:
             tokenizer: The HFTokenizer to write.
         """
-        with TemporaryDirectory() as temp_dir:
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
             tokenizer.save_pretrained(temp_dir)
             io_utils.copy_dir(
                 temp_dir,

zenml/integrations/langchain/__init__.py

@@ -26,7 +26,8 @@ class LangchainIntegration(Integration):
 
     NAME = LANGCHAIN
     REQUIREMENTS = [
-        "langchain==0.0.325",
+        "langchain==0.3.0",
+        "langchain-community",
         "pyyaml>=6.0.1",
         "tenacity!=8.4.0",  # https://github.com/jd/tenacity/issues/471
     ]

zenml/integrations/langchain/materializers/openai_embedding_materializer.py

@@ -24,11 +24,37 @@ from zenml.materializers.cloudpickle_materializer import (
 if TYPE_CHECKING and sys.version_info < (3, 8):
     OpenAIEmbeddings = Any
 else:
-    from langchain.embeddings import OpenAIEmbeddings
+    from langchain_community.embeddings import (
+        OpenAIEmbeddings,
+    )
 
 
 class LangchainOpenaiEmbeddingMaterializer(CloudpickleMaterializer):
-    """Handle langchain OpenAI embedding objects."""
+    """Materializer for Langchain OpenAI Embeddings."""
 
     ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.MODEL
     ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = (OpenAIEmbeddings,)
+
+    def save(self, embeddings: Any) -> None:
+        """Saves the embeddings model after clearing non-picklable clients.
+
+        Args:
+            embeddings: The embeddings model to save.
+        """
+        # Clear the clients which will be recreated on load
+        embeddings.client = None
+        embeddings.async_client = None
+
+        # Use the parent class's save implementation which uses cloudpickle
+        super().save(embeddings)
+
+    def load(self, data_type: Type[Any]) -> Any:
+        """Loads the embeddings model and lets it recreate clients when needed.
+
+        Args:
+            data_type: The type of the data to load.
+
+        Returns:
+            The loaded embeddings model.
+        """
+        return super().load(data_type)

zenml/integrations/lightgbm/materializers/lightgbm_booster_materializer.py

@@ -14,7 +14,6 @@
 """Implementation of the LightGBM booster materializer."""
 
 import os
-import tempfile
 from typing import Any, ClassVar, Tuple, Type
 
 import lightgbm as lgb
@@ -42,18 +41,13 @@ class LightGBMBoosterMaterializer(BaseMaterializer):
             A lightgbm Booster object.
         """
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
 
-        # Create a temporary folder
-        temp_dir = tempfile.mkdtemp(prefix="zenml-temp-")
-        temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
-
-        # Copy from artifact store to temporary file
-        fileio.copy(filepath, temp_file)
-        booster = lgb.Booster(model_file=temp_file)
-
-        # Cleanup and return
-        fileio.rmtree(temp_dir)
-        return booster
+            # Copy from artifact store to temporary file
+            fileio.copy(filepath, temp_file)
+            booster = lgb.Booster(model_file=temp_file)
+            return booster
 
     def save(self, booster: lgb.Booster) -> None:
         """Creates a JSON serialization for a lightgbm Booster model.
@@ -62,8 +56,7 @@ class LightGBMBoosterMaterializer(BaseMaterializer):
             booster: A lightgbm Booster model.
         """
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
-
-        with tempfile.TemporaryDirectory() as tmp_dir:
-            tmp_path = os.path.join(tmp_dir, "model.txt")
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            tmp_path = os.path.join(temp_dir, "model.txt")
             booster.save_model(tmp_path)
             fileio.copy(tmp_path, filepath)

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/lightning/orchestrators/lightning_orchestrator.py

@@ -390,10 +390,13 @@ class LightningOrchestrator(WheeledOrchestrator):
             studio.run(
                 f"tar -xvzf /teamspace/studios/this_studio/zenml_codes/{filename} -C /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
             )
-            studio.upload_file(
-                env_file_path,
-                remote_path="/teamspace/studios/this_studio/.lightning_studio/.studiorc",
+            studio.upload_file(env_file_path)
+            time.sleep(6)
+            studio.run(
+                f"cp {env_file_path.split('/')[-1]} ./.lightning_studio/.studiorc"
             )
+            studio.run(f"rm {env_file_path.split('/')[-1]}")
+
             studio.run("pip install uv")
             logger.info(
                 f"Installing requirements: {pipeline_requirements_to_string}"
@@ -481,16 +484,25 @@ class LightningOrchestrator(WheeledOrchestrator):
                 code_path,
                 remote_path=f"/teamspace/studios/this_studio/zenml_codes/{filename}",
             )
+            time.sleep(6)
             studio.run(
                 f"tar -xvzf /teamspace/studios/this_studio/zenml_codes/{filename} -C /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
             )
             logger.info(
                 "Uploading wheel package and installing dependencies on main studio"
             )
-            studio.upload_file(
-                env_file_path,
-                remote_path="/teamspace/studios/this_studio/.lightning_studio/.studiorc",
+            # for some reason uploading file directly to /teamspace/studios/this_studio/.lightning_studio/.studiorc
+            # doesn't work, while other file names can be uploaded just fine.
+            # below is a workaround to copy the file to the correct location.
+            # we first upload it under a different name and then copy it to the
+            # right location.
+            studio.upload_file(env_file_path)
+            time.sleep(6)
+            studio.run(
+                f"cp {env_file_path.split('/')[-1]} ./.lightning_studio/.studiorc"
             )
+            studio.run(f"rm {env_file_path.split('/')[-1]}")
+
             studio.run("pip install uv")
             studio.run(f"uv pip install {requirements}")
             studio.run("pip install zenml")
@@ -557,13 +569,21 @@ class LightningOrchestrator(WheeledOrchestrator):
         studio.run(
             f"mkdir -p /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
         )
-        studio.upload_file(code_path, remote_path=f"/zenml_codes/{filename}")
+        studio.upload_file(
+            code_path,
+            remote_path=f"/teamspace/studios/this_studio/zenml_codes/{filename}",
+        )
+        time.sleep(6)
         studio.run(
             f"tar -xvzf /teamspace/studios/this_studio/zenml_codes/{filename} -C /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
         )
-        studio.upload_file(
-            env_file_path, remote_path=".lightning_studio/.studiorc"
+        studio.upload_file(env_file_path)
+        time.sleep(6)
+        studio.run(
+            f"cp {env_file_path.split('/')[-1]} ./.lightning_studio/.studiorc"
         )
+        studio.run(f"rm {env_file_path.split('/')[-1]}")
+
         studio.run("pip install uv")
         studio.run(f"uv pip install {details['requirements']}")
         studio.run("pip install zenml")

zenml/integrations/openai/__init__.py

@@ -21,7 +21,7 @@ class OpenAIIntegration(Integration):
     """Definition of OpenAI integration for ZenML."""
 
     NAME = OPEN_AI
-    REQUIREMENTS = ["openai>=0.27.0,<1.0.0"]
+    REQUIREMENTS = ["openai>=1.0.0"]
 
 
 OpenAIIntegration.check_installation()

zenml/integrations/openai/hooks/open_ai_failure_hook.py

@@ -15,8 +15,9 @@
 
 import io
 import sys
+from typing import Optional
 
-import openai
+from openai import OpenAI
 from rich.console import Console
 
 from zenml import get_step_context
@@ -38,6 +39,8 @@ def openai_alerter_failure_hook_helper(
     Args:
         exception: The exception that was raised.
         model_name: The OpenAI model to use for the chatbot.
+
+    This implementation uses the OpenAI v1 SDK with automatic retries and backoff.
     """
     client = Client()
     context = get_step_context()
@@ -47,12 +50,15 @@ def openai_alerter_failure_hook_helper(
         openai_secret = client.get_secret(
             "openai", allow_partial_name_match=False
         )
-        openai_api_key = openai_secret.secret_values.get("api_key")
+        openai_api_key: Optional[str] = openai_secret.secret_values.get(
+            "api_key"
+        )
     except (KeyError, NotImplementedError):
         openai_api_key = None
 
     alerter = client.active_stack.alerter
     if alerter and openai_api_key:
+        # Capture rich traceback
         output_captured = io.StringIO()
         original_stdout = sys.stdout
         sys.stdout = output_captured
@@ -62,25 +68,44 @@ def openai_alerter_failure_hook_helper(
         sys.stdout = original_stdout
         rich_traceback = output_captured.getvalue()
 
-        response = openai.ChatCompletion.create(  # type: ignore
+        # Initialize OpenAI client with timeout and retry settings
+        openai_client = OpenAI(
+            api_key=openai_api_key,
+            max_retries=3,  # Will retry 3 times with exponential backoff
+            timeout=60.0,  # 60 second timeout
+        )
+
+        # Create chat completion using the new client pattern
+        response = openai_client.chat.completions.create(
             model=model_name,
             messages=[
                 {
                     "role": "user",
-                    "content": f"This is an error message (following an exception of type '{type(exception)}') I encountered while executing a ZenML step. Please suggest ways I might fix the problem. Feel free to give code snippets as examples, and note that your response will be piped to a Slack bot so make sure the formatting is appropriate: {exception} -- {rich_traceback}. Thank you!",
+                    "content": f"This is an error message (following an exception of type '{type(exception)}') "
+                    f"I encountered while executing a ZenML step. Please suggest ways I might fix the problem. "
+                    f"Feel free to give code snippets as examples, and note that your response will be piped "
+                    f"to a Slack bot so make sure the formatting is appropriate: {exception} -- {rich_traceback}. "
+                    f"Thank you!",
                 }
             ],
         )
-        suggestion = response["choices"][0]["message"]["content"]
-        message = "*Failure Hook Notification! Step failed!*" + "\n\n"
-        message += f"Run name: `{context.pipeline_run.name}`" + "\n"
-        message += f"Step name: `{context.step_run.name}`" + "\n"
-        message += f"Parameters: `{context.step_run.config.parameters}`" + "\n"
-        message += f"Exception: `({type(exception)}) {exception}`" + "\n\n"
-        message += (
-            f"*OpenAI ChatGPT's suggestion (model = `{model_name}`) on how to fix it:*\n `{suggestion}`"
-            + "\n"
+
+        suggestion = response.choices[0].message.content
+
+        # Format the alert message
+        message = "\n".join(
+            [
+                "*Failure Hook Notification! Step failed!*",
+                "",
+                f"Run name: `{context.pipeline_run.name}`",
+                f"Step name: `{context.step_run.name}`",
+                f"Parameters: `{context.step_run.config.parameters}`",
+                f"Exception: `({type(exception)}) {exception}`",
+                "",
+                f"*OpenAI ChatGPT's suggestion (model = `{model_name}`) on how to fix it:*\n `{suggestion}`",
+            ]
         )
+
         alerter.post(message)
     elif not openai_api_key:
         logger.warning(
@@ -111,4 +136,4 @@ def openai_gpt4_alerter_failure_hook(
     Args:
         exception: The exception that was raised.
     """
-    openai_alerter_failure_hook_helper(exception, "gpt-4")
+    openai_alerter_failure_hook_helper(exception, "gpt-4o")

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)