zenml-nightly 0.58.2.dev20240618__py3-none-any.whl → 0.58.2.dev20240620__py3-none-any.whl

zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py

@@ -58,7 +58,7 @@ class KubernetesOrchestratorSettings(BaseSettings):
     orchestrator_pod_settings: Optional[KubernetesPodSettings] = None
 
 
-class KubernetesOrchestratorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class KubernetesOrchestratorConfig(
     BaseOrchestratorConfig, KubernetesOrchestratorSettings
 ):
     """Configuration for the Kubernetes orchestrator.

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator_entrypoint.py

@@ -100,8 +100,17 @@ def main() -> None:
         )
 
         step_config = deployment_config.step_configurations[step_name].config
-        settings = KubernetesOrchestratorSettings.parse_obj(
-            step_config.settings.get("orchestrator.kubernetes", {})
+
+        kubernetes_settings = step_config.settings.get(
+            "orchestrator.kubernetes", None
+        )
+
+        orchestrator_settings = {}
+        if kubernetes_settings is not None:
+            orchestrator_settings = kubernetes_settings.model_dump()
+
+        settings = KubernetesOrchestratorSettings.model_validate(
+            orchestrator_settings
         )
 
         env = get_config_environment_vars()

zenml/integrations/kubernetes/pod_settings.py

@@ -13,22 +13,13 @@
 #  permissions and limitations under the License.
 """Kubernetes pod settings."""
 
-from typing import TYPE_CHECKING, Any, Dict, List, Union
+from typing import Any, Dict, List
 
-from pydantic import validator
+from pydantic import field_validator
 
 from zenml.config.base_settings import BaseSettings
 from zenml.integrations.kubernetes import serialization_utils
 
-if TYPE_CHECKING:
-    from kubernetes.client.models import (
-        V1Affinity,
-        V1ResourceRequirements,
-        V1Toleration,
-        V1Volume,
-        V1VolumeMount,
-    )
-
 
 class KubernetesPodSettings(BaseSettings):
     """Kubernetes Pod settings.
@@ -53,10 +44,9 @@ class KubernetesPodSettings(BaseSettings):
     volume_mounts: List[Dict[str, Any]] = []
     host_ipc: bool = False
 
-    @validator("volumes", pre=True)
-    def _convert_volumes(
-        cls, value: List[Union[Dict[str, Any], "V1Volume"]]
-    ) -> List[Dict[str, Any]]:
+    @field_validator("volumes", mode="before")
+    @classmethod
+    def _convert_volumes(cls, value: Any) -> Any:
         """Converts Kubernetes volumes to dicts.
 
         Args:
@@ -78,10 +68,9 @@ class KubernetesPodSettings(BaseSettings):
 
         return result
 
-    @validator("volume_mounts", pre=True)
-    def _convert_volume_mounts(
-        cls, value: List[Union[Dict[str, Any], "V1VolumeMount"]]
-    ) -> List[Dict[str, Any]]:
+    @field_validator("volume_mounts", mode="before")
+    @classmethod
+    def _convert_volume_mounts(cls, value: Any) -> Any:
         """Converts Kubernetes volume mounts to dicts.
 
         Args:
@@ -103,10 +92,9 @@ class KubernetesPodSettings(BaseSettings):
 
         return result
 
-    @validator("affinity", pre=True)
-    def _convert_affinity(
-        cls, value: Union[Dict[str, Any], "V1Affinity"]
-    ) -> Dict[str, Any]:
+    @field_validator("affinity", mode="before")
+    @classmethod
+    def _convert_affinity(cls, value: Any) -> Any:
         """Converts Kubernetes affinity to a dict.
 
         Args:
@@ -122,10 +110,9 @@ class KubernetesPodSettings(BaseSettings):
         else:
             return value
 
-    @validator("tolerations", pre=True)
-    def _convert_tolerations(
-        cls, value: List[Union[Dict[str, Any], "V1Toleration"]]
-    ) -> List[Dict[str, Any]]:
+    @field_validator("tolerations", mode="before")
+    @classmethod
+    def _convert_tolerations(cls, value: Any) -> Any:
         """Converts Kubernetes tolerations to dicts.
 
         Args:
@@ -147,10 +134,9 @@ class KubernetesPodSettings(BaseSettings):
 
         return result
 
