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

zenml/models/v2/core/schedule.py

@@ -14,13 +14,14 @@
 """Models representing schedules."""
 
 import datetime
-from typing import Any, Dict, Optional, Union
+from typing import Optional, Union
 from uuid import UUID
 
-from pydantic import Field, root_validator
+from pydantic import Field, model_validator
 
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.logger import get_logger
+from zenml.models.v2.base.base import BaseUpdate
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedFilter,
     WorkspaceScopedRequest,
@@ -29,7 +30,6 @@ from zenml.models.v2.base.scoped import (
     WorkspaceScopedResponseMetadata,
     WorkspaceScopedResponseResources,
 )
-from zenml.models.v2.base.update import update_model
 
 logger = get_logger(__name__)
 
@@ -53,15 +53,12 @@ class ScheduleRequest(WorkspaceScopedRequest):
     orchestrator_id: Optional[UUID]
     pipeline_id: Optional[UUID]
 
-    @root_validator
+    @model_validator(mode="after")
     def _ensure_cron_or_periodic_schedule_configured(
-        cls, values: Dict[str, Any]
-    ) -> Dict[str, Any]:
+        self,
+    ) -> "ScheduleRequest":
         """Ensures that the cron expression or start time + interval are set.
 
-        Args:
-            values: All attributes of the schedule.
-
         Returns:
             All schedule attributes.
 
@@ -69,11 +66,9 @@ class ScheduleRequest(WorkspaceScopedRequest):
             ValueError: If no cron expression or start time + interval were
                 provided.
         """
-        cron_expression = values.get("cron_expression")
-        periodic_schedule = values.get("start_time") and values.get(
-            "interval_second"
-        )
-        run_once_starts_at = values.get("run_once_start_time")
+        cron_expression = self.cron_expression
+        periodic_schedule = self.start_time and self.interval_second
+        run_once_starts_at = self.run_once_start_time
 
         if cron_expression and periodic_schedule:
             logger.warning(
@@ -83,7 +78,7 @@ class ScheduleRequest(WorkspaceScopedRequest):
                 "but will usually ignore the interval and use the cron "
                 "expression."
             )
-            return values
+            return self
         elif cron_expression and run_once_starts_at:
             logger.warning(
                 "This schedule was created with a cron expression as well as "
@@ -91,9 +86,9 @@ class ScheduleRequest(WorkspaceScopedRequest):
                 "depends on the concrete orchestrator implementation but will "
                 "usually ignore the `run_once_start_time`."
             )
-            return values
+            return self
         elif cron_expression or periodic_schedule or run_once_starts_at:
-            return values
+            return self
         else:
             raise ValueError(
                 "Either a cron expression, a start time and interval seconds "
@@ -105,10 +100,20 @@ class ScheduleRequest(WorkspaceScopedRequest):
 # ------------------ Update Model ------------------
 
 
-@update_model
-class ScheduleUpdate(ScheduleRequest):
+class ScheduleUpdate(BaseUpdate):
     """Update model for schedules."""
 
+    name: Optional[str] = None
+    active: Optional[bool] = None
+    cron_expression: Optional[str] = None
+    start_time: Optional[datetime.datetime] = None
+    end_time: Optional[datetime.datetime] = None
+    interval_second: Optional[datetime.timedelta] = None
+    catchup: Optional[bool] = None
+    run_once_start_time: Optional[datetime.datetime] = None
+    orchestrator_id: Optional[UUID] = None
+    pipeline_id: Optional[UUID] = None
+
 
 # ------------------ Response Model ------------------
 
@@ -275,17 +280,24 @@ class ScheduleFilter(WorkspaceScopedFilter):
     """Model to enable advanced filtering of all Users."""
 
     workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Workspace scope of the schedule."
+        default=None,
+        description="Workspace scope of the schedule.",
+        union_mode="left_to_right",
     )
     user_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="User that created the schedule"
+        default=None,
+        description="User that created the schedule",
+        union_mode="left_to_right",
     )
     pipeline_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Pipeline that the schedule is attached to."
+        default=None,
+        description="Pipeline that the schedule is attached to.",
+        union_mode="left_to_right",
     )
     orchestrator_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="Orchestrator that the schedule is attached to.",
