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

zenml/integrations/neptune/experiment_trackers/neptune_experiment_tracker.py

@@ -77,10 +77,12 @@ class NeptuneExperimentTracker(BaseExperimentTracker):
             NeptuneExperimentTrackerSettings, self.get_settings(info)
         )
 
-        self.run_state.token = self.config.api_token
-        self.run_state.project = self.config.project
-        self.run_state.run_name = info.run_name
-        self.run_state.tags = list(settings.tags)
+        self.run_state.initialize(
+            project=self.config.project,
+            token=self.config.api_token,
+            run_name=info.run_name,
+            tags=list(settings.tags),
+        )
 
     def get_step_run_metadata(
         self, info: "StepRunInfo"
@@ -107,4 +109,4 @@ class NeptuneExperimentTracker(BaseExperimentTracker):
         """
         self.run_state.active_run.sync()
         self.run_state.active_run.stop()
-        self.run_state.reset_active_run()
+        self.run_state.reset()

zenml/integrations/neptune/experiment_trackers/run_state.py

@@ -16,11 +16,10 @@
 from hashlib import md5
 from typing import TYPE_CHECKING, Any, List, Optional
 
-import neptune  # type: ignore[import-untyped]
+import neptune
 
 import zenml
 from zenml.client import Client
-from zenml.integrations.constants import NEPTUNE
 from zenml.utils.singleton import SingletonMetaClass
 
 if TYPE_CHECKING:
@@ -29,20 +28,38 @@ if TYPE_CHECKING:
 _INTEGRATION_VERSION_KEY = "source_code/integrations/zenml"
 
 
-class InvalidExperimentTrackerSelected(Exception):
-    """Raised if a Neptune run is fetched while using a different experiment tracker."""
-
-
 class RunProvider(metaclass=SingletonMetaClass):
     """Singleton object used to store and persist a Neptune run state across the pipeline."""
 
     def __init__(self) -> None:
         """Initialize RunProvider. Called with no arguments."""
-        self._active_run = None
-        self._project: Optional[str]
-        self._run_name: Optional[str]
-        self._token: Optional[str]
-        self._tags: Optional[List[str]]
+        self._active_run: Optional["Run"] = None
+        self._project: Optional[str] = None
+        self._run_name: Optional[str] = None
+        self._token: Optional[str] = None
+        self._tags: Optional[List[str]] = None
+        self._initialized = False
+
+    def initialize(
+        self,
+        project: Optional[str] = None,
+        token: Optional[str] = None,
+        run_name: Optional[str] = None,
+        tags: Optional[List[str]] = None,
+    ) -> None:
+        """Initialize the run state.
+
+        Args:
+            project: The neptune project.
+            token: The neptune token.
+            run_name: The neptune run name.
+            tags: Tags for the neptune run.
+        """
+        self._project = project
+        self._token = token
+        self._run_name = run_name
+        self._tags = tags
+        self._initialized = True
 
     @property
     def project(self) -> Optional[Any]:
@@ -53,15 +70,6 @@ class RunProvider(metaclass=SingletonMetaClass):
         """
         return self._project
 
-    @project.setter
-    def project(self, project: str) -> None:
-        """Setter for project name.
-
-        Args:
-            project: Neptune project name
-        """
-        self._project = project
-
     @property
     def token(self) -> Optional[Any]:
         """Getter for API token.
@@ -71,15 +79,6 @@ class RunProvider(metaclass=SingletonMetaClass):
         """
         return self._token
 
-    @token.setter
-    def token(self, token: str) -> None:
-        """Setter for API token.
-
-        Args:
-            token: Neptune API token
-        """
-        self._token = token
-
     @property
     def run_name(self) -> Optional[Any]:
         """Getter for run name.
@@ -89,15 +88,6 @@ class RunProvider(metaclass=SingletonMetaClass):
         """
         return self._run_name
 
-    @run_name.setter
-    def run_name(self, run_name: str) -> None:
-        """Setter for run name.
-
-        Args:
-            run_name: name of the pipeline run
-        """
-        self._run_name = run_name
-
     @property
     def tags(self) -> Optional[Any]:
         """Getter for run tags.
@@ -107,14 +97,14 @@ class RunProvider(metaclass=SingletonMetaClass):
         """
         return self._tags
 
-    @tags.setter
-    def tags(self, tags: List[str]) -> None:
-        """Setter for run tags.
+    @property
+    def initialized(self) -> bool:
+        """If the run state is initialized.
 
-        Args:
-            tags: list of tags associated with a Neptune run
+        Returns:
+            If the run state is initialized.
         """
-        self._tags = tags
+        return self._initialized
 
     @property
     def active_run(self) -> "Run":
@@ -137,9 +127,14 @@ class RunProvider(metaclass=SingletonMetaClass):
             self._active_run = run
         return self._active_run
 
-    def reset_active_run(self) -> None:
-        """Resets the active run state to None."""
+    def reset(self) -> None:
+        """Reset the run state."""
         self._active_run = None
+        self._project = None
+        self._run_name = None
+        self._token = None
+        self._tags = None
+        self._initialized = False
 
 
 def get_neptune_run() -> "Run":
@@ -149,14 +144,35 @@ def get_neptune_run() -> "Run":
         Neptune run object
 
     Raises:
-        InvalidExperimentTrackerSelected: when called while using an experiment tracker other than Neptune
+        RuntimeError: When unable to fetch the active neptune run.
     """
-    client = Client()
-    experiment_tracker = client.active_stack.experiment_tracker
-    if experiment_tracker.flavor == NEPTUNE:  # type: ignore
-        return experiment_tracker.run_state.active_run  # type: ignore
-    raise InvalidExperimentTrackerSelected(
-        "Fetching a Neptune run works only with the 'neptune' flavor of "
-        "the experiment tracker. The flavor currently selected is %s"
-        % experiment_tracker.flavor  # type: ignore
+    from zenml.integrations.neptune.experiment_trackers import (
+        NeptuneExperimentTracker,
     )
+
+    experiment_tracker = Client().active_stack.experiment_tracker
+
+    if not experiment_tracker:
+        raise RuntimeError(
+            "Unable to get neptune run: Missing experiment tracker in the "
+            "active stack."
+        )
+
+    if not isinstance(experiment_tracker, NeptuneExperimentTracker):
+        raise RuntimeError(
+            "Unable to get neptune run: Experiment tracker in the active "
+            f"stack ({experiment_tracker.flavor}) is not a neptune experiment "
+            "tracker."
+        )
+
+    run_state = experiment_tracker.run_state
+    if not run_state.initialized:
+        raise RuntimeError(
+            "Unable to get neptune run: The experiment tracker has not been "
+            "initialized. To solve this, make sure you use the experiment "
+            "tracker in your step. See "
+            "https://docs.zenml.io/stack-components/experiment-trackers/neptune#how-do-you-use-it "
+            "for more information."
+        )
+
+    return experiment_tracker.run_state.active_run

zenml/integrations/neptune/flavors/neptune_experiment_tracker_flavor.py

@@ -47,7 +47,7 @@ class NeptuneExperimentTrackerConfig(BaseExperimentTrackerConfig):
     """
 
     project: Optional[str] = None
-    api_token: Optional[str] = SecretField()
+    api_token: Optional[str] = SecretField(default=None)
 
 
 class NeptuneExperimentTrackerSettings(BaseSettings):

zenml/integrations/neural_prophet/__init__.py

@@ -21,7 +21,12 @@ class NeuralProphetIntegration(Integration):
     """Definition of NeuralProphet integration for ZenML."""
 
     NAME = NEURAL_PROPHET
-    REQUIREMENTS = ["neuralprophet>=0.3.2,<0.5.0", "holidays>=0.4.1,<0.25.0"]
+    REQUIREMENTS = [
+        "neuralprophet>=0.3.2,<0.5.0",
+        "holidays>=0.4.1,<0.25.0",
+        "tenacity!=8.4.0",  # https://github.com/jd/tenacity/issues/471
+    ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["tenacity"]
 
     @classmethod
     def activate(cls) -> None:

zenml/integrations/numpy/__init__.py

@@ -0,0 +1,32 @@
+#  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 Numpy integration."""
+
+from zenml.integrations.constants import NUMPY
+from zenml.integrations.integration import Integration
+
+
+class NumpyIntegration(Integration):
+    """Definition of Numpy integration for ZenML."""
+
+    NAME = NUMPY
+    REQUIREMENTS = ["numpy<2.0.0"]
+
+    @classmethod
+    def activate(cls) -> None:
+        """Activates the integration."""
+        from zenml.integrations.numpy import materializers  # noqa
+
+
+NumpyIntegration.check_installation()

zenml/integrations/{kserve/constants.py → numpy/materializers/__init__.py}

@@ -1,4 +1,4 @@
-#  Copyright (c) ZenML GmbH 2022. All Rights Reserved.
+#  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.
@@ -11,7 +11,8 @@
 #  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.
-"""KServe constants."""
+"""Initialization of the Numpy materializer."""
 
-KSERVE_DOCKER_IMAGE_KEY = "kserve_model_deployer"
-KSERVE_CUSTOM_DEPLOYMENT = "kserve_custom_deployment"
+from zenml.integrations.numpy.materializers.numpy_materializer import (  # noqa
+    NumpyMaterializer,
+)

zenml/integrations/numpy/materializers/numpy_materializer.py

@@ -0,0 +1,246 @@
+#  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.
+"""Implementation of the ZenML NumPy materializer."""
+
+import os
+from collections import Counter
+from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type
+
+import numpy as np
+
+from zenml.enums import ArtifactType, VisualizationType
+from zenml.logger import get_logger
+from zenml.materializers.base_materializer import BaseMaterializer
+from zenml.metadata.metadata_types import DType, MetadataType
+
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
+logger = get_logger(__name__)
+
+
+NUMPY_FILENAME = "data.npy"
+
+DATA_FILENAME = "data.parquet"
+SHAPE_FILENAME = "shape.json"
+DATA_VAR = "data_var"
+
+
+class NumpyMaterializer(BaseMaterializer):
+    """Materializer to read data to and from pandas."""
+
+    ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = (np.ndarray,)
+    ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.DATA
+
+    def load(self, data_type: Type[Any]) -> "Any":
+        """Reads a numpy array from a `.npy` file.
+
+        Args:
+            data_type: The type of the data to read.
+
+
+        Raises:
+            ImportError: If pyarrow is not installed.
+
+        Returns:
+            The numpy array.
+        """
+        numpy_file = os.path.join(self.uri, NUMPY_FILENAME)
+
+        if self.artifact_store.exists(numpy_file):
+            with self.artifact_store.open(numpy_file, "rb") as f:
+                return np.load(f, allow_pickle=True)
+        elif self.artifact_store.exists(os.path.join(self.uri, DATA_FILENAME)):
+            logger.warning(
+                "A legacy artifact was found. "
+                "This artifact was created with an older version of "
+                "ZenML. You can still use it, but it will be "
+                "converted to the new format on the next materialization."
+            )
+            try:
+                # Import old materializer dependencies
+                import pyarrow as pa  # type: ignore
+                import pyarrow.parquet as pq  # type: ignore
+
+                from zenml.utils import yaml_utils
+
+                # Read numpy array from parquet file
+                shape_dict = yaml_utils.read_json(
+                    os.path.join(self.uri, SHAPE_FILENAME)
+                )
+                shape_tuple = tuple(shape_dict.values())
+                with self.artifact_store.open(
+                    os.path.join(self.uri, DATA_FILENAME), "rb"
+                ) as f:
+                    input_stream = pa.input_stream(f)
+                    data = pq.read_table(input_stream)
+                vals = getattr(data.to_pandas(), DATA_VAR).values
+                return np.reshape(vals, shape_tuple)
+            except ImportError:
+                raise ImportError(
+                    "You have an old version of a `NumpyMaterializer` ",
+                    "data artifact stored in the artifact store ",
+                    "as a `.parquet` file, which requires `pyarrow` for reading. ",
+                    "You can install `pyarrow` by running `pip install pyarrow`.",
+                )
+
+    def save(self, arr: "NDArray[Any]") -> None:
+        """Writes a np.ndarray to the artifact store as a `.npy` file.
+
+        Args:
+            arr: The numpy array to write.
+        """
+        with self.artifact_store.open(
+            os.path.join(self.uri, NUMPY_FILENAME), "wb"
+        ) as f:
+            np.save(f, arr)
+
+    def save_visualizations(
+        self, arr: "NDArray[Any]"
+    ) -> Dict[str, VisualizationType]:
+        """Saves visualizations for a numpy array.
+
+        If the array is 1D, a histogram is saved. If the array is 2D or 3D with
+        3 or 4 channels, an image is saved.
+
+        Args:
+            arr: The numpy array to visualize.
+
+        Returns:
+            A dictionary of visualization URIs and their types.
+        """
+        if not np.issubdtype(arr.dtype, np.number):
+            return {}
+
+        try:
+            # Save histogram for 1D arrays
+            if len(arr.shape) == 1:
+                histogram_path = os.path.join(self.uri, "histogram.png")
+                histogram_path = histogram_path.replace("\\", "/")
+                self._save_histogram(histogram_path, arr)
+                return {histogram_path: VisualizationType.IMAGE}
+
+            # Save as image for 3D arrays with 3 or 4 channels
+            if len(arr.shape) == 3 and arr.shape[2] in [3, 4]:
+                image_path = os.path.join(self.uri, "image.png")
+                image_path = image_path.replace("\\", "/")
+                self._save_image(image_path, arr)
+                return {image_path: VisualizationType.IMAGE}
+
+        except ImportError:
+            logger.info(
+                "Skipping visualization of numpy array because matplotlib "
+                "is not installed. To install matplotlib, run "
+                "`pip install matplotlib`."
+            )
+
+        return {}
+
+    def _save_histogram(self, output_path: str, arr: "NDArray[Any]") -> None:
+        """Saves a histogram of a numpy array.
+
+        Args:
+            output_path: The path to save the histogram to.
+            arr: The numpy array of which to save the histogram.
+        """
+        import matplotlib.pyplot as plt
+
+        plt.hist(arr)
+        with self.artifact_store.open(output_path, "wb") as f:
+            plt.savefig(f)
+        plt.close()
+
+    def _save_image(self, output_path: str, arr: "NDArray[Any]") -> None:
+        """Saves a numpy array as an image.
+
+        Args:
+            output_path: The path to save the image to.
+            arr: The numpy array to save.
+        """
+        from matplotlib.image import imsave
+
+        with self.artifact_store.open(output_path, "wb") as f:
+            imsave(f, arr)
+
+    def extract_metadata(
+        self, arr: "NDArray[Any]"
+    ) -> Dict[str, "MetadataType"]:
+        """Extract metadata from the given numpy array.
+
+        Args:
+            arr: The numpy array to extract metadata from.
+
+        Returns:
+            The extracted metadata as a dictionary.
+        """
+        if np.issubdtype(arr.dtype, np.number):
+            return self._extract_numeric_metadata(arr)
+        elif np.issubdtype(arr.dtype, np.unicode_) or np.issubdtype(
+            arr.dtype, np.object_
+        ):
+            return self._extract_text_metadata(arr)
+        else:
+            return {}
+
+    def _extract_numeric_metadata(
+        self, arr: "NDArray[Any]"
+    ) -> Dict[str, "MetadataType"]:
+        """Extracts numeric metadata from a numpy array.
+
+        Args:
+            arr: The numpy array to extract metadata from.
+
+        Returns:
+            A dictionary of metadata.
+        """
+        min_val = np.min(arr).item()
+        max_val = np.max(arr).item()
+
+        numpy_metadata: Dict[str, "MetadataType"] = {
+            "shape": tuple(arr.shape),
+            "dtype": DType(arr.dtype.type),
+            "mean": np.mean(arr).item(),
+            "std": np.std(arr).item(),
+            "min": min_val,
+            "max": max_val,
+        }
+        return numpy_metadata
+
+    def _extract_text_metadata(
+        self, arr: "NDArray[Any]"
+    ) -> Dict[str, "MetadataType"]:
+        """Extracts text metadata from a numpy array.
+
+        Args:
+            arr: The numpy array to extract metadata from.
+
+        Returns:
+            A dictionary of metadata.
+        """
+        text = " ".join(arr)
+        words = text.split()
+        word_counts = Counter(words)
+        unique_words = len(word_counts)
+        total_words = len(words)
+        most_common_word, most_common_count = word_counts.most_common(1)[0]
+
+        text_metadata: Dict[str, "MetadataType"] = {
+            "shape": tuple(arr.shape),
+            "dtype": DType(arr.dtype.type),
+            "unique_words": unique_words,
+            "total_words": total_words,
+            "most_common_word": most_common_word,
+            "most_common_count": most_common_count,
+        }
+        return text_metadata

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/{steps/external_artifact.py → integrations/pandas/__init__.py}

@@ -1,4 +1,4 @@
-#  Copyright (c) ZenML GmbH 2023. All Rights Reserved.
+#  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.
@@ -11,16 +11,22 @@
 #  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.
-"""Backward compatibility for the `ExternalArtifact` class."""
+"""Initialization of the Pandas integration."""
 
-from zenml.artifacts.external_artifact import ExternalArtifact
-from zenml.logger import get_logger
+from zenml.integrations.constants import PANDAS
+from zenml.integrations.integration import Integration
 
-logger = get_logger(__name__)
-logger.warning(
-    "The `ExternalArtifact` object was moved to `zenml.artifacts.external_artifact`. "
-    "This import path is going to be deprecated. Please consider updating your "
-    "code to use `from zenml.artifacts.external_artifact import ExternalArtifact`."
-)
 
-__all__ = ["ExternalArtifact"]
+class PandasIntegration(Integration):
+    """Definition of Pandas integration for ZenML."""
+
+    NAME = PANDAS
+    REQUIREMENTS = ["pandas>=2.0.0"]
+
+    @classmethod
+    def activate(cls) -> None:
+        """Activates the integration."""
+        from zenml.integrations.pandas import materializers  # noqa
+
+
+PandasIntegration.check_installation()