-    @validator("resources", pre=True)
-    def _convert_resources(
-        cls, value: Union[Dict[str, Any], "V1ResourceRequirements"]
-    ) -> Dict[str, Any]:
+    @field_validator("resources", mode="before")
+    @classmethod
+    def _convert_resources(cls, value: Any) -> Any:
         """Converts Kubernetes resource requirements to a dict.
 
         Args:

zenml/integrations/kubernetes/service_connectors/kubernetes_service_connector.py

@@ -25,7 +25,7 @@ from typing import Any, List, Optional
 
 from kubernetes import client as k8s_client
 from kubernetes import config as k8s_config
-from pydantic import Field, SecretStr
+from pydantic import Field
 
 from zenml.constants import KUBERNETES_CLUSTER_RESOURCE_TYPE
 from zenml.exceptions import AuthorizationException
@@ -40,6 +40,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__)
 
@@ -47,7 +48,7 @@ logger = get_logger(__name__)
 class KubernetesServerCredentials(AuthenticationConfig):
     """Kubernetes server authentication config."""
 
-    certificate_authority: Optional[SecretStr] = Field(
+    certificate_authority: Optional[PlainSerializedSecretStr] = Field(
         default=None,
         title="Kubernetes CA Certificate (base64 encoded)",
     )
@@ -68,10 +69,10 @@ class KubernetesServerConfig(KubernetesServerCredentials):
 class KubernetesUserPasswordCredentials(AuthenticationConfig):
     """Kubernetes user/pass authentication config."""
 
-    username: SecretStr = Field(
+    username: PlainSerializedSecretStr = Field(
         title="Kubernetes Username",
     )
-    password: SecretStr = Field(
+    password: PlainSerializedSecretStr = Field(
         title="Kubernetes Password",
     )
 
@@ -94,15 +95,15 @@ class KubernetesUserPasswordConfig(
 class KubernetesTokenCredentials(AuthenticationConfig):
     """Kubernetes token authentication config."""
 
-    client_certificate: Optional[SecretStr] = Field(
+    client_certificate: Optional[PlainSerializedSecretStr] = Field(
         default=None,
         title="Kubernetes Client Certificate (base64 encoded)",
     )
-    client_key: Optional[SecretStr] = Field(
+    client_key: Optional[PlainSerializedSecretStr] = Field(
         default=None,
         title="Kubernetes Client Key (base64 encoded)",
     )
-    token: Optional[SecretStr] = Field(
+    token: Optional[PlainSerializedSecretStr] = Field(
         default=None,
         title="Kubernetes Token",
     )

zenml/integrations/label_studio/__init__.py

@@ -14,7 +14,6 @@
 """Initialization of the Label Studio integration."""
 from typing import List, Type
 
-from zenml.enums import StackComponentType
 from zenml.integrations.constants import LABEL_STUDIO
 from zenml.integrations.integration import Integration
 from zenml.stack import Flavor
@@ -27,8 +26,7 @@ class LabelStudioIntegration(Integration):
 
     NAME = LABEL_STUDIO
     REQUIREMENTS = [
-        "label-studio>=1.6.0,<=1.9.1",
-        "label-studio-sdk>=0.0.17,<=0.0.32",
+        "label-studio-sdk>=1.0.0",
     ]
 
     @classmethod

zenml/integrations/label_studio/annotators/label_studio_annotator.py

@@ -178,13 +178,12 @@ class LabelStudioAnnotator(BaseAnnotator, AuthenticationMixin):
             else:
                 api_key = settings.api_key
         except RuntimeError:
-            # Try to get from secret
-            secret = self.get_authentication_secret()
-            if not secret:
+            if secret := self.get_authentication_secret():
+                api_key = secret.secret_values.get("api_key", "")
+            else:
                 raise ValueError(
                     "Unable to access predefined secret to access Label Studio API key."
                 )
-            api_key = secret.secret_values.get("api_key", "")
         if not api_key:
             raise ValueError(
                 "Unable to access Label Studio API key from secret."

zenml/integrations/label_studio/flavors/label_studio_annotator_flavor.py

@@ -43,10 +43,10 @@ class LabelStudioAnnotatorSettings(BaseSettings):
 
     instance_url: str = DEFAULT_LOCAL_INSTANCE_URL
     port: Optional[int] = DEFAULT_LOCAL_LABEL_STUDIO_PORT
-    api_key: Optional[str] = SecretField()
+    api_key: Optional[str] = SecretField(default=None)
 
 
-class LabelStudioAnnotatorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class LabelStudioAnnotatorConfig(
     BaseAnnotatorConfig,
     LabelStudioAnnotatorSettings,
     AuthenticationConfigMixin,

zenml/integrations/langchain/materializers/document_materializer.py

@@ -13,20 +13,56 @@
 #  permissions and limitations under the License.
 """Implementation of ZenML's Langchain Document materializer."""
 
