chore(ci): Introduce pydantic mypy plugin (#3212)

Backported-from: main (24.12) Backported-to: 24.03 Backport-of: 3212
lablup · Dec 6, 2024 · c916ce0 · c916ce0
1 parent 4cbdc0f
commit c916ce0
Show file tree

Hide file tree

Showing 7 changed files with 176 additions and 23 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -101,10 +101,15 @@ line-length = 100
 preview = true
 
 [tool.mypy]
+plugins = ["pydantic.mypy"]
 ignore_missing_imports = true
-implicit_optional = true  # FIXME: remove after adding https://github.com/hauntsaninja/no_implicit_optional to fmt
 mypy_path = "stubs:src:tools/pants-plugins"
 namespace_packages = true
 explicit_package_bases = true
 python_executable = "dist/export/python/virtualenvs/python-default/3.12.6/bin/python"
 disable_error_code = ["typeddict-unknown-key"]
+
+[tool.pydantic-mypy]
+init_forbid_extra = true
+init_typed = true
+warn_required_dynamic_aliases = true
diff --git a/src/ai/backend/client/cli/pretty.py b/src/ai/backend/client/cli/pretty.py
@@ -1,6 +1,7 @@
 import asyncio
 import enum
 import functools
+import json
 import sys
 import textwrap
 import traceback
@@ -135,12 +136,10 @@ def format_error(exc: Exception):
             general_error_msg = exc.data.get("msg", None)
             if general_error_msg is not None:
                 yield f"\n- {general_error_msg}"
-            per_field_errors = exc.data.get("data", {})
-            if isinstance(per_field_errors, dict):
-                for k, v in per_field_errors.items():
-                    yield f'\n- "{k}": {v}'
-            else:
-                yield f"\n- {per_field_errors}"
+            per_field_errors = exc.data.get("data", None)
+            if per_field_errors:
+                yield "\n"
+                yield json.dumps(per_field_errors, indent=2)
         else:
             if exc.data["type"].endswith("/graphql-error"):
                 yield "\n\u279c Message:\n"

diff --git a/src/ai/backend/manager/api/service.py b/src/ai/backend/manager/api/service.py
@@ -5,7 +5,15 @@
 import uuid
 from dataclasses import dataclass
 from datetime import datetime
-from typing import TYPE_CHECKING, Annotated, Any, Iterable, Sequence, Tuple
+from typing import (
+    TYPE_CHECKING,
+    Annotated,
+    Any,
+    Iterable,
+    Self,
+    Sequence,
+    Tuple,
+)
 
 import aiohttp
 import aiohttp_cors
@@ -22,6 +30,7 @@
     HttpUrl,
     NonNegativeFloat,
     NonNegativeInt,
+    model_validator,
 )
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import selectinload
@@ -1018,11 +1027,31 @@ async def delete_route(request: web.Request) -> SuccessResponseModel:
 
 
 class TokenRequestModel(BaseModel):
-    duration: tv.TimeDuration = Field(default=None, description="duration of the token.")
+    duration: tv.TimeDuration | None = Field(
+        default=None, description="The lifetime duration of the token."
+    )
     valid_until: int | None = Field(
-        default=None, description="Absolute token expiry date, expressed in Unix epoch format."
+        default=None,
+        description="The absolute token expiry date expressed in the Unix epoch format.",
+    )
+    expires_at: int = Field(
+        default=-1,
+        description="The expiration timestamp computed from duration or valid_until.",
     )
 
+    @model_validator(mode="after")
+    def check_lifetime(self) -> Self:
+        now = datetime.now()
+        if self.valid_until is not None:
+            self.expires_at = self.valid_until
+        elif self.duration is not None:
+            self.expires_at = int((now + self.duration).timestamp())
+        else:
+            raise ValueError("Either valid_until or duration must be specified.")
+        if now.timestamp() > self.expires_at:
+            raise ValueError("The expiration time cannot be in the past.")
+        return self
+
 
 class TokenResponseModel(BaseResponseModel):
     token: str
@@ -1066,15 +1095,7 @@ async def generate_token(request: web.Request, params: TokenRequestModel) -> Tok
 
     await get_user_uuid_scopes(request, {"owner_uuid": endpoint.session_owner})
 
