Add the skeleton interface of VFolder CRUD APIs using the new layered architecture in storage-proxy

[HyeockJinKim]

Motivation  

• The current CRUD handler interface in storage_proxy suffers from code duplication and readability issues, making maintenance challenging.
• To improve code quality, the interface needs to be refactored for better modularity and scalability, enabling easier integration of new features and updates.

Required Features

• Define CRUD handler class and methods for storage-proxy
• Apply handler’s method to new routes in storage/api/manager.py, client.py

Implementation Detail (interface)

• Use typing.Protocol to define a formal interface for CRUD handlers, enabling better type checking and enforcing consistency across handler implementations.
• Refactor duplicated CRUD logic into standalone utility functions to improve reusability and reduce redundancy across handlers.
• Use dependency injection to pass RootContext to handlers dynamically, ensuring they receive the necessary context during initialization.

1. storage/api/vfolder/types.py

import uuid
from pathlib import Path, PurePath, PurePosixPath
from typing import Any, Mapping, TypeAlias

from pydantic import AliasChoices, Field, model_validator
from pydantic import BaseModel as PydanticBaseModel

from ai.backend.common.types import BinarySize, QuotaConfig, QuotaScopeID, VFolderID
from ai.backend.storage.types import CapacityUsage, TreeUsage

__all__ = (
    "VolumeIDModel",
    "VolumeInfoModel",
    "VolumeInfoListModel",
    "VFolderIDModel",
    "VFolderInfoRequestModel",
    "VFolderInfoModel",
    "VFolderCloneModel",
    "QuotaIDModel",
    "QuotaScopeInfoModel",
    "QuotaConfigModel",
)


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(
        "vid", "volumeid", "volume_id", "VolumeID", "Volume_Id", "Volumeid"
    ),
)
VFOLDER_ID_FIELD = Field(
    ...,
    validation_alias=AliasChoices(
        "vfid", "vfolderid", "vfolder_id", "VFolderID", "VFolder_Id", "VFolderid"
    ),
)
QUOTA_SCOPE_ID_FIELD = Field(
    ...,
    validation_alias=AliasChoices(
        "qsid",
        "quotascopeid",
        "quota_scope_id",
        "QuotaScopeID",
        "Quota_Scope_Id",
        "QuotaScopeid",
        "Quota_ScopeID",
        "Quota_Scopeid",
        "quotaScopeID",
        "quotaScopeid",
    ),
)


class VolumeIDModel(BaseModel):
    volume_id: VolumeID = VOLUME_ID_FIELD

    @model_validator(mode="before")
    def validate_all_fields(cls, values: dict[str, Any]) -> dict[str, Any]:
        if "volume_id" in values and not isinstance(values["volume_id"], uuid.UUID):
            raise ValueError("volume_id must be a UUID")
        return values


class VolumeInfoModel(BaseModel):
    """For `get_volume`, `get_volumes` requests"""

    volume_id: VolumeID = VOLUME_ID_FIELD
    backend: str = Field(...)
    path: Path = Field(...)
    fsprefix: PurePath | None
    capabilities: list[str] = Field(...)
    options: Mapping[str, Any] | None

    @model_validator(mode="before")
    def validate_all_fields(cls, values: dict[str, Any]) -> dict[str, Any]:
        if "volume_id" in values and not isinstance(values["volume_id"], uuid.UUID):
            raise ValueError("volume_id must be a UUID")
        if "backend" in values and not isinstance(values["backend"], str):
            raise ValueError("backend must be a string")
        if "path" in values and not isinstance(values["path"], Path):
            raise ValueError("path must be a Path object")
        if values.get("fsprefix") is not None and not isinstance(values["fsprefix"], PurePath):
            raise ValueError("fsprefix must be a PurePath or None")
        if "capabilities" in values and not isinstance(values["capabilities"], list):
            raise ValueError("capabilities must be a list of strings")
        if values.get("options") is not None and not isinstance(values["options"], Mapping):
            raise ValueError("options must be a mapping or None")
        return values


class VolumeInfoListModel(BaseModel):
    """For `get_volumes` response"""

    volumes: dict[VolumeID, VolumeInfoModel] = Field(...)

    @model_validator(mode="before")
    def validate_all_fields(cls, values: dict[str, Any]) -> dict[str, Any]:
        if "volumes" in values and not isinstance(values["volumes"], dict):
            raise ValueError("volumes must be a dictionary")
        for k, v in values.get("volumes", {}).items():
            if not isinstance(k, uuid.UUID):
                raise ValueError("keys in volumes must be UUIDs")
            if not isinstance(v, VolumeInfoModel):
                raise ValueError("values in volumes must be VolumeInfoModel instances")
        return values


