zenml-nightly 0.58.2.dev20240617__py3-none-any.whl → 0.58.2.dev20240619__py3-none-any.whl

zenml/integrations/aws/flavors/sagemaker_step_operator_flavor.py

@@ -15,6 +15,8 @@
 
 from typing import TYPE_CHECKING, Any, Dict, Optional, Type, Union
 
+from pydantic import Field
+
 from zenml.config.base_settings import BaseSettings
 from zenml.integrations.aws import (
     AWS_RESOURCE_TYPE,
@@ -58,7 +60,9 @@ class SagemakerStepOperatorSettings(BaseSettings):
 
     instance_type: Optional[str] = None
     experiment_name: Optional[str] = None
-    input_data_s3_uri: Optional[Union[str, Dict[str, str]]] = None
+    input_data_s3_uri: Optional[Union[str, Dict[str, str]]] = Field(
+        default=None, union_mode="left_to_right"
+    )
     estimator_args: Dict[str, Any] = {}
 
     _deprecation_validator = deprecation_utils.deprecate_pydantic_attributes(
@@ -66,7 +70,7 @@ class SagemakerStepOperatorSettings(BaseSettings):
     )
 
 
-class SagemakerStepOperatorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class SagemakerStepOperatorConfig(
     BaseStepOperatorConfig, SagemakerStepOperatorSettings
 ):
     """Config for the Sagemaker step operator.

zenml/integrations/aws/service_connectors/aws_service_connector.py

@@ -37,7 +37,7 @@ from aws_profile_manager import Common  # type: ignore[import-untyped]
 from botocore.client import BaseClient
 from botocore.exceptions import BotoCoreError, ClientError
 from botocore.signers import RequestSigner
-from pydantic import Field, SecretStr
+from pydantic import Field
 
 from zenml.constants import (
     DOCKER_REGISTRY_RESOURCE_TYPE,
@@ -65,6 +65,7 @@ from zenml.service_connectors.service_connector import (
     ServiceConnector,
 )
 from zenml.utils.enum_utils import StrEnum
+from zenml.utils.secret_utils import PlainSerializedSecretStr
 
 logger = get_logger(__name__)
 
@@ -77,11 +78,11 @@ BOTO3_SESSION_EXPIRATION_BUFFER = 15  # 15 minutes
 class AWSSecretKey(AuthenticationConfig):
     """AWS secret key credentials."""
 
-    aws_access_key_id: SecretStr = Field(
+    aws_access_key_id: PlainSerializedSecretStr = Field(
         title="AWS Access Key ID",
         description="An AWS access key ID associated with an AWS account or IAM user.",
     )
-    aws_secret_access_key: SecretStr = Field(
+    aws_secret_access_key: PlainSerializedSecretStr = Field(
         title="AWS Secret Access Key",
     )
 
@@ -89,7 +90,7 @@ class AWSSecretKey(AuthenticationConfig):
 class STSToken(AWSSecretKey):
     """AWS STS token."""
 
-    aws_session_token: SecretStr = Field(
+    aws_session_token: PlainSerializedSecretStr = Field(
         title="AWS Session Token",
     )
 

zenml/integrations/aws/step_operators/sagemaker_step_operator.py

@@ -261,7 +261,7 @@ class SagemakerStepOperator(BaseStepOperator):
                 "ExperimentName": settings.experiment_name,
                 "TrialName": sanitized_training_job_name,
             }
-
+        info.force_write_logs()
         estimator.fit(
             wait=True,
             inputs=inputs,

zenml/integrations/azure/flavors/azureml_step_operator_flavor.py

@@ -38,7 +38,7 @@ class AzureMLStepOperatorSettings(BaseSettings):
     environment_name: Optional[str] = None
 
 
-class AzureMLStepOperatorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class AzureMLStepOperatorConfig(
     BaseStepOperatorConfig, AzureMLStepOperatorSettings
 ):
     """Config for the AzureML step operator.
@@ -64,9 +64,9 @@ class AzureMLStepOperatorConfig(  # type: ignore[misc] # https://github.com/pyda
 
     # Service principal authentication
     # https://docs.microsoft.com/en-us/azure/machine-learning/how-to-setup-authentication#configure-a-service-principal
-    tenant_id: Optional[str] = SecretField()
-    service_principal_id: Optional[str] = SecretField()
-    service_principal_password: Optional[str] = SecretField()
+    tenant_id: Optional[str] = SecretField(default=None)
+    service_principal_id: Optional[str] = SecretField(default=None)
+    service_principal_password: Optional[str] = SecretField(default=None)
 
     @property
     def is_remote(self) -> bool:

zenml/integrations/azure/service_connectors/azure_service_connector.py

@@ -28,7 +28,7 @@ from azure.mgmt.containerservice import ContainerServiceClient
 from azure.mgmt.resource import SubscriptionClient
 from azure.mgmt.storage import StorageManagementClient
 from azure.storage.blob import BlobServiceClient
-from pydantic import Field, SecretStr
+from pydantic import Field
 
 from zenml.constants import (
     DOCKER_REGISTRY_RESOURCE_TYPE,
@@ -56,6 +56,7 @@ from zenml.service_connectors.service_connector import (
     ServiceConnector,
 )
 from zenml.utils.enum_utils import StrEnum
+from zenml.utils.secret_utils import PlainSerializedSecretStr
 
 logger = get_logger(__name__)
 
@@ -103,7 +104,7 @@ class AzureBaseConfig(AuthenticationConfig):
 class AzureClientSecret(AuthenticationConfig):
     """Azure client secret credentials."""
 
-    client_secret: SecretStr = Field(
+    client_secret: PlainSerializedSecretStr = Field(
         title="Service principal client secret",
         description="The client secret of the service principal",
     )
@@ -126,7 +127,7 @@ class AzureClientConfig(AzureBaseConfig):
 class AzureAccessToken(AuthenticationConfig):
     """Azure access token credentials."""
 
-    token: SecretStr = Field(
+    token: PlainSerializedSecretStr = Field(
         title="Azure Access Token",
         description="The Azure access token to use for authentication",
     )

zenml/integrations/azure/step_operators/azureml_step_operator.py

@@ -252,7 +252,7 @@ class AzureMLStepOperator(BaseStepOperator):
             "source_files",
         ]
         docker_settings = info.config.docker_settings
-        ignored_docker_fields = docker_settings.__fields_set__.intersection(
+        ignored_docker_fields = docker_settings.model_fields_set.intersection(
             unused_docker_fields
         )
 
@@ -297,4 +297,5 @@ class AzureMLStepOperator(BaseStepOperator):
         run = experiment.submit(config=run_config)
 
         run.display_name = info.run_name
+        info.force_write_logs()
         run.wait_for_completion(show_output=True)

zenml/integrations/bentoml/steps/bentoml_deployer.py

@@ -129,7 +129,7 @@ def bentoml_model_deployer_step(
 
     # fetch existing services with same pipeline name, step name and model name
     existing_services = model_deployer.find_model_server(
-        config=predictor_cfg.dict(),
+        config=predictor_cfg.model_dump(),
         service_type=BentoMLDeploymentService.SERVICE_TYPE,
     )
 

zenml/integrations/bitbucket/plugins/event_sources/bitbucket_webhook_event_source.py

@@ -16,7 +16,7 @@
 from typing import Any, Dict, List, Optional, Type, Union
 from uuid import UUID
 
-from pydantic import BaseModel, Extra, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 from zenml.enums import SecretScope
 from zenml.event_sources.base_event import (
@@ -59,9 +59,9 @@ class BitbucketEventType(StrEnum):
 class User(BaseModel):
     """Bitbucket User."""
 
-    name: Optional[str]
-    email: Optional[str]
-    username: Optional[str]
+    name: Optional[str] = None
+    email: Optional[str] = None
+    username: Optional[str] = None
 
 
 class Commit(BaseModel):
@@ -85,8 +85,8 @@ class Repository(BaseModel):
 class PushChange(BaseModel):
     """Bitbucket Push Change."""
 
-    new: Optional[Dict[str, Any]]
-    old: Optional[Dict[str, Any]]
+    new: Optional[Dict[str, Any]] = None
+    old: Optional[Dict[str, Any]] = None
     commits: List[Commit]
 
 
@@ -102,11 +102,7 @@ class BitbucketEvent(BaseEvent):
     actor: User
     repository: Repository
     push: Push
-
-    class Config:
-        """Pydantic configuration class."""
-
-        extra = Extra.allow
+    model_config = ConfigDict(extra="allow")
 
     @property
     def branch(self) -> Optional[str]:
@@ -418,7 +414,7 @@ class BitbucketWebhookEventSourceHandler(BaseWebhookEventSourceHandler):
         if config.rotate_secret:
             # In case the secret is being rotated
             secret_key_value = random_str(12)
-            webhook_secret = SecretUpdate(  # type: ignore[call-arg]
+            webhook_secret = SecretUpdate(
                 values={"webhook_secret": secret_key_value}
             )
             self.zen_store.update_secret(

zenml/integrations/comet/flavors/comet_experiment_tracker_flavor.py

@@ -50,7 +50,7 @@ class CometExperimentTrackerSettings(BaseSettings):
     settings: Dict[str, Any] = {}
 
 
-class CometExperimentTrackerConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class CometExperimentTrackerConfig(
     BaseExperimentTrackerConfig, CometExperimentTrackerSettings
 ):
     """Config for the Comet experiment tracker.

zenml/integrations/constants.py

@@ -30,7 +30,6 @@ FEAST = "feast"
 GCP = "gcp"
 GITHUB = "github"
 GITLAB = "gitlab"
-GRAPHVIZ = "graphviz"
 HUGGINGFACE = "huggingface"
 HYPERAI = "hyperai"
 GREAT_EXPECTATIONS = "great_expectations"

zenml/integrations/deepchecks/__init__.py

@@ -41,6 +41,7 @@ class DeepchecksIntegration(Integration):
         "pandas<2.0.0",
         "opencv-python==4.5.5.64",  # pin to same version
         "opencv-python-headless==4.5.5.64",  # pin to same version
+        "tenacity!=8.4.0",  # https://github.com/jd/tenacity/issues/471
     ]
     APT_PACKAGES = ["ffmpeg", "libsm6", "libxext6"]
 

zenml/integrations/evidently/__init__.py

@@ -31,10 +31,9 @@ from zenml.integrations.constants import EVIDENTLY
 from zenml.integrations.integration import Integration
 from zenml.stack import Flavor
 
-
 # Fix numba errors in Docker and suppress logs and deprecation warning spam
 try:
-    from numba.core.errors import (  # type: ignore[import-untyped]
+    from numba.core.errors import (  # type: ignore[import-not-found]
         NumbaDeprecationWarning,
         NumbaPendingDeprecationWarning,
     )
@@ -54,7 +53,10 @@ class EvidentlyIntegration(Integration):
     """[Evidently](https://github.com/evidentlyai/evidently) integration for ZenML."""
 
     NAME = EVIDENTLY
-    REQUIREMENTS = ["evidently>0.2.6,<0.4.5"]  # supports pyyaml 6
+    REQUIREMENTS = [
+        "evidently>=0.4.16,<=0.4.22",
+        "tenacity!=8.4.0",  # https://github.com/jd/tenacity/issues/471
+    ]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/evidently/column_mapping.py

@@ -16,7 +16,7 @@
 from typing import List, Optional, Sequence, Union
 
 from evidently import ColumnMapping  # type: ignore[import-untyped]
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict, Field
 
 
 class EvidentlyColumnMapping(BaseModel):
@@ -41,7 +41,9 @@ class EvidentlyColumnMapping(BaseModel):
     """
 
     target: Optional[str] = None
-    prediction: Optional[Union[str, Sequence[str]]] = "prediction"
+    prediction: Optional[Union[str, Sequence[str]]] = Field(
+        default="prediction", union_mode="left_to_right"
+    )
     datetime: Optional[str] = None
     id: Optional[str] = None
     numerical_features: Optional[List[str]] = None
@@ -49,9 +51,15 @@ class EvidentlyColumnMapping(BaseModel):
     datetime_features: Optional[List[str]] = None
     target_names: Optional[List[str]] = None
     task: Optional[str] = None
-    pos_label: Optional[Union[str, int]] = 1
+    pos_label: Optional[Union[str, int]] = Field(
+        default=1, union_mode="left_to_right"
+    )
     text_features: Optional[List[str]] = None
 
+    model_config = ConfigDict(
+        validate_assignment=True,
+    )
+
     def to_evidently_column_mapping(self) -> ColumnMapping:
         """Convert this Pydantic object to an Evidently ColumnMapping object.
 

zenml/integrations/evidently/data_validators/evidently_data_validator.py

@@ -118,14 +118,16 @@ class EvidentlyDataValidator(BaseDataValidator):
         return options
 
     @staticmethod
-    def _download_nltk_data() -> None:
-        """Download NLTK data for text metrics and tests.
+    def _set_nltk_data_path() -> None:
+        """Set the NLTK data path to the current working directory.
+
+        This is necessary because the default download directory is not writable
+        in some Docker containers.
 
         Raises:
             ImportError: if NLTK is not installed.
         """
         try:
-            import nltk  # type: ignore[import-untyped]
             from nltk.data import (  # type: ignore[import-untyped]
                 path as nltk_path,
             )
@@ -138,8 +140,23 @@ class EvidentlyDataValidator(BaseDataValidator):
         # Configure NLTK to use the current working directory to download and
         # lookup data. This is necessary because the default download directory
         # is not writable in some Docker containers.
+        os.environ["NLTK_DATA"] = os.getcwd()
         nltk_path.append(os.getcwd())
 
+    @staticmethod
+    def _download_nltk_data() -> None:
+        """Download NLTK data for text metrics and tests.
+
+        Raises:
+            ImportError: if NLTK is not installed.
+        """
+        try:
+            import nltk  # type: ignore[import-untyped]
+        except ImportError:
+            raise ImportError(
+                "NLTK is not installed. Please install NLTK to use "
+                "Evidently text metrics and tests."
+            )
         # Download NLTK data. We need this later on for the Evidently text report.
         nltk.download("words", download_dir=os.getcwd())
         nltk.download("wordnet", download_dir=os.getcwd())
@@ -196,6 +213,7 @@ class EvidentlyDataValidator(BaseDataValidator):
         Returns:
             The Evidently Report as JSON object and as HTML.
         """
+        self._set_nltk_data_path()
         if download_nltk_data:
             self._download_nltk_data()
 

zenml/integrations/evidently/metrics.py

@@ -33,7 +33,7 @@ from evidently.metrics.base_metric import (  # type: ignore[import-untyped]
 from evidently.utils.generators import (  # type: ignore[import-untyped]
     BaseGenerator,
 )
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 from zenml.logger import get_logger
 from zenml.utils import source_utils
@@ -75,7 +75,9 @@ class EvidentlyMetricConfig(BaseModel):
     class_path: str
     parameters: Dict[str, Any] = Field(default_factory=dict)
     is_generator: bool = False
-    columns: Optional[Union[str, List[str]]] = None
+    columns: Optional[Union[str, List[str]]] = Field(
+        default=None, union_mode="left_to_right"
+    )
     skip_id_column: bool = False
 
     @staticmethod
@@ -359,7 +361,4 @@ class EvidentlyMetricConfig(BaseModel):
                 f"`{self.class_path}`: {str(e)}"
             )
 
-    class Config:
-        """Pydantic config class."""
-
-        extra = "forbid"
+    model_config = ConfigDict(extra="forbid")

zenml/integrations/evidently/tests.py

@@ -33,7 +33,7 @@ from evidently.tests.base_test import (  # type: ignore[import-untyped]
 from evidently.utils.generators import (  # type: ignore[import-untyped]
     BaseGenerator,
 )
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 from zenml.logger import get_logger
 from zenml.utils import source_utils
@@ -73,7 +73,9 @@ class EvidentlyTestConfig(BaseModel):
     class_path: str
     parameters: Dict[str, Any] = Field(default_factory=dict)
     is_generator: bool = False
-    columns: Optional[Union[str, List[str]]] = None
+    columns: Optional[Union[str, List[str]]] = Field(
+        default=None, union_mode="left_to_right"
+    )
 
     @staticmethod
     def get_test_class(test_name: str) -> Union[Test, TestPreset]:
@@ -343,7 +345,4 @@ class EvidentlyTestConfig(BaseModel):
                 f"`{self.class_path}`: {str(e)}"
             )
 
-    class Config:
-        """Pydantic config class."""
-
-        extra = "forbid"
+    model_config = ConfigDict(extra="forbid")

zenml/integrations/facets/models.py

@@ -16,7 +16,7 @@
 from typing import Dict, List, Union
 
 import pandas as pd
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 
 class FacetsComparison(BaseModel):
@@ -31,8 +31,4 @@ class FacetsComparison(BaseModel):
     """
 
     datasets: List[Dict[str, Union[str, pd.DataFrame]]]
-
-    class Config:
-        """Pydantic config."""
-
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(arbitrary_types_allowed=True)

zenml/integrations/feast/__init__.py

@@ -31,7 +31,9 @@ class FeastIntegration(Integration):
     """Definition of Feast integration for ZenML."""
 
     NAME = FEAST
-    REQUIREMENTS = ["feast[redis]>=0.26.0", "redis-server>=6.0.9"]
+    # click is added to keep the feast click version
+    # in sync with ZenML's click
+    REQUIREMENTS = ["feast<=0.37.1", "click>=8.0.1,<8.1.4"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/feast/feature_stores/feast_feature_store.py

@@ -16,7 +16,6 @@
 from typing import Any, Dict, List, Union, cast
 
 import pandas as pd
-import redis
 from feast import FeatureStore  # type: ignore
 from feast.infra.registry.base_registry import BaseRegistry  # type: ignore
 
@@ -41,23 +40,6 @@ class FeastFeatureStore(BaseFeatureStore):
         """
         return cast(FeastFeatureStoreConfig, self._config)
 
-    def _validate_connection(self) -> None:
-        """Validates the connection to the feature store.
-
-        Raises:
-            ConnectionError: If the online component (Redis) is not available.
-        """
-        client = redis.Redis(
-            host=self.config.online_host, port=self.config.online_port
-        )
-        try:
-            client.ping()
-        except redis.exceptions.ConnectionError as e:
-            raise redis.exceptions.ConnectionError(
-                "Could not connect to feature store's online component. "
-                "Please make sure that Redis is running."
-            ) from e
-
     def get_historical_features(
         self,
         entity_df: Union[pd.DataFrame, str],
@@ -104,7 +86,6 @@ class FeastFeatureStore(BaseFeatureStore):
         Returns:
             The latest online feature data as a dictionary.
         """
-        self._validate_connection()
         fs = FeatureStore(repo_path=self.config.feast_repo)
 
         return fs.get_online_features(  # type: ignore[no-any-return]
@@ -122,7 +103,6 @@ class FeastFeatureStore(BaseFeatureStore):
         Returns:
             The data sources' names.
         """
-        self._validate_connection()
         fs = FeatureStore(repo_path=self.config.feast_repo)
         return [ds.name for ds in fs.list_data_sources()]
 
@@ -135,7 +115,6 @@ class FeastFeatureStore(BaseFeatureStore):
         Returns:
             The entity names.
         """
-        self._validate_connection()
         fs = FeatureStore(repo_path=self.config.feast_repo)
         return [ds.name for ds in fs.list_entities()]
 
@@ -148,7 +127,6 @@ class FeastFeatureStore(BaseFeatureStore):
         Returns:
             The feature service names.
         """
-        self._validate_connection()
         fs = FeatureStore(repo_path=self.config.feast_repo)
         return [ds.name for ds in fs.list_feature_services()]
 
@@ -161,7 +139,6 @@ class FeastFeatureStore(BaseFeatureStore):
         Returns:
             The feature view names.
         """
-        self._validate_connection()
         fs = FeatureStore(repo_path=self.config.feast_repo)
         return [ds.name for ds in fs.list_feature_views()]
 

zenml/integrations/gcp/__init__.py

@@ -43,7 +43,7 @@ class GcpIntegration(Integration):
 
     NAME = GCP
     REQUIREMENTS = [
-        "kfp==1.8.22",  # Only 1.x version that supports pyyaml 6
+        "kfp>=1.8.22",
         "gcsfs",
         "google-cloud-secret-manager",
         "google-cloud-container>=2.21.0",

zenml/integrations/gcp/flavors/vertex_orchestrator_flavor.py

@@ -68,7 +68,7 @@ class VertexOrchestratorSettings(BaseSettings):
     )
 
 
-class VertexOrchestratorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class VertexOrchestratorConfig(
     BaseOrchestratorConfig,
     GoogleCredentialsConfigMixin,
     VertexOrchestratorSettings,

zenml/integrations/gcp/flavors/vertex_step_operator_flavor.py

@@ -61,7 +61,7 @@ class VertexStepOperatorSettings(BaseSettings):
     boot_disk_type: str = "pd-ssd"
 
 
-class VertexStepOperatorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class VertexStepOperatorConfig(
     BaseStepOperatorConfig,
     GoogleCredentialsConfigMixin,
     VertexStepOperatorSettings,