zenml-nightly 0.58.2.dev20240614__py3-none-any.whl → 0.58.2.dev20240622__py3-none-any.whl

zenml/integrations/seldon/seldon_client.py

@@ -17,11 +17,11 @@ import base64
 import json
 import re
 import time
-from typing import Any, Dict, Generator, List, Optional
+from typing import Any, Dict, Generator, List, Literal, Optional
 
 from kubernetes import client as k8s_client
 from kubernetes import config as k8s_config
-from pydantic import BaseModel, Field, ValidationError
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
 
 from zenml.logger import get_logger
 from zenml.utils.enum_utils import StrEnum
@@ -30,9 +30,6 @@ logger = get_logger(__name__)
 
 api = k8s_client.ApiClient()
 
-SELDON_DEPLOYMENT_KIND = "SeldonDeployment"
-SELDON_DEPLOYMENT_API_VERSION = "machinelearning.seldon.io/v1"
-
 
 class SeldonDeploymentPredictorParameter(BaseModel):
     """Parameter for Seldon Deployment predictor.
@@ -46,14 +43,12 @@ class SeldonDeploymentPredictorParameter(BaseModel):
     name: str = ""
     type: str = ""
     value: str = ""
-
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # validate attribute assignments
-        validate_assignment = True
+        validate_assignment=True,
         # Ignore extra attributes from the CRD that are not reflected here
-        extra = "ignore"
+        extra="ignore",
+    )
 
 
 class SeldonResourceRequirements(BaseModel):
@@ -82,14 +77,12 @@ class SeldonDeploymentMetadata(BaseModel):
     labels: Dict[str, str] = Field(default_factory=dict)
     annotations: Dict[str, str] = Field(default_factory=dict)
     creationTimestamp: Optional[str] = None
-
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # validate attribute assignments
-        validate_assignment = True
+        validate_assignment=True,
         # Ignore extra attributes from the CRD that are not reflected here
-        extra = "ignore"
+        extra="ignore",
+    )
 
 
 class SeldonDeploymentPredictiveUnitType(StrEnum):
@@ -130,14 +123,12 @@ class SeldonDeploymentPredictiveUnit(BaseModel):
     serviceAccountName: Optional[str] = None
     envSecretRefName: Optional[str] = None
     children: Optional[List["SeldonDeploymentPredictiveUnit"]] = None
-
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # validate attribute assignments
-        validate_assignment = True
+        validate_assignment=True,
         # Ignore extra attributes from the CRD that are not reflected here
-        extra = "ignore"
+        extra="ignore",
+    )
 
 
 class SeldonDeploymentComponentSpecs(BaseModel):
@@ -147,16 +138,13 @@ class SeldonDeploymentComponentSpecs(BaseModel):
         spec: the component spec.
     """
 
-    spec: Optional[Dict[str, Any]]
-    # TODO [HIGH]: Add graph field to ComponentSpecs. graph: Optional[SeldonDeploymentPredictiveUnit]
-
-    class Config:
-        """Pydantic configuration class."""
-
+    spec: Optional[Dict[str, Any]] = None
+    model_config = ConfigDict(
         # validate attribute assignments
-        validate_assignment = True
+        validate_assignment=True,
         # Ignore extra attributes from the CRD that are not reflected here
-        extra = "ignore"
+        extra="ignore",
+    )
 
 
 class SeldonDeploymentPredictor(BaseModel):
@@ -175,14 +163,12 @@ class SeldonDeploymentPredictor(BaseModel):
         default_factory=SeldonResourceRequirements
     )
     componentSpecs: Optional[List[SeldonDeploymentComponentSpecs]] = None
-
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # validate attribute assignments
-        validate_assignment = True
+        validate_assignment=True,
         # Ignore extra attributes from the CRD that are not reflected here
-        extra = "ignore"
+        extra="ignore",
+    )
 
 
 class SeldonDeploymentSpec(BaseModel):