class VFolderIDModel(BaseModel):
    volume_id: VolumeID = VOLUME_ID_FIELD
    vfolder_id: VFolderID = VFOLDER_ID_FIELD

    @model_validator(mode="before")
    def validate_all_fields(cls, values: dict[str, Any]) -> dict[str, Any]:
        if "volume_id" in values and not isinstance(values["volume_id"], uuid.UUID):
            raise ValueError("volume_id must be a UUID")
        if "vfolder_id" in values and not isinstance(values["vfolder_id"], VFolderID):
            raise ValueError("vfolder_id must be a VFolderID")
        return values


class VFolderInfoRequestModel(BaseModel):
    """For `get_vfolder_info` request"""

    volume_id: VolumeID = VOLUME_ID_FIELD
    vfolder_id: VFolderID = VFOLDER_ID_FIELD
    subpath: PurePosixPath = Field(...)

    @model_validator(mode="before")
    def validate_all_fields(cls, values: dict[str, Any]) -> dict[str, Any]:
        if "volume_id" in values and not isinstance(values["volume_id"], uuid.UUID):
            raise ValueError("volume_id must be a UUID")
        if "vfolder_id" in values and not isinstance(values["vfolder_id"], VFolderID):
            raise ValueError("vfolder_id must be a VFolderID")
        if "subpath" in values and not isinstance(values["subpath"], PurePosixPath):
            raise ValueError("subpath must be a PurePosixPath")
        return values


class VFolderInfoModel(BaseModel):
    """For `get_vfolder_info` response"""

    vfolder_mount: Path = Field(...)
    vfolder_metadata: bytes = Field(...)  # 실제로 쓰이는지 확인 필요
    vfolder_usage: TreeUsage = Field(...)
    vfolder_used_bytes: BinarySize = Field(...)
    vfolder_fs_usage: CapacityUsage = Field(...)

    @model_validator(mode="before")
    def validate_all_fields(cls, values: dict[str, Any]) -> dict[str, Any]:
        if "vfolder_mount" in values and not isinstance(values["vfolder_mount"], Path):
            raise ValueError("vfolder_mount must be a Path object")
        if "vfolder_metadata" in values and not isinstance(values["vfolder_metadata"], bytes):
            raise ValueError("vfolder_metadata must be bytes")
        if "vfolder_usage" in values and not isinstance(values["vfolder_usage"], TreeUsage):
            raise ValueError("vfolder_usage must be a TreeUsage object")
        if "vfolder_used_bytes" in values and not isinstance(
            values["vfolder_used_bytes"], BinarySize
        ):
            raise ValueError("vfolder_used_bytes must be a BinarySize object")
        if "vfolder_fs_usage" in values and not isinstance(
            values["vfolder_fs_usage"], CapacityUsage
        ):
            raise ValueError("vfolder_fs_usage must be a CapacityUsage object")
        return values


class VFolderCloneModel(BaseModel):
    volume_id: VolumeID = VOLUME_ID_FIELD  # source volume
    src_vfolder_id: VFolderID = Field(
        ...,
        validation_alias=AliasChoices(
            "src_vfid",
            "src_vfolderid",
            "src_vfolder_id",
            "source",
            "src",
            "src_vfolderid",
            "source_vfid",
            "source_vfolderid",
            "source_vfolder_id",
            "SrcVfid",
            "SrcVfolderid",
            "Source",
            "Src",
            "SrcVfolderid",
            "SourceVfid",
            "SourceVfolderid",
        ),
    )
    dst_vfolder_id: VFolderID = Field(
        ...,
        validation_alias=AliasChoices(
            "dst_vfid",
            "dst_vfolderid",
            "destination",
            "dst",
            "dst_vfolderid",
            "dst_vfolder_id",
            "destination_vfid",
            "destination_vfolderid",
            "destination_vfolder_id",
            "DstVfid",
            "DstVfolderid",
            "Destination",
            "Dst",
            "DstVfolderid",
            "DestinationVfid",
            "DestinationVfolderid",
        ),
    )

    @model_validator(mode="before")
    def validate_all_fields(cls, values: dict[str, Any]) -> dict[str, Any]:
        if "volume_id" in values and not isinstance(values["volume_id"], uuid.UUID):
            raise ValueError("volume_id must be a UUID")
        if "src_vfolder_id" in values and not isinstance(values["src_vfolder_id"], VFolderID):
            raise ValueError("src_vfolder_id must be a VFolderID")
        if "dst_vfolder_id" in values and not isinstance(values["dst_vfolder_id"], VFolderID):
            raise ValueError("dst_vfolder_id must be a VFolderID")
        return values