-import sys
-from typing import TYPE_CHECKING, Any, ClassVar, Tuple, Type
+import os
+from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type
+
+from langchain.docstore.document import Document
 
 from zenml.enums import ArtifactType
-from zenml.materializers.pydantic_materializer import PydanticMaterializer
+from zenml.materializers.base_materializer import BaseMaterializer
+from zenml.utils import yaml_utils
+
+if TYPE_CHECKING:
+    from zenml.metadata.metadata_types import MetadataType
 
-if TYPE_CHECKING and sys.version_info < (3, 8):
-    Document = Any
-else:
-    from langchain.docstore.document import Document
+DEFAULT_FILENAME = "data.json"
 
 
-class LangchainDocumentMaterializer(PydanticMaterializer):
+class LangchainDocumentMaterializer(BaseMaterializer):
     """Handle Langchain Document objects."""
 
     ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.DATA
     ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = (Document,)
+
+    def load(self, data_type: Type["Document"]) -> Any:
+        """Reads BaseModel from JSON.
+
+        Args:
+            data_type: The type of the data to read.
+
+        Returns:
+            The data read.
+        """
+        data_path = os.path.join(self.uri, DEFAULT_FILENAME)
+        contents = yaml_utils.read_json(data_path)
+        return data_type.parse_raw(contents)
+
+    def save(self, data: "Document") -> None:
+        """Serialize a BaseModel to JSON.
+
+        Args:
+            data: The data to store.
+        """
+        data_path = os.path.join(self.uri, DEFAULT_FILENAME)
+        yaml_utils.write_json(data_path, data.json())
+
+    def extract_metadata(self, data: Document) -> Dict[str, "MetadataType"]:
+        """Extract metadata from the given BaseModel object.
+
+        Args:
+            data: The BaseModel object to extract metadata from.
+
+        Returns:
+            The extracted metadata as a dictionary.
+        """
+        return {"schema": data.schema()}

zenml/integrations/mlflow/__init__.py

@@ -31,15 +31,21 @@ class MlflowIntegration(Integration):
     """Definition of MLflow integration for ZenML."""
 
     NAME = MLFLOW
-    # We need to pin protobuf to a version <=4 here, as this mlflow release
-    # does not pin it. They fixed this in a later version, so we can probably
-    # remove this once we update the mlflow version.
+
     REQUIREMENTS = [
         "mlflow>=2.1.1,<=2.12.2",
         "mlserver>=1.3.3",
         "mlserver-mlflow>=1.3.3",
         # TODO: remove this requirement once rapidjson is fixed
         "python-rapidjson<1.15",
+        # When you do:
+        # pip install zenml
+        # You get all our required dependencies. However, if you follow it with:
+        # zenml integration install mlflow
+        # This downgrades pydantic to v1 even though mlflow does not have
+        # any issues with v2. This is why we have to pin it here so a downgrade
+        # will not happen.
+        "pydantic>=2.7.0,<2.8.0"
     ]
 
     @classmethod

zenml/integrations/mlflow/experiment_trackers/mlflow_experiment_tracker.py

@@ -173,7 +173,7 @@ class MLFlowExperimentTracker(BaseExperimentTracker):
         if as_plain_text:
             tracking_uri = self.config.tracking_uri
         else:
-            tracking_uri = self.config.dict()["tracking_uri"]
+            tracking_uri = self.config.model_dump()["tracking_uri"]
         return tracking_uri or self._local_mlflow_backend()
 
     def prepare_step_run(self, info: "StepRunInfo") -> None:

zenml/integrations/mlflow/flavors/mlflow_experiment_tracker_flavor.py

@@ -15,7 +15,7 @@
 
 from typing import TYPE_CHECKING, Any, Dict, Optional, Type
 