-    if params.valid_until:
-        exp = params.valid_until
-    elif params.duration:
-        exp = int((datetime.now() + params.duration).timestamp())
-    else:
-        raise InvalidAPIParameters("valid_until and duration can't be both unspecified")
-    if datetime.now().timestamp() > exp:
-        raise InvalidAPIParameters("valid_until is older than now")
-    body = {"user_uuid": str(endpoint.session_owner), "exp": exp}
+    body = {"user_uuid": str(endpoint.session_owner), "exp": params.expires_at}
     async with aiohttp.ClientSession() as session:
         async with session.post(
             f"{wsproxy_addr}/v2/endpoints/{endpoint.id}/token",

diff --git a/src/ai/backend/manager/api/utils.py b/src/ai/backend/manager/api/utils.py
@@ -298,8 +298,22 @@ async def wrapped(
                     kwargs["query"] = query_params
             except (json.decoder.JSONDecodeError, yaml.YAMLError, yaml.MarkedYAMLError):
                 raise InvalidAPIParameters("Malformed body")
-            except ValidationError as e:
-                raise InvalidAPIParameters("Input validation error", extra_data=e.errors())
+            except ValidationError as ex:
+                first_error = ex.errors()[0]
+                # Format the first validation error as the message
+                # The client may refer extra_data to access the full validation errors.
+                metadata = {
+                    "input": first_error["input"],
+                }
+                if loc := first_error["loc"]:
+                    metadata["loc"] = loc[0]
+                metadata_formatted_items = [
+                    f"type={first_error["type"]}",  # format as symbol
+                    *(f"{k}={v!r}" for k, v in metadata.items()),
+                ]
+                msg = f"{first_error["msg"]} [{", ".join(metadata_formatted_items)}]"
+                # To reuse the json serialization provided by pydantic, we call ex.json() and re-parse it.
+                raise InvalidAPIParameters(msg, extra_data=json.loads(ex.json()))
             result = await handler(request, checked_params, *args, **kwargs)
             return ensure_stream_response_type(result)
 

diff --git a/src/ai/backend/manager/server.py b/src/ai/backend/manager/server.py
@@ -244,6 +244,7 @@ async def exception_middleware(
     try:
         await stats_monitor.report_metric(INCREMENT, "ai.backend.manager.api.requests")
         resp = await handler(request)
+    # NOTE: pydantic.ValidationError is handled in utils.pydantic_params_api_handler()
     except InvalidArgument as ex:
         if len(ex.args) > 1:
             raise InvalidAPIParameters(f"{ex.args[0]}: {', '.join(map(str, ex.args[1:]))}")

diff --git a/tools/mypy-requirements.txt b/tools/mypy-requirements.txt
@@ -1 +1,2 @@
 mypy==1.13.0
+pydantic~=2.9.2
diff --git a/tools/mypy.lock b/tools/mypy.lock
@@ -9,7 +9,8 @@
 //     "CPython==3.12.6"
 //   ],
 //   "generated_with_requirements": [
-//     "mypy==1.13.0"
+//     "mypy==1.13.0",
+//     "pydantic~=2.9.2"
 //   ],
 //   "manylinux": "manylinux2014",
 //   "requirement_constraints": [],
@@ -28,6 +29,26 @@
   "locked_resolves": [
     {
       "locked_requirements": [
+        {
+          "artifacts": [
+            {
+              "algorithm": "sha256",
+              "hash": "1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53",
+              "url": "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "aff07c09a53a08bc8cfccb9c85b05f1aa9a2a6f23728d790723543408344ce89",
+              "url": "https://files.pythonhosted.org/packages/ee/67/531ea369ba64dcff5ec9c3402f9f51bf748cec26dde048a2f973a4eea7f5/annotated_types-0.7.0.tar.gz"
+            }
+          ],
+          "project_name": "annotated-types",
+          "requires_dists": [
+            "typing-extensions>=4.0.0; python_version < \"3.9\""
+          ],
+          "requires_python": ">=3.8",
+          "version": "0.7.0"
+        },
         {
           "artifacts": [
             {
@@ -103,6 +124,96 @@
           "requires_python": ">=3.5",
           "version": "1.0.0"
         },
+        {
+          "artifacts": [
+            {
+              "algorithm": "sha256",
+              "hash": "f048cec7b26778210e28a0459867920654d48e5e62db0958433636cde4254f12",
+              "url": "https://files.pythonhosted.org/packages/df/e4/ba44652d562cbf0bf320e0f3810206149c8a4e99cdbf66da82e97ab53a15/pydantic-2.9.2-py3-none-any.whl"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "d155cef71265d1e9807ed1c32b4c8deec042a44a50a4188b25ac67ecd81a9c0f",
+              "url": "https://files.pythonhosted.org/packages/a9/b7/d9e3f12af310e1120c21603644a1cd86f59060e040ec5c3a80b8f05fae30/pydantic-2.9.2.tar.gz"
+            }
+          ],
+          "project_name": "pydantic",
+          "requires_dists": [
+            "annotated-types>=0.6.0",
+            "email-validator>=2.0.0; extra == \"email\"",
+            "pydantic-core==2.23.4",
+            "typing-extensions>=4.12.2; python_version >= \"3.13\"",
+            "typing-extensions>=4.6.1; python_version < \"3.13\"",
+            "tzdata; (python_version >= \"3.9\" and sys_platform == \"win32\") and extra == \"timezone\""
+          ],
+          "requires_python": ">=3.8",
+          "version": "2.9.2"
+        },
+        {
+          "artifacts": [
+            {
+              "algorithm": "sha256",
+              "hash": "9261d3ce84fa1d38ed649c3638feefeae23d32ba9182963e465d58d62203bd24",
+              "url": "https://files.pythonhosted.org/packages/e3/b9/41f7efe80f6ce2ed3ee3c2dcfe10ab7adc1172f778cc9659509a79518c43/pydantic_core-2.23.4-cp312-cp312-musllinux_1_1_x86_64.whl"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "128585782e5bfa515c590ccee4b727fb76925dd04a98864182b22e89a4e6ed36",
+              "url": "https://files.pythonhosted.org/packages/06/c8/7d4b708f8d05a5cbfda3243aad468052c6e99de7d0937c9146c24d9f12e9/pydantic_core-2.23.4-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "f69a8e0b033b747bb3e36a44e7732f0c99f7edd5cea723d45bc0d6e95377ffee",
+              "url": "https://files.pythonhosted.org/packages/14/de/866bdce10ed808323d437612aca1ec9971b981e1c52e5e42ad9b8e17a6f6/pydantic_core-2.23.4-cp312-cp312-macosx_11_0_arm64.whl"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "33e3d65a85a2a4a0dc3b092b938a4062b1a05f3a9abde65ea93b233bca0e03f2",
+              "url": "https://files.pythonhosted.org/packages/2d/ad/b5f0fe9e6cfee915dd144edbd10b6e9c9c9c9d7a56b69256d124b8ac682e/pydantic_core-2.23.4-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "f3e0da4ebaef65158d4dfd7d3678aad692f7666877df0002b8a522cdf088f231",
+              "url": "https://files.pythonhosted.org/packages/74/7b/8e315f80666194b354966ec84b7d567da77ad927ed6323db4006cf915f3f/pydantic_core-2.23.4-cp312-cp312-macosx_10_12_x86_64.whl"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "bb2802e667b7051a1bebbfe93684841cc9351004e2badbd6411bf357ab8d5ac8",
+              "url": "https://files.pythonhosted.org/packages/80/33/9c24334e3af796ce80d2274940aae38dd4e5676298b4398eff103a79e02d/pydantic_core-2.23.4-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "68665f4c17edcceecc112dfed5dbe6f92261fb9d6054b47d01bf6371a6196126",
+              "url": "https://files.pythonhosted.org/packages/89/4d/3079d00c47f22c9a9a8220db088b309ad6e600a73d7a69473e3a8e5e3ea3/pydantic_core-2.23.4-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "d18ca8148bebe1b0a382a27a8ee60350091a6ddaf475fa05ef50dc35b5df6327",
+              "url": "https://files.pythonhosted.org/packages/a5/6f/e9567fd90104b79b101ca9d120219644d3314962caa7948dd8b965e9f83e/pydantic_core-2.23.4-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "723314c1d51722ab28bfcd5240d858512ffd3116449c557a1336cbe3919beb87",
+              "url": "https://files.pythonhosted.org/packages/dc/69/8edd5c3cd48bb833a3f7ef9b81d7666ccddd3c9a635225214e044b6e8281/pydantic_core-2.23.4-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "2584f7cf844ac4d970fba483a717dbe10c1c1c96a969bf65d61ffe94df1b2863",
+              "url": "https://files.pythonhosted.org/packages/e2/aa/6b6a9b9f8537b872f552ddd46dd3da230367754b6f707b8e1e963f515ea3/pydantic_core-2.23.4.tar.gz"
+            },
+            {
+              "algorithm": "sha256",
+              "hash": "20152074317d9bed6b7a95ade3b7d6054845d70584216160860425f4fbd5ee9e",
+              "url": "https://files.pythonhosted.org/packages/e9/88/9df5b7ce880a4703fcc2d76c8c2d8eb9f861f79d0c56f4b8f5f2607ccec8/pydantic_core-2.23.4-cp312-cp312-musllinux_1_1_aarch64.whl"
+            }
+          ],
+          "project_name": "pydantic-core",
+          "requires_dists": [
+            "typing-extensions!=4.7.0,>=4.6.0"
+          ],
+          "requires_python": ">=3.8",
+          "version": "2.23.4"
+        },
         {
           "artifacts": [
             {
@@ -133,7 +244,8 @@
   "pip_version": "24.1.2",
   "prefer_older_binary": false,
   "requirements": [
-    "mypy==1.13.0"
+    "mypy==1.13.0",
+    "pydantic~=2.9.2"
   ],
   "requires_python": [
     "==3.12.6"