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

zenml/models/v2/core/user.py

@@ -27,7 +27,7 @@ from typing import (
 )
 from uuid import UUID
 
-from pydantic import BaseModel, Field, root_validator
+from pydantic import BaseModel, ConfigDict, Field, model_validator
 
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.models.v2.base.base import (
@@ -171,15 +171,12 @@ class UserRequest(UserBase, BaseRequest):
     )
     active: bool = Field(default=False, title="Whether the account is active.")
 
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # Validate attributes when assigning them
-        validate_assignment = True
-
+        validate_assignment=True,
         # Forbid extra attributes to prevent unexpected behavior
-        extra = "forbid"
-        underscore_attrs_are_private = True
+        extra="forbid",
+    )
 
 
 # ------------------ Update Model ------------------
@@ -213,13 +210,10 @@ class UserUpdate(UserBase, BaseZenModel):
         max_length=STR_FIELD_MAX_LENGTH,
     )
 
-    @root_validator
-    def user_email_updates(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="after")
+    def user_email_updates(self) -> "UserUpdate":
         """Validate that the UserUpdateModel conforms to the email-opt-in-flow.
 
-        Args:
-            values: The values to validate.
-
         Returns:
             The validated values.
 
@@ -229,18 +223,18 @@ class UserUpdate(UserBase, BaseZenModel):
         """
         # When someone sets the email, or updates the email and hasn't
         #  before explicitly opted out, they are opted in
-        if values["email"] is not None:
-            if values["email_opted_in"] is None:
-                values["email_opted_in"] = True
+        if self.email is not None:
+            if self.email_opted_in is None:
+                self.email_opted_in = True
 
         # It should not be possible to do opt in without an email
-        if values["email_opted_in"] is True:
-            if values["email"] is None:
+        if self.email_opted_in is True:
+            if self.email is None:
                 raise ValueError(
                     "Please provide an email, when you are opting-in with "
                     "your email."
                 )
