WIP: ⬆️ Upgrade Director v2 service (Pydantic v2) (ITISFoundation#6619)

Co-authored-by: Andrei Neagu <[email protected]> Co-authored-by: Andrei Neagu <[email protected]> Co-authored-by: Sylvain <[email protected]>
sanderegg · Nov 14, 2024 · f335bca · f335bca
1 parent a00d55f
commit f335bca
Show file tree

Hide file tree

Showing 122 changed files with 1,649 additions and 991 deletions.
diff --git a/packages/aws-library/src/aws_library/s3/_client.py b/packages/aws-library/src/aws_library/s3/_client.py
@@ -73,7 +73,7 @@ async def create(
         session = aioboto3.Session()
         session_client = session.client(
             "s3",
-            endpoint_url=settings.S3_ENDPOINT,
+            endpoint_url=f"{settings.S3_ENDPOINT}",
             aws_access_key_id=settings.S3_ACCESS_KEY,
             aws_secret_access_key=settings.S3_SECRET_KEY,
             region_name=settings.S3_REGION,

diff --git a/packages/common-library/src/common_library/pydantic_validators.py b/packages/common-library/src/common_library/pydantic_validators.py
@@ -1,12 +1,11 @@
-import datetime
+import datetime as dt
 import re
 import warnings
-from datetime import timedelta
 
 from pydantic import TypeAdapter, field_validator
 
 
-def _validate_legacy_timedelta_str(time_str: str | timedelta) -> str | timedelta:
+def _validate_legacy_timedelta_str(time_str: str | dt.timedelta) -> str | dt.timedelta:
     if not isinstance(time_str, str):
         return time_str
 
@@ -34,14 +33,14 @@ def validate_numeric_string_as_timedelta(field: str):
     """Transforms a float/int number into a valid datetime as it used to work in the past"""
 
     def _numeric_string_as_timedelta(
-        v: datetime.timedelta | str | float,
-    ) -> datetime.timedelta | str | float:
+        v: dt.timedelta | str | float,
+    ) -> dt.timedelta | str | float:
         if isinstance(v, str):
             try:
                 converted_value = float(v)
 
-                iso8601_format = TypeAdapter(timedelta).dump_python(
-                    timedelta(seconds=converted_value), mode="json"
+                iso8601_format = TypeAdapter(dt.timedelta).dump_python(
+                    dt.timedelta(seconds=converted_value), mode="json"
                 )
                 warnings.warn(
                     f"{field}='{v}' -should be set to-> {field}='{iso8601_format}' (ISO8601 datetime format). "

diff --git a/packages/common-library/src/common_library/serialization.py b/packages/common-library/src/common_library/serialization.py
@@ -21,10 +21,9 @@ def model_dump_with_secrets(
             data[field_name] = field_data.total_seconds()
 
         elif isinstance(field_data, SecretStr):
-            if show_secrets:
-                data[field_name] = field_data.get_secret_value()
-            else:
-                data[field_name] = str(field_data)
+            data[field_name] = (
+                field_data.get_secret_value() if show_secrets else str(field_data)
+            )
 
         elif isinstance(field_data, Url):
             data[field_name] = str(field_data)

diff --git a/packages/common-library/tests/test_serialization.py b/packages/common-library/tests/test_serialization.py
@@ -4,22 +4,31 @@
 
 
 class Credentials(BaseModel):
-    USERNAME: str | None = None
-    PASSWORD: SecretStr | None = None
+    username: str
+    password: SecretStr
+
+
+class Access(BaseModel):
+    credentials: Credentials
 
 
 @pytest.mark.parametrize(
     "expected,show_secrets",
     [
         (
-            {"USERNAME": "DeepThought", "PASSWORD": "42"},
+            {"credentials": {"username": "DeepThought", "password": "42"}},
             True,
         ),
         (
-            {"USERNAME": "DeepThought", "PASSWORD": "**********"},
+            {"credentials": {"username": "DeepThought", "password": "**********"}},
             False,  # hide secrets
         ),
     ],
 )
 def test_model_dump_with_secrets(expected: dict, show_secrets: bool):
-    assert expected == model_dump_with_secrets(Credentials(USERNAME="DeepThought", PASSWORD=SecretStr("42")), show_secrets=show_secrets)
+    assert expected == model_dump_with_secrets(
+        Access(
+            credentials=Credentials(username="DeepThought", password=SecretStr("42"))
+        ),
+        show_secrets=show_secrets,
+    )
diff --git a/packages/models-library/src/models_library/aiodocker_api.py b/packages/models-library/src/models_library/aiodocker_api.py
@@ -11,12 +11,13 @@
 
 
 class AioDockerContainerSpec(ContainerSpec):
-    Env: dict[str, str | None] | None = Field(
+    env: dict[str, str | None] | None = Field(  # type: ignore[assignment]
         default=None,
-        description="aiodocker expects here a dictionary and re-convert it back internally`.\n",
+        alias="Env",
+        description="aiodocker expects here a dictionary and re-convert it back internally",
     )
 
-    @field_validator("Env", mode="before")
+    @field_validator("env", mode="before")
     @classmethod
     def convert_list_to_dict(cls, v):
         if v is not None and isinstance(v, list):
@@ -33,25 +34,22 @@ def convert_list_to_dict(cls, v):
 class AioDockerResources1(Resources1):
     # NOTE: The Docker REST API documentation is wrong!!!
     # Do not set that back to singular Reservation.
-    Reservation: ResourceObject | None = Field(
+    reservation: ResourceObject | None = Field(
         None, description="Define resources reservation.", alias="Reservations"
     )
 
     model_config = ConfigDict(populate_by_name=True)
 
 
 class AioDockerTaskSpec(TaskSpec):
-    ContainerSpec: AioDockerContainerSpec | None = Field(
-        None,
+    container_spec: AioDockerContainerSpec | None = Field(
+        default=None, alias="ContainerSpec"
     )
 
-    Resources: AioDockerResources1 | None = Field(
-        None,
-        description="Resource requirements which apply to each individual container created\nas part of the service.\n",
-    )
+    resources: AioDockerResources1 | None = Field(default=None, alias="Resources")
 
 
 class AioDockerServiceSpec(ServiceSpec):
-    TaskTemplate: AioDockerTaskSpec | None = None
+    task_template: AioDockerTaskSpec | None = Field(default=None, alias="TaskTemplate")
 
     model_config = ConfigDict(populate_by_name=True, alias_generator=camel_to_snake)
diff --git a/packages/models-library/src/models_library/api_schemas_directorv2/clusters.py b/packages/models-library/src/models_library/api_schemas_directorv2/clusters.py
@@ -1,4 +1,4 @@
-from typing import TypeAlias
+from typing import Any, TypeAlias
 
 from pydantic import (
     AnyHttpUrl,
@@ -48,13 +48,12 @@ class WorkerMetrics(BaseModel):
 class UsedResources(DictModel[str, NonNegativeFloat]):
     @model_validator(mode="before")
     @classmethod
-    def ensure_negative_value_is_zero(cls, values):
+    def ensure_negative_value_is_zero(cls, values: dict[str, Any]):
         # dasks adds/remove resource values and sometimes
         # they end up being negative instead of 0
-        if v := values.get("__root__", {}):
-            for res_key, res_value in v.items():
-                if res_value < 0:
-                    v[res_key] = 0
+        for res_key, res_value in values.items():
+            if res_value < 0:
+                values[res_key] = 0
         return values
 
 

diff --git a/packages/models-library/src/models_library/api_schemas_directorv2/comp_tasks.py b/packages/models-library/src/models_library/api_schemas_directorv2/comp_tasks.py
@@ -1,7 +1,15 @@
 from typing import Any, TypeAlias
 
 from models_library.basic_types import IDStr
-from pydantic import AnyHttpUrl, AnyUrl, BaseModel, ConfigDict, Field, field_validator
+from pydantic import (
+    AnyHttpUrl,
+    AnyUrl,
+    BaseModel,
+    ConfigDict,
+    Field,
+    ValidationInfo,
+    field_validator,
+)
 
 from ..clusters import ClusterID
 from ..projects import ProjectID
@@ -63,16 +71,18 @@ class ComputationCreate(BaseModel):
 
     @field_validator("product_name")
     @classmethod
-    def ensure_product_name_defined_if_computation_starts(cls, v, values):
-        if "start_pipeline" in values and values["start_pipeline"] and v is None:
+    def _ensure_product_name_defined_if_computation_starts(
+        cls, v, info: ValidationInfo
+    ):
+        if info.data.get("start_pipeline") and v is None:
             msg = "product_name must be set if computation shall start!"
             raise ValueError(msg)
         return v
 
     @field_validator("use_on_demand_clusters")
     @classmethod
-    def ensure_expected_options(cls, v, values):
-        if v is True and ("cluster_id" in values and values["cluster_id"] is not None):
+    def _ensure_expected_options(cls, v, info: ValidationInfo):
+        if v and info.data.get("cluster_id") is not None:
             msg = "cluster_id cannot be set if use_on_demand_clusters is set"
             raise ValueError(msg)
         return v

diff --git a/packages/models-library/src/models_library/clusters.py b/packages/models-library/src/models_library/clusters.py
@@ -2,6 +2,7 @@
 from pathlib import Path
 from typing import Final, Literal, TypeAlias
 
+from models_library.utils._original_fastapi_encoders import jsonable_encoder
 from pydantic import (
     AnyUrl,
     BaseModel,
@@ -224,18 +225,24 @@ class Cluster(BaseCluster):
     @model_validator(mode="before")
     @classmethod
     def check_owner_has_access_rights(cls, values):
+        values = jsonable_encoder(values)
+
         is_default_cluster = bool(values["id"] == DEFAULT_CLUSTER_ID)
         owner_gid = values["owner"]
 
         # check owner is in the access rights, if not add it
         access_rights = values.get("access_rights", values.get("accessRights", {}))
         if owner_gid not in access_rights:
             access_rights[owner_gid] = (
-                CLUSTER_USER_RIGHTS if is_default_cluster else CLUSTER_ADMIN_RIGHTS
+                CLUSTER_USER_RIGHTS.model_dump()
+                if is_default_cluster
+                else CLUSTER_ADMIN_RIGHTS.model_dump()
             )
         # check owner has the expected access
         if access_rights[owner_gid] != (
-            CLUSTER_USER_RIGHTS if is_default_cluster else CLUSTER_ADMIN_RIGHTS
+            CLUSTER_USER_RIGHTS.model_dump()
+            if is_default_cluster
+            else CLUSTER_ADMIN_RIGHTS.model_dump()
         ):
             msg = f"the cluster owner access rights are incorrectly set: {access_rights[owner_gid]}"
             raise ValueError(msg)

diff --git a/packages/models-library/src/models_library/docker.py b/packages/models-library/src/models_library/docker.py
@@ -37,13 +37,7 @@ def from_key(cls, key: str) -> "DockerLabelKey":
     str, StringConstraints(pattern=DOCKER_GENERIC_TAG_KEY_RE)
 ]
 
-
-class DockerPlacementConstraint(ConstrainedStr):
-    strip_whitespace = True
-    regex = re.compile(
-        r"^(?!-)(?![.])(?!.*--)(?!.*[.][.])[a-zA-Z0-9.-]*(?<!-)(?<![.])(!=|==)[a-zA-Z0-9_. -]*$"
-    )
-
+DockerPlacementConstraint: TypeAlias = Annotated[str, StringConstraints(strip_whitespace = True, pattern = re.compile(r"^(?!-)(?![.])(?!.*--)(?!.*[.][.])[a-zA-Z0-9.-]*(?<!-)(?<![.])(!=|==)[a-zA-Z0-9_. -]*$"))]
 
 _SIMCORE_RUNTIME_DOCKER_LABEL_PREFIX: Final[str] = "io.simcore.runtime."
 _BACKWARDS_COMPATIBILITY_SIMCORE_RUNTIME_DOCKER_LABELS_MAP: Final[dict[str, str]] = {

diff --git a/packages/models-library/src/models_library/errors.py b/packages/models-library/src/models_library/errors.py
@@ -1,4 +1,8 @@
-from typing import Any, TypedDict
+from typing import Any
+
+from typing_extensions import (  # https://docs.pydantic.dev/latest/api/standard_library_types/#typeddict
+    TypedDict,
+)
 
 Loc = tuple[int | str, ...]