@@ -196,17 +182,15 @@ class SeldonDeploymentSpec(BaseModel):
     """
 
     name: str
-    protocol: Optional[str]
+    protocol: Optional[str] = None
     predictors: List[SeldonDeploymentPredictor]
     replicas: int = 1
-
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # validate attribute assignments
-        validate_assignment = True
+        validate_assignment=True,
         # Ignore extra attributes from the CRD that are not reflected here
-        extra = "ignore"
+        extra="ignore",
+    )
 
 
 class SeldonDeploymentStatusState(StrEnum):
@@ -242,8 +226,8 @@ class SeldonDeploymentStatusCondition(BaseModel):
 
     type: str
     status: bool
-    reason: Optional[str]
-    message: Optional[str]
+    reason: Optional[str] = None
+    message: Optional[str] = None
 
 
 class SeldonDeploymentStatus(BaseModel):
@@ -258,18 +242,16 @@ class SeldonDeploymentStatus(BaseModel):
     """
 
     state: SeldonDeploymentStatusState = SeldonDeploymentStatusState.UNKNOWN
-    description: Optional[str]
-    replicas: Optional[int]
-    address: Optional[SeldonDeploymentStatusAddress]
+    description: Optional[str] = None
+    replicas: Optional[int] = None
+    address: Optional[SeldonDeploymentStatusAddress] = None
     conditions: List[SeldonDeploymentStatusCondition]
-
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # validate attribute assignments
-        validate_assignment = True
+        validate_assignment=True,
         # Ignore extra attributes from the CRD that are not reflected here
-        extra = "ignore"
+        extra="ignore",
+    )
 
 
 class SeldonDeployment(BaseModel):
@@ -292,8 +274,10 @@ class SeldonDeployment(BaseModel):
         status: Seldon Deployment status.
     """
 
-    kind: str = Field(SELDON_DEPLOYMENT_KIND, const=True)
-    apiVersion: str = Field(SELDON_DEPLOYMENT_API_VERSION, const=True)
+    kind: Literal["SeldonDeployment"] = "SeldonDeployment"
+    apiVersion: Literal["machinelearning.seldon.io/v1"] = (
+        "machinelearning.seldon.io/v1"
+    )
     metadata: SeldonDeploymentMetadata
     spec: SeldonDeploymentSpec
     status: Optional[SeldonDeploymentStatus] = None
@@ -304,7 +288,7 @@ class SeldonDeployment(BaseModel):
         Returns:
             A string representation of the Seldon Deployment.
         """
-        return json.dumps(self.dict(exclude_none=True), indent=4)
+        return json.dumps(self.model_dump(exclude_none=True), indent=4)
 
     @classmethod
     def build(
@@ -339,8 +323,10 @@ class SeldonDeployment(BaseModel):
             annotations: A dictionary of annotations to apply to the Seldon
                 Deployment.
             spec: A Kubernetes pod spec to use for the Seldon Deployment.
-            is_custom_deployment: Whether the Seldon Deployment is a custom or a built-in one.
-            serviceAccountName: The name of the service account to associate with the predictive unit container.
+            is_custom_deployment: Whether the Seldon Deployment is a custom
+                or a built-in one.
+            serviceAccountName: The name of the service account to associate
+                with the predictive unit container.
 
         Returns:
             A minimal SeldonDeployment object built from the provided
@@ -495,13 +481,12 @@ class SeldonDeployment(BaseModel):
             return None
         return ready_condition_message[0]
 
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # validate attribute assignments
-        validate_assignment = True
+        validate_assignment=True,
         # Ignore extra attributes from the CRD that are not reflected here
-        extra = "ignore"
+        extra="ignore",
+    )
 
 
 class SeldonClientError(Exception):
@@ -651,7 +636,7 @@ class SeldonClient:
             # are not
             deployment.mark_as_managed_by_zenml()
 
