refactor(BA-589): Apply pydantic handling decorator to VFolder APIs in storage-proxy (#3562)

changes/3565.enhance.md

@@ -0,0 +1 @@
+Apply pydantic handling decorator to VFolder APIs in storage-proxy

src/ai/backend/storage/api/vfolder/BUILD

@@ -0,0 +1 @@
+python_sources()

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

@@ -0,0 +1,144 @@
+from typing import Protocol
+
+from ai.backend.common.api_handlers import ApiResponse, BodyParam, api_handler
+from ai.backend.storage.api.vfolder.response_model import (
+    GetVolumeResponse,
+    QuotaScopeResponse,
+    VFolderMetadataResponse,
+    VolumeMetadataResponse,
+)
+from ai.backend.storage.api.vfolder.types import (
+    QuotaScopeIdData,
+    QuotaScopeMetadata,
+    VFolderIdData,
+    VFolderMetadata,
+    VolumeIdData,
+    VolumeMetadataList,
+)
+
+
+class VFolderServiceProtocol(Protocol):
+    async def get_volume(self, volume_data: VolumeIdData) -> VolumeMetadataList: ...
+
+    async def get_volumes(self) -> VolumeMetadataList: ...
+
+    async def create_quota_scope(self, quota_data: QuotaScopeIdData) -> None: ...
+
+    async def get_quota_scope(self, quota_data: QuotaScopeIdData) -> QuotaScopeMetadata: ...
+
+    async def update_quota_scope(self, quota_data: QuotaScopeIdData) -> None: ...
+
+    async def delete_quota_scope(self, quota_data: QuotaScopeIdData) -> None: ...
+
+    async def create_vfolder(self, vfolder_data: VFolderIdData) -> VFolderIdData: ...
+
+    async def clone_vfolder(self, vfolder_data: VFolderIdData) -> None: ...
+
+    async def get_vfolder_info(self, vfolder_data: VFolderIdData) -> VFolderMetadata: ...
+
+    async def delete_vfolder(self, vfolder_data: VFolderIdData) -> VFolderIdData: ...
+
+
+class VFolderHandler:
+    def __init__(self, storage_service: VFolderServiceProtocol) -> None:
+        self.storage_service = storage_service
+
+    @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
+                ]
+            ),
+        )
+
+    @api_handler
+    async def get_volumes(self) -> ApiResponse:
+        volumes_data = await self.storage_service.get_volumes()
+        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
+                ]
+            ),
+        )
+
+    @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)
+
+    @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
+            ),
+        )
+
+    @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)
+
+    @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)
+
+    @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)
+
+    @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)
+
+    @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,
+            ),
+        )
+
+    @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

@@ -0,0 +1,39 @@
+from typing import List, Optional
+
+from pydantic import Field
+
+from ai.backend.common.pydantic_handlers import BaseResponseModel
+from ai.backend.common.types import BinarySize
+
+
+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.")
+
+
+class GetVolumeResponse(BaseResponseModel):
+    volumes: List[VolumeMetadataResponse] = Field(..., description="The list of volumes.")
+
+
+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.",
+    )
+
+
+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."
+    )

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

@@ -0,0 +1,133 @@
+import uuid
+from pathlib import Path, PurePath, PurePosixPath
+from typing import List, Optional, TypeAlias
+
+from pydantic import AliasChoices, Field
+from pydantic import BaseModel as PydanticBaseModel
+
+from ai.backend.common.types import BinarySize, QuotaConfig, QuotaScopeID, VFolderID
+
+
+class BaseModel(PydanticBaseModel):
+    """Base model for all models in this module"""
+
+    model_config = {"arbitrary_types_allowed": True}
+
+
+VolumeID: TypeAlias = uuid.UUID
+
+
+# Common fields for VolumeID and VFolderID
+VOLUME_ID_FIELD = Field(
+    ...,
+    validation_alias=AliasChoices(
+        "volume",
+        "volumeid",
+        "volume_id",
+        "volumeId",
+    ),
+    description="A unique identifier for the volume.",
+)
+VFOLDER_ID_FIELD = Field(
+    ...,
+    validation_alias=AliasChoices(
+        "vfid",
+        "folderid",
+        "folder_id",
+        "folderId",
+        "vfolderid",
+        "vfolder_id",
+        "vfolderId",
+    ),
+    description="A unique identifier for the virtual folder.",
+)
+QUOTA_SCOPE_ID_FIELD = Field(
+    ...,
+    validation_alias=AliasChoices(
+        "qsid",
+        "quotascopeid",
+        "quota_scope_id",
+        "quotaScopeId",
+    ),
+    description="A unique identifier for the quota scope.",
+)
+
+
+class VolumeIdData(BaseModel):
+    volume_id: VolumeID = VOLUME_ID_FIELD
+
+
+class VolumeMetadata(BaseModel):
+    """For `get_volume`, `get_volumes`"""
+
+    volume_id: VolumeID = Field(..., description="The unique identifier for the volume.")
+    backend: str = Field(
+        ..., description="The backend storage type for the volume (e.g., CephFS, GPFS)."
+    )
+    path: Path = Field(..., description="The path where the volume is mounted.")
+    fsprefix: Optional[PurePath] = Field(
+        default=None, description="The filesystem prefix for the volume, or None if not applicable."
+    )
+    capabilities: list[str] = Field(
+        ..., description="A list of capabilities supported by the volume."
+    )
+
+
+class VolumeMetadataList(BaseModel):
+    volumes: List[VolumeMetadata] = Field(..., description="A list of volume information.")
+
+
+class VFolderIdData(BaseModel):
+    volume_id: VolumeID = VOLUME_ID_FIELD
+    vfolder_id: VFolderID = VFOLDER_ID_FIELD
+    # For `get_vfolder_info`: mount
+    subpath: Optional[PurePosixPath] = Field(
+        default=None,
+        description="For `get_vfolder_info`\n\
+            The subpath inside the virtual folder to be queried.",
+    )
+    # For `clone_vfolder`
+    # You can use volume_id and vfolder_id as src_volume and src_vfolder_id.
+    dst_vfolder_id: Optional[VFolderID] = Field(
+        default=None,
+        validation_alias=AliasChoices(
+            "dst_vfid",
+            "dstvfolderid",
+            "dst_vfolder_id",
+            "dstVfolderId",
+        ),
+        description="For `clone_vfolder`\n\
+            The destination virtual folder ID to clone to.",
+    )
+
+
+class VFolderMetadata(BaseModel):
+    """For `get_vfolder_info`"""
+
+    mount_path: Path = 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 amount of used bytes in the virtual folder."
+    )
+
+
+class QuotaScopeIdData(BaseModel):
+    volume_id: VolumeID = VOLUME_ID_FIELD
+    quota_scope_id: QuotaScopeID = QUOTA_SCOPE_ID_FIELD
+    options: Optional[QuotaConfig] = Field(
+        default=None, description="Optional configuration settings for the quota."
+    )
+
+
+class QuotaScopeMetadata(BaseModel):
+    """For `get_quota_scope`"""
+
+    used_bytes: Optional[int] = Field(
+        default=0, description="The number of bytes currently used in the quota scope."
+    )
+    limit_bytes: Optional[int] = Field(
+        default=0, description="The maximum number of bytes allowed in the quota scope."
+    )

tests/storage-proxy/vfolder/BUILD

@@ -0,0 +1,3 @@
+python_test_utils()
+
+python_tests(name="tests")