zenml-nightly 0.64.0.dev20240811__py3-none-any.whl → 0.66.0.dev20240910__py3-none-any.whl

zenml/materializers/__init__.py

@@ -29,8 +29,6 @@ from zenml.materializers.cloudpickle_materializer import (
 from zenml.materializers.structured_string_materializer import (
     StructuredStringMaterializer,
 )
-from zenml.materializers.numpy_materializer import NumpyMaterializer
-from zenml.materializers.pandas_materializer import PandasMaterializer
 from zenml.materializers.pydantic_materializer import PydanticMaterializer
 from zenml.materializers.service_materializer import ServiceMaterializer
 
@@ -40,8 +38,6 @@ __all__ = [
     "BytesMaterializer",
     "CloudpickleMaterializer",
     "StructuredStringMaterializer",
-    "NumpyMaterializer",
-    "PandasMaterializer",
     "PydanticMaterializer",
     "ServiceMaterializer",
 ]

zenml/materializers/base_materializer.py

@@ -98,6 +98,10 @@ class BaseMaterializerMeta(type):
                     associated_type, cls
                 )
 
+            from zenml.utils import notebook_utils
+
+            notebook_utils.try_to_save_notebook_cell_code(cls)
+
         return cls
 
 

zenml/materializers/numpy_materializer.py

@@ -1,4 +1,4 @@
-#  Copyright (c) ZenML GmbH 2021. 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,236 +11,25 @@
 #  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
+"""Placeholder for importing the Numpy Materializer from its former path."""
+
+# With the code below, we keep a reference to the NumpyMaterializer in its
+# original spot. So, if someone has an artifact version that was created with
+# the NumpyMaterializer, (assuming that they have it installed), they will still
+# be able to use the artifact version. However, if this file is to be removed,
+# you have to write a DB migration script to change the materializer source
+# of the old artifact version entries.
+
+try:
+    from zenml.integrations.numpy.materializers.numpy_materializer import (  # noqa
+        NumpyMaterializer,
+    )
+except (ImportError, ModuleNotFoundError) as e:
+    raise ImportError(
+        "The ZenML built-in Numpy materializer has been moved to an "
+        "integration. Before you use it, make sure you that the `numpy` "
+        "integration is installed with `zenml integration install numpy`, "
+        "and if you need to import the materializer explicitly, please use: "
+        "'from zenml.integrations.numpy.materializers import "
+        f"NumpyMaterializer': {e}"
+    )

zenml/materializers/pandas_materializer.py

@@ -1,4 +1,4 @@
-#  Copyright (c) ZenML GmbH 2021. 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,182 +11,25 @@
 #  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.
