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

zenml/integrations/tensorboard/__init__.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Initialization for TensorBoard integration."""
 
-import sys
 from typing import List, Optional
 from zenml.integrations.constants import TENSORBOARD
 from zenml.integrations.integration import Integration
@@ -35,18 +34,7 @@ class TensorBoardIntegration(Integration):
         Returns:
             A list of requirements.
         """
-        if sys.version_info > (3, 11):
-            tf_version = "2.13"
-        else:
-            # Capping tensorflow to 2.11 for Python 3.10 and below because it
-            # is not compatible with Pytorch
-            # (see https://github.com/pytorch/pytorch/issues/99637).
-            tf_version = "2.11"
-
-        requirements = [
-            f"tensorboard=={tf_version}",
-            "protobuf>=3.6.0,<4.0.0",
-        ]
+        requirements = ["tensorboard>=2.12,<2.15"]
         return requirements
 
     @classmethod

zenml/integrations/tensorboard/services/tensorboard_service.py

@@ -13,12 +13,10 @@
 #  permissions and limitations under the License.
 """Implementation of the TensorBoard service."""
 
+import uuid
 from typing import Any, Dict, Union
 
 from tensorboard import default, program  # type: ignore [import-untyped]
-from tensorboard.uploader import (  # type: ignore [import-untyped]
-    uploader_subcommand,
-)
 
 from zenml.logger import get_logger
 from zenml.services import (
@@ -103,6 +101,8 @@ class TensorboardService(LocalDaemonService):
                 ),
             )
             attrs["endpoint"] = endpoint
+        if "uuid" not in attrs:
+            attrs["uuid"] = uuid.uuid4()
         super().__init__(config=config, **attrs)
 
     def run(self) -> None:
@@ -117,7 +117,6 @@ class TensorboardService(LocalDaemonService):
         try:
             tensorboard = program.TensorBoard(
                 plugins=default.get_plugins(),
-                subcommands=[uploader_subcommand.UploaderSubcommand()],
             )
             tensorboard.configure(
                 logdir=self.config.logdir,

zenml/integrations/tensorboard/visualizers/tensorboard_visualizer.py

@@ -87,40 +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,
+                    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,
@@ -153,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":
@@ -193,19 +200,19 @@ def get_step(pipeline_name: str, step_name: str) -> "StepRunResponse":
     Raises:
         RuntimeError: If the step is not found.
     """
-    pipeline = Client().get_pipeline(pipeline_name)
-    if pipeline is None:
+    runs = Client().list_pipeline_runs(pipeline=pipeline_name)
+    if runs.total == 0:
         raise RuntimeError(
-            f"No pipeline with name `{pipeline_name}` was found"
+            f"No pipeline runs for pipeline `{pipeline_name}` were found"
         )
 
-    last_run = pipeline.runs[0]
-    step = last_run.steps[step_name]
-    if step is None:
+    last_run = runs[0]
+    if step_name not in last_run.steps:
         raise RuntimeError(
             f"No pipeline step with name `{step_name}` was found in "
             f"pipeline `{pipeline_name}`"
         )
+    step = last_run.steps[step_name]
     return step
 
 

zenml/integrations/tensorflow/__init__.py

@@ -18,6 +18,9 @@ import sys
 from typing import List, Optional
 from zenml.integrations.constants import TENSORFLOW
 from zenml.integrations.integration import Integration
+from zenml.logger import get_logger
+
+logger = get_logger(__name__)
 
 
 class TensorflowIntegration(Integration):
@@ -25,17 +28,15 @@ class TensorflowIntegration(Integration):
 
     NAME = TENSORFLOW
     REQUIREMENTS = []
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["typing-extensions"]
 
     @classmethod
     def activate(cls) -> None:
         """Activates the integration."""
         # need to import this explicitly to load the Tensorflow file IO support
         # for S3 and other file systems
-        if (
-            not platform.system() == "Darwin"
-            or not platform.machine() == "arm64"
-        ):
-            import tensorflow_io  # type: ignore [import-untyped]
+        if not platform.system() == "Darwin" or not platform.machine() == "arm64":
+            import tensorflow_io  # type: ignore
 
         from zenml.integrations.tensorflow import materializers  # noqa
 
@@ -49,24 +50,18 @@ class TensorflowIntegration(Integration):
         Returns:
             A list of requirements.
         """
-        if sys.version_info > (3, 11):
-            tf_version = "2.13"
-        else:
-            # Capping tensorflow to 2.11 for Python 3.10 and below because it
-            # is not compatible with Pytorch
-            # (see https://github.com/pytorch/pytorch/issues/99637).
-            tf_version = "2.11"
         target_os = target_os or platform.system()
         if target_os == "Darwin" and platform.machine() == "arm64":
             requirements = [
-                f"tensorflow-macos=={tf_version}",
+                "tensorflow-macos>=2.12,<2.15",
             ]
         else:
             requirements = [
-                f"tensorflow=={tf_version}",
+                "tensorflow>=2.12,<2.15",
                 "tensorflow_io>=0.24.0",
-                "protobuf>=3.6.0,<4.0.0",
             ]
+        if sys.version_info.minor == 8:
+            requirements.append("typing-extensions>=4.6.1")
         return requirements
 
 

zenml/integrations/tensorflow/materializers/keras_materializer.py

@@ -13,14 +13,14 @@
 #  permissions and limitations under the License.
 """Implementation of the TensorFlow Keras materializer."""
 
-import tempfile
+import os
 from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type
 
-from tensorflow import keras
+import tensorflow as tf
+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
 
@@ -31,48 +31,45 @@ if TYPE_CHECKING:
 class KerasMaterializer(BaseMaterializer):
     """Materializer to read/write Keras models."""
 
-    ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = (keras.Model,)
+    ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = (
+        tf.keras.Model,
+        tf_keras.Model,
+    )
     ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.MODEL
+    MODEL_FILE_NAME = "model.keras"
 
-    def load(self, data_type: Type[Any]) -> keras.Model:
+    def load(self, data_type: Type[Any]) -> tf_keras.Model:
         """Reads and returns a Keras model after copying it to temporary path.
 
         Args:
             data_type: The type of the data to read.
 
         Returns:
-            A tf.keras.Model model.
+            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
-        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 = keras.models.load_model(temp_dir.name)
+            return model
 
-        # Cleanup and return
-        fileio.rmtree(temp_dir.name)
-
-        return model
-
-    def save(self, model: keras.Model) -> None:
+    def save(self, model: tf_keras.Model) -> None:
         """Writes a keras model to the artifact store.
 
         Args:
-            model: A tf.keras.Model model.
+            model: A keras.Model model.
         """
-        # Create a temporary directory to store the model
-        temp_dir = tempfile.TemporaryDirectory()
-        model.save(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:
+            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: keras.Model
+        self, model: tf_keras.Model
     ) -> Dict[str, "MetadataType"]:
         """Extract metadata from the given `Model` object.
 

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.experimental.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:
-            tf.data.experimental.save(
+        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"]