-            body_deploy = deployment.dict(exclude_none=True)
+            body_deploy = deployment.model_dump(exclude_none=True)
             response = (
                 self._custom_objects_api.create_namespaced_custom_object(
                     group="machinelearning.seldon.io",
@@ -788,7 +773,7 @@ class SeldonClient:
                 namespace=self._namespace,
                 plural="seldondeployments",
                 name=deployment.name,
-                body=deployment.dict(exclude_none=True),
+                body=deployment.model_dump(exclude_none=True),
                 _request_timeout=poll_timeout or None,
             )
             logger.debug("Seldon Core API response: %s", response)

zenml/integrations/seldon/services/seldon_deployment.py

@@ -113,7 +113,7 @@ class SeldonDeploymentConfig(ServiceConfig):
             The annotations for the Seldon Core deployment.
         """
         annotations = {
-            "zenml.service_config": self.json(),
+            "zenml.service_config": self.model_dump_json(),
             "zenml.version": __version__,
         }
         return annotations
@@ -133,7 +133,7 @@ class SeldonDeploymentConfig(ServiceConfig):
 
         Raises:
             ValueError: if the given deployment resource does not contain
-                the expected annotations or it contains an invalid or
+                the expected annotations, or it contains an invalid or
                 incompatible Seldon Core service configuration.
         """
         config_data = deployment.metadata.annotations.get(
@@ -145,7 +145,7 @@ class SeldonDeploymentConfig(ServiceConfig):
                 f"'zenml.service_config' annotation: {deployment}"
             )
         try:
-            service_config = cls.parse_raw(config_data)
+            service_config = cls.model_validate_json(config_data)
         except ValidationError as e:
             raise ValueError(
                 f"The loaded Seldon Core deployment resource contains an "

zenml/integrations/seldon/steps/seldon_deployer.py

@@ -76,7 +76,7 @@ def seldon_model_deployer_step(
     step_name = context.step_run.name
 
     # update the step configuration with the real pipeline runtime information
-    service_config = service_config.copy()
+    service_config = service_config.model_copy()
     service_config.pipeline_name = pipeline_name
     service_config.pipeline_step_name = step_name
 
@@ -141,7 +141,7 @@ def seldon_model_deployer_step(
     # fetch existing services with same pipeline name, step name and
     # model name
     existing_services = model_deployer.find_model_server(
-        config=service_config.dict()
+        config=service_config.model_dump()
     )
 
     # even when the deploy decision is negative, if an existing model server
@@ -240,7 +240,7 @@ def seldon_custom_model_deployer_step(
     # fetch existing services with the same pipeline name, step name and
     # model name
     existing_services = model_deployer.find_model_server(
-        config=service_config.dict()
+        config=service_config.model_dump()
     )
     # even when the deploy decision is negative if an existing model server
     # is not running for this pipeline/step, we still have to serve the
@@ -305,7 +305,7 @@ def seldon_custom_model_deployer_step(
     )
 
     # prepare the service configuration for the deployment
-    service_config = service_config.copy()
+    service_config = service_config.model_copy()
     service_config.model_uri = served_model_uri
 
     # create the specification for the custom deployment

zenml/integrations/skypilot/flavors/skypilot_orchestrator_base_vm_config.py

@@ -15,6 +15,8 @@
 
 from typing import Dict, List, Literal, Optional, Union
 
+from pydantic import Field
+
 from zenml.config.base_settings import BaseSettings
 from zenml.logger import get_logger
 from zenml.orchestrators import BaseOrchestratorConfig
@@ -90,15 +92,23 @@ class SkypilotBaseOrchestratorSettings(BaseSettings):
 
     # Resources
     instance_type: Optional[str] = None
-    cpus: Union[None, int, float, str] = None
-    memory: Union[None, int, float, str] = None
-    accelerators: Union[None, str, Dict[str, int]] = None
+    cpus: Union[None, int, float, str] = Field(
+        default=None, union_mode="left_to_right"
+    )
+    memory: Union[None, int, float, str] = Field(
+        default=None, union_mode="left_to_right"
+    )
+    accelerators: Union[None, str, Dict[str, int]] = Field(
+        default=None, union_mode="left_to_right"
+    )
     accelerator_args: Optional[Dict[str, str]] = None
     use_spot: Optional[bool] = None
     spot_recovery: Optional[str] = None
     region: Optional[str] = None
     zone: Optional[str] = None
-    image_id: Union[Dict[str, str], str, None] = None
+    image_id: Union[Dict[str, str], str, None] = Field(
+        default=None, union_mode="left_to_right"
+    )
     disk_size: Optional[int] = None
     disk_tier: Optional[Literal["high", "medium", "low"]] = None
 
@@ -112,7 +122,7 @@ class SkypilotBaseOrchestratorSettings(BaseSettings):
     docker_run_args: List[str] = []
 
 
-class SkypilotBaseOrchestratorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class SkypilotBaseOrchestratorConfig(
     BaseOrchestratorConfig, SkypilotBaseOrchestratorSettings
 ):
     """Skypilot orchestrator base config.

zenml/integrations/skypilot_aws/__init__.py

@@ -32,7 +32,7 @@ class SkypilotAWSIntegration(Integration):
 
     NAME = SKYPILOT_AWS
     REQUIREMENTS = ["skypilot[aws]<=0.5.0"]
-    APT_PACKAGES = ["openssh-client","rsync"]
+    APT_PACKAGES = ["openssh-client", "rsync"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/skypilot_aws/flavors/skypilot_orchestrator_aws_vm_flavor.py

@@ -37,7 +37,7 @@ class SkypilotAWSOrchestratorSettings(SkypilotBaseOrchestratorSettings):
     """Skypilot orchestrator settings."""
 
 
-class SkypilotAWSOrchestratorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class SkypilotAWSOrchestratorConfig(
     SkypilotBaseOrchestratorConfig, SkypilotAWSOrchestratorSettings
 ):
     """Skypilot orchestrator config."""

zenml/integrations/skypilot_azure/__init__.py

@@ -34,7 +34,7 @@ class SkypilotAzureIntegration(Integration):
 
     NAME = SKYPILOT_AZURE
     REQUIREMENTS = ["skypilot[azure]<=0.5.0"]
-    APT_PACKAGES = ["openssh-client","rsync"]
+    APT_PACKAGES = ["openssh-client", "rsync"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/skypilot_azure/flavors/skypilot_orchestrator_azure_vm_flavor.py

@@ -39,7 +39,7 @@ class SkypilotAzureOrchestratorSettings(SkypilotBaseOrchestratorSettings):
     """Skypilot orchestrator settings for Azure."""
 
 
-class SkypilotAzureOrchestratorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class SkypilotAzureOrchestratorConfig(
     SkypilotBaseOrchestratorConfig, SkypilotAzureOrchestratorSettings
 ):
     """Skypilot orchestrator config for Azure."""

zenml/integrations/skypilot_gcp/__init__.py

@@ -26,12 +26,13 @@ from zenml.stack import Flavor
 
 SKYPILOT_GCP_ORCHESTRATOR_FLAVOR = "vm_gcp"
 
+
 class SkypilotGCPIntegration(Integration):
     """Definition of Skypilot (GCP) Integration for ZenML."""
 
     NAME = SKYPILOT_GCP
     REQUIREMENTS = ["skypilot[gcp]<=0.5.0"]
-    APT_PACKAGES = ["openssh-client","rsync"]
+    APT_PACKAGES = ["openssh-client", "rsync"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/skypilot_gcp/flavors/skypilot_orchestrator_gcp_vm_flavor.py

@@ -40,7 +40,7 @@ class SkypilotGCPOrchestratorSettings(SkypilotBaseOrchestratorSettings):
     """Skypilot orchestrator settings for GCP."""
 
 
-class SkypilotGCPOrchestratorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class SkypilotGCPOrchestratorConfig(
     SkypilotBaseOrchestratorConfig,
     GoogleCredentialsConfigMixin,
     SkypilotGCPOrchestratorSettings,

zenml/integrations/skypilot_lambda/flavors/skypilot_orchestrator_lambda_vm_flavor.py

@@ -60,12 +60,12 @@ class SkypilotLambdaOrchestratorSettings(SkypilotBaseOrchestratorSettings):
         super().__setattr__(name, value)
 
 
-class SkypilotLambdaOrchestratorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class SkypilotLambdaOrchestratorConfig(
     SkypilotBaseOrchestratorConfig, SkypilotLambdaOrchestratorSettings
 ):
     """Skypilot orchestrator config."""
 
-    api_key: Optional[str] = SecretField()
+    api_key: Optional[str] = SecretField(default=None)
 
 
 class SkypilotLambdaOrchestratorFlavor(BaseOrchestratorFlavor):

zenml/integrations/spark/flavors/spark_step_operator_flavor.py

@@ -43,7 +43,7 @@ class SparkStepOperatorSettings(BaseSettings):
     submit_kwargs: Optional[Dict[str, Any]] = None
 
 
-class SparkStepOperatorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class SparkStepOperatorConfig(
     BaseStepOperatorConfig, SparkStepOperatorSettings
 ):
     """Spark step operator config.

zenml/integrations/spark/step_operators/spark_step_operator.py

@@ -303,6 +303,8 @@ class SparkStepOperator(BaseStepOperator):
         # Add any additional configuration given by the user.
         self._additional_configuration(spark_config=conf, settings=settings)
 
+        info.force_write_logs()
+
         # Generate a spark-submit command given the configuration
         self._launch_spark_job(
             spark_config=conf,

zenml/integrations/tekton/__init__.py

@@ -31,7 +31,7 @@ class TektonIntegration(Integration):
     """Definition of Tekton Integration for ZenML."""
 
     NAME = TEKTON
-    REQUIREMENTS = ["kfp-tekton==1.7.1"]
+    REQUIREMENTS = ["kfp>=2.6.0", "kfp-kubernetes>=1.1.0"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/tekton/flavors/tekton_orchestrator_flavor.py

@@ -13,14 +13,16 @@
 #  permissions and limitations under the License.
 """Tekton orchestrator flavor."""
 
-from typing import TYPE_CHECKING, Any, Dict, Optional, Type
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type
 
-from pydantic import root_validator
+from pydantic import model_validator
 
 from zenml.constants import KUBERNETES_CLUSTER_RESOURCE_TYPE
 from zenml.integrations.tekton import TEKTON_ORCHESTRATOR_FLAVOR
 from zenml.models import ServiceConnectorRequirements
 from zenml.orchestrators import BaseOrchestratorConfig, BaseOrchestratorFlavor
+from zenml.utils.pydantic_utils import before_validator_handler
+from zenml.utils.secret_utils import SecretField
 
 if TYPE_CHECKING:
     from zenml.integrations.tekton.orchestrators import TektonOrchestrator
@@ -36,42 +38,83 @@ class TektonOrchestratorSettings(BaseSettings):
     """Settings for the Tekton orchestrator.
 
     Attributes:
-        pod_settings: Pod settings to apply.
+        synchronous: If `True`, the client running a pipeline using this
+            orchestrator waits until all steps finish running. If `False`,
+            the client returns immediately and the pipeline is executed
+            asynchronously. Defaults to `True`. This setting only
+            has an effect when specified on the pipeline and will be ignored if
+            specified on steps.
+        timeout: How many seconds to wait for synchronous runs.
+        client_args: Arguments to pass when initializing the KFP client.
+        client_username: Username to generate a session cookie for the kubeflow client. Both `client_username`
+        and `client_password` need to be set together.
+        client_password: Password to generate a session cookie for the kubeflow client. Both `client_username`
+        and `client_password` need to be set together.
+        user_namespace: The user namespace to use when creating experiments
+            and runs.
     """
 
+    synchronous: bool = True
+    timeout: int = 1200
+
+    client_args: Dict[str, Any] = {}
+    client_username: Optional[str] = SecretField(default=None)
+    client_password: Optional[str] = SecretField(default=None)
+    user_namespace: Optional[str] = None
+    node_selectors: Dict[str, str] = {}
+    node_affinity: Dict[str, List[str]] = {}
     pod_settings: Optional[KubernetesPodSettings] = None
 
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def _validate_and_migrate_pod_settings(
+        cls, data: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Validates settings and migrates pod settings from older version.
+
+        Args:
+            data: Dict representing user-specified runtime settings.
+
+        Returns:
+            Validated settings.
+
+        Raises:
+            ValueError: If username and password are not specified together.
+        """
+        # Validate username and password for auth cookie logic
+        username = data.get("client_username")
+        password = data.get("client_password")
+        client_creds_error = "`client_username` and `client_password` both need to be set together."
+        if (username and password is None) or (password and username is None):
+            raise ValueError(client_creds_error)
+        return data
+
 
-class TektonOrchestratorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class TektonOrchestratorConfig(
     BaseOrchestratorConfig, TektonOrchestratorSettings
 ):
     """Configuration for the Tekton orchestrator.
 
     Attributes:
+        tekton_hostname: Hostname of the Tekton server.
         kubernetes_context: Name of a kubernetes context to run
             pipelines in. If the stack component is linked to a Kubernetes
             service connector, this field is ignored. Otherwise, it is
             mandatory.
         kubernetes_namespace: Name of the kubernetes namespace in which the
             pods that run the pipeline steps should be running.
-        local: If `True`, the orchestrator will assume it is connected to a
-            local kubernetes cluster and will perform additional validations and
-            operations to allow using the orchestrator in combination with other
-            local stack components that store data in the local filesystem
-            (i.e. it will mount the local stores directory into the pipeline
-            containers).
-        skip_local_validations: If `True`, the local validations will be
-            skipped.
     """
 
+    tekton_hostname: Optional[str] = None
     kubernetes_context: Optional[str] = None
-    kubernetes_namespace: str = "zenml"
-    local: bool = False
-    skip_local_validations: bool = False
+    kubernetes_namespace: str = "kubeflow"
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
     def _validate_deprecated_attrs(
-        cls, values: Dict[str, Any]
+        cls, data: Dict[str, Any]
     ) -> Dict[str, Any]:
         """Pydantic root_validator for deprecated attributes.
 
@@ -80,7 +123,7 @@ class TektonOrchestratorConfig(  # type: ignore[misc] # https://github.com/pydan
         mandatory in the meantime.
 
         Args:
-            values: Values passed to the object constructor
+            data: Values passed to the object constructor
 
         Returns:
             Values passed to the object constructor
@@ -93,10 +136,10 @@ class TektonOrchestratorConfig(  # type: ignore[misc] # https://github.com/pydan
 
         # remove deprecated attributes from values dict
         for attr in provisioning_attrs:
-            if attr in values:
-                del values[attr]
+            if attr in data:
+                del data[attr]
 
-        return values
+        return data
 
     @property
     def is_remote(self) -> bool:
@@ -109,7 +152,7 @@ class TektonOrchestratorConfig(  # type: ignore[misc] # https://github.com/pydan
         Returns:
             True if this config is for a remote component, False otherwise.
         """
-        return not self.local
+        return True
 
     @property
     def is_local(self) -> bool:
@@ -118,7 +161,7 @@ class TektonOrchestratorConfig(  # type: ignore[misc] # https://github.com/pydan
         Returns:
             True if this config is for a local component, False otherwise.
         """
-        return self.local
+        return False
 
 
 class TektonOrchestratorFlavor(BaseOrchestratorFlavor):