-from pydantic import root_validator
+from pydantic import model_validator
 
 from zenml.config.base_settings import BaseSettings
 from zenml.experiment_trackers.base_experiment_tracker import (
@@ -72,7 +72,7 @@ class MLFlowExperimentTrackerSettings(BaseSettings):
     tags: Dict[str, Any] = {}
 
 
-class MLFlowExperimentTrackerConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class MLFlowExperimentTrackerConfig(
     BaseExperimentTrackerConfig, MLFlowExperimentTrackerSettings
 ):
     """Config for the MLflow experiment tracker.
@@ -101,69 +101,61 @@ class MLFlowExperimentTrackerConfig(  # type: ignore[misc] # https://github.com/
     """
 
     tracking_uri: Optional[str] = None
-    tracking_username: Optional[str] = SecretField()
-    tracking_password: Optional[str] = SecretField()
-    tracking_token: Optional[str] = SecretField()
+    tracking_username: Optional[str] = SecretField(default=None)
+    tracking_password: Optional[str] = SecretField(default=None)
+    tracking_token: Optional[str] = SecretField(default=None)
     tracking_insecure_tls: bool = False
     databricks_host: Optional[str] = None
 
-    @root_validator(skip_on_failure=True)
+    @model_validator(mode="after")
     def _ensure_authentication_if_necessary(
-        cls, values: Dict[str, Any]
-    ) -> Dict[str, Any]:
+        self,
+    ) -> "MLFlowExperimentTrackerConfig":
         """Ensures that credentials or a token for authentication exist.
 
         We make this check when running MLflow tracking with a remote backend.
 
-        Args:
-            values: The values to validate.
-
         Returns:
             The validated values.
 
         Raises:
-            ValueError: If neither credentials nor a token are provided.
+             ValueError: If neither credentials nor a token are provided.
         """
-        tracking_uri = values.get("tracking_uri")
-
-        if tracking_uri:
-            if is_databricks_tracking_uri(tracking_uri):
-                # If the tracking uri is "databricks", then we need the databricks
-                # host to be set.
-                databricks_host = values.get("databricks_host")
-
-                if not databricks_host:
+        if self.tracking_uri:
+            if is_databricks_tracking_uri(self.tracking_uri):
+                # If the tracking uri is "databricks", then we need the
+                # databricks host to be set.
+                if not self.databricks_host:
                     raise ValueError(
                         "MLflow experiment tracking with a Databricks MLflow "
-                        "managed tracking server requires the `databricks_host` "
-                        "to be set in your stack component. To update your "
-                        "component, run `zenml experiment-tracker update "
+                        "managed tracking server requires the "
+                        "`databricks_host` to be set in your stack component. "
+                        "To update your component, run "
+                        "`zenml experiment-tracker update "
                         "<NAME> --databricks_host=DATABRICKS_HOST` "
                         "and specify the hostname of your Databricks workspace."
                     )
 
-            if is_remote_mlflow_tracking_uri(tracking_uri):
-                # we need either username + password or a token to authenticate to
-                # the remote backend
-                basic_auth = values.get("tracking_username") and values.get(
-                    "tracking_password"
-                )
-                token_auth = values.get("tracking_token")
+            if is_remote_mlflow_tracking_uri(self.tracking_uri):
+                # we need either username + password or a token to authenticate
+                # to the remote backend
+                basic_auth = self.tracking_username and self.tracking_password
 
-                if not (basic_auth or token_auth):
+                if not (basic_auth or self.tracking_token):
                     raise ValueError(
                         f"MLflow experiment tracking with a remote backend "
-                        f"{tracking_uri} is only possible when specifying either "
-                        f"username and password or an authentication token in your "
-                        f"stack component. To update your component, run the "
-                        f"following command: `zenml experiment-tracker update "
+                        f"{self.tracking_uri} is only possible when specifying "
+                        f"either username and password or an authentication "
+                        f"token in your stack component. To update your "
+                        f"component, run the following command: "
+                        f"`zenml experiment-tracker update "
                         f"<NAME> --tracking_username=MY_USERNAME "
                         f"--tracking_password=MY_PASSWORD "
                         f"--tracking_token=MY_TOKEN` and specify either your "
                         f"username and password or token."
                     )
 
-        return values
+        return self
 
     @property
     def is_local(self) -> bool:

zenml/integrations/mlflow/model_registries/mlflow_model_registry.py

@@ -385,7 +385,7 @@ class MLFlowModelRegistry(BaseModelRegistry):
                     f"Registering a new version for the model `'{name}'` "
                     f"a version will be assigned automatically."
                 )
