zenml-nightly 0.68.1.dev20241102__py3-none-any.whl → 0.68.1.dev20241106__py3-none-any.whl

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/__init__.py

@@ -0,0 +1,50 @@
+#  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.
+"""Initialization for the ZenML vLLM integration."""
+from typing import List, Type
+from zenml.integrations.integration import Integration
+from zenml.stack import Flavor
+from zenml.logger import get_logger
+from zenml.integrations.constants import VLLM
+
+VLLM_MODEL_DEPLOYER = "vllm"
+
+logger = get_logger(__name__)
+
+
+class VLLMIntegration(Integration):
+    """Definition of vLLM integration for ZenML."""
+
+    NAME = VLLM
+
+    REQUIREMENTS = ["vllm>=0.6.0,<0.7.0", "openai>=1.0.0"]
+
+    @classmethod
+    def activate(cls) -> None:
+        """Activates the integration."""
+        from zenml.integrations.vllm import services
+
+    @classmethod
+    def flavors(cls) -> List[Type[Flavor]]:
+        """Declare the stack component flavors for the vLLM integration.
+
+        Returns:
+            List of stack component flavors for this integration.
+        """
+        from zenml.integrations.vllm.flavors import VLLMModelDeployerFlavor
+
+        return [VLLMModelDeployerFlavor]
+
+
+VLLMIntegration.check_installation()

zenml/integrations/vllm/flavors/__init__.py

@@ -0,0 +1,21 @@
+#  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.
+"""vLLM integration flavors."""
+
+from zenml.integrations.vllm.flavors.vllm_model_deployer_flavor import (  # noqa
+    VLLMModelDeployerConfig,
+    VLLMModelDeployerFlavor,
+)
+
+__all__ = ["VLLMModelDeployerConfig", "VLLMModelDeployerFlavor"]

zenml/integrations/vllm/flavors/vllm_model_deployer_flavor.py

@@ -0,0 +1,91 @@
+#  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.
+"""vLLM model deployer flavor."""
+
+from typing import TYPE_CHECKING, Optional, Type
+
+from zenml.integrations.vllm import VLLM_MODEL_DEPLOYER
+from zenml.model_deployers.base_model_deployer import (
+    BaseModelDeployerConfig,
+    BaseModelDeployerFlavor,
+)
+
+if TYPE_CHECKING:
+    from zenml.integrations.vllm.model_deployers import VLLMModelDeployer
+
+
+class VLLMModelDeployerConfig(BaseModelDeployerConfig):
+    """Configuration for vLLM Inference model deployer."""
+
+    service_path: str = ""
+
+
+class VLLMModelDeployerFlavor(BaseModelDeployerFlavor):
+    """vLLM model deployer flavor."""
+
+    @property
+    def name(self) -> str:
+        """Name of the flavor.
+
+        Returns:
+            The name of the flavor.
+        """
+        return VLLM_MODEL_DEPLOYER
+
+    @property
+    def docs_url(self) -> Optional[str]:
+        """A url to point at docs explaining this flavor.
+
+        Returns:
+            A flavor docs url.
+        """
+        return self.generate_default_docs_url()
+
+    @property
+    def sdk_docs_url(self) -> Optional[str]:
+        """A url to point at SDK docs explaining this flavor.
+
+        Returns:
+            A flavor SDK docs url.
+        """
+        return self.generate_default_sdk_docs_url()
+
+    @property
+    def logo_url(self) -> str:
+        """A url to represent the flavor in the dashboard.
+
+        Returns:
+            The flavor logo.
+        """
+        return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/model_deployer/vllm.png"
+
+    @property
+    def config_class(self) -> Type[VLLMModelDeployerConfig]:
+        """Returns `VLLMModelDeployerConfig` config class.
+
+        Returns:
+            The config class.
+        """
+        return VLLMModelDeployerConfig
+
+    @property
+    def implementation_class(self) -> Type["VLLMModelDeployer"]:
+        """Implementation class for this flavor.
+
+        Returns:
+            The implementation class.
+        """
+        from zenml.integrations.vllm.model_deployers import VLLMModelDeployer
+
+        return VLLMModelDeployer

zenml/integrations/vllm/model_deployers/__init__.py

@@ -0,0 +1,19 @@
+#  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.
+"""Initialization of the vLLM model deployers."""
+from zenml.integrations.vllm.model_deployers.vllm_model_deployer import (  # noqa
+    VLLMModelDeployer,
+)
+
+__all__ = ["VLLMModelDeployer"]