-"""Materializer for Pandas."""
-
-import os
-from typing import Any, ClassVar, Dict, Optional, Tuple, Type, Union
-
-import pandas as pd
-
-from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
-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
-
-logger = get_logger(__name__)
-
-PARQUET_FILENAME = "df.parquet.gzip"
-COMPRESSION_TYPE = "gzip"
-
-CSV_FILENAME = "df.csv"
-
-
-class PandasMaterializer(BaseMaterializer):
-    """Materializer to read data to and from pandas."""
-
-    ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = (
-        pd.DataFrame,
-        pd.Series,
+"""Placeholder for importing the Pandas Materializer from its former path."""
+
+# With the code below, we keep a reference to the PandasMaterializer in its
+# original spot. So, if someone has an artifact version that was created with
+# the NumpyMaterializer, (assuming that they have it installed), they will still
+# be able to use the artifact version. However, if this file is to be removed,
+# you have to write a DB migration script to change the materializer source
+# of the old artifact version entries.
+
+try:
+    from zenml.integrations.pandas.materializers.pandas_materializer import (  # noqa
+        PandasMaterializer,
+    )
+except (ImportError, ModuleNotFoundError) as e:
+    raise ImportError(
+        "The ZenML built-in Pandas materializer has been moved to an "
+        "integration. Before you use it, make sure you that the `pandas` "
+        "integration is installed with `zenml integration install pandas`, "
+        "and if you need to import the materializer explicitly, please use: "
+        "'from zenml.integrations.pandas.materializers import "
+        f"PandasMaterializer': {e}"
     )
-    ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.DATA
-
-    def __init__(
-        self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
-    ):
-        """Define `self.data_path`.
-
-        Args:
-            uri: The URI where the artifact data is stored.
-            artifact_store: The artifact store where the artifact data is stored.
-        """
-        super().__init__(uri, artifact_store)
-        try:
-            import pyarrow  # type: ignore # noqa
-
-            self.pyarrow_exists = True
-        except ImportError:
-            self.pyarrow_exists = False
-            logger.warning(
-                "By default, the `PandasMaterializer` stores data as a "
-                "`.csv` file. If you want to store data more efficiently, "
-                "you can install `pyarrow` by running "
-                "'`pip install pyarrow`'. This will allow `PandasMaterializer` "
-                "to automatically store the data as a `.parquet` file instead."
-            )
-        finally:
-            self.parquet_path = os.path.join(self.uri, PARQUET_FILENAME)
-            self.csv_path = os.path.join(self.uri, CSV_FILENAME)
-
-    def load(self, data_type: Type[Any]) -> Union[pd.DataFrame, pd.Series]:
-        """Reads `pd.DataFrame` or `pd.Series` from a `.parquet` or `.csv` file.
-
-        Args:
-            data_type: The type of the data to read.
-
-        Raises:
-            ImportError: If pyarrow or fastparquet is not installed.
-
-        Returns:
-            The pandas dataframe or series.
-        """
-        if self.artifact_store.exists(self.parquet_path):
-            if self.pyarrow_exists:
-                with self.artifact_store.open(
-                    self.parquet_path, mode="rb"
-                ) as f:
-                    df = pd.read_parquet(f)
-            else:
-                raise ImportError(
-                    "You have an old version of a `PandasMaterializer` "
-                    "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 fastparquet`'."
-                )
-        else:
-            with self.artifact_store.open(self.csv_path, mode="rb") as f:
-                df = pd.read_csv(f, index_col=0, parse_dates=True)
-
-        # validate the type of the data.
-        def is_dataframe_or_series(
-            df: Union[pd.DataFrame, pd.Series],
-        ) -> Union[pd.DataFrame, pd.Series]:
-            """Checks if the data is a `pd.DataFrame` or `pd.Series`.
-
-            Args:
-                df: The data to check.
-
-            Returns:
-                The data if it is a `pd.DataFrame` or `pd.Series`.
-            """
-            if issubclass(data_type, pd.Series):
-                # Taking the first column if its a series as the assumption
-                # is that there will only be one
-                assert len(df.columns) == 1
-                df = df[df.columns[0]]
-                return df
-            else:
-                return df
-
-        return is_dataframe_or_series(df)
-
-    def save(self, df: Union[pd.DataFrame, pd.Series]) -> None:
-        """Writes a pandas dataframe or series to the specified filename.
-
-        Args:
-            df: The pandas dataframe or series to write.
-        """
-        if isinstance(df, pd.Series):
-            df = df.to_frame(name="series")
-
-        if self.pyarrow_exists:
-            with self.artifact_store.open(self.parquet_path, mode="wb") as f:
-                df.to_parquet(f, compression=COMPRESSION_TYPE)
-        else:
-            with self.artifact_store.open(self.csv_path, mode="wb") as f:
-                df.to_csv(f, index=True)
-
-    def save_visualizations(
-        self, df: Union[pd.DataFrame, pd.Series]
-    ) -> Dict[str, VisualizationType]:
-        """Save visualizations of the given pandas dataframe or series.
-
-        Args:
-            df: The pandas dataframe or series to visualize.
-
-        Returns:
-            A dictionary of visualization URIs and their types.
-        """
-        describe_uri = os.path.join(self.uri, "describe.csv")
-        describe_uri = describe_uri.replace("\\", "/")
-        with self.artifact_store.open(describe_uri, mode="wb") as f:
-            df.describe().to_csv(f)
-        return {describe_uri: VisualizationType.CSV}
-
-    def extract_metadata(
-        self, df: Union[pd.DataFrame, pd.Series]
-    ) -> Dict[str, "MetadataType"]:
-        """Extract metadata from the given pandas dataframe or series.
-
-        Args:
-            df: The pandas dataframe or series to extract metadata from.
-
-        Returns:
-            The extracted metadata as a dictionary.
-        """
-        pandas_metadata: Dict[str, "MetadataType"] = {"shape": df.shape}
-
-        if isinstance(df, pd.Series):
-            pandas_metadata["dtype"] = DType(df.dtype.type)
-            pandas_metadata["mean"] = float(df.mean().item())
-            pandas_metadata["std"] = float(df.std().item())
-            pandas_metadata["min"] = float(df.min().item())
-            pandas_metadata["max"] = float(df.max().item())
-
-        else:
-            pandas_metadata["dtype"] = {
-                str(key): DType(value.type) for key, value in df.dtypes.items()
-            }
-            for stat_name, stat in {
-                "mean": df.mean,
-                "std": df.std,
-                "min": df.min,
-                "max": df.max,
-            }.items():
-                pandas_metadata[stat_name] = {
-                    str(key): float(value)
-                    for key, value in stat(numeric_only=True).to_dict().items()
-                }
-
-        return pandas_metadata

zenml/model/model.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """Model user facing interface to pass into pipeline or step."""
 
+import datetime
 import time
 from typing import (
     TYPE_CHECKING,
@@ -31,6 +32,7 @@ from zenml.enums import MetadataResourceTypes, ModelStages
 from zenml.exceptions import EntityExistsError
 from zenml.logger import get_logger
 from zenml.utils.pydantic_utils import before_validator_handler
+from zenml.utils.string_utils import format_name_template
 
 if TYPE_CHECKING:
     from zenml.metadata.metadata_types import MetadataType
@@ -40,6 +42,7 @@ if TYPE_CHECKING:
         ModelVersionResponse,
         PipelineRunResponse,
         RunMetadataResponse,
+        StepRunResponse,
     )
 
 logger = get_logger(__name__)
@@ -508,8 +511,12 @@ class Model(BaseModel):
         Returns:
             Dict of validated values.
         """
-        suppress_class_validation_warnings = data.get(
-            "suppress_class_validation_warnings", False
+        suppress_class_validation_warnings = (
+            data.get(
+                "suppress_class_validation_warnings",
+                False,
+            )
+            or data.get("model_version_id", None) is not None
         )
         version = data.get("version", None)
 
@@ -677,6 +684,10 @@ class Model(BaseModel):
 
         model = self._get_or_create_model()
 
+        # backup logic, if the Model class is used directly from the code
+        if isinstance(self.version, str):
+            self.version = format_name_template(self.version)
+
         zenml_client = Client()
         model_version_request = ModelVersionRequest(
             user=zenml_client.active_user.id,
@@ -819,3 +830,81 @@ class Model(BaseModel):
                 )
             )
         )
+
+    def _prepare_model_version_before_step_launch(
+        self,
+        pipeline_run: "PipelineRunResponse",
+        step_run: Optional["StepRunResponse"],
+        return_logs: bool,
+    ) -> str:
+        """Prepares model version inside pipeline run.
+
+        Args:
+            pipeline_run: pipeline run
+            step_run: step run (passed only if model version is defined in a step explicitly)
+            return_logs: whether to return logs or not
+
+        Returns:
+            Logs related to the Dashboard URL to show later.
+        """
+        from zenml.client import Client
+        from zenml.models import PipelineRunUpdate, StepRunUpdate
+
+        logs = ""
+
+        # copy Model instance to prevent corrupting configs of the
+        # subsequent runs, if they share the same config object
+        self_copy = self.model_copy()
+
+        # in case request is within the step and no self-configuration is provided
+        # try reuse what's in the pipeline run first
+        if step_run is None and pipeline_run.model_version is not None:
+            self_copy.version = pipeline_run.model_version.name
+            self_copy.model_version_id = pipeline_run.model_version.id
+        # otherwise try to fill the templated name, if needed
+        elif isinstance(self_copy.version, str):
+            if pipeline_run.start_time:
+                start_time = pipeline_run.start_time
+            else:
+                start_time = datetime.datetime.now(datetime.timezone.utc)
+            self_copy.version = format_name_template(
+                self_copy.version,
+                date=start_time.strftime("%Y_%m_%d"),
+                time=start_time.strftime("%H_%M_%S_%f"),
+            )
+
+        # if exact model not yet defined - try to get/create and update it
+        # back to the run accordingly
+        if self_copy.model_version_id is None:
+            model_version_response = self_copy._get_or_create_model_version()
+
+            # update the configured model version id in runs accordingly
+            if step_run:
+                Client().zen_store.update_run_step(
+                    step_run_id=step_run.id,
+                    step_run_update=StepRunUpdate(
+                        model_version_id=model_version_response.id
+                    ),
+                )
+            else:
+                Client().zen_store.update_run(
+                    run_id=pipeline_run.id,
+                    run_update=PipelineRunUpdate(
+                        model_version_id=model_version_response.id
+                    ),
+                )
+
+            if return_logs:
+                from zenml.utils.cloud_utils import try_get_model_version_url
+
+                if logs_to_show := try_get_model_version_url(
+                    model_version_response
+                ):
+                    logs = logs_to_show
+                else:
+                    logs = (
+                        "Models can be viewed in the dashboard using ZenML Pro. Sign up "
+                        "for a free trial at https://www.zenml.io/pro/"
+                    )
+        self.model_version_id = self_copy.model_version_id
+        return logs

zenml/model/utils.py

@@ -61,13 +61,8 @@ def link_step_artifacts_to_model(
             artifact_name
         ).artifact_config
 
-        # Implicit linking
         if artifact_config is None and model is not None:
             artifact_config = ArtifactConfig(name=artifact_name)
-            logger.info(
-                f"Implicitly linking artifact `{artifact_name}` to model "
-                f"`{model.name}` version `{model.version}`."
-            )
 
         if artifact_config:
             link_artifact_config_to_model(
@@ -101,6 +96,11 @@ def link_artifact_config_to_model(
         )
 
     if model:
+        logger.debug(
+            f"Linking artifact `{artifact_config.name}` to model "
+            f"`{model.name}` version `{model.version}` using config "
+            f"`{artifact_config}`."
+        )
         request = ModelVersionArtifactRequest(
             user=client.active_user.id,
             workspace=client.active_workspace.id,