-            metadata_dict = metadata.dict() if metadata else {}
+            metadata_dict = metadata.model_dump() if metadata else {}
             # Set the run ID and link.
             run_id = metadata_dict.get("mlflow_run_id", None)
             run_link = metadata_dict.get("mlflow_run_link", None)
@@ -477,7 +477,7 @@ class MLFlowModelRegistry(BaseModelRegistry):
         # Update the model tags.
         if metadata:
             try:
-                for key, value in metadata.dict().items():
+                for key, value in metadata.model_dump().items():
                     self.mlflow_client.set_model_version_tag(
                         name=name,
                         version=version,
@@ -593,7 +593,7 @@ class MLFlowModelRegistry(BaseModelRegistry):
                 filter_string += " AND "
             filter_string += f"run_id='{kwargs['mlflow_run_id']}'"
         if metadata:
-            for tag, value in metadata.dict().items():
+            for tag, value in metadata.model_dump().items():
                 if value:
                     if filter_string:
                         filter_string += " AND "
@@ -737,7 +737,7 @@ class MLFlowModelRegistry(BaseModelRegistry):
             registered_model=RegisteredModel(name=mlflow_model_version.name),
             model_format=MLFLOW_MODEL_FORMAT,
             model_library=model_library,
-            version=mlflow_model_version.version,
+            version=str(mlflow_model_version.version),
             created_at=datetime.fromtimestamp(
                 int(mlflow_model_version.creation_timestamp) / 1e3
             ),

zenml/integrations/mlflow/steps/mlflow_deployer.py

@@ -130,7 +130,7 @@ def mlflow_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(),
     )
 
     # Check whether to deploy a new service

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/pigeon/flavors/pigeon_annotator_flavor.py

@@ -31,7 +31,7 @@ class PigeonAnnotatorSettings(BaseSettings):
     """Settings for the Pigeon annotator."""
 
 
-class PigeonAnnotatorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class PigeonAnnotatorConfig(
     BaseAnnotatorConfig, PigeonAnnotatorSettings, AuthenticationConfigMixin
 ):
     """Config for the Pigeon annotator.

zenml/integrations/s3/flavors/s3_artifact_store_flavor.py

@@ -24,7 +24,7 @@ from typing import (
     Type,
 )
 
-from pydantic import validator
+from pydantic import field_validator
 
 from zenml.artifact_stores import (
     BaseArtifactStoreConfig,
@@ -63,14 +63,15 @@ class S3ArtifactStoreConfig(
 
     SUPPORTED_SCHEMES: ClassVar[Set[str]] = {"s3://"}
 
-    key: Optional[str] = SecretField()
-    secret: Optional[str] = SecretField()
-    token: Optional[str] = SecretField()
+    key: Optional[str] = SecretField(default=None)
+    secret: Optional[str] = SecretField(default=None)
+    token: Optional[str] = SecretField(default=None)
     client_kwargs: Optional[Dict[str, Any]] = None
     config_kwargs: Optional[Dict[str, Any]] = None
     s3_additional_kwargs: Optional[Dict[str, Any]] = None
 
-    @validator("client_kwargs")
+    @field_validator("client_kwargs")
+    @classmethod
     def _validate_client_kwargs(
         cls, value: Optional[Dict[str, Any]]
     ) -> Optional[Dict[str, Any]]:
@@ -138,7 +139,7 @@ class S3ArtifactStoreFlavor(BaseArtifactStoreFlavor):
 
     @property
     def docs_url(self) -> Optional[str]:
-        """A url to point at docs explaining this flavor.
+        """A URL to point at docs explaining this flavor.
 
         Returns:
             A flavor docs url.
@@ -147,7 +148,7 @@ class S3ArtifactStoreFlavor(BaseArtifactStoreFlavor):
 
     @property
     def sdk_docs_url(self) -> Optional[str]:
-        """A url to point at SDK docs explaining this flavor.
+        """A URL to point at SDK docs explaining this flavor.
 
         Returns:
             A flavor SDK docs url.
@@ -156,7 +157,7 @@ class S3ArtifactStoreFlavor(BaseArtifactStoreFlavor):
 
     @property
     def logo_url(self) -> str:
-        """A url to represent the flavor in the dashboard.
+        """A URL to represent the flavor in the dashboard.
 
         Returns:
             The flavor logo.