class QuotaIDModel(BaseModel):
    volume_id: VolumeID = VOLUME_ID_FIELD
    quota_scope_id: QuotaScopeID = QUOTA_SCOPE_ID_FIELD

    @model_validator(mode="before")
    def validate_all_fields(cls, values: dict[str, Any]) -> dict[str, Any]:
        if "volume_id" in values and not isinstance(values["volume_id"], uuid.UUID):
            raise ValueError("volume_id must be a UUID")
        if "quota_scope_id" in values and not isinstance(values["quota_scope_id"], QuotaScopeID):
            raise ValueError("quota_scope_id must be a QuotaScopeID")
        return values


class QuotaScopeInfoModel(BaseModel):
    used_bytes: int | None
    limit_bytes: int | None

    @model_validator(mode="before")
    def validate_all_fields(cls, values: dict[str, Any]) -> dict[str, Any]:
        for field in ["used_bytes", "limit_bytes"]:
            value = values.get(field)
            if value is None:
                values[field] = 0
            elif not isinstance(value, int):
                raise ValueError(f"{field} must be an integer or None")
        return values


class QuotaConfigModel(BaseModel):
    volume_id: VolumeID = VOLUME_ID_FIELD
    quota_scope_id: QuotaScopeID = QUOTA_SCOPE_ID_FIELD
    options: QuotaConfig | None

    @model_validator(mode="before")
    def validate_all_fields(cls, values: dict[str, Any]) -> dict[str, Any]:
        if "volume_id" in values and not isinstance(values["volume_id"], uuid.UUID):
            raise ValueError("volume_id must be a UUID")
        if "quota_scope_id" in values and not isinstance(values["quota_scope_id"], QuotaScopeID):
            raise ValueError("quota_scope_id must be a QuotaScopeID")
        if values.get("options") is not None and not isinstance(values["options"], QuotaConfig):
            raise ValueError("options must be a QuotaConfig or None")
        return values

1. storage/api/vfolder/manager_service.py

from pathlib import Path, PurePath
from typing import Protocol

from ai.backend.common.types import BinarySize
from ai.backend.storage.api.vfolder.types import (
    QuotaConfigModel,
    QuotaIDModel,
    QuotaScopeInfoModel,
    VFolderCloneModel,
    VFolderIDModel,
    VFolderInfoModel,
    VFolderInfoRequestModel,
    VolumeIDModel,
    VolumeInfoListModel,
    VolumeInfoModel,
)
from ai.backend.storage.types import CapacityUsage, TreeUsage


class VFolderServiceProtocol(Protocol):
    async def get_volume(self, volume_data: VolumeIDModel) -> VolumeInfoModel:
        """by volume_id"""
        ...

    async def get_volumes(self) -> VolumeInfoListModel: ...

    async def create_quota_scope(self, quota_config_data: QuotaConfigModel) -> None: ...

    async def get_quota_scope(self, quota_data: QuotaIDModel) -> QuotaScopeInfoModel: ...

    async def update_quota_scope(self, quota_config_data: QuotaConfigModel) -> None: ...

    async def delete_quota_scope(self, quota_data: QuotaIDModel) -> None:
        """Previous: unset_quota"""
        ...

    async def create_vfolder(self, vfolder_data: VFolderIDModel) -> None: ...

    async def clone_vfolder(self, vfolder_clone_data: VFolderCloneModel) -> None: ...

    async def get_vfolder_info(self, vfolder_info: VFolderInfoRequestModel) -> VFolderInfoModel:
        # Integration: vfolder_mount, metadata, vfolder_usage, vfolder_used_bytes, vfolder_fs_usage
        ...

    async def delete_vfolder(self, vfolder_data: VFolderIDModel) -> None: ...


class VFolderService:
    async def get_volume(self, volume_data: VolumeIDModel) -> VolumeInfoModel:
        return VolumeInfoModel(
            volume_id=volume_data.volume_id,
            backend="default-backend",
            path=Path("/default/path"),
            fsprefix=PurePath("/fsprefix"),
            capabilities=["read", "write"],
            options={"option1": "value1"},
        )

    async def get_volumes(self) -> VolumeInfoListModel:
        return VolumeInfoListModel(volumes={})

    async def create_quota_scope(self, quota_config_data: QuotaConfigModel) -> None:
        return None

    async def get_quota_scope(self, quota_data: QuotaIDModel) -> QuotaScopeInfoModel:
        return QuotaScopeInfoModel(used_bytes=0, limit_bytes=0)

    async def update_quota_scope(self, quota_config_data: QuotaConfigModel) -> None:
        return None

    async def delete_quota_scope(self, quota_data: QuotaIDModel) -> None:
        return None

    async def create_vfolder(self, vfolder_data: VFolderIDModel) -> None:
        return None

    async def clone_vfolder(self, vfolder_clone_data: VFolderCloneModel) -> None:
        return None

    async def get_vfolder_info(self, vfolder_info: VFolderInfoRequestModel) -> VFolderInfoModel:
        return VFolderInfoModel(
            vfolder_mount=Path("/mount/point"),
            vfolder_metadata=b"",
            vfolder_usage=TreeUsage(file_count=0, used_bytes=0),
            vfolder_used_bytes=BinarySize(0),
            vfolder_fs_usage=CapacityUsage(used_bytes=0, capacity_bytes=0),
        )

    async def delete_vfolder(self, vfolder_data: VFolderIDModel) -> None:
        return None

