refactor(BA-6003): bulk-migrate all pydantic BaseModel subclasses to BackendAISchema (#11554)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: jopemachine

changes/11554.enhance.md

@@ -0,0 +1 @@
+Migrate every remaining pydantic `BaseModel` subclass across `src/ai/backend/` to `BackendAISchema`, so any `model_validate()` failure auto-converts to a `BackendAISchemaValidationFailed` (HTTP 400) instead of leaking as raw `pydantic.ValidationError`.

src/ai/backend/account_manager/api/utils.py

@@ -13,6 +13,8 @@
 from pydantic import BaseModel, ConfigDict, ValidationError
 
 from ai.backend.account_manager.exceptions import AuthorizationFailed, InvalidAPIParameters
+from ai.backend.common.exception import BackendAISchemaValidationFailed
+from ai.backend.common.types import BackendAISchema
 from ai.backend.logging import BraceStyleAdapter
 
 
@@ -69,7 +71,7 @@ def mask_sensitive_keys(data: Mapping[str, Any]) -> Mapping[str, Any]:
 TBaseModel = TypeVar("TBaseModel", bound=BaseModel)
 
 
-class RequestData(BaseModel):
+class RequestData(BackendAISchema):
     model_config = ConfigDict(
         extra="allow",
     )
@@ -173,8 +175,13 @@ async def wrapped(
                     kwargs["query"] = query_params
             except (json.decoder.JSONDecodeError, yaml.YAMLError, yaml.MarkedYAMLError) as e:
                 raise InvalidAPIParameters("Malformed body") from e
-            except ValidationError as e:
-                raise InvalidAPIParameters("Input validation error", extra_data=e.errors()) from e
+            except (BackendAISchemaValidationFailed, ValidationError) as e:
+                # ``ValidationError`` covers plain ``BaseModel`` subclasses that
+                # skip the ``BackendAISchema`` auto-conversion override.
+                raise InvalidAPIParameters(
+                    "Input validation error",
+                    extra_data={"errors": e.errors()},
+                ) from e
             result = await handler(request, checked_params, *args, **kwargs)
             return ensure_stream_response_type(result)
 

src/ai/backend/account_manager/config.py

@@ -12,7 +12,6 @@
 
 import click
 from pydantic import (
-    BaseModel,
     ConfigDict,
     Field,
     GetCoreSchemaHandler,
@@ -23,6 +22,7 @@
 from pydantic_core import PydanticUndefined, core_schema
 
 from ai.backend.common import config
+from ai.backend.common.types import BackendAISchema
 from ai.backend.logging import LogLevel
 
 from .types import EventLoopType
@@ -46,7 +46,7 @@ class TransactionIsolationLevel(enum.StrEnum):
     SERIALIZABLE = "SERIALIZABLE"
 
 
-class BaseSchema(BaseModel):
+class BaseSchema(BackendAISchema):
     model_config = ConfigDict(
         validate_by_name=True,
         from_attributes=True,

src/ai/backend/agent/kernel_registry/types.py

@@ -3,19 +3,19 @@
 from collections.abc import Mapping
 from typing import TYPE_CHECKING, Self
 
-from pydantic import BaseModel, Field
+from pydantic import Field
 
 from ai.backend.agent.kernel import KernelOwnershipData
 from ai.backend.agent.proxy import DomainSocketPathPair
 from ai.backend.agent.resources import KernelResourceSpec
 from ai.backend.common.docker import ImageRef
-from ai.backend.common.types import AgentId, KernelId, ServicePort, SessionTypes
+from ai.backend.common.types import AgentId, BackendAISchema, KernelId, ServicePort, SessionTypes
 
 if TYPE_CHECKING:
     from ai.backend.agent.docker.kernel import DockerKernel
 
 
-class KernelRecoveryData(BaseModel):
+class KernelRecoveryData(BackendAISchema):
     """
     Data required for recovering a Kernel.
     Agent should load and write Kernel data using this structure

src/ai/backend/agent/proxy.py

@@ -4,9 +4,9 @@
 from typing import Any
 
 import attrs
-from pydantic import BaseModel
 
 from ai.backend.common.asyncio import current_loop
+from ai.backend.common.types import BackendAISchema
 
 
 @attrs.define(auto_attribs=True, slots=True)
@@ -16,7 +16,7 @@ class DomainSocketProxy:
     proxy_server: asyncio.AbstractServer
 
 
-class DomainSocketPathPair(BaseModel):
+class DomainSocketPathPair(BackendAISchema):
     host_sock_path: Path
     host_proxy_path: Path
 

src/ai/backend/agent/scratch/types.py

@@ -1,17 +1,15 @@
 from collections.abc import Mapping
 from typing import Self
 
-from pydantic import BaseModel
-
 from ai.backend.agent.kernel import KernelOwnershipData
 from ai.backend.agent.kernel_registry.types import KernelRecoveryData
 from ai.backend.agent.proxy import DomainSocketPathPair
 from ai.backend.agent.resources import KernelResourceSpec
 from ai.backend.common.docker import ImageRef
-from ai.backend.common.types import AgentId, KernelId, ServicePort, SessionTypes
+from ai.backend.common.types import AgentId, BackendAISchema, KernelId, ServicePort, SessionTypes
 
 
-class KernelRecoveryScratchData(BaseModel):
+class KernelRecoveryScratchData(BackendAISchema):
     """
     Serializable subset of KernelRecoveryData for scratch storage.
     Excludes `resource_spec` and `environ` which are loaded separately.

src/ai/backend/agent/server.py

@@ -44,7 +44,6 @@
 from callosum.ordering import ExitOrderedAsyncScheduler
 from callosum.rpc import Peer, RPCMessage
 from etcd_client import WatchEventType
-from pydantic import ValidationError
 from setproctitle import setproctitle
 from zmq.auth.certs import load_certificate
 
@@ -87,7 +86,7 @@
     KernelTerminatedBroadcastEvent,
 )
 from ai.backend.common.events.event_types.kernel.types import KernelLifecycleEventReason
-from ai.backend.common.exception import ConfigurationError
+from ai.backend.common.exception import BackendAISchemaValidationFailed, ConfigurationError
 from ai.backend.common.health_checker.checkers.etcd import EtcdHealthChecker
 from ai.backend.common.health_checker.checkers.valkey import ValkeyHealthChecker
 from ai.backend.common.health_checker.probe import HealthProbe, HealthProbeOptions
@@ -1740,7 +1739,7 @@ def main(
         if server_config.debug.enabled:
             print("== Agent configuration ==")
             pprint(server_config.model_dump(by_alias=True))
-    except ValidationError as e:
+    except BackendAISchemaValidationFailed as e:
         print(
             "ConfigurationError: Agent local config failed validation checks:",
             file=sys.stderr,

src/ai/backend/appproxy/common/config.py

@@ -17,6 +17,7 @@
 from pydantic_core import PydanticUndefined, core_schema
 
 from ai.backend.common.meta import BackendAIConfigMeta, CompositeType, ConfigExample
+from ai.backend.common.types import BackendAISchema
 
 from .errors import (
     GroupNotFoundError,
@@ -30,7 +31,7 @@
 # FIXME: merge majority of common definitions to ai.backend.common when ready
 
 
-class BaseSchema(BaseModel):
+class BaseSchema(BackendAISchema):
     model_config = ConfigDict(
         populate_by_name=True,
         from_attributes=True,
@@ -723,7 +724,7 @@ class UnsupportedTypeError(RuntimeError):
 
 
 def generate_example_json(
-    schema: type[BaseModel] | types.GenericAlias | types.UnionType,
+    schema: type[BackendAISchema] | types.GenericAlias | types.UnionType,
     parent: list[str] | None = None,
 ) -> dict[str, Any] | list[Any]:
     if parent is None:

src/ai/backend/appproxy/common/openapi.py

@@ -8,6 +8,7 @@
 from pydantic import BaseModel
 
 from ai.backend.appproxy.common.types import PydanticResponse
+from ai.backend.common.types import BackendAISchema
 
 from . import __version__
 
@@ -131,7 +132,7 @@ def generate_openapi(
             route_def["description"] = "\n".join(description)
             type_hints = get_type_hints(route.handler)
 
-            def _parse_schema(model_cls: type[BaseModel]) -> dict[str, Any]:
+            def _parse_schema(model_cls: type[BackendAISchema]) -> dict[str, Any]:
                 if not issubclass(model_cls, BaseModel):
                     raise RuntimeError(f"{model_cls} not considered as a valid response type")
 

src/ai/backend/appproxy/common/types.py

@@ -16,7 +16,7 @@
 from aiohttp import web
 from pydantic import AliasChoices, AnyUrl, BaseModel, ConfigDict, Field
 
-from ai.backend.common.types import ModelServiceStatus, RuntimeVariant
+from ai.backend.common.types import BackendAISchema, ModelServiceStatus, RuntimeVariant
 
 # FIXME: merge majority of common definitions to ai.backend.common when ready
 
@@ -84,7 +84,7 @@ class DigestModType(enum.StrEnum):
 ]
 
 
-class RouteInfo(BaseModel):
+class RouteInfo(BackendAISchema):
     """Information about a route within a circuit.
 
     Routes describe a kernel endpoint (``kernel_host`` + ``kernel_port``)
@@ -157,7 +157,7 @@ def current_kernel_host(self) -> str:
         return self.kernel_host or "localhost"
 
 
-class SerializableCircuit(BaseModel):
+class SerializableCircuit(BackendAISchema):
     """
     Serializable representation of `ai.backend.appproxy.coordinator.models.Circuit`
     """
@@ -241,7 +241,7 @@ def traefik_router_name(self) -> str:
         return f"bai_router_{self.id}@etcd"
 
 
-class SerializableToken(BaseModel):
+class SerializableToken(BackendAISchema):
     login_session_token: str | None
     kernel_host: str
     kernel_port: int
@@ -252,7 +252,7 @@ class SerializableToken(BaseModel):
     domain_name: str
 
 
-class SessionConfig(BaseModel):
+class SessionConfig(BackendAISchema):
     model_config = ConfigDict(populate_by_name=True)
 
     id: Annotated[UUID | None, Field(default=None)]
@@ -265,13 +265,13 @@ class SessionConfig(BaseModel):
     domain_name: str
 
 
-class EndpointConfig(BaseModel):
+class EndpointConfig(BackendAISchema):
     id: UUID
     runtime_variant: Annotated[RuntimeVariant | None, Field(default=None)]
     existing_url: AnyUrl | None
 
 
-class HealthCheckState(BaseModel):
+class HealthCheckState(BackendAISchema):
     """
     Runtime health check state
     """