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

zenml/integrations/great_expectations/__init__.py

@@ -17,9 +17,8 @@ The Great Expectations integration enables you to use Great Expectations as a
 way of profiling and validating your data.
 """
 
-from typing import List, Type
+from typing import List, Type, Optional
 
-from zenml.enums import StackComponentType
 from zenml.integrations.constants import GREAT_EXPECTATIONS
 from zenml.integrations.integration import Integration
 from zenml.stack import Flavor
@@ -31,14 +30,12 @@ class GreatExpectationsIntegration(Integration):
     """Definition of Great Expectations integration for ZenML."""
 
     NAME = GREAT_EXPECTATIONS
-    REQUIREMENTS = [
-        "great-expectations>=0.15.0,<=0.15.47",
-        # typing_extensions 4.6.0 and above doesn't work with GE
-        "typing_extensions<4.6.0",
-    ]
-
-    @staticmethod
-    def activate() -> None:
+    REQUIREMENTS = ["great-expectations>=0.17.15,<1.0"]
+
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["pandas"]
+
+    @classmethod
+    def activate(cls) -> None:
         """Activate the Great Expectations integration."""
         from zenml.integrations.great_expectations import materializers  # noqa
 
@@ -55,5 +52,20 @@ class GreatExpectationsIntegration(Integration):
 
         return [GreatExpectationsDataValidatorFlavor]
 
+    @classmethod
+    def get_requirements(cls, target_os: Optional[str] = None) -> List[str]:
+        """Method to get the requirements for the integration.
+
+        Args:
+            target_os: The target operating system to get the requirements for.
+
+        Returns:
+            A list of requirements.
+        """
+        from zenml.integrations.pandas import PandasIntegration
+
+        return cls.REQUIREMENTS + \
+            PandasIntegration.get_requirements(target_os=target_os)
+
 
 GreatExpectationsIntegration.check_installation()

zenml/integrations/great_expectations/data_validators/ge_data_validator.py

@@ -17,21 +17,25 @@ import os
 from typing import Any, ClassVar, Dict, List, Optional, Sequence, Type, cast
 
 import pandas as pd
-import yaml
 from great_expectations.checkpoint.types.checkpoint_result import (  # type: ignore[import-untyped]
     CheckpointResult,
 )
 from great_expectations.core import (  # type: ignore[import-untyped]
     ExpectationSuite,
 )
-from great_expectations.data_context.data_context import (  # type: ignore[import-untyped]
-    BaseDataContext,
-    DataContext,
+from great_expectations.data_context.data_context.abstract_data_context import (
+    AbstractDataContext,
 )
-from great_expectations.data_context.types.base import (  # type: ignore[import-untyped]
+from great_expectations.data_context.data_context.context_factory import (
+    get_context,
+)
+from great_expectations.data_context.data_context.ephemeral_data_context import (
+    EphemeralDataContext,
+)
+from great_expectations.data_context.types.base import (
     DataContextConfig,
 )
-from great_expectations.data_context.types.resource_identifiers import (  # type: ignore[import-untyped]
+from great_expectations.data_context.types.resource_identifiers import (
     ExpectationSuiteIdentifier,
 )
 from great_expectations.profile.user_configurable_profiler import (  # type: ignore[import-untyped]
@@ -61,12 +65,12 @@ class GreatExpectationsDataValidator(BaseDataValidator):
     """Great Expectations data validator stack component."""
 
     NAME: ClassVar[str] = "Great Expectations"
-    FLAVOR: ClassVar[
-        Type[BaseDataValidatorFlavor]
-    ] = GreatExpectationsDataValidatorFlavor
+    FLAVOR: ClassVar[Type[BaseDataValidatorFlavor]] = (
+        GreatExpectationsDataValidatorFlavor
+    )
 
-    _context: BaseDataContext = None
-    _context_config: Optional[Dict[str, Any]] = None
+    _context: Optional[AbstractDataContext] = None
+    _context_config: Optional[DataContextConfig] = None
 
     @property
     def config(self) -> GreatExpectationsDataValidatorConfig:
@@ -78,7 +82,7 @@ class GreatExpectationsDataValidator(BaseDataValidator):
         return cast(GreatExpectationsDataValidatorConfig, self._config)
 
     @classmethod
-    def get_data_context(cls) -> BaseDataContext:
+    def get_data_context(cls) -> AbstractDataContext:
         """Get the Great Expectations data context managed by ZenML.
 
         Call this method to retrieve the data context managed by ZenML