+        union_mode="left_to_right",
     )
     active: Optional[bool] = Field(
         default=None,
@@ -296,10 +308,10 @@ class ScheduleFilter(WorkspaceScopedFilter):
         description="The cron expression, describing the schedule",
     )
     start_time: Optional[Union[datetime.datetime, str]] = Field(
-        default=None, description="Start time"
+        default=None, description="Start time", union_mode="left_to_right"
     )
     end_time: Optional[Union[datetime.datetime, str]] = Field(
-        default=None, description="End time"
+        default=None, description="End time", union_mode="left_to_right"
     )
     interval_second: Optional[Optional[float]] = Field(
         default=None,
@@ -317,4 +329,5 @@ class ScheduleFilter(WorkspaceScopedFilter):
     run_once_start_time: Optional[Union[datetime.datetime, str]] = Field(
         default=None,
         description="The time at which the schedule should run once",
+        union_mode="left_to_right",
     )

zenml/models/v2/core/secret.py

@@ -26,6 +26,7 @@ from zenml.enums import (
     SecretScope,
     SorterOps,
 )
+from zenml.models.v2.base.base import BaseUpdate
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedFilter,
     WorkspaceScopedRequest,
@@ -34,7 +35,7 @@ from zenml.models.v2.base.scoped import (
     WorkspaceScopedResponseMetadata,
     WorkspaceScopedResponseResources,
 )
-from zenml.models.v2.base.update import update_model
+from zenml.utils.secret_utils import PlainSerializedSecretStr
 
 # ------------------ Request Model ------------------
 
@@ -51,7 +52,7 @@ class SecretRequest(WorkspaceScopedRequest):
     scope: SecretScope = Field(
         SecretScope.WORKSPACE, title="The scope of the secret."
     )
-    values: Dict[str, Optional[SecretStr]] = Field(
+    values: Dict[str, Optional[PlainSerializedSecretStr]] = Field(
         default_factory=dict, title="The values stored in this secret."
     )
 
@@ -77,13 +78,23 @@ class SecretRequest(WorkspaceScopedRequest):
 # ------------------ Update Model ------------------
 
 
-@update_model
-class SecretUpdate(SecretRequest):
+class SecretUpdate(BaseUpdate):
     """Secret update model."""
 
-    scope: Optional[SecretScope] = Field(  # type: ignore[assignment]
+    ANALYTICS_FIELDS: ClassVar[List[str]] = ["scope"]
+
+    name: Optional[str] = Field(
+        title="The name of the secret.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    scope: Optional[SecretScope] = Field(
         default=None, title="The scope of the secret."
     )
+    values: Optional[Dict[str, Optional[PlainSerializedSecretStr]]] = Field(
+        title="The values stored in this secret.",
+        default=None,
+    )
 
     def get_secret_values_update(self) -> Dict[str, Optional[str]]:
         """Returns a dictionary with the secret values to update.
@@ -91,10 +102,13 @@ class SecretUpdate(SecretRequest):
         Returns:
             A dictionary with the secret values to update.
         """
-        return {
-            k: v.get_secret_value() if v is not None else None
-            for k, v in self.values.items()
-        }
+        if self.values is not None:
+            return {
+                k: v.get_secret_value() if v is not None else None
+                for k, v in self.values.items()
+            }
+
+        return {}
 
 
 # ------------------ Response Model ------------------
@@ -106,7 +120,7 @@ class SecretResponseBody(WorkspaceScopedResponseBody):
     scope: SecretScope = Field(
         SecretScope.WORKSPACE, title="The scope of the secret."
     )
-    values: Dict[str, Optional[SecretStr]] = Field(
+    values: Dict[str, Optional[PlainSerializedSecretStr]] = Field(
         default_factory=dict, title="The values stored in this secret."
     )
 
@@ -243,14 +257,19 @@ class SecretFilter(WorkspaceScopedFilter):
     scope: Optional[Union[SecretScope, str]] = Field(
         default=None,
         description="Scope in which to filter secrets",
+        union_mode="left_to_right",
     )
 
     workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Workspace of the Secret"
+        default=None,
+        description="Workspace of the Secret",
+        union_mode="left_to_right",
     )
 
     user_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="User that created the Secret"
+        default=None,
+        description="User that created the Secret",
+        union_mode="left_to_right",
     )
 
     @staticmethod

zenml/models/v2/core/service.py

@@ -15,7 +15,6 @@
 
 from datetime import datetime
 from typing import (
-    TYPE_CHECKING,
     Any,
     ClassVar,
     Dict,
@@ -26,8 +25,8 @@ from typing import (
 )
 from uuid import UUID
 
-from pydantic import BaseModel, Field
-from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList
+from pydantic import BaseModel, ConfigDict, Field
+from sqlalchemy.sql.elements import ColumnElement
 from sqlmodel import SQLModel
 
 from zenml.constants import STR_FIELD_MAX_LENGTH
@@ -43,9 +42,6 @@ from zenml.models.v2.base.scoped import (
 from zenml.services.service_status import ServiceState
 from zenml.services.service_type import ServiceType
 
-if TYPE_CHECKING:
-    pass
-
 # ------------------ Request Model ------------------
 
 
@@ -56,50 +52,46 @@ class ServiceRequest(WorkspaceScopedRequest):
         title="The name of the service.",
         max_length=STR_FIELD_MAX_LENGTH,
     )
-
     service_type: ServiceType = Field(
         title="The type of the service.",
     )
-
     service_source: Optional[str] = Field(
         title="The class of the service.",
-        description="The fully qualified class name of the service implementation.",
+        description="The fully qualified class name of the service "
+        "implementation.",
+        default=None,
     )
-
     admin_state: Optional[ServiceState] = Field(
         title="The admin state of the service.",
-        description="The administrative state of the service, e.g., ACTIVE, INACTIVE.",
+        description="The administrative state of the service, e.g., ACTIVE, "
+        "INACTIVE.",
+        default=None,
     )
-
     config: Dict[str, Any] = Field(
         title="The service config.",
-        description="A dictionary containing configuration parameters for the service.",
+        description="A dictionary containing configuration parameters for the "
+        "service.",
     )
-
     labels: Optional[Dict[str, str]] = Field(
         default=None,
         title="The service labels.",
     )
-
     status: Optional[Dict[str, Any]] = Field(
+        default=None,
         title="The status of the service.",
     )
-
     endpoint: Optional[Dict[str, Any]] = Field(
         default=None,
         title="The service endpoint.",
     )
-
     prediction_url: Optional[str] = Field(
         default=None,
         title="The service endpoint URL.",
     )
-
     health_check_url: Optional[str] = Field(
         default=None,
         title="The service health check URL.",
     )
-
     model_version_id: Optional[UUID] = Field(
         default=None,
         title="The model version id linked to the service.",
@@ -107,8 +99,17 @@ class ServiceRequest(WorkspaceScopedRequest):
     pipeline_run_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="By the event source this trigger is attached to.",
+        union_mode="left_to_right",
     )
 
+    # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
+    #  fields defined under base models. If not handled, this raises a warning.
+    #  It is possible to suppress this warning message with the following
+    #  configuration, however the ultimate solution is to rename these fields.
+    #  Even though they do not cause any problems right now, if we are not
+    #  careful we might overwrite some fields protected by pydantic.
+    model_config = ConfigDict(protected_namespaces=())
+
 
 # ------------------ Update Model ------------------
 
@@ -117,46 +118,55 @@ class ServiceUpdate(BaseModel):
     """Update model for stack components."""
 
     name: Optional[str] = Field(
+        None,
         title="The name of the service.",
         max_length=STR_FIELD_MAX_LENGTH,
     )
-
     admin_state: Optional[ServiceState] = Field(
+        None,
         title="The admin state of the service.",
-        description="The administrative state of the service, e.g., ACTIVE, INACTIVE.",
+        description="The administrative state of the service, e.g., ACTIVE, "
+        "INACTIVE.",
     )
-
     service_source: Optional[str] = Field(
+        None,
         title="The class of the service.",
-        description="The fully qualified class name of the service implementation.",
+        description="The fully qualified class name of the service "
+        "implementation.",
     )
-
     status: Optional[Dict[str, Any]] = Field(
+        None,
         title="The status of the service.",
     )
-
     endpoint: Optional[Dict[str, Any]] = Field(
+        None,
         title="The service endpoint.",
     )
-
     prediction_url: Optional[str] = Field(
+        None,
         title="The service endpoint URL.",
     )
-
     health_check_url: Optional[str] = Field(
+        None,
         title="The service health check URL.",
     )
-
     labels: Optional[Dict[str, str]] = Field(
         default=None,
         title="The service labels.",
     )
-
     model_version_id: Optional[UUID] = Field(
         default=None,
         title="The model version id linked to the service.",
     )
 
+    # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
+    #  fields defined under base models. If not handled, this raises a warning.
+    #  It is possible to suppress this warning message with the following
+    #  configuration, however the ultimate solution is to rename these fields.
+    #  Even though they do not cause any problems right now, if we are not
+    #  careful we might overwrite some fields protected by pydantic.
+    model_config = ConfigDict(protected_namespaces=())
+
 
 # ------------------ Response Model ------------------
 
@@ -362,13 +372,19 @@ class ServiceFilter(WorkspaceScopedFilter):
     """
 
     name: Optional[str] = Field(
-        description="Name of the service. Use this to filter services by their name.",
+        default=None,
+        description="Name of the service. Use this to filter services by "
+        "their name.",
     )
     workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Workspace of the service"
+        default=None,
+        description="Workspace of the service",
+        union_mode="left_to_right",
     )
     user_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="User of the service"
+        default=None,
+        description="User of the service",
+        union_mode="left_to_right",
     )
     type: Optional[str] = Field(
         default=None,
@@ -376,11 +392,13 @@ class ServiceFilter(WorkspaceScopedFilter):
     )
     flavor: Optional[str] = Field(
         default=None,
-        description="Flavor of the service. Use this to filter services by their flavor.",
+        description="Flavor of the service. Use this to filter services by "
+        "their flavor.",
     )
     config: Optional[bytes] = Field(
         default=None,
-        description="Config of the service. Use this to filter services by their config.",
+        description="Config of the service. Use this to filter services by "
+        "their config.",
     )
     pipeline_name: Optional[str] = Field(
         default=None,
@@ -396,12 +414,22 @@ class ServiceFilter(WorkspaceScopedFilter):
     model_version_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="By the model version this service is attached to.",
+        union_mode="left_to_right",
     )
     pipeline_run_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="By the pipeline run this service is attached to.",
+        union_mode="left_to_right",
     )
 
+    # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
+    #  fields defined under base models. If not handled, this raises a warning.
+    #  It is possible to suppress this warning message with the following
+    #  configuration, however the ultimate solution is to rename these fields.
+    #  Even though they do not cause any problems right now, if we are not
+    #  careful we might overwrite some fields protected by pydantic.
+    model_config = ConfigDict(protected_namespaces=())
+
     def set_type(self, type: str) -> None:
         """Set the type of the service.
 
@@ -441,7 +469,7 @@ class ServiceFilter(WorkspaceScopedFilter):
 
     def generate_filter(
         self, table: Type["SQLModel"]
-    ) -> Union["BinaryExpression[Any]", "BooleanClauseList[Any]"]:
+    ) -> Union["ColumnElement[bool]"]:
         """Generate the filter for the query.
 
         Services can be scoped by type to narrow the search.
@@ -452,7 +480,7 @@ class ServiceFilter(WorkspaceScopedFilter):
         Returns:
             The filter expression for the query.
         """
-        from sqlalchemy import and_
+        from sqlmodel import and_
 
         base_filter = super().generate_filter(table)
 

zenml/models/v2/core/service_account.py

@@ -15,7 +15,7 @@
 
 from typing import TYPE_CHECKING, ClassVar, List, Optional, Type, Union
 
-from pydantic import Field
+from pydantic import ConfigDict, Field
 
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
 from zenml.models.v2.base.base import (
@@ -24,9 +24,9 @@ from zenml.models.v2.base.base import (
     BaseRequest,
     BaseResponseMetadata,
     BaseResponseResources,
+    BaseUpdate,
 )
 from zenml.models.v2.base.filter import AnyQuery, BaseFilter
-from zenml.models.v2.base.update import update_model
 
 if TYPE_CHECKING:
     from zenml.models.v2.base.filter import AnySchema
@@ -52,22 +52,34 @@ class ServiceAccountRequest(BaseRequest):
         max_length=TEXT_FIELD_MAX_LENGTH,
     )
     active: bool = Field(title="Whether the service account is active or not.")
-
-    class Config:
-        """Pydantic configuration class."""
-
-        # Validate attributes when assigning them
-        validate_assignment = True
-        extra = "ignore"
+    model_config = ConfigDict(validate_assignment=True, extra="ignore")
 
 
 # ------------------ Update Model ------------------
 
 
-@update_model
-class ServiceAccountUpdate(ServiceAccountRequest):
+class ServiceAccountUpdate(BaseUpdate):
     """Update model for service accounts."""
 
+    ANALYTICS_FIELDS: ClassVar[List[str]] = ["name", "active"]
+
+    name: Optional[str] = Field(
+        title="The unique name for the service account.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    description: Optional[str] = Field(
+        title="A description of the service account.",
+        max_length=TEXT_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    active: Optional[bool] = Field(
+        title="Whether the service account is active or not.",
+        default=None,
+    )
+
+    model_config = ConfigDict(validate_assignment=True)
+
 
 # ------------------ Response Model ------------------
 
@@ -189,6 +201,7 @@ class ServiceAccountFilter(BaseFilter):
     active: Optional[Union[bool, str]] = Field(
         default=None,
         description="Whether the user is active",
+        union_mode="left_to_right",
     )
 
     def apply_filter(