1. storage/api/vfolder/manager_handler.py

import uuid

from aiohttp import web

from ai.backend.storage.api.vfolder.manager_service import VFolderService
from ai.backend.storage.api.vfolder.types import (
    QuotaConfigModel,
    QuotaIDModel,
    VFolderCloneModel,
    VFolderIDModel,
    VFolderInfoRequestModel,
    VolumeIDModel,
)


class VFolderHandler:
    def __init__(self, storage_service: VFolderService) -> None:
        self.storage_service = storage_service

    async def get_volume(self, request: web.Request) -> web.Response:
        data = await request.json()
        data["volume_id"] = uuid.UUID(data["volume_id"])
        req = VolumeIDModel(**data)
        result = await self.storage_service.get_volume(req)
        return web.json_response(result)

    async def get_volumes(self, request: web.Request) -> web.Response:
        result = await self.storage_service.get_volumes()
        # Assume that the volume_dict is a dictionary of VolumeInfoModel objects
        volumes_dict = result.volumes
        volumes_dict = {k: v for k, v in volumes_dict.items()}
        return web.json_response(volumes_dict)

    async def create_quota_scope(self, request: web.Request) -> web.Response:
        data = await request.json()
        data["volume_id"] = uuid.UUID(data["volume_id"])
        req = QuotaConfigModel(**data)
        await self.storage_service.create_quota_scope(req)
        return web.Response(status=204)

    async def get_quota_scope(self, request: web.Request) -> web.Response:
        data = await request.json()
        data["volume_id"] = uuid.UUID(data["volume_id"])
        req = QuotaIDModel(**data)
        result = await self.storage_service.get_quota_scope(req)
        return web.json_response(result)

    async def update_quota_scope(self, request: web.Request) -> web.Response:
        data = await request.json()
        data["volume_id"] = uuid.UUID(data["volume_id"])
        req = QuotaConfigModel(**data)
        await self.storage_service.update_quota_scope(req)
        return web.Response(status=204)

    async def delete_quota_scope(self, request: web.Request) -> web.Response:
        data = await request.json()
        data["volume_id"] = uuid.UUID(data["volume_id"])
        req = QuotaIDModel(**data)
        await self.storage_service.delete_quota_scope(req)
        return web.Response(status=204)

    async def create_vfolder(self, request: web.Request) -> web.Response:
        data = await request.json()
        data["volume_id"] = uuid.UUID(data["volume_id"])
        req = VFolderIDModel(**data)
        await self.storage_service.create_vfolder(req)
        return web.Response(status=204)

    async def clone_vfolder(self, request: web.Request) -> web.Response:
        data = await request.json()
        data["volume_id"] = uuid.UUID(data["volume_id"])
        req = VFolderCloneModel(**data)
        await self.storage_service.clone_vfolder(req)
        return web.Response(status=204)

    async def get_vfolder_info(self, request: web.Request) -> web.Response:
        data = await request.json()
        data["volume_id"] = uuid.UUID(data["volume_id"])
        req = VFolderInfoRequestModel(**data)
        result = await self.storage_service.get_vfolder_info(req)
        return web.json_response(result)

    async def delete_vfolder(self, request: web.Request) -> web.Response:
        data = await request.json()
        data["volume_id"] = uuid.UUID(data["volume_id"])
        req = VFolderIDModel(**data)
        await self.storage_service.delete_vfolder(req)
        return web.Response(status=202)

Testing Scenarios  

• CREATE
  • A valid request successfully creates a vfolder.
  • A request with a duplicate name returns an error response.
  • A request with invalid fields (e.g., empty name) returns an appropriate error message.
• READ
  • Retrieving an existing vfolder returns the correct information.
  • Retrieving a non-existent vfolder results in a 404 Not Found response.
  • Requests from users without access permissions return a 403 Forbidden response.
• UPDATE
  • Valid update requests successfully apply changes.
  • Requests from unauthorized users return a 403 Forbidden response.
  • Requests with invalid data return an appropriate error message.
• DELETE
  • Valid delete requests successfully remove the vfolder.
  • Requests to delete an already deleted vfolder return an appropriate status message.
  • Requests from unauthorized users return a 403 Forbidden response.