@@ -94,15 +98,11 @@ class GreatExpectationsDataValidator(BaseDataValidator):
         return data_validator.data_context
 
     @property
-    def context_config(self) -> Optional[Dict[str, Any]]:
+    def context_config(self) -> Optional[DataContextConfig]:
         """Get the Great Expectations data context configuration.
 
-        The first time the context config is loaded from the stack component
-        config, it is converted from JSON/YAML string format to a dict.
-
         Raises:
-            ValueError: If the context_config value is not a valid JSON/YAML or
-                if the GE configuration extracted from it fails GE validation.
+            ValueError: In case there is an invalid context_config value
 
         Returns:
             A dictionary with the GE data context configuration.
@@ -111,31 +111,18 @@ class GreatExpectationsDataValidator(BaseDataValidator):
         if self._context_config is not None:
             return self._context_config
 
-        # Otherwise, load it from the stack component config
-        context_config = self.config.context_config
-        if context_config is None:
+        # Otherwise, use the configuration from the stack component config, if
+        # set
+        context_config_dict = self.config.context_config
+        if context_config_dict is None:
             return None
-        if isinstance(context_config, dict):
-            self._context_config = context_config
-            return self._context_config
-
-        # If the context config is a string, try to parse it as JSON/YAML
-        try:
-            context_config_dict = yaml.safe_load(context_config)
-        except yaml.parser.ParserError as e:
-            raise ValueError(
-                f"Malformed `context_config` value. Only JSON and YAML "
-                f"formats are supported: {str(e)}"
-            )
 
         # Validate that the context config is a valid GE config
         try:
-            context_config = DataContextConfig(**context_config_dict)
-            BaseDataContext(project_config=context_config)
+            self._context_config = DataContextConfig(**context_config_dict)
         except Exception as e:
             raise ValueError(f"Invalid `context_config` value: {str(e)}")
 
-        self._context_config = cast(Dict[str, Any], context_config_dict)
         return self._context_config
 
     @property
@@ -203,7 +190,7 @@ class GreatExpectationsDataValidator(BaseDataValidator):
         }
 
     @property
-    def data_context(self) -> BaseDataContext:
+    def data_context(self) -> AbstractDataContext:
         """Returns the Great Expectations data context configured for this component.
 
         Returns:
@@ -216,7 +203,9 @@ class GreatExpectationsDataValidator(BaseDataValidator):
             profiler_store_name = "zenml_profiler_store"
             evaluation_parameter_store_name = "evaluation_parameter_store"
 
