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

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/tensorboard/visualizers/tensorboard_visualizer.py

@@ -87,41 +87,44 @@ class TensorboardVisualizer:
             *args: Additional arguments.
             **kwargs: Additional keyword arguments.
         """
-        for _, artifact_view in object.outputs.items():
-            # filter out anything but model artifacts
-            if artifact_view.type == ArtifactType.MODEL:
-                logdir = os.path.dirname(artifact_view.uri)
-
-                # first check if a TensorBoard server is already running for
-                # the same logdir location and use that one
-                running_server = self.find_running_tensorboard_server(logdir)
-                if running_server:
-                    self.visualize_tensorboard(running_server.port, height)
-                    return
-
-                if sys.platform == "win32":
-                    # Daemon service functionality is currently not supported
-                    # on Windows
-                    print(
-                        "You can run:\n"
-                        f"[italic green]    tensorboard --logdir {logdir}"
-                        "[/italic green]\n"
-                        "...to visualize the TensorBoard logs for your trained model."
+        for output in object.outputs.values():
+            for artifact_view in output:
+                # filter out anything but model artifacts
+                if artifact_view.type == ArtifactType.MODEL:
+                    logdir = os.path.dirname(artifact_view.uri)
+
+                    # first check if a TensorBoard server is already running for
+                    # the same logdir location and use that one
+                    running_server = self.find_running_tensorboard_server(
+                        logdir
                     )
-                else:
-                    # start a new TensorBoard server
-                    service = TensorboardService(
-                        TensorboardServiceConfig(
-                            logdir=logdir,
-                            name=f"zenml-tensorboard-{logdir}",
+                    if running_server:
+                        self.visualize_tensorboard(running_server.port, height)
+                        return
+
+                    if sys.platform == "win32":
+                        # Daemon service functionality is currently not supported
+                        # on Windows
+                        print(
+                            "You can run:\n"
+                            f"[italic green]    tensorboard --logdir {logdir}"
+                            "[/italic green]\n"
+                            "...to visualize the TensorBoard logs for your trained model."
                         )
-                    )
-                    service.start(timeout=60)
-                    if service.endpoint.status.port:
-                        self.visualize_tensorboard(
-                            service.endpoint.status.port, height
+                    else:
+                        # start a new TensorBoard server
+                        service = TensorboardService(
+                            TensorboardServiceConfig(
+                                logdir=logdir,
+                                name=f"zenml-tensorboard-{logdir}",
+                            )
                         )
-                return
+                        service.start(timeout=60)
+                        if service.endpoint.status.port:
+                            self.visualize_tensorboard(
+                                service.endpoint.status.port, height
+                            )
+                    return
 
     def visualize_tensorboard(
         self,
@@ -154,31 +157,34 @@ class TensorboardVisualizer:
         Args:
             object: StepRunResponseModel fetched from get_step().
         """
-        for _, artifact_view in object.outputs.items():
-            # filter out anything but model artifacts
-            if artifact_view.type == ArtifactType.MODEL:
-                logdir = os.path.dirname(artifact_view.uri)
-
-                # first check if a TensorBoard server is already running for
-                # the same logdir location and use that one
-                running_server = self.find_running_tensorboard_server(logdir)
-                if not running_server:
-                    return
+        for output in object.outputs.values():
+            for artifact_view in output:
+                # filter out anything but model artifacts
+                if artifact_view.type == ArtifactType.MODEL:
+                    logdir = os.path.dirname(artifact_view.uri)
+
+                    # first check if a TensorBoard server is already running for
+                    # the same logdir location and use that one
+                    running_server = self.find_running_tensorboard_server(
+                        logdir
+                    )
+                    if not running_server:
+                        return
 
-                logger.debug(
-                    "Stopping tensorboard server with PID '%d' ...",
-                    running_server.pid,
-                )
-                try:
-                    p = psutil.Process(running_server.pid)
-                except psutil.Error:
-                    logger.error(
-                        "Could not find process for PID '%d' ...",
+                    logger.debug(
+                        "Stopping tensorboard server with PID '%d' ...",
                         running_server.pid,
                     )
-                    continue
-                p.kill()
-                return
+                    try:
+                        p = psutil.Process(running_server.pid)
+                    except psutil.Error:
+                        logger.error(
+                            "Could not find process for PID '%d' ...",
+                            running_server.pid,
+                        )
+                        continue
+                    p.kill()
+                    return
 
 
 def get_step(pipeline_name: str, step_name: str) -> "StepRunResponse":

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"]