refactor(BA-575): Apply pydantic handling decorator to VFolder APIs i…

…n storage-proxy (#3562)

src/ai/backend/storage/api/vfolder/manager_handler.py

@@ -1,156 +1,150 @@
 from typing import Protocol
 
-from aiohttp import web
-
+from ai.backend.common.pydantic_handlers import ApiResponse, BodyParam, pydantic_api_handler
 from ai.backend.storage.api.vfolder.response_model import (
-    GetVolumeResponseModel,
-    NoContentResponseModel,
-    ProcessingResponseModel,
-    QuotaScopeResponseModel,
-    VFolderMetadataResponseModel,
-    VolumeMetadataResponseModel,
+    GetVolumeResponse,
+    QuotaScopeResponse,
+    VFolderMetadataResponse,
+    VolumeMetadataResponse,
 )
 from ai.backend.storage.api.vfolder.types import (
-    QuotaScopeIDModel,
-    QuotaScopeMetadataModel,
-    VFolderIDModel,
-    VFolderMetadataModel,
-    VolumeIDModel,
-    VolumeMetadataListModel,
+    QuotaScopeIdData,
+    QuotaScopeMetadata,
+    VFolderIdData,
+    VFolderMetadata,
+    VolumeIdData,
+    VolumeMetadataList,
 )
 
 
 class VFolderServiceProtocol(Protocol):
-    async def get_volume(self, volume_data: VolumeIDModel) -> VolumeMetadataListModel:
+    async def get_volume(self, volume_data: VolumeIdData) -> VolumeMetadataList:
         """by volume_id"""
         ...
 
-    async def get_volumes(self) -> VolumeMetadataListModel: ...
+    async def get_volumes(self) -> VolumeMetadataList: ...
 
-    async def create_quota_scope(self, quota_data: QuotaScopeIDModel) -> None: ...
+    async def create_quota_scope(self, quota_data: QuotaScopeIdData) -> None: ...
 
-    async def get_quota_scope(self, quota_data: QuotaScopeIDModel) -> QuotaScopeMetadataModel: ...
+    async def get_quota_scope(self, quota_data: QuotaScopeIdData) -> QuotaScopeMetadata: ...
 
-    async def update_quota_scope(self, quota_data: QuotaScopeIDModel) -> None: ...
+    async def update_quota_scope(self, quota_data: QuotaScopeIdData) -> None: ...
 
-    async def delete_quota_scope(self, quota_data: QuotaScopeIDModel) -> None:
+    async def delete_quota_scope(self, quota_data: QuotaScopeIdData) -> None:
         """Previous: unset_quota"""
         ...
 
-    async def create_vfolder(self, vfolder_data: VFolderIDModel) -> VFolderIDModel: ...
+    async def create_vfolder(self, vfolder_data: VFolderIdData) -> VFolderIdData: ...
 
-    async def clone_vfolder(self, vfolder_data: VFolderIDModel) -> None: ...
+    async def clone_vfolder(self, vfolder_data: VFolderIdData) -> None: ...
 
-    async def get_vfolder_info(self, vfolder_data: VFolderIDModel) -> VFolderMetadataModel:
+    async def get_vfolder_info(self, vfolder_data: VFolderIdData) -> VFolderMetadata:
         # Integration: vfolder_mount, metadata, vfolder_usage, vfolder_used_bytes, vfolder_fs_usage
         ...
 
-    async def delete_vfolder(self, vfolder_data: VFolderIDModel) -> VFolderIDModel: ...
+    async def delete_vfolder(self, vfolder_data: VFolderIdData) -> VFolderIdData: ...
 
 
 class VFolderHandler:
     def __init__(self, storage_service: VFolderServiceProtocol) -> None:
         self.storage_service = storage_service
 
-    async def get_volume(self, request: web.Request) -> GetVolumeResponseModel:
-        data = await request.json()
-        params = VolumeIDModel(volume_id=data["volume_id"])
-        volume_data = await self.storage_service.get_volume(params)
-        return GetVolumeResponseModel(
-            volumes=[
-                VolumeMetadataResponseModel(
-                    volume_id=str(volume.volume_id),
-                    backend=str(volume.backend),
-                    path=str(volume.path),
-                    fsprefix=str(volume.fsprefix) if volume.fsprefix else None,
-                    capabilities=[str(cap) for cap in volume.capabilities],
-                )
-                for volume in volume_data.volumes
-            ]
+    @pydantic_api_handler
+    async def get_volume(self, body: BodyParam[VolumeIdData]) -> ApiResponse:
+        volume_params = body.parsed
+        volume_data = await self.storage_service.get_volume(volume_params)
+        return ApiResponse.build(
+            status_code=200,
+            response_model=GetVolumeResponse(
+                volumes=[
+                    VolumeMetadataResponse(
+                        volume_id=str(volume.volume_id),
+                        backend=str(volume.backend),
+                        path=str(volume.path),
+                        fsprefix=str(volume.fsprefix) if volume.fsprefix else None,
+                        capabilities=[str(cap) for cap in volume.capabilities],
+                    )
+                    for volume in volume_data.volumes
+                ]
+            ),
         )
 
-    async def get_volumes(self, request: web.Request) -> GetVolumeResponseModel:
+    @pydantic_api_handler
+    async def get_volumes(self) -> ApiResponse:
         volumes_data = await self.storage_service.get_volumes()
-        return GetVolumeResponseModel(
-            volumes=[
-                VolumeMetadataResponseModel(
-                    volume_id=str(volume.volume_id),
-                    backend=str(volume.backend),
-                    path=str(volume.path),
-                    fsprefix=str(volume.fsprefix) if volume.fsprefix else None,
-                    capabilities=[str(cap) for cap in volume.capabilities],
-                )
-                for volume in volumes_data.volumes
-            ]
-        )
-
-    async def create_quota_scope(self, request: web.Request) -> NoContentResponseModel:
-        data = await request.json()
-        params = QuotaScopeIDModel(
-            volume_id=data["volume_id"],
-            quota_scope_id=data["quota_scope_id"],
-            options=data.get("options"),
-        )
-        await self.storage_service.create_quota_scope(params)
-        return NoContentResponseModel()
-
-    async def get_quota_scope(self, request: web.Request) -> QuotaScopeResponseModel:
-        data = await request.json()
-        params = QuotaScopeIDModel(
-            volume_id=data["volume_id"], quota_scope_id=data["quota_scope_id"]
-        )
-        quota_scope = await self.storage_service.get_quota_scope(params)
-        return QuotaScopeResponseModel(
-            used_bytes=quota_scope.used_bytes, limit_bytes=quota_scope.limit_bytes
+        return ApiResponse.build(
+            status_code=200,
+            response_model=GetVolumeResponse(
+                volumes=[
+                    VolumeMetadataResponse(
+                        volume_id=str(volume.volume_id),
+                        backend=str(volume.backend),
+                        path=str(volume.path),
+                        fsprefix=str(volume.fsprefix) if volume.fsprefix else None,
+                        capabilities=[str(cap) for cap in volume.capabilities],
+                    )
+                    for volume in volumes_data.volumes
+                ]
+            ),
         )
 
-    async def update_quota_scope(self, request: web.Request) -> NoContentResponseModel:
-        data = await request.json()
-        params = QuotaScopeIDModel(
-            volume_id=data["volume_id"],
-            quota_scope_id=data["quota_scope_id"],
-            options=data.get("options"),
+    @pydantic_api_handler
+    async def create_quota_scope(self, body: BodyParam[QuotaScopeIdData]) -> ApiResponse:
+        quota_params = body.parsed
+        await self.storage_service.create_quota_scope(quota_params)
+        return ApiResponse.no_content(status_code=201)
+
+    @pydantic_api_handler
+    async def get_quota_scope(self, body: BodyParam[QuotaScopeIdData]) -> ApiResponse:
+        quota_params = body.parsed
+        quota_scope = await self.storage_service.get_quota_scope(quota_params)
+        return ApiResponse.build(
+            status_code=204,
+            response_model=QuotaScopeResponse(
+                used_bytes=quota_scope.used_bytes, limit_bytes=quota_scope.limit_bytes
+            ),
         )
-        await self.storage_service.update_quota_scope(params)
-        return NoContentResponseModel()
 
-    async def delete_quota_scope(self, request: web.Request) -> NoContentResponseModel:
-        data = await request.json()
-        params = QuotaScopeIDModel(
-            volume_id=data["volume_id"], quota_scope_id=data["quota_scope_id"]
-        )
-        await self.storage_service.delete_quota_scope(params)
-        return NoContentResponseModel()
-
-    async def create_vfolder(self, request: web.Request) -> NoContentResponseModel:
-        data = await request.json()
-        params = VFolderIDModel(volume_id=data["volume_id"], vfolder_id=data["vfolder_id"])
-        await self.storage_service.create_vfolder(params)
-        return NoContentResponseModel()
-
-    async def clone_vfolder(self, request: web.Request) -> NoContentResponseModel:
-        data = await request.json()
-        params = VFolderIDModel(
-            volume_id=data["volume_id"],
-            vfolder_id=data["vfolder_id"],
-            dst_vfolder_id=data["dst_vfolder_id"],
-        )
-        await self.storage_service.clone_vfolder(params)
-        return NoContentResponseModel()
-
-    async def get_vfolder_info(self, request: web.Request) -> VFolderMetadataResponseModel:
-        data = await request.json()
-        params = VFolderIDModel(**data)
-        metadata = await self.storage_service.get_vfolder_info(params)
-        return VFolderMetadataResponseModel(
-            mount_path=str(metadata.mount_path),
-            file_count=metadata.file_count,
-            capacity_bytes=metadata.capacity_bytes,
-            used_bytes=metadata.used_bytes,
+    @pydantic_api_handler
+    async def update_quota_scope(self, body: BodyParam[QuotaScopeIdData]) -> ApiResponse:
+        quota_params = body.parsed
+        await self.storage_service.update_quota_scope(quota_params)
+        return ApiResponse.no_content(status_code=204)
+
+    @pydantic_api_handler
+    async def delete_quota_scope(self, body: BodyParam[QuotaScopeIdData]) -> ApiResponse:
+        quota_params = body.parsed
+        await self.storage_service.delete_quota_scope(quota_params)
+        return ApiResponse.no_content(status_code=204)
+
+    @pydantic_api_handler
+    async def create_vfolder(self, body: BodyParam[VFolderIdData]) -> ApiResponse:
+        vfolder_params = body.parsed
+        await self.storage_service.create_vfolder(vfolder_params)
+        return ApiResponse.no_content(status_code=201)
+
+    @pydantic_api_handler
+    async def clone_vfolder(self, body: BodyParam[VFolderIdData]) -> ApiResponse:
+        vfolder_params = body.parsed
+        await self.storage_service.clone_vfolder(vfolder_params)
+        return ApiResponse.no_content(status_code=204)
+
+    @pydantic_api_handler
+    async def get_vfolder_info(self, body: BodyParam[VFolderIdData]) -> ApiResponse:
+        vfolder_params = body.parsed
+        metadata = await self.storage_service.get_vfolder_info(vfolder_params)
+        return ApiResponse.build(
+            status_code=200,
+            response_model=VFolderMetadataResponse(
+                mount_path=str(metadata.mount_path),
+                file_count=metadata.file_count,
+                capacity_bytes=metadata.capacity_bytes,
+                used_bytes=metadata.used_bytes,
+            ),
         )
 
-    async def delete_vfolder(self, request: web.Request) -> ProcessingResponseModel:
-        data = await request.json()
-        params = VFolderIDModel(volume_id=data["volume_id"], vfolder_id=data["vfolder_id"])
-        await self.storage_service.delete_vfolder(params)
-        return ProcessingResponseModel()
+    @pydantic_api_handler
+    async def delete_vfolder(self, body: BodyParam[VFolderIdData]) -> ApiResponse:
+        vfolder_params = body.parsed
+        await self.storage_service.delete_vfolder(vfolder_params)
+        return ApiResponse.no_content(status_code=202)

src/ai/backend/storage/api/vfolder/response_model.py

@@ -1,55 +1,39 @@
-from dataclasses import dataclass
-from typing import Annotated, List, Optional
+from typing import List, Optional
 
-from pydantic import BaseModel as PydanticBaseModel
 from pydantic import Field
 
+from ai.backend.common.pydantic_handlers import BaseResponseModel
 from ai.backend.common.types import BinarySize
 
 
-class BaseModel(PydanticBaseModel):
-    """Base model for all models in this module"""
+class VolumeMetadataResponse(BaseResponseModel):
+    volume_id: str = Field(..., description="A unique identifier for the volume.")
+    backend: str = Field(..., description="The backend name.")
+    path: str = Field(..., description="The path to the volume.")
+    fsprefix: Optional[str] = Field(default=None, description="The prefix for the filesystem.")
+    capabilities: List[str] = Field(..., description="The capabilities of the volume.")
 
-    model_config = {"arbitrary_types_allowed": True}
 
+class GetVolumeResponse(BaseResponseModel):
+    volumes: List[VolumeMetadataResponse] = Field(..., description="The list of volumes.")
 
-@dataclass
-class ResponseModel:
-    user_model: Optional[BaseModel] = None
-    status: Annotated[int, Field(strict=True, exclude=True, ge=100, lt=600)] = 200
 
+class QuotaScopeResponse(BaseResponseModel):
+    used_bytes: Optional[int] = Field(
+        default=0, description="The number of bytes currently used within the quota scope."
+    )
+    limit_bytes: Optional[int] = Field(
+        default=0,
+        description="The maximum number of bytes that can be used within the quota scope.",
+    )
 
-@dataclass
-class ProcessingResponseModel(ResponseModel):
-    user_model: Optional[BaseModel] = None
-    status: int = 202
 
-
-@dataclass
-class NoContentResponseModel(ResponseModel):
-    user_model: Optional[BaseModel] = None
-    status: int = 204
-
-
-class VolumeMetadataResponseModel(BaseModel):
-    volume_id: str
-    backend: str
-    path: str
-    fsprefix: Optional[str] = None
-    capabilities: List[str]
-
-
-class GetVolumeResponseModel(BaseModel):
-    volumes: List[VolumeMetadataResponseModel]
-
-
-class QuotaScopeResponseModel(BaseModel):
-    used_bytes: Optional[int] = 0
-    limit_bytes: Optional[int] = 0
-
-
-class VFolderMetadataResponseModel(BaseModel):
-    mount_path: str
-    file_count: int
-    capacity_bytes: int
-    used_bytes: BinarySize
+class VFolderMetadataResponse(BaseResponseModel):
+    mount_path: str = Field(..., description="The path where the virtual folder is mounted.")
+    file_count: int = Field(..., description="The number of files in the virtual folder.")
+    capacity_bytes: int = Field(
+        ..., description="The total capacity in bytes of the virtual folder."
+    )
+    used_bytes: BinarySize = Field(
+        ..., description="The used capacity in bytes of the virtual folder."
+    )