-            zenml_context_config = dict(
+            # Define default configuration options that plug the GX stores
+            # in the active ZenML artifact store
+            zenml_context_config: Dict[str, Any] = dict(
                 stores={
                     expectations_store_name: self.get_store_config(
                         "ExpectationsStore", "expectations"
@@ -250,18 +239,29 @@ class GreatExpectationsDataValidator(BaseDataValidator):
             if self.config.context_root_dir:
                 # initialize the local data context, if a local path was
                 # configured
-                self._context = DataContext(self.config.context_root_dir)
+                self._context = get_context(
+                    context_root_dir=self.config.context_root_dir
+                )
+
             else:
-                # create an in-memory data context configuration that is not
-                # backed by a local YAML file (see https://docs.greatexpectations.io/docs/guides/setup/configuring_data_contexts/how_to_instantiate_a_data_context_without_a_yml_file/).
+                # create an ephemeral in-memory data context that is not
+                # backed by a local YAML file (see https://docs.greatexpectations.io/docs/oss/guides/setup/configuring_data_contexts/instantiating_data_contexts/instantiate_data_context/).
                 if self.context_config:
-                    context_config = DataContextConfig(**self.context_config)
+                    # Use the data context configuration provided in the stack
+                    # component configuration
+                    context_config = self.context_config
                 else:
+                    # Initialize the data context with the default ZenML
+                    # configuration options effectively plugging the GX stores
+                    # into the ZenML artifact store
                     context_config = DataContextConfig(**zenml_context_config)
                     # skip adding the stores after initialization, as they are
                     # already baked in the initial configuration
                     configure_zenml_stores = False
-                self._context = BaseDataContext(project_config=context_config)
+
+                self._context = EphemeralDataContext(
+                    project_config=context_config
+                )
 
             if configure_zenml_stores:
                 self._context.config.expectations_store_name = (
@@ -277,27 +277,31 @@ class GreatExpectationsDataValidator(BaseDataValidator):
                 self._context.config.evaluation_parameter_store_name = (
                     evaluation_parameter_store_name
                 )
-                for store_name, store_config in zenml_context_config[  # type: ignore[attr-defined]
+                for store_name, store_config in zenml_context_config[
                     "stores"
                 ].items():
                     self._context.add_store(
                         store_name=store_name,
                         store_config=store_config,
                     )
-                for site_name, site_config in zenml_context_config[  # type: ignore[attr-defined]
-                    "data_docs_sites"
-                ].items():
-                    self._context.config.data_docs_sites[
-                        site_name
-                    ] = site_config
-
-            if self.config.configure_local_docs:
+                if self._context.config.data_docs_sites is not None:
+                    for site_name, site_config in zenml_context_config[
+                        "data_docs_sites"
+                    ].items():
+                        self._context.config.data_docs_sites[site_name] = (
+                            site_config
+                        )
+
+            if (
+                self.config.configure_local_docs
+                and self._context.config.data_docs_sites is not None
+            ):
                 client = Client()
                 artifact_store = client.active_stack.artifact_store
                 if artifact_store.flavor != "local":
-                    self._context.config.data_docs_sites[
-                        "zenml_local"
-                    ] = self.get_data_docs_config("data_docs", local=True)
+                    self._context.config.data_docs_sites["zenml_local"] = (
+                        self.get_data_docs_config("data_docs", local=True)
+                    )
 
         return self._context
 
@@ -509,7 +513,7 @@ class GreatExpectationsDataValidator(BaseDataValidator):
             },
         ]
 
-        checkpoint_config = {
+        checkpoint_config: Dict[str, Any] = {
             "name": checkpoint_name,
             "run_name_template": run_name,
             "config_version": 1,
@@ -517,7 +521,7 @@ class GreatExpectationsDataValidator(BaseDataValidator):
             "expectation_suite_name": expectation_suite_name,
             "action_list": action_list,
         }
-        context.add_checkpoint(**checkpoint_config)
+        context.add_checkpoint(**checkpoint_config)  # type: ignore[has-type]
 
         try:
             results = context.run_checkpoint(

zenml/integrations/great_expectations/flavors/great_expectations_data_validator_flavor.py

@@ -16,7 +16,9 @@
 import os
 from typing import TYPE_CHECKING, Any, Dict, Optional, Type
 
-from pydantic import validator
+import yaml
+from pydantic import field_validator, model_validator
+from yaml.parser import ParserError
 
 from zenml.data_validators.base_data_validator import (
     BaseDataValidatorConfig,
@@ -26,6 +28,7 @@ from zenml.integrations.great_expectations import (
     GREAT_EXPECTATIONS_DATA_VALIDATOR_FLAVOR,
 )
 from zenml.io import fileio
+from zenml.utils.pydantic_utils import before_validator_handler
 
 if TYPE_CHECKING:
     from zenml.integrations.great_expectations.data_validators import (
@@ -41,6 +44,8 @@ class GreatExpectationsDataValidatorConfig(BaseDataValidatorConfig):
             data context. If configured, the data validator will only be usable
             with local orchestrators.
         context_config: in-line Great Expectations data context configuration.
+            If the `context_root_dir` attribute is also set, this configuration
+            will be ignored.
         configure_zenml_stores: if set, ZenML will automatically configure
             stores that use the Artifact Store as a backend. If neither
             `context_root_dir` nor `context_config` are set, this is the default
@@ -54,7 +59,8 @@ class GreatExpectationsDataValidatorConfig(BaseDataValidatorConfig):
     configure_zenml_stores: bool = False
     configure_local_docs: bool = True
 
-    @validator("context_root_dir")
+    @field_validator("context_root_dir")
+    @classmethod
     def _ensure_valid_context_root_dir(
         cls, context_root_dir: Optional[str] = None
     ) -> Optional[str]:
@@ -78,6 +84,33 @@ class GreatExpectationsDataValidatorConfig(BaseDataValidatorConfig):
                 )
         return context_root_dir
 
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def validate_context_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Convert the context configuration if given in JSON/YAML format.
+
+        Args:
+            data: The configuration values.
+
+        Returns:
+            The validated configuration values.
+
+        Raises:
+            ValueError: If the context configuration is not a valid
+                JSON/YAML object.
+        """
+        if isinstance(data.get("context_config"), str):
+            try:
+                data["context_config"] = yaml.safe_load(data["context_config"])
+            except ParserError as e:
+                raise ValueError(
+                    f"Malformed `context_config` value. Only JSON and YAML "
+                    f"formats are supported: {str(e)}"
+                )
+
+        return data
+
     @property
     def is_local(self) -> bool:
         """Checks if this stack component is running locally.

zenml/integrations/great_expectations/ge_store_backend.py

@@ -17,14 +17,16 @@ import os
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple, cast
 
-from great_expectations.data_context.store.tuple_store_backend import (  # type: ignore[import-untyped]
+from great_expectations.data_context.store.tuple_store_backend import (
     TupleStoreBackend,
-    filter_properties_dict,
 )
 from great_expectations.exceptions import (  # type: ignore[import-untyped]
     InvalidKeyError,
     StoreBackendError,
 )
+from great_expectations.util import (  # type: ignore[import-untyped]
+    filter_properties_dict,
+)
 
 from zenml.client import Client
 from zenml.io import fileio
@@ -34,7 +36,7 @@ from zenml.utils import io_utils
 logger = get_logger(__name__)
 
 
-class ZenMLArtifactStoreBackend(TupleStoreBackend):  # type: ignore[misc]
+class ZenMLArtifactStoreBackend(TupleStoreBackend):
     """Great Expectations store backend that uses the active ZenML Artifact Store as a store."""
 
     def __init__(
@@ -105,7 +107,7 @@ class ZenMLArtifactStoreBackend(TupleStoreBackend):  # type: ignore[misc]
         if not isinstance(key, tuple):
             key = key.to_tuple()
         if not is_prefix:
-            object_relative_path = self._convert_key_to_filepath(key)
+            object_relative_path = self._convert_key_to_filepath(key)  # type: ignore[no-untyped-call]
         elif key:
             object_relative_path = os.path.join(*key)
         else:
@@ -116,7 +118,7 @@ class ZenMLArtifactStoreBackend(TupleStoreBackend):  # type: ignore[misc]
             object_key = object_relative_path
         return os.path.join(self.root_path, object_key)
 
-    def _get(self, key: Tuple[str, ...]) -> str:
+    def _get(self, key: Tuple[str, ...]) -> str:  # type: ignore[override]
         """Get the value of an object from the store.
 
         Args:
@@ -140,7 +142,18 @@ class ZenMLArtifactStoreBackend(TupleStoreBackend):  # type: ignore[misc]
             )
         return contents
 
-    def _set(self, key: Tuple[str, ...], value: str, **kwargs: Any) -> str:
+    def _get_all(self) -> List[Any]:
+        """Get all objects in the store.
+
+        Raises:
+            NotImplementedError: if the method is not implemented for this store
+                backend.
+        """
+        raise NotImplementedError(
+            "Method `_get_all` is not implemented for this store backend."
+        )
+
+    def _set(self, key: Tuple[str, ...], value: str, **kwargs: Any) -> str:  # type: ignore[override]
         """Set the value of an object in the store.
 
         Args:
@@ -212,12 +225,12 @@ class ZenMLArtifactStoreBackend(TupleStoreBackend):  # type: ignore[misc]
                     self.filepath_suffix
                 ):
                     continue
-                key = self._convert_filepath_to_key(filepath)
-                if key and not self.is_ignored_key(key):
+                key = self._convert_filepath_to_key(filepath)  # type: ignore[no-untyped-call]
+                if key and not self.is_ignored_key(key):  # type: ignore[no-untyped-call]
                     key_list.append(key)
         return key_list
 
-    def remove_key(self, key: Tuple[str, ...]) -> bool:
+    def remove_key(self, key: Tuple[str, ...]) -> bool:  # type: ignore[override]
         """Delete an object from the store.
 
         Args:
@@ -250,7 +263,7 @@ class ZenMLArtifactStoreBackend(TupleStoreBackend):  # type: ignore[misc]
         result = fileio.exists(filepath)
         return result
 
-    def get_url_for_key(
+    def get_url_for_key(  # type: ignore[override]
         self, key: Tuple[str, ...], protocol: Optional[str] = None
     ) -> str:
         """Get the URL of an object in the store.
@@ -292,7 +305,7 @@ class ZenMLArtifactStoreBackend(TupleStoreBackend):  # type: ignore[misc]
                 f"requested but `base_public_path` was not configured for the "
                 f"{self.__class__.__name__}"
             )
-        filepath = self._convert_key_to_filepath(key)
+        filepath = self._convert_key_to_filepath(key)  # type: ignore[no-untyped-call]
         public_url = self.base_public_path + filepath.replace(self.proto, "")
         return cast(str, public_url)
 

zenml/integrations/great_expectations/materializers/ge_materializer.py

@@ -25,10 +25,10 @@ from great_expectations.core import (  # type: ignore[import-untyped]
 from great_expectations.core.expectation_validation_result import (  # type: ignore[import-untyped]
     ExpectationSuiteValidationResult,
 )
-from great_expectations.data_context.types.base import (  # type: ignore[import-untyped]
+from great_expectations.data_context.types.base import (
     CheckpointConfig,
 )
-from great_expectations.data_context.types.resource_identifiers import (  # type: ignore[import-untyped]
+from great_expectations.data_context.types.resource_identifiers import (
     ExpectationSuiteIdentifier,
     ValidationResultIdentifier,
 )
@@ -56,9 +56,9 @@ class GreatExpectationsMaterializer(BaseMaterializer):
         ExpectationSuite,
         CheckpointResult,
     )
-    ASSOCIATED_ARTIFACT_TYPE: ClassVar[
-        ArtifactType
-    ] = ArtifactType.DATA_ANALYSIS
+    ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = (
+        ArtifactType.DATA_ANALYSIS
+    )
 
     @staticmethod
     def preprocess_checkpoint_result_dict(
@@ -86,7 +86,7 @@ class GreatExpectationsMaterializer(BaseMaterializer):
         validation_dict = {}
         for result_ident, results in artifact_dict["run_results"].items():
             validation_ident = (
-                ValidationResultIdentifier.from_fixed_length_tuple(
+                ValidationResultIdentifier.from_fixed_length_tuple(  # type: ignore[no-untyped-call]
                     result_ident.split("::")[1].split("/")
                 )
             )
@@ -124,9 +124,9 @@ class GreatExpectationsMaterializer(BaseMaterializer):
         filepath = os.path.join(self.uri, ARTIFACT_FILENAME)
         artifact_dict = obj.to_json_dict()
         artifact_type = type(obj)
-        artifact_dict[
-            "data_type"
-        ] = f"{artifact_type.__module__}.{artifact_type.__name__}"
+        artifact_dict["data_type"] = (
+            f"{artifact_type.__module__}.{artifact_type.__name__}"
+        )
         yaml_utils.write_json(filepath, artifact_dict)
 
     def save_visualizations(

zenml/integrations/great_expectations/utils.py

@@ -13,14 +13,14 @@
 #  permissions and limitations under the License.
 """Great Expectations data profiling standard step."""
 
-from typing import Optional
+from typing import Any, Dict, Optional
 
 import pandas as pd
 from great_expectations.core.batch import (  # type: ignore[import-untyped]
     RuntimeBatchRequest,
 )
-from great_expectations.data_context.data_context import (  # type: ignore[import-untyped]
-    BaseDataContext,
+from great_expectations.data_context.data_context.abstract_data_context import (
+    AbstractDataContext,
 )
 
 from zenml import get_step_context
@@ -31,7 +31,7 @@ logger = get_logger(__name__)
 
 
 def create_batch_request(
-    context: BaseDataContext,
+    context: AbstractDataContext,
     dataset: pd.DataFrame,
     data_asset_name: Optional[str],
 ) -> RuntimeBatchRequest:
@@ -62,7 +62,7 @@ def create_batch_request(
     data_asset_name = data_asset_name or f"{pipeline_name}_{step_name}"
     batch_identifier = "default"
 
-    datasource_config = {
+    datasource_config: Dict[str, Any] = {
         "name": datasource_name,
         "class_name": "Datasource",
         "module_name": "great_expectations.datasource",

zenml/integrations/huggingface/__init__.py

@@ -12,21 +12,72 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Initialization of the Huggingface integration."""
+import sys
+from typing import List, Type, Optional
 
 from zenml.integrations.constants import HUGGINGFACE
 from zenml.integrations.integration import Integration
+from zenml.stack import Flavor
+
+HUGGINGFACE_MODEL_DEPLOYER_FLAVOR = "huggingface"
+HUGGINGFACE_SERVICE_ARTIFACT = "hf_deployment_service"
 
 
 class HuggingfaceIntegration(Integration):
     """Definition of Huggingface integration for ZenML."""
 
     NAME = HUGGINGFACE
-    REQUIREMENTS = ["transformers<=4.31", "datasets"]
+
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["fsspec", "pandas"]
 
     @classmethod
     def activate(cls) -> None:
         """Activates the integration."""
         from zenml.integrations.huggingface import materializers  # noqa
+        from zenml.integrations.huggingface import services
+
+    @classmethod
+    def get_requirements(cls, target_os: Optional[str] = None) -> List[str]:
+        """Defines platform specific requirements for the integration.
+
+        Args:
+            target_os: The target operating system.
+
+        Returns:
+            A list of requirements.
+        """
+        requirements = [
+            "datasets",
+            "huggingface_hub>0.19.0",
+            "accelerate",
+            "bitsandbytes>=0.41.3",
+            "peft",
+            # temporary fix for CI issue similar to:
+            # - https://github.com/huggingface/datasets/issues/6737
+            # - https://github.com/huggingface/datasets/issues/6697
+            # TODO try relaxing it back going forward
+            "fsspec<=2023.12.0",
+            "transformers",
+        ]
+
+        # Add the pandas integration requirements
+        from zenml.integrations.pandas import PandasIntegration
+
+        return requirements + \
+            PandasIntegration.get_requirements(target_os=target_os)
+
+    @classmethod
+    def flavors(cls) -> List[Type[Flavor]]:
+        """Declare the stack component flavors for the Huggingface integration.
+
+        Returns:
+            List of stack component flavors for this integration.
+        """
+        from zenml.integrations.huggingface.flavors import (
+            HuggingFaceModelDeployerFlavor,
+        )
+
+        return [HuggingFaceModelDeployerFlavor]
 
 
 HuggingfaceIntegration.check_installation()

zenml/integrations/huggingface/flavors/__init__.py

@@ -0,0 +1,26 @@
+#  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.
+"""Hugging Face integration flavors."""
+
+from zenml.integrations.huggingface.flavors.huggingface_model_deployer_flavor import (  # noqa
+    HuggingFaceModelDeployerConfig,
+    HuggingFaceModelDeployerFlavor,
+    HuggingFaceBaseConfig,
+)
+
+__all__ = [
+    "HuggingFaceModelDeployerConfig",
+    "HuggingFaceModelDeployerFlavor",
+    "HuggingFaceBaseConfig",
+]