refactor(BA-620): Add Service Layer to Avoid Direct Volume and Vfolder Operations in Storage-Proxy Handler

changes/3588.enhance.md

@@ -0,0 +1 @@
+Add Service Layer to Avoid Direct Volume and Vfolder Operations in Storage-Proxy Handler

src/ai/backend/common/BUILD

@@ -27,6 +27,7 @@ python_distribution(
         ":src",
         "src/ai/backend/common/auth:src",
         "src/ai/backend/common/dto/manager:src",
+        "src/ai/backend/common/dto/storage:src",
         "src/ai/backend/common/metrics:src",
         "src/ai/backend/common/plugin:src",
         "src/ai/backend/common/web/session:src",  # not auto-inferred

src/ai/backend/common/api_handlers.py

@@ -22,7 +22,7 @@
 from aiohttp import web
 from aiohttp.web_urldispatcher import UrlMappingMatchInfo
 from multidict import CIMultiDictProxy, MultiMapping
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 from pydantic_core._pydantic_core import ValidationError
 
 from .exception import (
@@ -126,6 +126,10 @@ def from_request(cls, request: web.Request) -> Self:
         pass
 
 
+class BaseRequestModel(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+
 class BaseResponseModel(BaseModel):
     pass
 

src/ai/backend/common/dto/manager/request.py

@@ -1,23 +1,25 @@
 import uuid
 from typing import Optional
 
-from pydantic import AliasChoices, BaseModel, Field
+from pydantic import AliasChoices, Field
 
-from ai.backend.common import typed_validators as tv
-from ai.backend.common.dto.manager.dto import VFolderPermissionDTO
+from ai.backend.common.api_handlers import BaseRequestModel
 from ai.backend.common.types import VFolderUsageMode
 
+from ...typed_validators import VFolderName
+from .field import VFolderPermissionField
 
-class VFolderCreateReq(BaseModel):
-    name: tv.VFolderName = Field(
+
+class VFolderCreateReq(BaseRequestModel):
+    name: VFolderName = Field(
         description="Name of the vfolder",
     )
     folder_host: Optional[str] = Field(
         validation_alias=AliasChoices("host", "folder_host"),
         default=None,
     )
     usage_mode: VFolderUsageMode = Field(default=VFolderUsageMode.GENERAL)
-    permission: VFolderPermissionDTO = Field(default=VFolderPermissionDTO.READ_WRITE)
+    permission: VFolderPermissionField = Field(default=VFolderPermissionField.READ_WRITE)
     unmanaged_path: Optional[str] = Field(
         validation_alias=AliasChoices("unmanaged_path", "unmanagedPath"),
         default=None,
@@ -31,7 +33,7 @@ class VFolderCreateReq(BaseModel):
     )
 
 
-class RenameVFolderReq(BaseModel):
-    new_name: tv.VFolderName = Field(
+class RenameVFolderReq(BaseRequestModel):
+    new_name: VFolderName = Field(
         description="Name of the vfolder",
     )

src/ai/backend/common/dto/storage/BUILD

@@ -0,0 +1 @@
+python_sources(name="src")

src/ai/backend/common/dto/storage/field.py

@@ -0,0 +1,21 @@
+from typing import Optional
+
+from pydantic import BaseModel
+
+from ...types import VolumeID
+
+
+class VolumeMetaField(BaseModel):
+    volume_id: VolumeID
+    backend: str
+    path: str
+    fsprefix: Optional[str]
+    capabilities: list[str]
+
+
+class VFolderMetaField(BaseModel):
+    mount_path: str
+    file_count: int
+    used_bytes: int
+    capacity_bytes: int
+    fs_used_bytes: int

src/ai/backend/common/dto/storage/path.py

@@ -0,0 +1,27 @@
+from uuid import UUID
+
+from pydantic import Field
+
+from ...api_handlers import BaseRequestModel
+from ...types import QuotaScopeType, VolumeID
+
+
+class VolumeIDPath(BaseRequestModel):
+    volume_id: VolumeID = Field(
+        description="A unique identifier for the volume.",
+    )
+
+
+class QuotaScopeKeyPath(VolumeIDPath):
+    scope_type: QuotaScopeType = Field(
+        description="The type of the quota scope.",
+    )
+    scope_uuid: UUID = Field(
+        description="A unique uuid for the quota scope.",
+    )
+
+
+class VFolderKeyPath(QuotaScopeKeyPath):
+    folder_uuid: UUID = Field(
+        description="A unique uuid for the virtual folder.",
+    )

src/ai/backend/common/dto/storage/request.py

@@ -0,0 +1,26 @@
+from typing import Optional
+
+from pydantic import AliasChoices, Field
+
+from ...api_handlers import BaseRequestModel
+from ...types import QuotaConfig, VFolderID
+
+
+class QuotaScopeReq(BaseRequestModel):
+    options: Optional[QuotaConfig] = Field(
+        default=None,
+        description="The options for the quota scope.",
+    )
+
+
+class GetVFolderMetaReq(BaseRequestModel):
+    subpath: str = Field(
+        description="The subpath of the virtual folder.",
+    )
+
+
+class CloneVFolderReq(BaseRequestModel):
+    dst_vfolder_id: VFolderID = Field(
+        description="The destination virtual folder ID.",
+        validation_alias=AliasChoices("dst_vfid", "dst_vfolder_id"),
+    )

src/ai/backend/common/dto/storage/response.py

@@ -0,0 +1,23 @@
+from typing import Optional
+
+from pydantic import Field
+
+from ai.backend.common.api_handlers import BaseResponseModel
+from ai.backend.common.dto.storage.field import VFolderMetaField, VolumeMetaField
+
+
+class GetVolumeResponse(BaseResponseModel):
+    item: VolumeMetaField
+
+
+class GetVolumesResponse(BaseResponseModel):
+    items: list[VolumeMetaField]
+
+
+class QuotaScopeResponse(BaseResponseModel):
+    used_bytes: Optional[int] = Field(default=0)
+    limit_bytes: Optional[int] = Field(default=0)
+
+
+class VFolderMetadataResponse(BaseResponseModel):
+    item: VFolderMetaField

src/ai/backend/common/types.py

@@ -7,6 +7,7 @@
 import math
 import numbers
 import textwrap
+import uuid
 from abc import ABCMeta, abstractmethod
 from collections import UserDict, defaultdict, namedtuple
 from collections.abc import Iterable
@@ -974,6 +975,9 @@ def as_trafaret(cls) -> t.Trafaret:
         raise NotImplementedError
 
 
+type VolumeID = uuid.UUID
+
+
 @attrs.define(slots=True, frozen=True)
 class QuotaScopeID:
     scope_type: QuotaScopeType

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

@@ -0,0 +1,142 @@
+from typing import Optional, Protocol
+
+from ai.backend.common.api_handlers import APIResponse, BodyParam, PathParam, api_handler
+from ai.backend.common.dto.storage.path import QuotaScopeKeyPath, VFolderKeyPath, VolumeIDPath
+from ai.backend.common.dto.storage.request import (
+    CloneVFolderReq,
+    GetVFolderMetaReq,
+    QuotaScopeReq,
+)
+from ai.backend.common.dto.storage.response import (
+    GetVolumeResponse,
+    GetVolumesResponse,
+    VFolderMetadataResponse,
+)
+from ai.backend.common.types import QuotaConfig, VFolderID, VolumeID
+
+from ...volumes.types import (
+    QuotaScopeKey,
+    QuotaScopeMeta,
+    VFolderKey,
+    VFolderMeta,
+    VolumeMeta,
+)
+
+
+class VFolderServiceProtocol(Protocol):
+    async def get_volume(self, volume_id: VolumeID) -> VolumeMeta: ...
+
+    async def get_volumes(self) -> list[VolumeMeta]: ...
+
+    async def create_quota_scope(
+        self, quota_scope_key: QuotaScopeKey, options: Optional[QuotaConfig]
+    ) -> None: ...
+
+    async def get_quota_scope(self, quota_scope_key: QuotaScopeKey) -> QuotaScopeMeta: ...
+
+    async def update_quota_scope(
+        self, quota_scope_key: QuotaScopeKey, options: Optional[QuotaConfig]
+    ) -> None: ...
+
+    async def delete_quota_scope(self, quota_scope_key: QuotaScopeKey) -> None: ...
+
+    async def create_vfolder(self, vfolder_key: VFolderKey) -> None: ...
+
+    async def clone_vfolder(self, vfolder_key: VFolderKey, dst_vfolder_id: VFolderID) -> None: ...
+
+    async def get_vfolder_info(self, vfolder_key: VFolderKey, subpath: str) -> VFolderMeta: ...
+
+    async def delete_vfolder(self, vfolder_key: VFolderKey) -> None: ...
+
+
+class VFolderHandler:
+    _storage_service: VFolderServiceProtocol
+
+    def __init__(self, storage_service: VFolderServiceProtocol) -> None:
+        self._storage_service = storage_service
+
+    @api_handler
+    async def get_volume(self, path: PathParam[VolumeIDPath]) -> APIResponse:
+        volume_meta = await self._storage_service.get_volume(path.parsed.volume_id)
+        return APIResponse.build(
+            status_code=200,
+            response_model=GetVolumeResponse(
+                item=volume_meta.to_field(),
+            ),
+        )
+
+    @api_handler
+    async def get_volumes(self) -> APIResponse:
+        volume_meta_list = await self._storage_service.get_volumes()
+        return APIResponse.build(
+            status_code=200,
+            response_model=GetVolumesResponse(
+                items=[volume.to_field() for volume in volume_meta_list],
+            ),
+        )
+
+    @api_handler
+    async def create_quota_scope(
+        self, path: PathParam[QuotaScopeKeyPath], body: BodyParam[QuotaScopeReq]
+    ) -> APIResponse:
+        quota_scope_key = QuotaScopeKey.from_quota_scope_path(path.parsed)
+        await self._storage_service.create_quota_scope(quota_scope_key, body.parsed.options)
+        return APIResponse.no_content(status_code=204)
+
+    @api_handler
+    async def get_quota_scope(self, path: PathParam[QuotaScopeKeyPath]) -> APIResponse:
+        quota_scope_key = QuotaScopeKey.from_quota_scope_path(path.parsed)
+        quota_scope = await self._storage_service.get_quota_scope(quota_scope_key)
+        return APIResponse.build(
+            status_code=200,
+            response_model=quota_scope.to_response(),
+        )
+
+    @api_handler
+    async def update_quota_scope(
+        self, path: PathParam[QuotaScopeKeyPath], body: BodyParam[QuotaScopeReq]
+    ) -> APIResponse:
+        quota_scope_key = QuotaScopeKey.from_quota_scope_path(path.parsed)
+        await self._storage_service.update_quota_scope(quota_scope_key, body.parsed.options)
+        return APIResponse.no_content(status_code=204)
+
+    @api_handler
+    async def delete_quota_scope(self, path: PathParam[QuotaScopeKeyPath]) -> APIResponse:
+        quota_scope_key = QuotaScopeKey.from_quota_scope_path(path.parsed)
+        await self._storage_service.delete_quota_scope(quota_scope_key)
+        return APIResponse.no_content(status_code=204)
+
+    @api_handler
+    async def create_vfolder(self, path: PathParam[VFolderKeyPath]) -> APIResponse:
+        vfolder_key = VFolderKey.from_vfolder_path(path.parsed)
+        await self._storage_service.create_vfolder(vfolder_key)
+        return APIResponse.no_content(status_code=204)
+
+    @api_handler
+    async def clone_vfolder(
+        self, path: PathParam[VFolderKeyPath], body: BodyParam[CloneVFolderReq]
+    ) -> APIResponse:
+        vfolder_key = VFolderKey.from_vfolder_path(path.parsed)
+        await self._storage_service.clone_vfolder(vfolder_key, body.parsed.dst_vfolder_id)
+        return APIResponse.no_content(status_code=204)
+
+    @api_handler
+    async def get_vfolder_info(
+        self, path: PathParam[VFolderKeyPath], body: BodyParam[GetVFolderMetaReq]
+    ) -> APIResponse:
+        vfolder_key = VFolderKey.from_vfolder_path(path.parsed)
+        vfolder_meta = await self._storage_service.get_vfolder_info(
+            vfolder_key, body.parsed.subpath
+        )
+        return APIResponse.build(
+            status_code=200,
+            response_model=VFolderMetadataResponse(
+                item=vfolder_meta.to_field(),
+            ),
+        )
+
+    @api_handler
+    async def delete_vfolder(self, path: PathParam[VFolderKeyPath]) -> APIResponse:
+        vfolder_key = VFolderKey.from_vfolder_path(path.parsed)
+        await self._storage_service.delete_vfolder(vfolder_key)
+        return APIResponse.no_content(status_code=204)