-        return values
+        return self
 
     def create_copy(self, exclude: AbstractSet[str]) -> "UserUpdate":
         """Create a copy of the current instance.
@@ -252,9 +246,9 @@ class UserUpdate(UserBase, BaseZenModel):
             A copy of the current instance.
         """
         return UserUpdate(
-            **self.dict(
+            **self.model_dump(
+                exclude=set(exclude),
                 exclude_unset=True,
-                exclude=exclude,
             )
         )
 
@@ -483,14 +477,17 @@ class UserFilter(BaseFilter):
     active: Optional[Union[bool, str]] = Field(
         default=None,
         description="Whether the user is active",
+        union_mode="left_to_right",
     )
     email_opted_in: Optional[Union[bool, str]] = Field(
         default=None,
         description="Whether the user has opted in to emails",
+        union_mode="left_to_right",
     )
     external_user_id: Optional[Union[UUID, str]] = Field(
         default=None,
         title="The external user ID associated with the account.",
+        union_mode="left_to_right",
     )
 
     def apply_filter(

zenml/models/v2/core/workspace.py

@@ -24,9 +24,9 @@ from zenml.models.v2.base.base import (
     BaseRequest,
     BaseResponseMetadata,
     BaseResponseResources,
+    BaseUpdate,
 )
 from zenml.models.v2.base.filter import BaseFilter
-from zenml.models.v2.base.update import update_model
 
 # ------------------ Request Model ------------------
 
@@ -48,10 +48,20 @@ class WorkspaceRequest(BaseRequest):
 # ------------------ Update Model ------------------
 
 
-@update_model
-class WorkspaceUpdate(WorkspaceRequest):
+class WorkspaceUpdate(BaseUpdate):
     """Update model for workspaces."""
 
+    name: Optional[str] = Field(
+        title="The unique name of the workspace.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    description: Optional[str] = Field(
+        title="The description of the workspace.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+
 
 # ------------------ Response Model ------------------
 

zenml/models/v2/misc/build_item.py

@@ -33,13 +33,13 @@ class BuildItem(BaseModel):
 
     image: str = Field(title="The image name or digest.")
     dockerfile: Optional[str] = Field(
-        title="The dockerfile used to build the image."
+        default=None, title="The dockerfile used to build the image."
     )
     requirements: Optional[str] = Field(
-        title="The pip requirements installed in the image."
+        default=None, title="The pip requirements installed in the image."
     )
     settings_checksum: Optional[str] = Field(
-        title="The checksum of the build settings."
+        default=None, title="The checksum of the build settings."
     )
     contains_code: bool = Field(
         default=True, title="Whether the image contains user files."

zenml/models/v2/misc/external_user.py

@@ -16,7 +16,7 @@
 from typing import Optional
 from uuid import UUID
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 
 class ExternalUserModel(BaseModel):
@@ -27,8 +27,4 @@ class ExternalUserModel(BaseModel):
     name: Optional[str] = None
     is_admin: bool = False
 
-    class Config:
-        """Pydantic configuration."""
-
-        # ignore arbitrary fields
-        extra = "ignore"
+    model_config = ConfigDict(extra="ignore")

zenml/models/v2/misc/hub_plugin_models.py

@@ -42,22 +42,22 @@ class HubUserResponseModel(BaseModel):
 
     id: UUID
     email: str
-    username: Optional[str]
+    username: Optional[str] = None
 
 
 class HubPluginBaseModel(BaseModel):
     """Base model for a ZenML Hub plugin."""
 
     name: str
-    description: Optional[str]
-    version: Optional[str]
-    release_notes: Optional[str]
+    description: Optional[str] = None
+    version: Optional[str] = None
+    release_notes: Optional[str] = None
     repository_url: str
-    repository_subdirectory: Optional[str]
-    repository_branch: Optional[str]
-    repository_commit: Optional[str]
-    tags: Optional[List[str]]
-    logo_url: Optional[str]
+    repository_subdirectory: Optional[str] = None
+    repository_branch: Optional[str] = None
+    repository_commit: Optional[str] = None
+    tags: Optional[List[str]] = None
+    logo_url: Optional[str] = None
 
 
 class HubPluginRequestModel(HubPluginBaseModel):

zenml/models/v2/misc/loaded_visualization.py

@@ -15,7 +15,7 @@
 
 from typing import Union
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from zenml.enums import VisualizationType
 
@@ -24,4 +24,4 @@ class LoadedVisualization(BaseModel):
     """Model for loaded visualizations."""
 
     type: VisualizationType
-    value: Union[str, bytes]
+    value: Union[str, bytes] = Field(union_mode="left_to_right")

zenml/models/v2/misc/service_connector_type.py

@@ -13,11 +13,10 @@
 #  permissions and limitations under the License.
 """Model definitions for ZenML service connectors."""
 
-import json
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
 from uuid import UUID
 
-from pydantic import BaseModel, Field, validator
+from pydantic import BaseModel, Field, field_validator
 
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.logger import get_logger
@@ -141,7 +140,8 @@ class AuthenticationMethodModel(BaseModel):
             **values: The data to initialize the authentication method with.
         """
         if config_class:
-            values["config_schema"] = json.loads(config_class.schema_json())
+            values["config_schema"] = config_class.model_json_schema()
+
         super().__init__(**values)
         self._config_class = config_class
 
@@ -215,11 +215,6 @@ class AuthenticationMethodModel(BaseModel):
 
         return expiration_seconds
 
-    class Config:
-        """Pydantic config class."""
-
-        underscore_attrs_are_private = True
-
 
 class ServiceConnectorTypeModel(BaseModel):
     """Service connector type specification.
@@ -333,7 +328,8 @@ class ServiceConnectorTypeModel(BaseModel):
         """
         self._connector_class = connector_class
 
-    @validator("resource_types")
+    @field_validator("resource_types")
+    @classmethod
     def validate_resource_types(
         cls, values: List[ResourceTypeModel]
     ) -> List[ResourceTypeModel]:
@@ -360,7 +356,8 @@ class ServiceConnectorTypeModel(BaseModel):
 
         return values
 
-    @validator("auth_methods")
+    @field_validator("auth_methods")
+    @classmethod
     def validate_auth_methods(
         cls, values: List[AuthenticationMethodModel]
     ) -> List[AuthenticationMethodModel]:
@@ -472,11 +469,6 @@ class ServiceConnectorTypeModel(BaseModel):
 
         return auth_method_spec, resource_type_spec
 
-    class Config:
-        """Pydantic config class."""
-
-        underscore_attrs_are_private = True
-
 
 class ServiceConnectorRequirements(BaseModel):
     """Service connector requirements.
@@ -600,8 +592,7 @@ class ServiceConnectorResourcesModel(BaseModel):
     )
 
     connector_type: Union[str, "ServiceConnectorTypeModel"] = Field(
-        title="The type of service connector.",
-        max_length=STR_FIELD_MAX_LENGTH,
+        title="The type of service connector.", union_mode="left_to_right"
     )
 
     resources: List[ServiceConnectorTypedResourcesModel] = Field(

zenml/models/v2/misc/user_auth.py

@@ -22,6 +22,7 @@ from pydantic import Field, SecretStr
 
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.models.v2.base.base import BaseZenModel
+from zenml.utils.secret_utils import PlainSerializedSecretStr
 
 if TYPE_CHECKING:
     from passlib.context import CryptContext
@@ -47,8 +48,12 @@ class UserAuthModel(BaseZenModel):
         "account."
     )
 
-    activation_token: Optional[SecretStr] = Field(default=None, exclude=True)
-    password: Optional[SecretStr] = Field(default=None, exclude=True)
+    activation_token: Optional[PlainSerializedSecretStr] = Field(
+        default=None, exclude=True
+    )
+    password: Optional[PlainSerializedSecretStr] = Field(
+        default=None, exclude=True
+    )
     name: str = Field(
         title="The unique username for the account.",
         max_length=STR_FIELD_MAX_LENGTH,

zenml/new/pipelines/build_utils.py

@@ -62,7 +62,7 @@ def _create_deployment(
 
     Args:
         deployment: Base of the deployment to create.
-        pipeline_id: Pipeline ID to use for the deloyment.
+        pipeline_id: Pipeline ID to use for the deployment.
         code_repository: Code repository to use for the deployment.
 
     Returns:
@@ -93,7 +93,7 @@ def _create_deployment(
         stack=Client().active_stack.id,
         pipeline=pipeline_id,
         code_reference=code_reference,
-        **deployment.dict(),
+        **deployment.model_dump(),
     )
     return (
         Client().zen_store.create_deployment(deployment=deployment_request).id
@@ -184,7 +184,7 @@ def reuse_or_create_pipeline_build(
             workspace=Client().active_workspace.id,
             stack=Client().active_stack_model.id,
             pipeline=pipeline_id,
-            **build.dict(),
+            **build.model_dump(),
         )
         build_model = Client().zen_store.create_build(build=build_request)
 

zenml/new/pipelines/pipeline.py

@@ -38,7 +38,7 @@ from typing import (
 from uuid import UUID
 
 import yaml
-from pydantic import ValidationError
+from pydantic import ConfigDict, ValidationError
 
 from zenml import constants
 from zenml.analytics.enums import AnalyticsEvent
@@ -378,7 +378,7 @@ class Pipeline:
             to_be_reapplied = []
             for param_, value_ in values.items():
                 if (
-                    param_ in PipelineRunConfiguration.__fields__
+                    param_ in PipelineRunConfiguration.model_fields
                     and param_ in self._from_config_file
                     and value_ != self._from_config_file[param_]
                 ):
@@ -389,16 +389,17 @@ class Pipeline:
                 msg = ""
                 reapply_during_run_warning = (
                     "The value of parameter '{name}' has changed from "
-                    "'{file_value}' to '{new_value}' set in your configuration file.\n"
+                    "'{file_value}' to '{new_value}' set in your configuration "
+                    "file.\n"
                 )
                 for name, file_value, new_value in to_be_reapplied:
                     msg += reapply_during_run_warning.format(
                         name=name, file_value=file_value, new_value=new_value
                     )
                 msg += (
-                    "Configuration file value will be used during pipeline run, "
-                    "so you change will not be efficient. Consider updating your "
-                    "configuration file instead."
+                    "Configuration file value will be used during pipeline "
+                    "run, so you change will not be efficient. Consider "
+                    "updating your configuration file instead."
                 )
                 logger.warning(msg)
 
@@ -501,7 +502,9 @@ To avoid this consider setting pipeline parameters only in one place (config or
         self._prepare_if_possible()
         integration_registry.activate_integrations()
 
-        if self.configuration.dict(exclude_defaults=True, exclude={"name"}):
+        if self.configuration.model_dump(
+            exclude_defaults=True, exclude={"name"}
+        ):
             logger.warning(
                 f"The pipeline `{self.name}` that you're registering has "
                 "custom configurations applied to it. These will not be "
@@ -729,7 +732,7 @@ To avoid this consider setting pipeline parameters only in one place (config or
                 build=build_id,
                 schedule=schedule_id,
                 code_reference=code_reference,
-                **deployment.dict(),
+                **deployment.model_dump(),
             )
             deployment_model = Client().zen_store.create_deployment(
                 deployment=deployment_request
@@ -1006,13 +1009,13 @@ To avoid this consider setting pipeline parameters only in one place (config or
 
         self._parse_config_file(
             config_path=config_path,
-            matcher=list(PipelineRunConfiguration.__fields__.keys()),
+            matcher=list(PipelineRunConfiguration.model_fields.keys()),
         )
 
         run_config = PipelineRunConfiguration(**self._from_config_file)
 
         new_values = dict_utils.remove_none_values(run_configuration_args)
-        update = PipelineRunConfiguration.parse_obj(new_values)
+        update = PipelineRunConfiguration.model_validate(new_values)
 
         # Update with the values in code so they take precedence
         run_config = pydantic_utils.update_model(run_config, update=update)
@@ -1291,7 +1294,7 @@ To avoid this consider setting pipeline parameters only in one place (config or
                 else:
                     from zenml.model.model import Model
 
-                    _from_config_file["model"] = Model.parse_obj(
+                    _from_config_file["model"] = Model.model_validate(
                         _from_config_file["model"]
                     )
         self._from_config_file = _from_config_file
@@ -1410,14 +1413,15 @@ To avoid this consider setting pipeline parameters only in one place (config or
         try:
             validated_args = pydantic_utils.validate_function_args(
                 self.entrypoint,
-                {"arbitrary_types_allowed": False, "smart_union": True},
+                ConfigDict(arbitrary_types_allowed=False),
                 *args,
                 **kwargs,
             )
         except ValidationError as e:
             raise ValueError(
-                "Invalid or missing inputs for pipeline entrypoint function. "
+                "Invalid or missing pipeline function entrypoint arguments. "
                 "Only JSON serializable inputs are allowed as pipeline inputs."
+                "Check out the pydantic error above for more details."
             ) from e
 
         self._parameters = validated_args

zenml/new/pipelines/run_utils.py

@@ -1,5 +1,6 @@
 """Utility functions for running pipelines."""
 
+import time
 from collections import defaultdict
 from datetime import datetime
 from types import FunctionType
@@ -16,19 +17,23 @@ from uuid import UUID
 
 from zenml import constants
 from zenml.client import Client
+from zenml.config.pipeline_run_configuration import PipelineRunConfiguration
 from zenml.config.step_configurations import StepConfigurationUpdate
 from zenml.enums import ExecutionStatus, ModelStages
 from zenml.logger import get_logger
 from zenml.models import (
+    FlavorFilter,
     PipelineDeploymentBase,
     PipelineDeploymentResponse,
     PipelineRunRequest,
     PipelineRunResponse,
+    StackResponse,
 )
 from zenml.new.pipelines.model_utils import NewModelRequest
 from zenml.orchestrators.utils import get_run_name
-from zenml.stack import Stack
+from zenml.stack import Flavor, Stack
 from zenml.utils import cloud_utils
+from zenml.zen_stores.base_zen_store import BaseZenStore
 
 if TYPE_CHECKING:
     from zenml.config.source import Source
@@ -258,3 +263,100 @@ def _validate_new_version_requests(
             "Models can be viewed in the dashboard using ZenML Pro. Sign up "
             "for a free trial at https://www.zenml.io/cloud/"
         )
+
+
+def wait_for_pipeline_run_to_finish(run_id: UUID) -> "PipelineRunResponse":
+    """Waits until a pipeline run is finished.
+
+    Args:
+        run_id: ID of the run for which to wait.
+
+    Returns:
+        Model of the finished run.
+    """
+    sleep_interval = 1
+    max_sleep_interval = 64
+
+    while True:
+        run = Client().get_pipeline_run(run_id)
+
+        if run.status.is_finished:
+            return run
+
+        logger.info(
+            "Waiting for pipeline run with ID %s to finish (current status: %s)",
+            run_id,
+            run.status,
+        )
+        time.sleep(sleep_interval)
+        if sleep_interval < max_sleep_interval:
+            sleep_interval *= 2
+
+
+def validate_stack_is_runnable_from_server(
+    zen_store: BaseZenStore, stack: StackResponse
+) -> None:
+    """Validate if a stack model is runnable from the server.
+
+    Args:
+        zen_store: ZenStore to use for listing flavors.
+        stack: The stack to validate.
+
+    Raises:
+        ValueError: If the stack has components of a custom flavor or local
+            components.
+    """
+    for component_list in stack.components.values():
+        assert len(component_list) == 1
+        component = component_list[0]
+        flavors = zen_store.list_flavors(
+            FlavorFilter(name=component.flavor, type=component.type)
+        )
+        assert len(flavors) == 1
+        flavor_model = flavors[0]
+
+        if flavor_model.workspace is not None:
+            raise ValueError("No custom stack component flavors allowed.")
+
+        flavor = Flavor.from_model(flavor_model)
+        component_config = flavor.config_class(**component.configuration)
+
+        if component_config.is_local:
+            raise ValueError("No local stack components allowed.")
+
+
+def validate_run_config_is_runnable_from_server(
+    run_configuration: "PipelineRunConfiguration",
+) -> None:
+    """Validates that the run configuration can be used to run from the server.
+
+    Args:
+        run_configuration: The run configuration to validate.
+
+    Raises:
+        ValueError: If there are values in the run configuration that are not
+            allowed when running a pipeline from the server.
+    """
+    if run_configuration.parameters:
+        raise ValueError(
+            "Can't set parameters when running pipeline via Rest API."
+        )
+
+    if run_configuration.build:
+        raise ValueError("Can't set build when running pipeline via Rest API.")
+
+    if run_configuration.schedule:
+        raise ValueError(
+            "Can't set schedule when running pipeline via Rest API."
+        )
+
+    if run_configuration.settings.get("docker"):
+        raise ValueError(
+            "Can't set DockerSettings when running pipeline via Rest API."
+        )
+
+    for step_update in run_configuration.steps.values():
+        if step_update.settings.get("docker"):
+            raise ValueError(
+                "Can't set DockerSettings when running pipeline via Rest API."
+            )

zenml/orchestrators/base_orchestrator.py

@@ -16,13 +16,14 @@
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Dict, Optional, Type, cast
 
-from pydantic import root_validator
+from pydantic import model_validator
 
 from zenml.enums import StackComponentType
 from zenml.logger import get_logger
 from zenml.orchestrators.step_launcher import StepLauncher
 from zenml.orchestrators.utils import get_config_environment_vars
 from zenml.stack import Flavor, Stack, StackComponent, StackComponentConfig
+from zenml.utils.pydantic_utils import before_validator_handler
 
 if TYPE_CHECKING:
     from zenml.config.step_configurations import Step
@@ -34,18 +35,20 @@ logger = get_logger(__name__)
 class BaseOrchestratorConfig(StackComponentConfig):
     """Base orchestrator config."""
 
-    @root_validator(pre=True)
-    def _deprecations(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def _deprecations(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         """Validate and/or remove deprecated fields.
 
         Args:
-            values: The values to validate.
+            data: The values to validate.
 
         Returns:
             The validated values.
         """
-        if "custom_docker_base_image_name" in values:
-            image_name = values.pop("custom_docker_base_image_name", None)
+        if "custom_docker_base_image_name" in data:
+            image_name = data.pop("custom_docker_base_image_name", None)
             if image_name:
                 logger.warning(
                     "The 'custom_docker_base_image_name' field has been "
@@ -54,7 +57,7 @@ class BaseOrchestratorConfig(StackComponentConfig):
                     "pipeline (see https://docs.zenml.io/how-to/customize-docker-builds)."
                 )
 
-        return values
+        return data
 
     @property
     def is_synchronous(self) -> bool:

zenml/orchestrators/local_docker/local_docker_orchestrator.py

@@ -206,7 +206,7 @@ class LocalDockerOrchestratorSettings(BaseSettings):
     run_args: Dict[str, Any] = {}
 
 
-class LocalDockerOrchestratorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class LocalDockerOrchestratorConfig(
     BaseOrchestratorConfig, LocalDockerOrchestratorSettings
 ):
     """Local Docker orchestrator config."""