From b86731224452594bbca75ef39eeb71ed10b31fd6 Mon Sep 17 00:00:00 2001 From: Mincheol Kang Date: Tue, 11 Feb 2025 19:21:29 +0900 Subject: [PATCH 1/7] feat: Add Service Layer to Avoid Direct Volume and Vfolder Operations in Storage-Proxy Handler --- changes/3588.enhance.md | 1 + src/ai/backend/common/dto/BUILD | 1 + src/ai/backend/common/dto/identifiers.py | 4 + src/ai/backend/common/dto/storage/BUILD | 1 + src/ai/backend/common/dto/storage/request.py | 69 ++ src/ai/backend/common/dto/storage/response.py | 48 ++ src/ai/backend/storage/api/vfolder/handler.py | 247 +++++++ .../storage/api/vfolder/response_model.py | 39 -- src/ai/backend/storage/api/vfolder/types.py | 133 ---- src/ai/backend/storage/volumes/abc.py | 2 +- src/ai/backend/storage/volumes/pool.py | 11 +- src/ai/backend/storage/volumes/service.py | 286 ++++++++ src/ai/backend/storage/volumes/types.py | 105 +++ tests/storage-proxy/vfolder/test_handler.py | 653 ++++++++++-------- tests/storage-proxy/vfolder/test_pool.py | 77 +++ 15 files changed, 1217 insertions(+), 460 deletions(-) create mode 100644 changes/3588.enhance.md create mode 100644 src/ai/backend/common/dto/BUILD create mode 100644 src/ai/backend/common/dto/identifiers.py create mode 100644 src/ai/backend/common/dto/storage/BUILD create mode 100644 src/ai/backend/common/dto/storage/request.py create mode 100644 src/ai/backend/common/dto/storage/response.py create mode 100644 src/ai/backend/storage/api/vfolder/handler.py delete mode 100644 src/ai/backend/storage/api/vfolder/response_model.py delete mode 100644 src/ai/backend/storage/api/vfolder/types.py create mode 100644 src/ai/backend/storage/volumes/service.py create mode 100644 src/ai/backend/storage/volumes/types.py create mode 100644 tests/storage-proxy/vfolder/test_pool.py diff --git a/changes/3588.enhance.md b/changes/3588.enhance.md new file mode 100644 index 00000000000..de2c160331e --- /dev/null +++ b/changes/3588.enhance.md @@ -0,0 +1 @@ +Add Service Layer to Avoid Direct Volume and Vfolder Operations in Storage-Proxy Handler diff --git a/src/ai/backend/common/dto/BUILD b/src/ai/backend/common/dto/BUILD new file mode 100644 index 00000000000..c1ffc1a1410 --- /dev/null +++ b/src/ai/backend/common/dto/BUILD @@ -0,0 +1 @@ +python_sources(name="src") \ No newline at end of file diff --git a/src/ai/backend/common/dto/identifiers.py b/src/ai/backend/common/dto/identifiers.py new file mode 100644 index 00000000000..a29b8e13a94 --- /dev/null +++ b/src/ai/backend/common/dto/identifiers.py @@ -0,0 +1,4 @@ +import uuid +from typing import TypeAlias + +VolumeID: TypeAlias = uuid.UUID diff --git a/src/ai/backend/common/dto/storage/BUILD b/src/ai/backend/common/dto/storage/BUILD new file mode 100644 index 00000000000..c1ffc1a1410 --- /dev/null +++ b/src/ai/backend/common/dto/storage/BUILD @@ -0,0 +1 @@ +python_sources(name="src") \ No newline at end of file diff --git a/src/ai/backend/common/dto/storage/request.py b/src/ai/backend/common/dto/storage/request.py new file mode 100644 index 00000000000..50a5b7c45bb --- /dev/null +++ b/src/ai/backend/common/dto/storage/request.py @@ -0,0 +1,69 @@ +from pathlib import PurePosixPath +from typing import Optional + +from pydantic import AliasChoices, BaseModel, ConfigDict, Field + +from ai.backend.common.dto.identifiers import VolumeID +from ai.backend.common.types import QuotaConfig, QuotaScopeID, VFolderID + + +class _BaseModel(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + + +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 VolumeKeyDataParams(_BaseModel): + volume_id: VolumeID = VOLUME_ID_FIELD + + +class VFolderKeyDataParams(_BaseModel): + volume_id: VolumeID = VOLUME_ID_FIELD + vfolder_id: VFolderID = VFOLDER_ID_FIELD + subpath: Optional[PurePosixPath] = Field(default=None) + # 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", + ), + ) + + +class QuotaScopeKeyDataParams(_BaseModel): + volume_id: VolumeID = VOLUME_ID_FIELD + quota_scope_id: QuotaScopeID = QUOTA_SCOPE_ID_FIELD + options: Optional[QuotaConfig] = Field(default=None) diff --git a/src/ai/backend/common/dto/storage/response.py b/src/ai/backend/common/dto/storage/response.py new file mode 100644 index 00000000000..0cd5a898e61 --- /dev/null +++ b/src/ai/backend/common/dto/storage/response.py @@ -0,0 +1,48 @@ +from typing import List, Optional + +from pydantic import Field + +from ai.backend.common.api_handlers import BaseResponseModel + + +class VolumeMetadataResponse(BaseResponseModel): + volume_id: str + backend: str + path: str + fsprefix: Optional[str] = Field(default=None) + capabilities: List[str] + + +class GetVolumeResponse(BaseResponseModel): + volumes: List[VolumeMetadataResponse] + + +class QuotaScopeResponse(BaseResponseModel): + used_bytes: Optional[int] = Field(default=0) + limit_bytes: Optional[int] = Field(default=0) + + +class VFolderMetadataResponse(BaseResponseModel): + mount_path: str + file_count: int + used_bytes: int + capacity_bytes: int + fs_used_bytes: int + + +class VFolderMountResponse(BaseResponseModel): + mount_path: str + + +class VFolderUsageResponse(BaseResponseModel): + file_count: int + used_bytes: int + + +class VFolderUsedBytesResponse(BaseResponseModel): + used_bytes: int + + +class VFolderFSUsageResponse(BaseResponseModel): + capacity_bytes: int + fs_used_bytes: int diff --git a/src/ai/backend/storage/api/vfolder/handler.py b/src/ai/backend/storage/api/vfolder/handler.py new file mode 100644 index 00000000000..69bfcb3a1dd --- /dev/null +++ b/src/ai/backend/storage/api/vfolder/handler.py @@ -0,0 +1,247 @@ +from typing import Protocol + +from ai.backend.common.api_handlers import APIResponse, BodyParam, PathParam, api_handler +from ai.backend.common.dto.storage.request import ( + QuotaScopeKeyDataParams, + VFolderKeyDataParams, + VolumeKeyDataParams, +) +from ai.backend.common.dto.storage.response import ( + GetVolumeResponse, + QuotaScopeResponse, + VFolderFSUsageResponse, + VFolderMetadataResponse, + VFolderMountResponse, + VFolderUsageResponse, + VFolderUsedBytesResponse, + VolumeMetadataResponse, +) +from ai.backend.storage.volumes.types import ( + NewQuotaScopeCreated, + NewVFolderCreated, + QuotaScopeKeyData, + QuotaScopeMetadata, + VFolderKeyData, + VFolderMetadata, + VolumeKeyData, + VolumeMetadata, + VolumeMetadataList, +) + + +class VFolderServiceProtocol(Protocol): + async def get_volume(self, volume_data: VolumeKeyData) -> VolumeMetadata: ... + + async def get_volumes(self) -> VolumeMetadataList: ... + + async def create_quota_scope(self, quota_data: QuotaScopeKeyData) -> NewQuotaScopeCreated: ... + + async def get_quota_scope(self, quota_data: QuotaScopeKeyData) -> QuotaScopeMetadata: ... + + async def update_quota_scope(self, quota_data: QuotaScopeKeyData) -> None: ... + + async def delete_quota_scope(self, quota_data: QuotaScopeKeyData) -> None: ... + + async def create_vfolder(self, vfolder_data: VFolderKeyData) -> NewVFolderCreated: ... + + async def clone_vfolder(self, vfolder_data: VFolderKeyData) -> NewVFolderCreated: ... + + async def get_vfolder_info(self, vfolder_data: VFolderKeyData) -> VFolderMetadata: ... + + async def delete_vfolder(self, vfolder_data: VFolderKeyData) -> None: ... + + +class VFolderHandler: + def __init__(self, storage_service: VFolderServiceProtocol) -> None: + self.storage_service = storage_service + + @api_handler + async def get_volume(self, body: PathParam[VolumeKeyDataParams]) -> APIResponse: + volume_parsed = body.parsed + volume_params = VolumeKeyData(volume_id=volume_parsed.volume_id) + volume_data = await self.storage_service.get_volume(volume_params) + return APIResponse.build( + status_code=200, + response_model=VolumeMetadataResponse( + volume_id=str(volume_data.volume_id), + backend=volume_data.backend, + path=str(volume_data.path), + fsprefix=str(volume_data.fsprefix) if volume_data.fsprefix else None, + capabilities=[cap for cap in volume_data.capabilities], + ), + ) + + @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=volume.backend, + path=str(volume.path), + fsprefix=str(volume.fsprefix) if volume.fsprefix else None, + capabilities=[cap for cap in volume.capabilities], + ) + for volume in volumes_data.volumes + ] + ), + ) + + @api_handler + async def create_quota_scope(self, body: BodyParam[QuotaScopeKeyDataParams]) -> APIResponse: + quota_parsed = body.parsed + quota_params = QuotaScopeKeyData( + volume_id=quota_parsed.volume_id, + quota_scope_id=quota_parsed.quota_scope_id, + options=quota_parsed.options, + ) + 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: PathParam[QuotaScopeKeyDataParams]) -> APIResponse: + quota_parsed = body.parsed + quota_params = QuotaScopeKeyData( + volume_id=quota_parsed.volume_id, + quota_scope_id=quota_parsed.quota_scope_id, + ) + quota_scope = await self.storage_service.get_quota_scope(quota_params) + return APIResponse.build( + status_code=200, + 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[QuotaScopeKeyDataParams]) -> APIResponse: + quota_parsed = body.parsed + quota_params = QuotaScopeKeyData( + volume_id=quota_parsed.volume_id, + quota_scope_id=quota_parsed.quota_scope_id, + options=quota_parsed.options, + ) + 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: PathParam[QuotaScopeKeyDataParams]) -> APIResponse: + quota_parsed = body.parsed + quota_params = QuotaScopeKeyData( + volume_id=quota_parsed.volume_id, + quota_scope_id=quota_parsed.quota_scope_id, + ) + 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[VFolderKeyDataParams]) -> APIResponse: + vfolder_parsed = body.parsed + vfolder_params = VFolderKeyData( + volume_id=vfolder_parsed.volume_id, + vfolder_id=vfolder_parsed.vfolder_id, + ) + await self.storage_service.create_vfolder(vfolder_params) + return APIResponse.no_content(status_code=201) + + @api_handler + async def clone_vfolder(self, body: BodyParam[VFolderKeyDataParams]) -> APIResponse: + vfolder_parsed = body.parsed + vfolder_params = VFolderKeyData( + volume_id=vfolder_parsed.volume_id, + vfolder_id=vfolder_parsed.vfolder_id, + dst_vfolder_id=vfolder_parsed.dst_vfolder_id, + ) + await self.storage_service.clone_vfolder(vfolder_params) + return APIResponse.no_content(status_code=201) + + @api_handler + async def get_vfolder_info(self, body: PathParam[VFolderKeyDataParams]) -> APIResponse: + vfolder_parsed = body.parsed + vfolder_params = VFolderKeyData( + volume_id=vfolder_parsed.volume_id, + vfolder_id=vfolder_parsed.vfolder_id, + subpath=vfolder_parsed.subpath, + ) + 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, + used_bytes=metadata.used_bytes, + capacity_bytes=metadata.capacity_bytes, + fs_used_bytes=metadata.fs_used_bytes, + ), + ) + + @api_handler + async def get_vfolder_mount(self, body: PathParam[VFolderKeyDataParams]) -> APIResponse: + vfolder_parsed = body.parsed + vfolder_params = VFolderKeyData( + volume_id=vfolder_parsed.volume_id, + vfolder_id=vfolder_parsed.vfolder_id, + subpath=vfolder_parsed.subpath, + ) + metadata = await self.storage_service.get_vfolder_info(vfolder_params) + return APIResponse.build( + status_code=200, + response_model=VFolderMountResponse(mount_path=str(metadata.mount_path)), + ) + + @api_handler + async def get_vfolder_usage(self, body: PathParam[VFolderKeyDataParams]) -> APIResponse: + vfolder_parsed = body.parsed + vfolder_params = VFolderKeyData( + volume_id=vfolder_parsed.volume_id, + vfolder_id=vfolder_parsed.vfolder_id, + ) + metadata = await self.storage_service.get_vfolder_info(vfolder_params) + return APIResponse.build( + status_code=200, + response_model=VFolderUsageResponse( + file_count=metadata.file_count, used_bytes=metadata.used_bytes + ), + ) + + @api_handler + async def get_vfolder_used_bytes(self, body: PathParam[VFolderKeyDataParams]) -> APIResponse: + vfolder_parsed = body.parsed + vfolder_params = VFolderKeyData( + volume_id=vfolder_parsed.volume_id, + vfolder_id=vfolder_parsed.vfolder_id, + ) + metadata = await self.storage_service.get_vfolder_info(vfolder_params) + return APIResponse.build( + status_code=200, + response_model=VFolderUsedBytesResponse(used_bytes=metadata.used_bytes), + ) + + @api_handler + async def get_vfolder_fs_usage(self, body: PathParam[VFolderKeyDataParams]) -> APIResponse: + vfolder_parsed = body.parsed + vfolder_params = VFolderKeyData( + volume_id=vfolder_parsed.volume_id, + vfolder_id=vfolder_parsed.vfolder_id, + ) + metadata = await self.storage_service.get_vfolder_info(vfolder_params) + return APIResponse.build( + status_code=200, + response_model=VFolderFSUsageResponse( + capacity_bytes=metadata.capacity_bytes, + fs_used_bytes=metadata.fs_used_bytes, + ), + ) + + @api_handler + async def delete_vfolder(self, body: PathParam[VFolderKeyDataParams]) -> APIResponse: + vfolder_parsed = body.parsed + vfolder_params = VFolderKeyData( + volume_id=vfolder_parsed.volume_id, + vfolder_id=vfolder_parsed.vfolder_id, + ) + await self.storage_service.delete_vfolder(vfolder_params) + return APIResponse.no_content(status_code=202) diff --git a/src/ai/backend/storage/api/vfolder/response_model.py b/src/ai/backend/storage/api/vfolder/response_model.py deleted file mode 100644 index c088a1253e6..00000000000 --- a/src/ai/backend/storage/api/vfolder/response_model.py +++ /dev/null @@ -1,39 +0,0 @@ -from typing import List, Optional - -from pydantic import Field - -from ai.backend.common.api_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." - ) diff --git a/src/ai/backend/storage/api/vfolder/types.py b/src/ai/backend/storage/api/vfolder/types.py deleted file mode 100644 index 3075e0fcdd5..00000000000 --- a/src/ai/backend/storage/api/vfolder/types.py +++ /dev/null @@ -1,133 +0,0 @@ -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." - ) diff --git a/src/ai/backend/storage/volumes/abc.py b/src/ai/backend/storage/volumes/abc.py index 1aef1b103bd..732d31b403a 100644 --- a/src/ai/backend/storage/volumes/abc.py +++ b/src/ai/backend/storage/volumes/abc.py @@ -20,6 +20,7 @@ from ai.backend.common.events import EventDispatcher, EventProducer from ai.backend.common.types import BinarySize, HardwareMetadata, QuotaScopeID from ai.backend.logging import BraceStyleAdapter +from ai.backend.storage.watcher import WatcherClient from ..exception import InvalidSubpathError, VFolderNotFoundError from ..types import ( @@ -31,7 +32,6 @@ TreeUsage, VFolderID, ) -from ..watcher import WatcherClient # Available capabilities of a volume implementation CAP_VFOLDER: Final = "vfolder" # ability to create vfolder diff --git a/src/ai/backend/storage/volumes/pool.py b/src/ai/backend/storage/volumes/pool.py index d33bfb7bbf9..60c899bf1ca 100644 --- a/src/ai/backend/storage/volumes/pool.py +++ b/src/ai/backend/storage/volumes/pool.py @@ -4,6 +4,7 @@ from pathlib import Path from typing import Any, AsyncIterator, Mapping, Type +from ai.backend.common.dto.identifiers import VolumeID from ai.backend.common.etcd import AsyncEtcd from ai.backend.common.events import EventDispatcher, EventProducer from ai.backend.storage.volumes.cephfs import CephFSVolume @@ -39,12 +40,11 @@ class VolumePool: - _volumes: dict[str, AbstractVolume] + _volumes: dict[VolumeID, AbstractVolume] _local_config: Mapping[str, Any] _etcd: AsyncEtcd _event_dispatcher: EventDispatcher _event_producer: EventProducer - _backends: dict[str, Type[AbstractVolume]] def __init__( self, @@ -68,8 +68,13 @@ def list_volumes(self) -> Mapping[str, VolumeInfo]: for volume_id, info in self._local_config["volume"].items() } + def get_volume_info(self, volume_id: VolumeID) -> VolumeInfo: + if volume_id not in self._local_config["volume"]: + raise InvalidVolumeError(volume_id) + return VolumeInfo(**self._local_config["volume"][volume_id]) + @actxmgr - async def get_volume(self, volume_id: str) -> AsyncIterator[AbstractVolume]: + async def get_volume(self, volume_id: VolumeID) -> AsyncIterator[AbstractVolume]: if volume_id in self._volumes: yield self._volumes[volume_id] else: diff --git a/src/ai/backend/storage/volumes/service.py b/src/ai/backend/storage/volumes/service.py new file mode 100644 index 00000000000..bb91eedaeb5 --- /dev/null +++ b/src/ai/backend/storage/volumes/service.py @@ -0,0 +1,286 @@ +import asyncio +import json +import logging +import uuid +import weakref +from contextlib import asynccontextmanager as actxmgr +from typing import AsyncIterator + +from aiohttp import web + +from ai.backend.common.dto.identifiers import VolumeID +from ai.backend.common.events import VFolderDeletionFailureEvent, VFolderDeletionSuccessEvent +from ai.backend.common.types import VFolderID +from ai.backend.logging.utils import BraceStyleAdapter +from ai.backend.storage.exception import ( + ExternalError, + InvalidQuotaConfig, + InvalidSubpathError, + QuotaScopeAlreadyExists, + QuotaScopeNotFoundError, + VFolderNotFoundError, +) +from ai.backend.storage.utils import log_manager_api_entry +from ai.backend.storage.volumes.pool import VolumePool +from ai.backend.storage.volumes.types import ( + NewQuotaScopeCreated, + NewVFolderCreated, + QuotaScopeKeyData, + QuotaScopeMetadata, + VFolderKeyData, + VFolderMetadata, + VolumeKeyData, + VolumeMetadata, + VolumeMetadataList, +) + +log = BraceStyleAdapter(logging.getLogger(__spec__.name)) + + +class VolumeService: + _volume_pool: VolumePool + _deletion_tasks: weakref.WeakValueDictionary[VFolderID, asyncio.Task] + + def __init__( + self, + volume_pool: VolumePool, + ) -> None: + self._volume_pool = volume_pool + self._deletion_tasks = weakref.WeakValueDictionary[VFolderID, asyncio.Task]() + + async def _get_capabilities(self, volume_id: VolumeID) -> list[str]: + async with self._volume_pool.get_volume(volume_id) as volume: + return [*await volume.get_capabilities()] + + @actxmgr + async def _handle_external_errors(self) -> AsyncIterator[None]: + try: + yield + except ExternalError as e: + log.exception("An external error occurred: %s", str(e)) + raise web.HTTPInternalServerError( + body=json.dumps({ + "msg": "An internal error has occurred.", + }), + content_type="application/json", + ) + + async def _delete_vfolder( + self, + vfolder_data: VFolderKeyData, + ) -> None: + volume_id = vfolder_data.volume_id + vfolder_id = vfolder_data.vfolder_id + + current_task = asyncio.current_task() + assert current_task is not None + self._deletion_tasks[vfolder_id] = current_task + + try: + async with self._volume_pool.get_volume(volume_id) as volume: + await volume.delete_vfolder(vfolder_id) + except OSError as e: + msg = str(e) if e.strerror is None else e.strerror + msg = f"{msg} (errno:{e.errno})" + log.exception(f"VFolder deletion task failed. (vfolder_id:{vfolder_id}, e:{msg})") + await self._volume_pool._event_producer.produce_event( + VFolderDeletionFailureEvent( + vfid=vfolder_id, + message=msg, + ) + ) + except Exception as e: + log.exception(f"VFolder deletion task failed. (vfolder_id:{vfolder_id}, e:{str(e)})") + await self._volume_pool._event_producer.produce_event( + VFolderDeletionFailureEvent( + vfid=vfolder_id, + message=str(e), + ) + ) + except asyncio.CancelledError: + log.warning(f"Vfolder deletion task cancelled. (vfolder_id:{vfolder_id})") + else: + log.info(f"VFolder deletion task successed. (vfolder_id:{vfolder_id})") + await self._volume_pool._event_producer.produce_event( + VFolderDeletionSuccessEvent(vfolder_id) + ) + + async def get_volume(self, volume_data: VolumeKeyData) -> VolumeMetadata: + volume_id = volume_data.volume_id + await log_manager_api_entry(log, "get_volume", volume_id) + volume = self._volume_pool.get_volume_info(volume_id) + return VolumeMetadata( + volume_id=volume_id, + backend=volume.backend, + path=volume.path, + fsprefix=volume.fsprefix, + capabilities=await self._get_capabilities(volume_id), + ) + + async def get_volumes(self) -> VolumeMetadataList: + await log_manager_api_entry(log, "get_volumes", params=None) + volumes = self._volume_pool.list_volumes() + return VolumeMetadataList( + volumes=[ + VolumeMetadata( + volume_id=uuid.UUID(volume_id), + backend=info.backend, + path=info.path, + fsprefix=info.fsprefix, + capabilities=await self._get_capabilities(uuid.UUID(volume_id)), + ) + for volume_id, info in volumes.items() + ] + ) + + async def create_quota_scope(self, quota_data: QuotaScopeKeyData) -> NewQuotaScopeCreated: + volume_id = quota_data.volume_id + quota_scope_id = quota_data.quota_scope_id + options = quota_data.options + + await log_manager_api_entry(log, "create_quota_scope", quota_data) + async with self._volume_pool.get_volume(volume_id) as volume: + try: + async with self._handle_external_errors(): + await volume.quota_model.create_quota_scope( + quota_scope_id=quota_scope_id, options=options, extra_args=None + ) + except QuotaScopeAlreadyExists: + raise web.HTTPConflict(reason="Volume already exists with given quota scope.") + return NewQuotaScopeCreated( + quota_scope_id=quota_scope_id, + quota_scope_path=volume.quota_model.mangle_qspath(quota_scope_id), + ) + + async def get_quota_scope(self, quota_data: QuotaScopeKeyData) -> QuotaScopeMetadata: + volume_id = quota_data.volume_id + quota_scope_id = quota_data.quota_scope_id + + await log_manager_api_entry(log, "get_quota_scope", quota_data) + async with self._volume_pool.get_volume(volume_id) as volume: + async with self._handle_external_errors(): + quota_usage = await volume.quota_model.describe_quota_scope(quota_scope_id) + if not quota_usage: + raise QuotaScopeNotFoundError + return QuotaScopeMetadata( + used_bytes=quota_usage.used_bytes, limit_bytes=quota_usage.limit_bytes + ) + + async def update_quota_scope(self, quota_data: QuotaScopeKeyData) -> None: + volume_id = quota_data.volume_id + quota_scope_id = quota_data.quota_scope_id + options = quota_data.options + + await log_manager_api_entry(log, "update_quota_scope", quota_data) + async with self._volume_pool.get_volume(volume_id) as volume: + async with self._handle_external_errors(): + quota_usage = await volume.quota_model.describe_quota_scope(quota_scope_id) + if not quota_usage: + await volume.quota_model.create_quota_scope( + quota_scope_id=quota_scope_id, options=options, extra_args=None + ) + else: + assert options is not None + try: + await volume.quota_model.update_quota_scope( + quota_scope_id=quota_scope_id, + config=options, + ) + except InvalidQuotaConfig: + raise web.HTTPBadRequest(reason="Invalid quota config option") + return None + + async def delete_quota_scope(self, quota_data: QuotaScopeKeyData) -> None: + volume_id = quota_data.volume_id + quota_scope_id = quota_data.quota_scope_id + + await log_manager_api_entry(log, "delete_quota_scope", quota_data) + async with self._volume_pool.get_volume(volume_id) as volume: + async with self._handle_external_errors(): + quota_usage = await volume.quota_model.describe_quota_scope(quota_scope_id) + if not quota_usage: + raise QuotaScopeNotFoundError + await volume.quota_model.unset_quota(quota_scope_id) + return None + + async def create_vfolder(self, vfolder_data: VFolderKeyData) -> NewVFolderCreated: + volume_id = vfolder_data.volume_id + vfolder_id = vfolder_data.vfolder_id + quota_scope_id = vfolder_id.quota_scope_id + + await log_manager_api_entry(log, "create_vfolder", vfolder_data) + assert quota_scope_id is not None + async with self._volume_pool.get_volume(volume_id) as volume: + try: + await volume.create_vfolder(vfolder_id) + except QuotaScopeNotFoundError: + assert quota_scope_id + await volume.quota_model.create_quota_scope(quota_scope_id) + try: + await volume.create_vfolder(vfolder_id) + except QuotaScopeNotFoundError: + raise ExternalError("Failed to create vfolder due to quota scope not found") + return NewVFolderCreated( + vfolder_id=vfolder_id, + quota_scope_path=volume.quota_model.mangle_qspath(quota_scope_id), + vfolder_path=await volume.get_vfolder_mount(vfolder_id, "."), + ) + + async def clone_vfolder(self, vfolder_data: VFolderKeyData) -> NewVFolderCreated: + volume_id = vfolder_data.volume_id + src_vfolder_id = vfolder_data.vfolder_id + dst_vfolder_id = vfolder_data.dst_vfolder_id + + if dst_vfolder_id is None: + raise ValueError("Destination vfolder ID cannot be None") + await log_manager_api_entry(log, "clone_vfolder", vfolder_data) + async with self._volume_pool.get_volume(volume_id) as volume: + await volume.clone_vfolder(src_vfolder_id, dst_vfolder_id) + return NewVFolderCreated( + vfolder_id=dst_vfolder_id, + quota_scope_path=volume.quota_model.mangle_qspath(dst_vfolder_id), + vfolder_path=await volume.get_vfolder_mount(dst_vfolder_id, "."), + ) + + async def get_vfolder_info(self, vfolder_data: VFolderKeyData) -> VFolderMetadata: + volume_id = vfolder_data.volume_id + vfolder_id = vfolder_data.vfolder_id + subpath = vfolder_data.subpath + + await log_manager_api_entry(log, "get_vfolder_info", vfolder_data) + async with self._volume_pool.get_volume(volume_id) as volume: + try: + mount_path = await volume.get_vfolder_mount(vfolder_id, str(subpath)) + usage = await volume.get_usage(vfolder_id) + fs_usage = await volume.get_fs_usage() + except VFolderNotFoundError: + raise web.HTTPGone(reason="VFolder not found") + except InvalidSubpathError: + raise web.HTTPBadRequest(reason="Invalid vfolder subpath") + + return VFolderMetadata( + mount_path=mount_path, + file_count=usage.file_count, + used_bytes=usage.used_bytes, + capacity_bytes=fs_usage.capacity_bytes, + fs_used_bytes=fs_usage.used_bytes, + ) + + async def delete_vfolder(self, vfolder_data: VFolderKeyData) -> None: + volume_id = vfolder_data.volume_id + vfolder_id = vfolder_data.vfolder_id + + await log_manager_api_entry(log, "delete_vfolder", vfolder_data) + try: + async with self._volume_pool.get_volume(volume_id) as volume: + await volume.get_vfolder_mount(vfolder_id, ".") + except VFolderNotFoundError: + ongoing_task = self._deletion_tasks.get(vfolder_id) + if ongoing_task is not None: + ongoing_task.cancel() + raise web.HTTPGone(reason="VFolder not found") + else: + ongoing_task = self._deletion_tasks.get(vfolder_id) + if ongoing_task is not None and ongoing_task.done(): + asyncio.create_task(self._delete_vfolder(vfolder_data)) + return None diff --git a/src/ai/backend/storage/volumes/types.py b/src/ai/backend/storage/volumes/types.py new file mode 100644 index 00000000000..594afa5f365 --- /dev/null +++ b/src/ai/backend/storage/volumes/types.py @@ -0,0 +1,105 @@ +from pathlib import Path, PurePath, PurePosixPath +from typing import List, Optional + +from pydantic import AliasChoices, BaseModel, ConfigDict, Field + +from ai.backend.common.dto.identifiers import VolumeID +from ai.backend.common.types import QuotaConfig, QuotaScopeID, VFolderID + + +class _BaseModel(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + + +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 VolumeKeyData(_BaseModel): + volume_id: VolumeID = VOLUME_ID_FIELD + + +class VolumeMetadata(_BaseModel): + volume_id: VolumeID + backend: str + path: Path + fsprefix: Optional[PurePath] = Field(default=None) + capabilities: list[str] + + +class VolumeMetadataList(_BaseModel): + volumes: List[VolumeMetadata] + + +class VFolderKeyData(_BaseModel): + volume_id: VolumeID = VOLUME_ID_FIELD + vfolder_id: VFolderID = VFOLDER_ID_FIELD + subpath: Optional[PurePosixPath] = Field(default=None) + # 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", + ), + ) + + +class VFolderMetadata(BaseModel): + mount_path: Path + file_count: int + used_bytes: int + capacity_bytes: int + fs_used_bytes: int + + +class NewVFolderCreated(_BaseModel): + vfolder_id: VFolderID = VFOLDER_ID_FIELD + quota_scope_path: Path + vfolder_path: Path + + +class QuotaScopeKeyData(_BaseModel): + volume_id: VolumeID = VOLUME_ID_FIELD + quota_scope_id: QuotaScopeID = QUOTA_SCOPE_ID_FIELD + options: Optional[QuotaConfig] = Field(default=None) + + +class QuotaScopeMetadata(BaseModel): + used_bytes: Optional[int] = Field(default=0) + limit_bytes: Optional[int] = Field(default=0) + + +class NewQuotaScopeCreated(_BaseModel): + quota_scope_id: QuotaScopeID = QUOTA_SCOPE_ID_FIELD + quota_scope_path: Path diff --git a/tests/storage-proxy/vfolder/test_handler.py b/tests/storage-proxy/vfolder/test_handler.py index 6941bffe066..9cc365d160a 100644 --- a/tests/storage-proxy/vfolder/test_handler.py +++ b/tests/storage-proxy/vfolder/test_handler.py @@ -1,291 +1,376 @@ -# import uuid -# from pathlib import Path -# from unittest.mock import AsyncMock - -# import pytest -# from aiohttp import web - -# from ai.backend.common.types import BinarySize, QuotaConfig, QuotaScopeID, QuotaScopeType, VFolderID -# from ai.backend.storage.api.vfolder.manager_handler import VFolderHandler -# from ai.backend.storage.api.vfolder.response_model import ( -# GetVolumeResponseModel, -# NoContentResponseModel, -# ProcessingResponseModel, -# QuotaScopeResponseModel, -# VFolderMetadataResponseModel, -# ) -# from ai.backend.storage.api.vfolder.types import ( -# QuotaScopeIDModel, -# QuotaScopeMetadataModel, -# VFolderIDModel, -# VFolderMetadataModel, -# VolumeIDModel, -# VolumeMetadataListModel, -# VolumeMetadataModel, -# ) - - -# @pytest.fixture -# def mock_vfolder_service(): -# class MockVFolderService: -# async def get_volume(self, volume_data: VolumeIDModel) -> VolumeMetadataListModel: -# return VolumeMetadataListModel( -# volumes=[ -# VolumeMetadataModel( -# volume_id=volume_data.volume_id, -# backend="mock-backend", -# path=Path("/mock/path"), -# fsprefix=None, -# capabilities=["read", "write"], -# ) -# ] -# ) - -# async def get_volumes(self) -> VolumeMetadataListModel: -# return VolumeMetadataListModel( -# volumes=[ -# VolumeMetadataModel( -# volume_id=uuid.UUID("123e4567-e89b-12d3-a456-426614174000"), -# backend="mock-backend", -# path=Path("/mock/path"), -# fsprefix=None, -# capabilities=["read", "write"], -# ) -# ] -# ) - -# async def create_quota_scope(self, quota_data: QuotaScopeIDModel) -> None: -# pass - -# async def get_quota_scope(self, quota_data: QuotaScopeIDModel) -> QuotaScopeMetadataModel: -# return QuotaScopeMetadataModel( -# used_bytes=1024, -# limit_bytes=2048, -# ) - -# async def update_quota_scope(self, quota_data: QuotaScopeIDModel) -> None: -# pass - -# async def delete_quota_scope(self, quota_data: QuotaScopeIDModel) -> None: -# pass - -# async def create_vfolder(self, vfolder_data: VFolderIDModel) -> None: -# pass - -# async def clone_vfolder(self, vfolder_data: VFolderIDModel) -> None: -# pass - -# async def get_vfolder_info(self, vfolder_data: VFolderIDModel) -> VFolderMetadataModel: -# return VFolderMetadataModel( -# mount_path=Path("/mock/mount/path"), -# file_count=100, -# capacity_bytes=1024 * 1024 * 1024, -# used_bytes=BinarySize(1024), -# ) - -# async def delete_vfolder(self, vfolder_data: VFolderIDModel) -> None: -# pass - -# return MockVFolderService() - - -# @pytest.mark.asyncio -# async def test_get_volume(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) - -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# request.json.return_value = {"volume_id": "123e4567-e89b-12d3-a456-426614174000"} -# return request - -# response = await handler.get_volume(await mock_request()) - -# assert isinstance(response, GetVolumeResponseModel) -# assert len(response.volumes) == 1 -# assert response.volumes[0].volume_id == "123e4567-e89b-12d3-a456-426614174000" - - -# @pytest.mark.asyncio -# async def test_get_volumes(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) - -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# return request - -# response = await handler.get_volumes(await mock_request()) - -# assert isinstance(response, GetVolumeResponseModel) -# assert len(response.volumes) == 1 - - -# @pytest.mark.asyncio -# async def test_create_quota_scope(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) - -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# quota_scope_id = QuotaScopeID( -# scope_type=QuotaScopeType.USER, -# scope_id=uuid.UUID("123e4567-e89b-12d3-a456-426614174000"), -# ) -# request.json.return_value = { -# "volume_id": "123e4567-e89b-12d3-a456-426614174000", -# "quota_scope_id": quota_scope_id, -# } -# return request - -# response = await handler.create_quota_scope(await mock_request()) - -# assert isinstance(response, NoContentResponseModel) - - -# @pytest.mark.asyncio -# async def test_get_quota_scope(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) - -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# quota_scope_id = QuotaScopeID( -# scope_type=QuotaScopeType.USER, -# scope_id=uuid.UUID("123e4567-e89b-12d3-a456-426614174000"), -# ) -# request.json.return_value = { -# "volume_id": "123e4567-e89b-12d3-a456-426614174000", -# "quota_scope_id": quota_scope_id, -# } -# return request +import json +import uuid +from pathlib import Path +from unittest.mock import AsyncMock + +import pytest +from aiohttp import web + +from ai.backend.common.types import QuotaConfig, QuotaScopeID, QuotaScopeType, VFolderID +from ai.backend.storage.api.vfolder.handler import VFolderHandler +from ai.backend.storage.volumes.types import ( + NewQuotaScopeCreated, + NewVFolderCreated, + QuotaScopeKeyData, + QuotaScopeMetadata, + VFolderKeyData, + VFolderMetadata, + VolumeKeyData, + VolumeMetadata, + VolumeMetadataList, +) + +UUID = uuid.UUID("123e4567-e89b-12d3-a456-426614174000") +QUOTA_SCOPE_ID = QuotaScopeID( + scope_type=QuotaScopeType.USER, + scope_id=UUID, +) +VFOLDER_SCOPE_ID = VFolderID( + quota_scope_id=QUOTA_SCOPE_ID, + folder_id=UUID, +) + + +@pytest.fixture +def mock_vfolder_service(): + class MockVFolderService: + async def get_volume(self, volume_data: VolumeKeyData) -> VolumeMetadata: + return VolumeMetadata( + volume_id=volume_data.volume_id, + backend="mock-backend", + path=Path("/mock/path"), + fsprefix=None, + capabilities=["read", "write"], + ) + + async def get_volumes(self) -> VolumeMetadataList: + return VolumeMetadataList( + volumes=[ + VolumeMetadata( + volume_id=UUID, + backend="mock-backend", + path=Path("/mock/path"), + fsprefix=None, + capabilities=["read", "write"], + ) + ] + ) + + async def create_quota_scope(self, quota_data: QuotaScopeKeyData) -> NewQuotaScopeCreated: + return NewQuotaScopeCreated( + quota_scope_id=QUOTA_SCOPE_ID, + quota_scope_path=Path("/mock/quota/scope/path"), + ) + + async def get_quota_scope(self, quota_data: QuotaScopeKeyData) -> QuotaScopeMetadata: + return QuotaScopeMetadata( + used_bytes=1024, + limit_bytes=2048, + ) + + async def update_quota_scope(self, quota_data: QuotaScopeKeyData) -> None: + pass + + async def delete_quota_scope(self, quota_data: QuotaScopeKeyData) -> None: + pass + + async def create_vfolder(self, vfolder_data: VFolderKeyData) -> NewVFolderCreated: + return NewVFolderCreated( + vfolder_id=VFOLDER_SCOPE_ID, + quota_scope_path=Path("/mock/quota/scope/path"), + vfolder_path=Path("/mock/vfolder/path"), + ) + + async def clone_vfolder(self, vfolder_data: VFolderKeyData) -> NewVFolderCreated: + return NewVFolderCreated( + vfolder_id=VFOLDER_SCOPE_ID, + quota_scope_path=Path("/mock/quota/scope/path"), + vfolder_path=Path("/mock/vfolder/path"), + ) + + async def get_vfolder_info(self, vfolder_data: VFolderKeyData) -> VFolderMetadata: + return VFolderMetadata( + mount_path=Path("/mock/mount/path"), + file_count=100, + capacity_bytes=1024 * 1024 * 1024, + used_bytes=1024, + fs_used_bytes=512000, + ) + + async def delete_vfolder(self, vfolder_data: VFolderKeyData) -> None: + pass + + return MockVFolderService() + + +@pytest.mark.asyncio +async def test_get_volume(mock_vfolder_service): + handler = VFolderHandler(storage_service=mock_vfolder_service) + + async def mock_request(): + request = AsyncMock(spec=web.Request) + request.match_info = {"volume_id": "123e4567-e89b-12d3-a456-426614174000"} + return request + + response = await handler.get_volume(await mock_request()) + + assert isinstance(response, web.Response) + assert response.status == 200 + assert response.content_type == "application/json" + response_volume_id = json.loads(response.text)["volume_id"] + assert response_volume_id == "123e4567-e89b-12d3-a456-426614174000" + + +@pytest.mark.asyncio +async def test_get_volumes(mock_vfolder_service): + handler = VFolderHandler(storage_service=mock_vfolder_service) + + async def mock_request(): + request = AsyncMock(spec=web.Request) + return request + + response = await handler.get_volumes(await mock_request()) + + assert isinstance(response, web.Response) + assert response.status == 200 + assert response.content_type == "application/json" + response_volumes = json.loads(response.text)["volumes"] + assert len(response_volumes) == 1 + + +@pytest.mark.asyncio +async def test_create_quota_scope(mock_vfolder_service): + handler = VFolderHandler(storage_service=mock_vfolder_service) + + async def mock_request(): + request = AsyncMock(spec=web.Request) + request.json.return_value = { + "volume_id": "123e4567-e89b-12d3-a456-426614174000", + "quota_scope_id": QUOTA_SCOPE_ID, + } + return request -# response = await handler.get_quota_scope(await mock_request()) + response = await handler.create_quota_scope(await mock_request()) -# assert isinstance(response, QuotaScopeResponseModel) -# assert response.used_bytes == 1024 -# assert response.limit_bytes == 2048 + assert isinstance(response, web.Response) + assert response.status == 201 -# @pytest.mark.asyncio -# async def test_update_quota_scope(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) +@pytest.mark.asyncio +async def test_get_quota_scope(mock_vfolder_service): + handler = VFolderHandler(storage_service=mock_vfolder_service) -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# quota_scope_id = QuotaScopeID( -# scope_type=QuotaScopeType.USER, -# scope_id=uuid.UUID("123e4567-e89b-12d3-a456-426614174000"), -# ) -# request.json.return_value = { -# "volume_id": "123e4567-e89b-12d3-a456-426614174000", -# "quota_scope_id": quota_scope_id, -# "options": QuotaConfig(limit_bytes=2048), # QuotaConfig 객체 사용 -# } -# return request + async def mock_request(): + request = AsyncMock(spec=web.Request) + request.match_info = { + "volume_id": "123e4567-e89b-12d3-a456-426614174000", + "quota_scope_id": QUOTA_SCOPE_ID, + } + return request -# response = await handler.update_quota_scope(await mock_request()) + response = await handler.get_quota_scope(await mock_request()) -# assert isinstance(response, NoContentResponseModel) - - -# @pytest.mark.asyncio -# async def test_delete_quota_scope(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) - -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# quota_scope_id = QuotaScopeID( -# scope_type=QuotaScopeType.USER, -# scope_id=uuid.UUID("123e4567-e89b-12d3-a456-426614174000"), -# ) -# request.json.return_value = { -# "volume_id": "123e4567-e89b-12d3-a456-426614174000", -# "quota_scope_id": quota_scope_id, -# } -# return request - -# response = await handler.delete_quota_scope(await mock_request()) - -# assert isinstance(response, NoContentResponseModel) - - -# @pytest.mark.asyncio -# async def test_create_vfolder(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) - -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# vfolder_id = VFolderID( -# folder_id=uuid.UUID("123e4567-e89b-12d3-a456-426614174000"), quota_scope_id=None -# ) -# request.json.return_value = { -# "volume_id": "123e4567-e89b-12d3-a456-426614174000", -# "vfolder_id": vfolder_id, -# } -# return request - -# response = await handler.create_vfolder(await mock_request()) - -# assert isinstance(response, NoContentResponseModel) - - -# @pytest.mark.asyncio -# async def test_clone_vfolder(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) - -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# vfolder_id = VFolderID( -# folder_id=uuid.UUID("123e4567-e89b-12d3-a456-426614174000"), quota_scope_id=None -# ) -# request.json.return_value = { -# "volume_id": "123e4567-e89b-12d3-a456-426614174000", -# "vfolder_id": vfolder_id, -# "dst_vfolder_id": vfolder_id, -# } -# return request - -# response = await handler.clone_vfolder(await mock_request()) - -# assert isinstance(response, NoContentResponseModel) - - -# @pytest.mark.asyncio -# async def test_get_vfolder_info(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) - -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# vfolder_id = VFolderID( -# folder_id=uuid.UUID("123e4567-e89b-12d3-a456-426614174000"), quota_scope_id=None -# ) -# request.json.return_value = { -# "volume_id": "123e4567-e89b-12d3-a456-426614174000", -# "vfolder_id": vfolder_id, -# } -# return request - -# response = await handler.get_vfolder_info(await mock_request()) - -# assert isinstance(response, VFolderMetadataResponseModel) -# assert response.mount_path == "/mock/mount/path" -# assert response.file_count == 100 -# assert response.capacity_bytes == 1024 * 1024 * 1024 -# assert response.used_bytes == 1024 - - -# @pytest.mark.asyncio -# async def test_delete_vfolder(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) - -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# vfolder_id = VFolderID( -# folder_id=uuid.UUID("123e4567-e89b-12d3-a456-426614174000"), quota_scope_id=None -# ) -# request.json.return_value = { -# "volume_id": "123e4567-e89b-12d3-a456-426614174000", -# "vfolder_id": vfolder_id, -# } -# return request - -# response = await handler.delete_vfolder(await mock_request()) - -# assert isinstance(response, ProcessingResponseModel) + assert isinstance(response, web.Response) + assert response.status == 200 + assert response.content_type == "application/json" + response_data = json.loads(response.text)["used_bytes"] + assert response_data == 1024 + + +@pytest.mark.asyncio +async def test_update_quota_scope(mock_vfolder_service): + handler = VFolderHandler(storage_service=mock_vfolder_service) + + async def mock_request(): + request = AsyncMock(spec=web.Request) + request.json.return_value = { + "volume_id": "123e4567-e89b-12d3-a456-426614174000", + "quota_scope_id": QUOTA_SCOPE_ID, + "options": QuotaConfig(limit_bytes=2048), # QuotaConfig 객체 사용 + } + return request + + response = await handler.update_quota_scope(await mock_request()) + + assert isinstance(response, web.Response) + assert response.status == 204 + + +@pytest.mark.asyncio +async def test_delete_quota_scope(mock_vfolder_service): + handler = VFolderHandler(storage_service=mock_vfolder_service) + + async def mock_request(): + request = AsyncMock(spec=web.Request) + request.match_info = { + "volume_id": "123e4567-e89b-12d3-a456-426614174000", + "quota_scope_id": QUOTA_SCOPE_ID, + } + return request + + response = await handler.delete_quota_scope(await mock_request()) + + assert isinstance(response, web.Response) + assert response.status == 204 + + +@pytest.mark.asyncio +async def test_create_vfolder(mock_vfolder_service): + handler = VFolderHandler(storage_service=mock_vfolder_service) + + async def mock_request(): + request = AsyncMock(spec=web.Request) + request.json.return_value = { + "volume_id": "123e4567-e89b-12d3-a456-426614174000", + "vfolder_id": VFOLDER_SCOPE_ID, + } + return request + + response = await handler.create_vfolder(await mock_request()) + + assert isinstance(response, web.Response) + assert response.status == 201 + + +@pytest.mark.asyncio +async def test_clone_vfolder(mock_vfolder_service): + handler = VFolderHandler(storage_service=mock_vfolder_service) + + async def mock_request(): + request = AsyncMock(spec=web.Request) + request.json.return_value = { + "volume_id": "123e4567-e89b-12d3-a456-426614174000", + "vfolder_id": VFOLDER_SCOPE_ID, + "dst_vfolder_id": VFOLDER_SCOPE_ID, + } + return request + + response = await handler.clone_vfolder(await mock_request()) + + assert isinstance(response, web.Response) + assert response.status == 201 + + +@pytest.mark.asyncio +async def test_get_vfolder_info(mock_vfolder_service): + handler = VFolderHandler(storage_service=mock_vfolder_service) + + async def mock_request(): + request = AsyncMock(spec=web.Request) + request.match_info = { + "volume_id": "123e4567-e89b-12d3-a456-426614174000", + "vfolder_id": VFOLDER_SCOPE_ID, + "subpath": "/mock/subpath", + } + return request + + response = await handler.get_vfolder_info(await mock_request()) + + assert isinstance(response, web.Response) + assert response.status == 200 + assert response.content_type == "application/json" + response_data = json.loads(response.text)["mount_path"] + assert response_data == "/mock/mount/path" + + +@pytest.mark.asyncio +async def test_get_vfolder_mount(mock_vfolder_service): + handler = VFolderHandler(storage_service=mock_vfolder_service) + + async def mock_request(): + request = AsyncMock(spec=web.Request) + request.match_info = { + "volume_id": "123e4567-e89b-12d3-a456-426614174000", + "vfolder_id": VFOLDER_SCOPE_ID, + "subpath": "/mock/subpath", + } + return request + + response = await handler.get_vfolder_mount(await mock_request()) + + assert isinstance(response, web.Response) + assert response.status == 200 + assert response.content_type == "application/json" + response_data = json.loads(response.text)["mount_path"] + assert response_data == "/mock/mount/path" + + +@pytest.mark.asyncio +async def test_get_vfolder_usage(mock_vfolder_service): + handler = VFolderHandler(storage_service=mock_vfolder_service) + + async def mock_request(): + request = AsyncMock(spec=web.Request) + request.match_info = { + "volume_id": "123e4567-e89b-12d3-a456-426614174000", + "vfolder_id": VFOLDER_SCOPE_ID, + } + return request + + response = await handler.get_vfolder_usage(await mock_request()) + + assert isinstance(response, web.Response) + assert response.status == 200 + assert response.content_type == "application/json" + response_data = json.loads(response.text)["file_count"] + assert response_data == 100 + + +@pytest.mark.asyncio +async def test_get_vfolder_used_bytes(mock_vfolder_service): + handler = VFolderHandler(storage_service=mock_vfolder_service) + + async def mock_request(): + request = AsyncMock(spec=web.Request) + request.match_info = { + "volume_id": "123e4567-e89b-12d3-a456-426614174000", + "vfolder_id": VFOLDER_SCOPE_ID, + } + return request + + response = await handler.get_vfolder_used_bytes(await mock_request()) + + assert isinstance(response, web.Response) + assert response.status == 200 + assert response.content_type == "application/json" + response_data = json.loads(response.text)["used_bytes"] + assert response_data == 1024 + + +@pytest.mark.asyncio +async def test_get_vfolder_fs_usage(mock_vfolder_service): + handler = VFolderHandler(storage_service=mock_vfolder_service) + + async def mock_request(): + request = AsyncMock(spec=web.Request) + request.match_info = { + "volume_id": "123e4567-e89b-12d3-a456-426614174000", + "vfolder_id": VFOLDER_SCOPE_ID, + } + return request + + response = await handler.get_vfolder_fs_usage(await mock_request()) + + assert isinstance(response, web.Response) + assert response.status == 200 + assert response.content_type == "application/json" + response_data = json.loads(response.text)["capacity_bytes"] + assert response_data == 1024 * 1024 * 1024 + + +@pytest.mark.asyncio +async def test_delete_vfolder(mock_vfolder_service): + handler = VFolderHandler(storage_service=mock_vfolder_service) + + async def mock_request(): + request = AsyncMock(spec=web.Request) + request.match_info = { + "volume_id": "123e4567-e89b-12d3-a456-426614174000", + "vfolder_id": VFOLDER_SCOPE_ID, + } + return request + + response = await handler.delete_vfolder(await mock_request()) + + assert isinstance(response, web.Response) + assert response.status == 202 diff --git a/tests/storage-proxy/vfolder/test_pool.py b/tests/storage-proxy/vfolder/test_pool.py new file mode 100644 index 00000000000..0e86e25e3e7 --- /dev/null +++ b/tests/storage-proxy/vfolder/test_pool.py @@ -0,0 +1,77 @@ +from typing import Mapping +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from ai.backend.common.dto.identifiers import VolumeID +from ai.backend.common.etcd import AsyncEtcd +from ai.backend.common.events import EventDispatcher, EventProducer +from ai.backend.storage.exception import InvalidVolumeError +from ai.backend.storage.types import VolumeInfo +from ai.backend.storage.volumes.pool import VolumePool + + +@pytest.fixture +def mock_etcd(): + return AsyncMock(spec=AsyncEtcd) + + +@pytest.fixture +def mock_event_dispatcher(): + return MagicMock(spec=EventDispatcher) + + +@pytest.fixture +def mock_event_producer(): + return MagicMock(spec=EventProducer) + + +def list_volumes(self) -> Mapping[str, VolumeInfo]: + return { + volume_id: VolumeInfo( + backend=info["backend"], + path=info["path"], + fsprefix=info.get("fsprefix", ""), + options=None, + ) + for volume_id, info in self._local_config["volume"].items() + } + + +def get_volume_info(self, volume_id: VolumeID) -> VolumeInfo: + if volume_id not in self._local_config["volume"]: + raise InvalidVolumeError(volume_id) + volume_config = self._local_config["volume"][volume_id] + return VolumeInfo( + backend=volume_config["backend"], + path=volume_config["path"], + fsprefix=volume_config.get("fsprefix", ""), + options=None, + ) + + +@pytest.mark.asyncio +async def test_get_volume(): + local_config = { + "volume": { + "test_volume": { + "backend": "vfs", + "path": "/mnt/test_volume", + "options": {}, + "fsprefix": "vfs-test", + } + }, + "storage-proxy": {"scandir-limit": 1000}, + } + + mock_etcd = AsyncMock() + mock_event_dispatcher = MagicMock() + mock_event_producer = MagicMock() + + pool = VolumePool(local_config, mock_etcd, mock_event_dispatcher, mock_event_producer) + await pool.__aenter__() + async with pool.get_volume("test_volume") as volume: + assert volume is not None + + async with pool.get_volume("test_volume") as volume2: + assert volume is volume2 From b200f07faf389bd00bd248d398f4b51bfd69a0fa Mon Sep 17 00:00:00 2001 From: HyeockJinKim Date: Wed, 12 Feb 2025 03:01:26 +0900 Subject: [PATCH 2/7] refactor: Organize new objects and refine structure --- src/ai/backend/common/BUILD | 1 + src/ai/backend/common/api_handlers.py | 4 + src/ai/backend/common/dto/BUILD | 1 - src/ai/backend/common/dto/identifiers.py | 4 - src/ai/backend/common/dto/manager/request.py | 18 +- src/ai/backend/common/dto/storage/field.py | 21 + src/ai/backend/common/dto/storage/path.py | 27 + src/ai/backend/common/dto/storage/request.py | 75 +- src/ai/backend/common/dto/storage/response.py | 39 +- src/ai/backend/common/types.py | 4 + src/ai/backend/storage/api/vfolder/handler.py | 255 ++----- .../storage/api/vfolder/manager_handler.py | 144 ---- src/ai/backend/storage/services/BUILD | 1 + .../storage/{volumes => services}/service.py | 185 ++--- src/ai/backend/storage/volumes/pool.py | 26 +- src/ai/backend/storage/volumes/types.py | 153 ++-- tests/storage-proxy/vfolder/test_handler.py | 666 +++++++++--------- tests/storage-proxy/vfolder/test_pool.py | 2 +- 18 files changed, 650 insertions(+), 976 deletions(-) delete mode 100644 src/ai/backend/common/dto/BUILD delete mode 100644 src/ai/backend/common/dto/identifiers.py create mode 100644 src/ai/backend/common/dto/storage/field.py create mode 100644 src/ai/backend/common/dto/storage/path.py delete mode 100644 src/ai/backend/storage/api/vfolder/manager_handler.py create mode 100644 src/ai/backend/storage/services/BUILD rename src/ai/backend/storage/{volumes => services}/service.py (61%) diff --git a/src/ai/backend/common/BUILD b/src/ai/backend/common/BUILD index 6bf7cd623e3..5068f572167 100644 --- a/src/ai/backend/common/BUILD +++ b/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 diff --git a/src/ai/backend/common/api_handlers.py b/src/ai/backend/common/api_handlers.py index d7892138a10..def733926da 100644 --- a/src/ai/backend/common/api_handlers.py +++ b/src/ai/backend/common/api_handlers.py @@ -126,6 +126,10 @@ def from_request(cls, request: web.Request) -> Self: pass +class BaseRequestModel(BaseModel): + pass + + class BaseResponseModel(BaseModel): pass diff --git a/src/ai/backend/common/dto/BUILD b/src/ai/backend/common/dto/BUILD deleted file mode 100644 index c1ffc1a1410..00000000000 --- a/src/ai/backend/common/dto/BUILD +++ /dev/null @@ -1 +0,0 @@ -python_sources(name="src") \ No newline at end of file diff --git a/src/ai/backend/common/dto/identifiers.py b/src/ai/backend/common/dto/identifiers.py deleted file mode 100644 index a29b8e13a94..00000000000 --- a/src/ai/backend/common/dto/identifiers.py +++ /dev/null @@ -1,4 +0,0 @@ -import uuid -from typing import TypeAlias - -VolumeID: TypeAlias = uuid.UUID diff --git a/src/ai/backend/common/dto/manager/request.py b/src/ai/backend/common/dto/manager/request.py index a9835adbc5d..e3c5b4bd737 100644 --- a/src/ai/backend/common/dto/manager/request.py +++ b/src/ai/backend/common/dto/manager/request.py @@ -1,15 +1,17 @@ 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( @@ -17,7 +19,7 @@ class VFolderCreateReq(BaseModel): 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", ) diff --git a/src/ai/backend/common/dto/storage/field.py b/src/ai/backend/common/dto/storage/field.py new file mode 100644 index 00000000000..a6c6cef558e --- /dev/null +++ b/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 diff --git a/src/ai/backend/common/dto/storage/path.py b/src/ai/backend/common/dto/storage/path.py new file mode 100644 index 00000000000..df0f762586a --- /dev/null +++ b/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.", + ) diff --git a/src/ai/backend/common/dto/storage/request.py b/src/ai/backend/common/dto/storage/request.py index 50a5b7c45bb..babd5dedae8 100644 --- a/src/ai/backend/common/dto/storage/request.py +++ b/src/ai/backend/common/dto/storage/request.py @@ -1,69 +1,26 @@ -from pathlib import PurePosixPath from typing import Optional -from pydantic import AliasChoices, BaseModel, ConfigDict, Field +from pydantic import AliasChoices, Field -from ai.backend.common.dto.identifiers import VolumeID -from ai.backend.common.types import QuotaConfig, QuotaScopeID, VFolderID +from ...api_handlers import BaseRequestModel +from ...types import QuotaConfig, VFolderID -class _BaseModel(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True) - - -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 VolumeKeyDataParams(_BaseModel): - volume_id: VolumeID = VOLUME_ID_FIELD +class QuotaScopeReq(BaseRequestModel): + options: Optional[QuotaConfig] = Field( + default=None, + description="The options for the quota scope.", + ) -class VFolderKeyDataParams(_BaseModel): - volume_id: VolumeID = VOLUME_ID_FIELD - vfolder_id: VFolderID = VFOLDER_ID_FIELD - subpath: Optional[PurePosixPath] = Field(default=None) - # 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", - ), +class GetVFolderMetaReq(BaseRequestModel): + subpath: str = Field( + description="The subpath of the virtual folder.", ) -class QuotaScopeKeyDataParams(_BaseModel): - volume_id: VolumeID = VOLUME_ID_FIELD - quota_scope_id: QuotaScopeID = QUOTA_SCOPE_ID_FIELD - options: Optional[QuotaConfig] = Field(default=None) +class CloneVFolderReq(BaseRequestModel): + dst_vfolder_id: VFolderID = Field( + description="The destination virtual folder ID.", + validation_alias=AliasChoices("dst_vfid", "dst_vfolder_id"), + ) diff --git a/src/ai/backend/common/dto/storage/response.py b/src/ai/backend/common/dto/storage/response.py index 0cd5a898e61..fc358a31b41 100644 --- a/src/ai/backend/common/dto/storage/response.py +++ b/src/ai/backend/common/dto/storage/response.py @@ -1,20 +1,17 @@ -from typing import List, Optional +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 VolumeMetadataResponse(BaseResponseModel): - volume_id: str - backend: str - path: str - fsprefix: Optional[str] = Field(default=None) - capabilities: List[str] +class GetVolumeResponse(BaseResponseModel): + item: VolumeMetaField -class GetVolumeResponse(BaseResponseModel): - volumes: List[VolumeMetadataResponse] +class GetVolumesResponse(BaseResponseModel): + items: list[VolumeMetaField] class QuotaScopeResponse(BaseResponseModel): @@ -23,26 +20,4 @@ class QuotaScopeResponse(BaseResponseModel): class VFolderMetadataResponse(BaseResponseModel): - mount_path: str - file_count: int - used_bytes: int - capacity_bytes: int - fs_used_bytes: int - - -class VFolderMountResponse(BaseResponseModel): - mount_path: str - - -class VFolderUsageResponse(BaseResponseModel): - file_count: int - used_bytes: int - - -class VFolderUsedBytesResponse(BaseResponseModel): - used_bytes: int - - -class VFolderFSUsageResponse(BaseResponseModel): - capacity_bytes: int - fs_used_bytes: int + item: VFolderMetaField diff --git a/src/ai/backend/common/types.py b/src/ai/backend/common/types.py index bada0fbad73..eec9591ea69 100644 --- a/src/ai/backend/common/types.py +++ b/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 diff --git a/src/ai/backend/storage/api/vfolder/handler.py b/src/ai/backend/storage/api/vfolder/handler.py index 69bfcb3a1dd..c0e5b3e1743 100644 --- a/src/ai/backend/storage/api/vfolder/handler.py +++ b/src/ai/backend/storage/api/vfolder/handler.py @@ -1,54 +1,52 @@ -from typing import Protocol +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 ( - QuotaScopeKeyDataParams, - VFolderKeyDataParams, - VolumeKeyDataParams, + CloneVFolderReq, + GetVFolderMetaReq, + QuotaScopeReq, ) from ai.backend.common.dto.storage.response import ( GetVolumeResponse, - QuotaScopeResponse, - VFolderFSUsageResponse, + GetVolumesResponse, VFolderMetadataResponse, - VFolderMountResponse, - VFolderUsageResponse, - VFolderUsedBytesResponse, - VolumeMetadataResponse, ) -from ai.backend.storage.volumes.types import ( - NewQuotaScopeCreated, - NewVFolderCreated, - QuotaScopeKeyData, - QuotaScopeMetadata, - VFolderKeyData, - VFolderMetadata, - VolumeKeyData, - VolumeMetadata, - VolumeMetadataList, +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_data: VolumeKeyData) -> VolumeMetadata: ... + async def get_volume(self, volume_id: VolumeID) -> VolumeMeta: ... - async def get_volumes(self) -> VolumeMetadataList: ... + async def get_volumes(self) -> list[VolumeMeta]: ... - async def create_quota_scope(self, quota_data: QuotaScopeKeyData) -> NewQuotaScopeCreated: ... + async def create_quota_scope( + self, quota_scope_key: QuotaScopeKey, options: Optional[QuotaConfig] + ) -> None: ... - async def get_quota_scope(self, quota_data: QuotaScopeKeyData) -> QuotaScopeMetadata: ... + async def get_quota_scope(self, quota_scope_key: QuotaScopeKey) -> QuotaScopeMeta: ... - async def update_quota_scope(self, quota_data: QuotaScopeKeyData) -> None: ... + async def update_quota_scope( + self, quota_scope_key: QuotaScopeKey, options: Optional[QuotaConfig] + ) -> None: ... - async def delete_quota_scope(self, quota_data: QuotaScopeKeyData) -> None: ... + async def delete_quota_scope(self, quota_scope_key: QuotaScopeKey) -> None: ... - async def create_vfolder(self, vfolder_data: VFolderKeyData) -> NewVFolderCreated: ... + async def create_vfolder(self, vfolder_key: VFolderKey) -> None: ... - async def clone_vfolder(self, vfolder_data: VFolderKeyData) -> NewVFolderCreated: ... + async def clone_vfolder(self, vfolder_key: VFolderKey, dst_vfolder_id: VFolderID) -> None: ... - async def get_vfolder_info(self, vfolder_data: VFolderKeyData) -> VFolderMetadata: ... + async def get_vfolder_info(self, vfolder_key: VFolderKey, subpath: str) -> VFolderMeta: ... - async def delete_vfolder(self, vfolder_data: VFolderKeyData) -> None: ... + async def delete_vfolder(self, vfolder_key: VFolderKey) -> None: ... class VFolderHandler: @@ -56,192 +54,85 @@ def __init__(self, storage_service: VFolderServiceProtocol) -> None: self.storage_service = storage_service @api_handler - async def get_volume(self, body: PathParam[VolumeKeyDataParams]) -> APIResponse: - volume_parsed = body.parsed - volume_params = VolumeKeyData(volume_id=volume_parsed.volume_id) - volume_data = await self.storage_service.get_volume(volume_params) + 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=VolumeMetadataResponse( - volume_id=str(volume_data.volume_id), - backend=volume_data.backend, - path=str(volume_data.path), - fsprefix=str(volume_data.fsprefix) if volume_data.fsprefix else None, - capabilities=[cap for cap in volume_data.capabilities], + response_model=GetVolumeResponse( + item=volume_meta.to_field(), ), ) @api_handler async def get_volumes(self) -> APIResponse: - volumes_data = await self.storage_service.get_volumes() + volume_meta_list = await self.storage_service.get_volumes() return APIResponse.build( status_code=200, - response_model=GetVolumeResponse( - volumes=[ - VolumeMetadataResponse( - volume_id=str(volume.volume_id), - backend=volume.backend, - path=str(volume.path), - fsprefix=str(volume.fsprefix) if volume.fsprefix else None, - capabilities=[cap for cap in volume.capabilities], - ) - for volume in volumes_data.volumes - ] + response_model=GetVolumesResponse( + items=[volume.to_field() for volume in volume_meta_list], ), ) @api_handler - async def create_quota_scope(self, body: BodyParam[QuotaScopeKeyDataParams]) -> APIResponse: - quota_parsed = body.parsed - quota_params = QuotaScopeKeyData( - volume_id=quota_parsed.volume_id, - quota_scope_id=quota_parsed.quota_scope_id, - options=quota_parsed.options, - ) - await self.storage_service.create_quota_scope(quota_params) - return APIResponse.no_content(status_code=201) + 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, body: PathParam[QuotaScopeKeyDataParams]) -> APIResponse: - quota_parsed = body.parsed - quota_params = QuotaScopeKeyData( - volume_id=quota_parsed.volume_id, - quota_scope_id=quota_parsed.quota_scope_id, - ) - quota_scope = await self.storage_service.get_quota_scope(quota_params) + 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=QuotaScopeResponse( - used_bytes=quota_scope.used_bytes, limit_bytes=quota_scope.limit_bytes - ), + response_model=quota_scope.to_response(), ) @api_handler - async def update_quota_scope(self, body: BodyParam[QuotaScopeKeyDataParams]) -> APIResponse: - quota_parsed = body.parsed - quota_params = QuotaScopeKeyData( - volume_id=quota_parsed.volume_id, - quota_scope_id=quota_parsed.quota_scope_id, - options=quota_parsed.options, - ) - await self.storage_service.update_quota_scope(quota_params) + 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, body: PathParam[QuotaScopeKeyDataParams]) -> APIResponse: - quota_parsed = body.parsed - quota_params = QuotaScopeKeyData( - volume_id=quota_parsed.volume_id, - quota_scope_id=quota_parsed.quota_scope_id, - ) - await self.storage_service.delete_quota_scope(quota_params) + 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, body: BodyParam[VFolderKeyDataParams]) -> APIResponse: - vfolder_parsed = body.parsed - vfolder_params = VFolderKeyData( - volume_id=vfolder_parsed.volume_id, - vfolder_id=vfolder_parsed.vfolder_id, - ) - await self.storage_service.create_vfolder(vfolder_params) - return APIResponse.no_content(status_code=201) + 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, body: BodyParam[VFolderKeyDataParams]) -> APIResponse: - vfolder_parsed = body.parsed - vfolder_params = VFolderKeyData( - volume_id=vfolder_parsed.volume_id, - vfolder_id=vfolder_parsed.vfolder_id, - dst_vfolder_id=vfolder_parsed.dst_vfolder_id, - ) - await self.storage_service.clone_vfolder(vfolder_params) - return APIResponse.no_content(status_code=201) + 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, body: PathParam[VFolderKeyDataParams]) -> APIResponse: - vfolder_parsed = body.parsed - vfolder_params = VFolderKeyData( - volume_id=vfolder_parsed.volume_id, - vfolder_id=vfolder_parsed.vfolder_id, - subpath=vfolder_parsed.subpath, - ) - metadata = await self.storage_service.get_vfolder_info(vfolder_params) + 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( - mount_path=str(metadata.mount_path), - file_count=metadata.file_count, - used_bytes=metadata.used_bytes, - capacity_bytes=metadata.capacity_bytes, - fs_used_bytes=metadata.fs_used_bytes, + item=vfolder_meta.to_field(), ), ) @api_handler - async def get_vfolder_mount(self, body: PathParam[VFolderKeyDataParams]) -> APIResponse: - vfolder_parsed = body.parsed - vfolder_params = VFolderKeyData( - volume_id=vfolder_parsed.volume_id, - vfolder_id=vfolder_parsed.vfolder_id, - subpath=vfolder_parsed.subpath, - ) - metadata = await self.storage_service.get_vfolder_info(vfolder_params) - return APIResponse.build( - status_code=200, - response_model=VFolderMountResponse(mount_path=str(metadata.mount_path)), - ) - - @api_handler - async def get_vfolder_usage(self, body: PathParam[VFolderKeyDataParams]) -> APIResponse: - vfolder_parsed = body.parsed - vfolder_params = VFolderKeyData( - volume_id=vfolder_parsed.volume_id, - vfolder_id=vfolder_parsed.vfolder_id, - ) - metadata = await self.storage_service.get_vfolder_info(vfolder_params) - return APIResponse.build( - status_code=200, - response_model=VFolderUsageResponse( - file_count=metadata.file_count, used_bytes=metadata.used_bytes - ), - ) - - @api_handler - async def get_vfolder_used_bytes(self, body: PathParam[VFolderKeyDataParams]) -> APIResponse: - vfolder_parsed = body.parsed - vfolder_params = VFolderKeyData( - volume_id=vfolder_parsed.volume_id, - vfolder_id=vfolder_parsed.vfolder_id, - ) - metadata = await self.storage_service.get_vfolder_info(vfolder_params) - return APIResponse.build( - status_code=200, - response_model=VFolderUsedBytesResponse(used_bytes=metadata.used_bytes), - ) - - @api_handler - async def get_vfolder_fs_usage(self, body: PathParam[VFolderKeyDataParams]) -> APIResponse: - vfolder_parsed = body.parsed - vfolder_params = VFolderKeyData( - volume_id=vfolder_parsed.volume_id, - vfolder_id=vfolder_parsed.vfolder_id, - ) - metadata = await self.storage_service.get_vfolder_info(vfolder_params) - return APIResponse.build( - status_code=200, - response_model=VFolderFSUsageResponse( - capacity_bytes=metadata.capacity_bytes, - fs_used_bytes=metadata.fs_used_bytes, - ), - ) - - @api_handler - async def delete_vfolder(self, body: PathParam[VFolderKeyDataParams]) -> APIResponse: - vfolder_parsed = body.parsed - vfolder_params = VFolderKeyData( - volume_id=vfolder_parsed.volume_id, - vfolder_id=vfolder_parsed.vfolder_id, - ) - await self.storage_service.delete_vfolder(vfolder_params) - return APIResponse.no_content(status_code=202) + 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) diff --git a/src/ai/backend/storage/api/vfolder/manager_handler.py b/src/ai/backend/storage/api/vfolder/manager_handler.py deleted file mode 100644 index 344c8d4d390..00000000000 --- a/src/ai/backend/storage/api/vfolder/manager_handler.py +++ /dev/null @@ -1,144 +0,0 @@ -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) diff --git a/src/ai/backend/storage/services/BUILD b/src/ai/backend/storage/services/BUILD new file mode 100644 index 00000000000..73574424040 --- /dev/null +++ b/src/ai/backend/storage/services/BUILD @@ -0,0 +1 @@ +python_sources(name="src") diff --git a/src/ai/backend/storage/volumes/service.py b/src/ai/backend/storage/services/service.py similarity index 61% rename from src/ai/backend/storage/volumes/service.py rename to src/ai/backend/storage/services/service.py index bb91eedaeb5..5b3815f320a 100644 --- a/src/ai/backend/storage/volumes/service.py +++ b/src/ai/backend/storage/services/service.py @@ -4,34 +4,31 @@ import uuid import weakref from contextlib import asynccontextmanager as actxmgr -from typing import AsyncIterator +from typing import AsyncIterator, Optional from aiohttp import web -from ai.backend.common.dto.identifiers import VolumeID from ai.backend.common.events import VFolderDeletionFailureEvent, VFolderDeletionSuccessEvent -from ai.backend.common.types import VFolderID +from ai.backend.common.types import QuotaConfig, VFolderID, VolumeID from ai.backend.logging.utils import BraceStyleAdapter -from ai.backend.storage.exception import ( + +from ..exception import ( ExternalError, InvalidQuotaConfig, + InvalidQuotaScopeError, InvalidSubpathError, QuotaScopeAlreadyExists, QuotaScopeNotFoundError, VFolderNotFoundError, ) -from ai.backend.storage.utils import log_manager_api_entry -from ai.backend.storage.volumes.pool import VolumePool -from ai.backend.storage.volumes.types import ( - NewQuotaScopeCreated, - NewVFolderCreated, - QuotaScopeKeyData, - QuotaScopeMetadata, - VFolderKeyData, - VFolderMetadata, - VolumeKeyData, - VolumeMetadata, - VolumeMetadataList, +from ..utils import log_manager_api_entry +from ..volumes.pool import VolumePool +from ..volumes.types import ( + QuotaScopeKey, + QuotaScopeMeta, + VFolderKey, + VFolderMeta, + VolumeMeta, ) log = BraceStyleAdapter(logging.getLogger(__spec__.name)) @@ -67,10 +64,10 @@ async def _handle_external_errors(self) -> AsyncIterator[None]: async def _delete_vfolder( self, - vfolder_data: VFolderKeyData, + vfolder_key: VFolderKey, ) -> None: - volume_id = vfolder_data.volume_id - vfolder_id = vfolder_data.vfolder_id + volume_id = vfolder_key.volume_id + vfolder_id = vfolder_key.vfolder_id current_task = asyncio.current_task() assert current_task is not None @@ -105,11 +102,10 @@ async def _delete_vfolder( VFolderDeletionSuccessEvent(vfolder_id) ) - async def get_volume(self, volume_data: VolumeKeyData) -> VolumeMetadata: - volume_id = volume_data.volume_id + async def get_volume(self, volume_id: VolumeID) -> VolumeMeta: await log_manager_api_entry(log, "get_volume", volume_id) volume = self._volume_pool.get_volume_info(volume_id) - return VolumeMetadata( + return VolumeMeta( volume_id=volume_id, backend=volume.backend, path=volume.path, @@ -117,29 +113,26 @@ async def get_volume(self, volume_data: VolumeKeyData) -> VolumeMetadata: capabilities=await self._get_capabilities(volume_id), ) - async def get_volumes(self) -> VolumeMetadataList: + async def get_volumes(self) -> list[VolumeMeta]: await log_manager_api_entry(log, "get_volumes", params=None) volumes = self._volume_pool.list_volumes() - return VolumeMetadataList( - volumes=[ - VolumeMetadata( - volume_id=uuid.UUID(volume_id), - backend=info.backend, - path=info.path, - fsprefix=info.fsprefix, - capabilities=await self._get_capabilities(uuid.UUID(volume_id)), - ) - for volume_id, info in volumes.items() - ] - ) - - async def create_quota_scope(self, quota_data: QuotaScopeKeyData) -> NewQuotaScopeCreated: - volume_id = quota_data.volume_id - quota_scope_id = quota_data.quota_scope_id - options = quota_data.options + return [ + VolumeMeta( + volume_id=uuid.UUID(volume_id), + backend=info.backend, + path=info.path, + fsprefix=info.fsprefix, + capabilities=await self._get_capabilities(uuid.UUID(volume_id)), + ) + for volume_id, info in volumes.items() + ] - await log_manager_api_entry(log, "create_quota_scope", quota_data) - async with self._volume_pool.get_volume(volume_id) as volume: + async def create_quota_scope( + self, quota_scope_key: QuotaScopeKey, options: Optional[QuotaConfig] + ) -> None: + quota_scope_id = quota_scope_key.quota_scope_id + await log_manager_api_entry(log, "create_quota_scope", quota_scope_key) + async with self._volume_pool.get_volume(quota_scope_key.volume_id) as volume: try: async with self._handle_external_errors(): await volume.quota_model.create_quota_scope( @@ -147,32 +140,26 @@ async def create_quota_scope(self, quota_data: QuotaScopeKeyData) -> NewQuotaSco ) except QuotaScopeAlreadyExists: raise web.HTTPConflict(reason="Volume already exists with given quota scope.") - return NewQuotaScopeCreated( - quota_scope_id=quota_scope_id, - quota_scope_path=volume.quota_model.mangle_qspath(quota_scope_id), - ) - - async def get_quota_scope(self, quota_data: QuotaScopeKeyData) -> QuotaScopeMetadata: - volume_id = quota_data.volume_id - quota_scope_id = quota_data.quota_scope_id - await log_manager_api_entry(log, "get_quota_scope", quota_data) - async with self._volume_pool.get_volume(volume_id) as volume: + async def get_quota_scope(self, quota_scope_key: QuotaScopeKey) -> QuotaScopeMeta: + await log_manager_api_entry(log, "get_quota_scope", quota_scope_key) + async with self._volume_pool.get_volume(quota_scope_key.volume_id) as volume: async with self._handle_external_errors(): - quota_usage = await volume.quota_model.describe_quota_scope(quota_scope_id) + quota_usage = await volume.quota_model.describe_quota_scope( + quota_scope_key.quota_scope_id + ) if not quota_usage: raise QuotaScopeNotFoundError - return QuotaScopeMetadata( + return QuotaScopeMeta( used_bytes=quota_usage.used_bytes, limit_bytes=quota_usage.limit_bytes ) - async def update_quota_scope(self, quota_data: QuotaScopeKeyData) -> None: - volume_id = quota_data.volume_id - quota_scope_id = quota_data.quota_scope_id - options = quota_data.options - - await log_manager_api_entry(log, "update_quota_scope", quota_data) - async with self._volume_pool.get_volume(volume_id) as volume: + async def update_quota_scope( + self, quota_scope_key: QuotaScopeKey, options: Optional[QuotaConfig] + ) -> None: + quota_scope_id = quota_scope_key.quota_scope_id + await log_manager_api_entry(log, "update_quota_scope", quota_scope_key) + async with self._volume_pool.get_volume(quota_scope_key.volume_id) as volume: async with self._handle_external_errors(): quota_usage = await volume.quota_model.describe_quota_scope(quota_scope_id) if not quota_usage: @@ -188,69 +175,45 @@ async def update_quota_scope(self, quota_data: QuotaScopeKeyData) -> None: ) except InvalidQuotaConfig: raise web.HTTPBadRequest(reason="Invalid quota config option") - return None - async def delete_quota_scope(self, quota_data: QuotaScopeKeyData) -> None: - volume_id = quota_data.volume_id - quota_scope_id = quota_data.quota_scope_id - - await log_manager_api_entry(log, "delete_quota_scope", quota_data) - async with self._volume_pool.get_volume(volume_id) as volume: + async def delete_quota_scope(self, quota_scope_key: QuotaScopeKey) -> None: + quota_scope_id = quota_scope_key.quota_scope_id + await log_manager_api_entry(log, "delete_quota_scope", quota_scope_key) + async with self._volume_pool.get_volume(quota_scope_key.volume_id) as volume: async with self._handle_external_errors(): quota_usage = await volume.quota_model.describe_quota_scope(quota_scope_id) if not quota_usage: raise QuotaScopeNotFoundError await volume.quota_model.unset_quota(quota_scope_id) - return None - async def create_vfolder(self, vfolder_data: VFolderKeyData) -> NewVFolderCreated: - volume_id = vfolder_data.volume_id - vfolder_id = vfolder_data.vfolder_id + async def create_vfolder(self, vfolder_key: VFolderKey) -> None: + vfolder_id = vfolder_key.vfolder_id quota_scope_id = vfolder_id.quota_scope_id - await log_manager_api_entry(log, "create_vfolder", vfolder_data) - assert quota_scope_id is not None - async with self._volume_pool.get_volume(volume_id) as volume: + await log_manager_api_entry(log, "create_vfolder", vfolder_key) + if quota_scope_id is None: + raise InvalidQuotaScopeError("Quota scope ID is not set in the vfolder key.") + async with self._volume_pool.get_volume(vfolder_key.volume_id) as volume: try: await volume.create_vfolder(vfolder_id) except QuotaScopeNotFoundError: - assert quota_scope_id await volume.quota_model.create_quota_scope(quota_scope_id) try: await volume.create_vfolder(vfolder_id) except QuotaScopeNotFoundError: raise ExternalError("Failed to create vfolder due to quota scope not found") - return NewVFolderCreated( - vfolder_id=vfolder_id, - quota_scope_path=volume.quota_model.mangle_qspath(quota_scope_id), - vfolder_path=await volume.get_vfolder_mount(vfolder_id, "."), - ) - async def clone_vfolder(self, vfolder_data: VFolderKeyData) -> NewVFolderCreated: - volume_id = vfolder_data.volume_id - src_vfolder_id = vfolder_data.vfolder_id - dst_vfolder_id = vfolder_data.dst_vfolder_id + async def clone_vfolder(self, vfolder_key: VFolderKey, dst_vfolder_id: VFolderID) -> None: + await log_manager_api_entry(log, "clone_vfolder", vfolder_key) + async with self._volume_pool.get_volume(vfolder_key.volume_id) as volume: + await volume.clone_vfolder(vfolder_key.vfolder_id, dst_vfolder_id) - if dst_vfolder_id is None: - raise ValueError("Destination vfolder ID cannot be None") - await log_manager_api_entry(log, "clone_vfolder", vfolder_data) - async with self._volume_pool.get_volume(volume_id) as volume: - await volume.clone_vfolder(src_vfolder_id, dst_vfolder_id) - return NewVFolderCreated( - vfolder_id=dst_vfolder_id, - quota_scope_path=volume.quota_model.mangle_qspath(dst_vfolder_id), - vfolder_path=await volume.get_vfolder_mount(dst_vfolder_id, "."), - ) - - async def get_vfolder_info(self, vfolder_data: VFolderKeyData) -> VFolderMetadata: - volume_id = vfolder_data.volume_id - vfolder_id = vfolder_data.vfolder_id - subpath = vfolder_data.subpath - - await log_manager_api_entry(log, "get_vfolder_info", vfolder_data) - async with self._volume_pool.get_volume(volume_id) as volume: + async def get_vfolder_info(self, vfolder_key: VFolderKey, subpath: str) -> VFolderMeta: + vfolder_id = vfolder_key.vfolder_id + await log_manager_api_entry(log, "get_vfolder_info", vfolder_key) + async with self._volume_pool.get_volume(vfolder_key.volume_id) as volume: try: - mount_path = await volume.get_vfolder_mount(vfolder_id, str(subpath)) + mount_path = await volume.get_vfolder_mount(vfolder_id, subpath) usage = await volume.get_usage(vfolder_id) fs_usage = await volume.get_fs_usage() except VFolderNotFoundError: @@ -258,7 +221,7 @@ async def get_vfolder_info(self, vfolder_data: VFolderKeyData) -> VFolderMetadat except InvalidSubpathError: raise web.HTTPBadRequest(reason="Invalid vfolder subpath") - return VFolderMetadata( + return VFolderMeta( mount_path=mount_path, file_count=usage.file_count, used_bytes=usage.used_bytes, @@ -266,13 +229,11 @@ async def get_vfolder_info(self, vfolder_data: VFolderKeyData) -> VFolderMetadat fs_used_bytes=fs_usage.used_bytes, ) - async def delete_vfolder(self, vfolder_data: VFolderKeyData) -> None: - volume_id = vfolder_data.volume_id - vfolder_id = vfolder_data.vfolder_id - - await log_manager_api_entry(log, "delete_vfolder", vfolder_data) + async def delete_vfolder(self, vfolder_key: VFolderKey) -> None: + vfolder_id = vfolder_key.vfolder_id + await log_manager_api_entry(log, "delete_vfolder", vfolder_key) try: - async with self._volume_pool.get_volume(volume_id) as volume: + async with self._volume_pool.get_volume(vfolder_key.volume_id) as volume: await volume.get_vfolder_mount(vfolder_id, ".") except VFolderNotFoundError: ongoing_task = self._deletion_tasks.get(vfolder_id) @@ -282,5 +243,5 @@ async def delete_vfolder(self, vfolder_data: VFolderKeyData) -> None: else: ongoing_task = self._deletion_tasks.get(vfolder_id) if ongoing_task is not None and ongoing_task.done(): - asyncio.create_task(self._delete_vfolder(vfolder_data)) + asyncio.create_task(self._delete_vfolder(vfolder_key)) return None diff --git a/src/ai/backend/storage/volumes/pool.py b/src/ai/backend/storage/volumes/pool.py index 60c899bf1ca..149234c4cf4 100644 --- a/src/ai/backend/storage/volumes/pool.py +++ b/src/ai/backend/storage/volumes/pool.py @@ -4,25 +4,25 @@ from pathlib import Path from typing import Any, AsyncIterator, Mapping, Type -from ai.backend.common.dto.identifiers import VolumeID from ai.backend.common.etcd import AsyncEtcd from ai.backend.common.events import EventDispatcher, EventProducer -from ai.backend.storage.volumes.cephfs import CephFSVolume -from ai.backend.storage.volumes.ddn import EXAScalerFSVolume -from ai.backend.storage.volumes.dellemc import DellEMCOneFSVolume -from ai.backend.storage.volumes.gpfs import GPFSVolume -from ai.backend.storage.volumes.netapp import NetAppVolume -from ai.backend.storage.volumes.purestorage import FlashBladeVolume -from ai.backend.storage.volumes.vast import VASTVolume -from ai.backend.storage.volumes.vfs import BaseVolume -from ai.backend.storage.volumes.weka import WekaVolume -from ai.backend.storage.volumes.xfs import XfsVolume +from ai.backend.common.types import VolumeID from ..exception import InvalidVolumeError from ..types import VolumeInfo from .abc import AbstractVolume +from .cephfs import CephFSVolume +from .ddn import EXAScalerFSVolume +from .dellemc import DellEMCOneFSVolume +from .gpfs import GPFSVolume +from .netapp import NetAppVolume +from .purestorage import FlashBladeVolume +from .vast import VASTVolume +from .vfs import BaseVolume +from .weka import WekaVolume +from .xfs import XfsVolume -DEFAULT_BACKENDS: Mapping[str, Type[AbstractVolume]] = { +_DEFAULT_BACKENDS: Mapping[str, Type[AbstractVolume]] = { FlashBladeVolume.name: FlashBladeVolume, BaseVolume.name: BaseVolume, XfsVolume.name: XfsVolume, @@ -60,7 +60,7 @@ def __init__( self._event_producer = event_producer async def __aenter__(self) -> None: - self._backends = {**DEFAULT_BACKENDS} + self._backends = {**_DEFAULT_BACKENDS} def list_volumes(self) -> Mapping[str, VolumeInfo]: return { diff --git a/src/ai/backend/storage/volumes/types.py b/src/ai/backend/storage/volumes/types.py index 594afa5f365..a93d8fe586a 100644 --- a/src/ai/backend/storage/volumes/types.py +++ b/src/ai/backend/storage/volumes/types.py @@ -1,105 +1,84 @@ -from pathlib import Path, PurePath, PurePosixPath -from typing import List, Optional - -from pydantic import AliasChoices, BaseModel, ConfigDict, Field - -from ai.backend.common.dto.identifiers import VolumeID -from ai.backend.common.types import QuotaConfig, QuotaScopeID, VFolderID - - -class _BaseModel(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True) - - -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 VolumeKeyData(_BaseModel): - volume_id: VolumeID = VOLUME_ID_FIELD - - -class VolumeMetadata(_BaseModel): +from dataclasses import dataclass +from pathlib import Path, PurePath +from typing import Optional, Self + +from pydantic import BaseModel, Field + +from ai.backend.common.dto.storage.field import VFolderMetaField, VolumeMetaField +from ai.backend.common.dto.storage.path import QuotaScopeKeyPath, VFolderKeyPath +from ai.backend.common.dto.storage.response import QuotaScopeResponse +from ai.backend.common.types import QuotaScopeID, VFolderID, VolumeID + + +@dataclass +class QuotaScopeKey: volume_id: VolumeID - backend: str - path: Path - fsprefix: Optional[PurePath] = Field(default=None) - capabilities: list[str] + quota_scope_id: QuotaScopeID + @classmethod + def from_quota_scope_path(cls, path: QuotaScopeKeyPath) -> Self: + return cls( + volume_id=path.volume_id, quota_scope_id=QuotaScopeID(path.scope_type, path.scope_uuid) + ) -class VolumeMetadataList(_BaseModel): - volumes: List[VolumeMetadata] +@dataclass +class VFolderKey: + volume_id: VolumeID + vfolder_id: VFolderID -class VFolderKeyData(_BaseModel): - volume_id: VolumeID = VOLUME_ID_FIELD - vfolder_id: VFolderID = VFOLDER_ID_FIELD - subpath: Optional[PurePosixPath] = Field(default=None) - # 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", - ), - ) + @classmethod + def from_vfolder_path(cls, path: VFolderKeyPath) -> Self: + quota_scope_id = QuotaScopeID(path.scope_type, path.scope_uuid) + return cls( + volume_id=path.volume_id, + vfolder_id=VFolderID(quota_scope_id, path.folder_uuid), + ) -class VFolderMetadata(BaseModel): +@dataclass +class VolumeMeta: + volume_id: VolumeID + backend: str + path: Path + fsprefix: Optional[PurePath] + capabilities: list[str] + + def to_field(self) -> VolumeMetaField: + return VolumeMetaField( + volume_id=self.volume_id, + backend=self.backend, + path=str(self.path), + fsprefix=str(self.fsprefix) if self.fsprefix is not None else None, + capabilities=self.capabilities, + ) + + +@dataclass +class VFolderMeta(BaseModel): mount_path: Path file_count: int used_bytes: int capacity_bytes: int fs_used_bytes: int + def to_field(self) -> VFolderMetaField: + return VFolderMetaField( + mount_path=str(self.mount_path), + file_count=self.file_count, + used_bytes=self.used_bytes, + capacity_bytes=self.capacity_bytes, + fs_used_bytes=self.fs_used_bytes, + ) -class NewVFolderCreated(_BaseModel): - vfolder_id: VFolderID = VFOLDER_ID_FIELD - quota_scope_path: Path - vfolder_path: Path - -class QuotaScopeKeyData(_BaseModel): - volume_id: VolumeID = VOLUME_ID_FIELD - quota_scope_id: QuotaScopeID = QUOTA_SCOPE_ID_FIELD - options: Optional[QuotaConfig] = Field(default=None) - - -class QuotaScopeMetadata(BaseModel): +@dataclass +class QuotaScopeMeta(BaseModel): used_bytes: Optional[int] = Field(default=0) limit_bytes: Optional[int] = Field(default=0) - -class NewQuotaScopeCreated(_BaseModel): - quota_scope_id: QuotaScopeID = QUOTA_SCOPE_ID_FIELD - quota_scope_path: Path + def to_response(self) -> QuotaScopeResponse: + return QuotaScopeResponse( + used_bytes=self.used_bytes, + limit_bytes=self.limit_bytes, + ) diff --git a/tests/storage-proxy/vfolder/test_handler.py b/tests/storage-proxy/vfolder/test_handler.py index 9cc365d160a..68bd6e91326 100644 --- a/tests/storage-proxy/vfolder/test_handler.py +++ b/tests/storage-proxy/vfolder/test_handler.py @@ -1,376 +1,376 @@ -import json -import uuid -from pathlib import Path -from unittest.mock import AsyncMock - -import pytest -from aiohttp import web - -from ai.backend.common.types import QuotaConfig, QuotaScopeID, QuotaScopeType, VFolderID -from ai.backend.storage.api.vfolder.handler import VFolderHandler -from ai.backend.storage.volumes.types import ( - NewQuotaScopeCreated, - NewVFolderCreated, - QuotaScopeKeyData, - QuotaScopeMetadata, - VFolderKeyData, - VFolderMetadata, - VolumeKeyData, - VolumeMetadata, - VolumeMetadataList, -) - -UUID = uuid.UUID("123e4567-e89b-12d3-a456-426614174000") -QUOTA_SCOPE_ID = QuotaScopeID( - scope_type=QuotaScopeType.USER, - scope_id=UUID, -) -VFOLDER_SCOPE_ID = VFolderID( - quota_scope_id=QUOTA_SCOPE_ID, - folder_id=UUID, -) - - -@pytest.fixture -def mock_vfolder_service(): - class MockVFolderService: - async def get_volume(self, volume_data: VolumeKeyData) -> VolumeMetadata: - return VolumeMetadata( - volume_id=volume_data.volume_id, - backend="mock-backend", - path=Path("/mock/path"), - fsprefix=None, - capabilities=["read", "write"], - ) - - async def get_volumes(self) -> VolumeMetadataList: - return VolumeMetadataList( - volumes=[ - VolumeMetadata( - volume_id=UUID, - backend="mock-backend", - path=Path("/mock/path"), - fsprefix=None, - capabilities=["read", "write"], - ) - ] - ) - - async def create_quota_scope(self, quota_data: QuotaScopeKeyData) -> NewQuotaScopeCreated: - return NewQuotaScopeCreated( - quota_scope_id=QUOTA_SCOPE_ID, - quota_scope_path=Path("/mock/quota/scope/path"), - ) - - async def get_quota_scope(self, quota_data: QuotaScopeKeyData) -> QuotaScopeMetadata: - return QuotaScopeMetadata( - used_bytes=1024, - limit_bytes=2048, - ) - - async def update_quota_scope(self, quota_data: QuotaScopeKeyData) -> None: - pass - - async def delete_quota_scope(self, quota_data: QuotaScopeKeyData) -> None: - pass - - async def create_vfolder(self, vfolder_data: VFolderKeyData) -> NewVFolderCreated: - return NewVFolderCreated( - vfolder_id=VFOLDER_SCOPE_ID, - quota_scope_path=Path("/mock/quota/scope/path"), - vfolder_path=Path("/mock/vfolder/path"), - ) - - async def clone_vfolder(self, vfolder_data: VFolderKeyData) -> NewVFolderCreated: - return NewVFolderCreated( - vfolder_id=VFOLDER_SCOPE_ID, - quota_scope_path=Path("/mock/quota/scope/path"), - vfolder_path=Path("/mock/vfolder/path"), - ) - - async def get_vfolder_info(self, vfolder_data: VFolderKeyData) -> VFolderMetadata: - return VFolderMetadata( - mount_path=Path("/mock/mount/path"), - file_count=100, - capacity_bytes=1024 * 1024 * 1024, - used_bytes=1024, - fs_used_bytes=512000, - ) - - async def delete_vfolder(self, vfolder_data: VFolderKeyData) -> None: - pass - - return MockVFolderService() - - -@pytest.mark.asyncio -async def test_get_volume(mock_vfolder_service): - handler = VFolderHandler(storage_service=mock_vfolder_service) - - async def mock_request(): - request = AsyncMock(spec=web.Request) - request.match_info = {"volume_id": "123e4567-e89b-12d3-a456-426614174000"} - return request - - response = await handler.get_volume(await mock_request()) - - assert isinstance(response, web.Response) - assert response.status == 200 - assert response.content_type == "application/json" - response_volume_id = json.loads(response.text)["volume_id"] - assert response_volume_id == "123e4567-e89b-12d3-a456-426614174000" - - -@pytest.mark.asyncio -async def test_get_volumes(mock_vfolder_service): - handler = VFolderHandler(storage_service=mock_vfolder_service) - - async def mock_request(): - request = AsyncMock(spec=web.Request) - return request - - response = await handler.get_volumes(await mock_request()) - - assert isinstance(response, web.Response) - assert response.status == 200 - assert response.content_type == "application/json" - response_volumes = json.loads(response.text)["volumes"] - assert len(response_volumes) == 1 - - -@pytest.mark.asyncio -async def test_create_quota_scope(mock_vfolder_service): - handler = VFolderHandler(storage_service=mock_vfolder_service) - - async def mock_request(): - request = AsyncMock(spec=web.Request) - request.json.return_value = { - "volume_id": "123e4567-e89b-12d3-a456-426614174000", - "quota_scope_id": QUOTA_SCOPE_ID, - } - return request +# import json +# import uuid +# from pathlib import Path +# from unittest.mock import AsyncMock + +# import pytest +# from aiohttp import web + +# from ai.backend.common.types import QuotaConfig, QuotaScopeID, QuotaScopeType, VFolderID +# from ai.backend.storage.api.vfolder.handler import VFolderHandler +# from ai.backend.storage.volumes.types import ( +# NewQuotaScopeCreated, +# NewVFolderCreated, +# QuotaScopeKeyData, +# QuotaScopeMetadata, +# VFolderKeyData, +# VFolderMetadata, +# VolumeKeyData, +# VolumeMetadata, +# VolumeMetadataList, +# ) + +# UUID = uuid.UUID("123e4567-e89b-12d3-a456-426614174000") +# QUOTA_SCOPE_ID = QuotaScopeID( +# scope_type=QuotaScopeType.USER, +# scope_id=UUID, +# ) +# VFOLDER_SCOPE_ID = VFolderID( +# quota_scope_id=QUOTA_SCOPE_ID, +# folder_id=UUID, +# ) + + +# @pytest.fixture +# def mock_vfolder_service(): +# class MockVFolderService: +# async def get_volume(self, volume_data: VolumeKeyData) -> VolumeMetadata: +# return VolumeMetadata( +# volume_id=volume_data.volume_id, +# backend="mock-backend", +# path=Path("/mock/path"), +# fsprefix=None, +# capabilities=["read", "write"], +# ) + +# async def get_volumes(self) -> VolumeMetadataList: +# return VolumeMetadataList( +# volumes=[ +# VolumeMetadata( +# volume_id=UUID, +# backend="mock-backend", +# path=Path("/mock/path"), +# fsprefix=None, +# capabilities=["read", "write"], +# ) +# ] +# ) + +# async def create_quota_scope(self, quota_data: QuotaScopeKeyData) -> NewQuotaScopeCreated: +# return NewQuotaScopeCreated( +# quota_scope_id=QUOTA_SCOPE_ID, +# quota_scope_path=Path("/mock/quota/scope/path"), +# ) + +# async def get_quota_scope(self, quota_data: QuotaScopeKeyData) -> QuotaScopeMetadata: +# return QuotaScopeMetadata( +# used_bytes=1024, +# limit_bytes=2048, +# ) + +# async def update_quota_scope(self, quota_data: QuotaScopeKeyData) -> None: +# pass + +# async def delete_quota_scope(self, quota_data: QuotaScopeKeyData) -> None: +# pass + +# async def create_vfolder(self, vfolder_data: VFolderKeyData) -> NewVFolderCreated: +# return NewVFolderCreated( +# vfolder_id=VFOLDER_SCOPE_ID, +# quota_scope_path=Path("/mock/quota/scope/path"), +# vfolder_path=Path("/mock/vfolder/path"), +# ) + +# async def clone_vfolder(self, vfolder_data: VFolderKeyData) -> NewVFolderCreated: +# return NewVFolderCreated( +# vfolder_id=VFOLDER_SCOPE_ID, +# quota_scope_path=Path("/mock/quota/scope/path"), +# vfolder_path=Path("/mock/vfolder/path"), +# ) + +# async def get_vfolder_info(self, vfolder_data: VFolderKeyData) -> VFolderMetadata: +# return VFolderMetadata( +# mount_path=Path("/mock/mount/path"), +# file_count=100, +# capacity_bytes=1024 * 1024 * 1024, +# used_bytes=1024, +# fs_used_bytes=512000, +# ) + +# async def delete_vfolder(self, vfolder_data: VFolderKeyData) -> None: +# pass + +# return MockVFolderService() + + +# @pytest.mark.asyncio +# async def test_get_volume(mock_vfolder_service): +# handler = VFolderHandler(storage_service=mock_vfolder_service) + +# async def mock_request(): +# request = AsyncMock(spec=web.Request) +# request.match_info = {"volume_id": "123e4567-e89b-12d3-a456-426614174000"} +# return request + +# response = await handler.get_volume(await mock_request()) + +# assert isinstance(response, web.Response) +# assert response.status == 200 +# assert response.content_type == "application/json" +# response_volume_id = json.loads(response.text)["volume_id"] +# assert response_volume_id == "123e4567-e89b-12d3-a456-426614174000" + + +# @pytest.mark.asyncio +# async def test_get_volumes(mock_vfolder_service): +# handler = VFolderHandler(storage_service=mock_vfolder_service) + +# async def mock_request(): +# request = AsyncMock(spec=web.Request) +# return request + +# response = await handler.get_volumes(await mock_request()) + +# assert isinstance(response, web.Response) +# assert response.status == 200 +# assert response.content_type == "application/json" +# response_volumes = json.loads(response.text)["volumes"] +# assert len(response_volumes) == 1 + + +# @pytest.mark.asyncio +# async def test_create_quota_scope(mock_vfolder_service): +# handler = VFolderHandler(storage_service=mock_vfolder_service) + +# async def mock_request(): +# request = AsyncMock(spec=web.Request) +# request.json.return_value = { +# "volume_id": "123e4567-e89b-12d3-a456-426614174000", +# "quota_scope_id": QUOTA_SCOPE_ID, +# } +# return request - response = await handler.create_quota_scope(await mock_request()) +# response = await handler.create_quota_scope(await mock_request()) - assert isinstance(response, web.Response) - assert response.status == 201 +# assert isinstance(response, web.Response) +# assert response.status == 201 -@pytest.mark.asyncio -async def test_get_quota_scope(mock_vfolder_service): - handler = VFolderHandler(storage_service=mock_vfolder_service) +# @pytest.mark.asyncio +# async def test_get_quota_scope(mock_vfolder_service): +# handler = VFolderHandler(storage_service=mock_vfolder_service) - async def mock_request(): - request = AsyncMock(spec=web.Request) - request.match_info = { - "volume_id": "123e4567-e89b-12d3-a456-426614174000", - "quota_scope_id": QUOTA_SCOPE_ID, - } - return request +# async def mock_request(): +# request = AsyncMock(spec=web.Request) +# request.match_info = { +# "volume_id": "123e4567-e89b-12d3-a456-426614174000", +# "quota_scope_id": QUOTA_SCOPE_ID, +# } +# return request - response = await handler.get_quota_scope(await mock_request()) +# response = await handler.get_quota_scope(await mock_request()) - assert isinstance(response, web.Response) - assert response.status == 200 - assert response.content_type == "application/json" - response_data = json.loads(response.text)["used_bytes"] - assert response_data == 1024 +# assert isinstance(response, web.Response) +# assert response.status == 200 +# assert response.content_type == "application/json" +# response_data = json.loads(response.text)["used_bytes"] +# assert response_data == 1024 -@pytest.mark.asyncio -async def test_update_quota_scope(mock_vfolder_service): - handler = VFolderHandler(storage_service=mock_vfolder_service) +# @pytest.mark.asyncio +# async def test_update_quota_scope(mock_vfolder_service): +# handler = VFolderHandler(storage_service=mock_vfolder_service) - async def mock_request(): - request = AsyncMock(spec=web.Request) - request.json.return_value = { - "volume_id": "123e4567-e89b-12d3-a456-426614174000", - "quota_scope_id": QUOTA_SCOPE_ID, - "options": QuotaConfig(limit_bytes=2048), # QuotaConfig 객체 사용 - } - return request +# async def mock_request(): +# request = AsyncMock(spec=web.Request) +# request.json.return_value = { +# "volume_id": "123e4567-e89b-12d3-a456-426614174000", +# "quota_scope_id": QUOTA_SCOPE_ID, +# "options": QuotaConfig(limit_bytes=2048), # QuotaConfig 객체 사용 +# } +# return request - response = await handler.update_quota_scope(await mock_request()) +# response = await handler.update_quota_scope(await mock_request()) - assert isinstance(response, web.Response) - assert response.status == 204 +# assert isinstance(response, web.Response) +# assert response.status == 204 -@pytest.mark.asyncio -async def test_delete_quota_scope(mock_vfolder_service): - handler = VFolderHandler(storage_service=mock_vfolder_service) +# @pytest.mark.asyncio +# async def test_delete_quota_scope(mock_vfolder_service): +# handler = VFolderHandler(storage_service=mock_vfolder_service) - async def mock_request(): - request = AsyncMock(spec=web.Request) - request.match_info = { - "volume_id": "123e4567-e89b-12d3-a456-426614174000", - "quota_scope_id": QUOTA_SCOPE_ID, - } - return request +# async def mock_request(): +# request = AsyncMock(spec=web.Request) +# request.match_info = { +# "volume_id": "123e4567-e89b-12d3-a456-426614174000", +# "quota_scope_id": QUOTA_SCOPE_ID, +# } +# return request - response = await handler.delete_quota_scope(await mock_request()) +# response = await handler.delete_quota_scope(await mock_request()) - assert isinstance(response, web.Response) - assert response.status == 204 +# assert isinstance(response, web.Response) +# assert response.status == 204 -@pytest.mark.asyncio -async def test_create_vfolder(mock_vfolder_service): - handler = VFolderHandler(storage_service=mock_vfolder_service) +# @pytest.mark.asyncio +# async def test_create_vfolder(mock_vfolder_service): +# handler = VFolderHandler(storage_service=mock_vfolder_service) - async def mock_request(): - request = AsyncMock(spec=web.Request) - request.json.return_value = { - "volume_id": "123e4567-e89b-12d3-a456-426614174000", - "vfolder_id": VFOLDER_SCOPE_ID, - } - return request +# async def mock_request(): +# request = AsyncMock(spec=web.Request) +# request.json.return_value = { +# "volume_id": "123e4567-e89b-12d3-a456-426614174000", +# "vfolder_id": VFOLDER_SCOPE_ID, +# } +# return request - response = await handler.create_vfolder(await mock_request()) +# response = await handler.create_vfolder(await mock_request()) - assert isinstance(response, web.Response) - assert response.status == 201 +# assert isinstance(response, web.Response) +# assert response.status == 201 -@pytest.mark.asyncio -async def test_clone_vfolder(mock_vfolder_service): - handler = VFolderHandler(storage_service=mock_vfolder_service) +# @pytest.mark.asyncio +# async def test_clone_vfolder(mock_vfolder_service): +# handler = VFolderHandler(storage_service=mock_vfolder_service) - async def mock_request(): - request = AsyncMock(spec=web.Request) - request.json.return_value = { - "volume_id": "123e4567-e89b-12d3-a456-426614174000", - "vfolder_id": VFOLDER_SCOPE_ID, - "dst_vfolder_id": VFOLDER_SCOPE_ID, - } - return request +# async def mock_request(): +# request = AsyncMock(spec=web.Request) +# request.json.return_value = { +# "volume_id": "123e4567-e89b-12d3-a456-426614174000", +# "vfolder_id": VFOLDER_SCOPE_ID, +# "dst_vfolder_id": VFOLDER_SCOPE_ID, +# } +# return request - response = await handler.clone_vfolder(await mock_request()) +# response = await handler.clone_vfolder(await mock_request()) - assert isinstance(response, web.Response) - assert response.status == 201 +# assert isinstance(response, web.Response) +# assert response.status == 201 -@pytest.mark.asyncio -async def test_get_vfolder_info(mock_vfolder_service): - handler = VFolderHandler(storage_service=mock_vfolder_service) +# @pytest.mark.asyncio +# async def test_get_vfolder_info(mock_vfolder_service): +# handler = VFolderHandler(storage_service=mock_vfolder_service) - async def mock_request(): - request = AsyncMock(spec=web.Request) - request.match_info = { - "volume_id": "123e4567-e89b-12d3-a456-426614174000", - "vfolder_id": VFOLDER_SCOPE_ID, - "subpath": "/mock/subpath", - } - return request +# async def mock_request(): +# request = AsyncMock(spec=web.Request) +# request.match_info = { +# "volume_id": "123e4567-e89b-12d3-a456-426614174000", +# "vfolder_id": VFOLDER_SCOPE_ID, +# "subpath": "/mock/subpath", +# } +# return request - response = await handler.get_vfolder_info(await mock_request()) +# response = await handler.get_vfolder_info(await mock_request()) - assert isinstance(response, web.Response) - assert response.status == 200 - assert response.content_type == "application/json" - response_data = json.loads(response.text)["mount_path"] - assert response_data == "/mock/mount/path" +# assert isinstance(response, web.Response) +# assert response.status == 200 +# assert response.content_type == "application/json" +# response_data = json.loads(response.text)["mount_path"] +# assert response_data == "/mock/mount/path" -@pytest.mark.asyncio -async def test_get_vfolder_mount(mock_vfolder_service): - handler = VFolderHandler(storage_service=mock_vfolder_service) +# @pytest.mark.asyncio +# async def test_get_vfolder_mount(mock_vfolder_service): +# handler = VFolderHandler(storage_service=mock_vfolder_service) - async def mock_request(): - request = AsyncMock(spec=web.Request) - request.match_info = { - "volume_id": "123e4567-e89b-12d3-a456-426614174000", - "vfolder_id": VFOLDER_SCOPE_ID, - "subpath": "/mock/subpath", - } - return request +# async def mock_request(): +# request = AsyncMock(spec=web.Request) +# request.match_info = { +# "volume_id": "123e4567-e89b-12d3-a456-426614174000", +# "vfolder_id": VFOLDER_SCOPE_ID, +# "subpath": "/mock/subpath", +# } +# return request - response = await handler.get_vfolder_mount(await mock_request()) +# response = await handler.get_vfolder_mount(await mock_request()) - assert isinstance(response, web.Response) - assert response.status == 200 - assert response.content_type == "application/json" - response_data = json.loads(response.text)["mount_path"] - assert response_data == "/mock/mount/path" +# assert isinstance(response, web.Response) +# assert response.status == 200 +# assert response.content_type == "application/json" +# response_data = json.loads(response.text)["mount_path"] +# assert response_data == "/mock/mount/path" -@pytest.mark.asyncio -async def test_get_vfolder_usage(mock_vfolder_service): - handler = VFolderHandler(storage_service=mock_vfolder_service) - - async def mock_request(): - request = AsyncMock(spec=web.Request) - request.match_info = { - "volume_id": "123e4567-e89b-12d3-a456-426614174000", - "vfolder_id": VFOLDER_SCOPE_ID, - } - return request +# @pytest.mark.asyncio +# async def test_get_vfolder_usage(mock_vfolder_service): +# handler = VFolderHandler(storage_service=mock_vfolder_service) + +# async def mock_request(): +# request = AsyncMock(spec=web.Request) +# request.match_info = { +# "volume_id": "123e4567-e89b-12d3-a456-426614174000", +# "vfolder_id": VFOLDER_SCOPE_ID, +# } +# return request - response = await handler.get_vfolder_usage(await mock_request()) +# response = await handler.get_vfolder_usage(await mock_request()) - assert isinstance(response, web.Response) - assert response.status == 200 - assert response.content_type == "application/json" - response_data = json.loads(response.text)["file_count"] - assert response_data == 100 - - -@pytest.mark.asyncio -async def test_get_vfolder_used_bytes(mock_vfolder_service): - handler = VFolderHandler(storage_service=mock_vfolder_service) - - async def mock_request(): - request = AsyncMock(spec=web.Request) - request.match_info = { - "volume_id": "123e4567-e89b-12d3-a456-426614174000", - "vfolder_id": VFOLDER_SCOPE_ID, - } - return request +# assert isinstance(response, web.Response) +# assert response.status == 200 +# assert response.content_type == "application/json" +# response_data = json.loads(response.text)["file_count"] +# assert response_data == 100 + + +# @pytest.mark.asyncio +# async def test_get_vfolder_used_bytes(mock_vfolder_service): +# handler = VFolderHandler(storage_service=mock_vfolder_service) + +# async def mock_request(): +# request = AsyncMock(spec=web.Request) +# request.match_info = { +# "volume_id": "123e4567-e89b-12d3-a456-426614174000", +# "vfolder_id": VFOLDER_SCOPE_ID, +# } +# return request - response = await handler.get_vfolder_used_bytes(await mock_request()) - - assert isinstance(response, web.Response) - assert response.status == 200 - assert response.content_type == "application/json" - response_data = json.loads(response.text)["used_bytes"] - assert response_data == 1024 - - -@pytest.mark.asyncio -async def test_get_vfolder_fs_usage(mock_vfolder_service): - handler = VFolderHandler(storage_service=mock_vfolder_service) - - async def mock_request(): - request = AsyncMock(spec=web.Request) - request.match_info = { - "volume_id": "123e4567-e89b-12d3-a456-426614174000", - "vfolder_id": VFOLDER_SCOPE_ID, - } - return request - - response = await handler.get_vfolder_fs_usage(await mock_request()) - - assert isinstance(response, web.Response) - assert response.status == 200 - assert response.content_type == "application/json" - response_data = json.loads(response.text)["capacity_bytes"] - assert response_data == 1024 * 1024 * 1024 - - -@pytest.mark.asyncio -async def test_delete_vfolder(mock_vfolder_service): - handler = VFolderHandler(storage_service=mock_vfolder_service) - - async def mock_request(): - request = AsyncMock(spec=web.Request) - request.match_info = { - "volume_id": "123e4567-e89b-12d3-a456-426614174000", - "vfolder_id": VFOLDER_SCOPE_ID, - } - return request - - response = await handler.delete_vfolder(await mock_request()) +# response = await handler.get_vfolder_used_bytes(await mock_request()) + +# assert isinstance(response, web.Response) +# assert response.status == 200 +# assert response.content_type == "application/json" +# response_data = json.loads(response.text)["used_bytes"] +# assert response_data == 1024 + + +# @pytest.mark.asyncio +# async def test_get_vfolder_fs_usage(mock_vfolder_service): +# handler = VFolderHandler(storage_service=mock_vfolder_service) + +# async def mock_request(): +# request = AsyncMock(spec=web.Request) +# request.match_info = { +# "volume_id": "123e4567-e89b-12d3-a456-426614174000", +# "vfolder_id": VFOLDER_SCOPE_ID, +# } +# return request + +# response = await handler.get_vfolder_fs_usage(await mock_request()) + +# assert isinstance(response, web.Response) +# assert response.status == 200 +# assert response.content_type == "application/json" +# response_data = json.loads(response.text)["capacity_bytes"] +# assert response_data == 1024 * 1024 * 1024 + + +# @pytest.mark.asyncio +# async def test_delete_vfolder(mock_vfolder_service): +# handler = VFolderHandler(storage_service=mock_vfolder_service) + +# async def mock_request(): +# request = AsyncMock(spec=web.Request) +# request.match_info = { +# "volume_id": "123e4567-e89b-12d3-a456-426614174000", +# "vfolder_id": VFOLDER_SCOPE_ID, +# } +# return request + +# response = await handler.delete_vfolder(await mock_request()) - assert isinstance(response, web.Response) - assert response.status == 202 +# assert isinstance(response, web.Response) +# assert response.status == 202 diff --git a/tests/storage-proxy/vfolder/test_pool.py b/tests/storage-proxy/vfolder/test_pool.py index 0e86e25e3e7..f7c95514eaf 100644 --- a/tests/storage-proxy/vfolder/test_pool.py +++ b/tests/storage-proxy/vfolder/test_pool.py @@ -3,9 +3,9 @@ import pytest -from ai.backend.common.dto.identifiers import VolumeID from ai.backend.common.etcd import AsyncEtcd from ai.backend.common.events import EventDispatcher, EventProducer +from ai.backend.common.types import VolumeID from ai.backend.storage.exception import InvalidVolumeError from ai.backend.storage.types import VolumeInfo from ai.backend.storage.volumes.pool import VolumePool From e68775027a00c7fa6ba93b9ed71c9c3c8f4c86c6 Mon Sep 17 00:00:00 2001 From: Mincheol Kang Date: Wed, 12 Feb 2025 13:10:32 +0900 Subject: [PATCH 3/7] refactor: Make class instance private --- src/ai/backend/storage/api/vfolder/handler.py | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/src/ai/backend/storage/api/vfolder/handler.py b/src/ai/backend/storage/api/vfolder/handler.py index c0e5b3e1743..d8e88e241b5 100644 --- a/src/ai/backend/storage/api/vfolder/handler.py +++ b/src/ai/backend/storage/api/vfolder/handler.py @@ -50,12 +50,14 @@ 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 + 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) + volume_meta = await self._storage_service.get_volume(path.parsed.volume_id) return APIResponse.build( status_code=200, response_model=GetVolumeResponse( @@ -65,7 +67,7 @@ async def get_volume(self, path: PathParam[VolumeIDPath]) -> APIResponse: @api_handler async def get_volumes(self) -> APIResponse: - volume_meta_list = await self.storage_service.get_volumes() + volume_meta_list = await self._storage_service.get_volumes() return APIResponse.build( status_code=200, response_model=GetVolumesResponse( @@ -78,13 +80,13 @@ 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) + 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) + quota_scope = await self._storage_service.get_quota_scope(quota_scope_key) return APIResponse.build( status_code=200, response_model=quota_scope.to_response(), @@ -95,19 +97,19 @@ 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) + 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) + 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) + await self._storage_service.create_vfolder(vfolder_key) return APIResponse.no_content(status_code=204) @api_handler @@ -115,7 +117,7 @@ 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) + await self._storage_service.clone_vfolder(vfolder_key, body.parsed.dst_vfolder_id) return APIResponse.no_content(status_code=204) @api_handler @@ -123,7 +125,9 @@ 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) + vfolder_meta = await self._storage_service.get_vfolder_info( + vfolder_key, body.parsed.subpath + ) return APIResponse.build( status_code=200, response_model=VFolderMetadataResponse( @@ -134,5 +138,5 @@ async def get_vfolder_info( @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) + await self._storage_service.delete_vfolder(vfolder_key) return APIResponse.no_content(status_code=204) From f0f2d0aa9b0a0be137d48959d2bb968fe0b1f0e9 Mon Sep 17 00:00:00 2001 From: Mincheol Kang Date: Fri, 14 Feb 2025 14:35:16 +0900 Subject: [PATCH 4/7] feat: handler test --- src/ai/backend/common/api_handlers.py | 4 +- tests/storage-proxy/vfolder/test_handler.py | 556 ++++++++------------ 2 files changed, 225 insertions(+), 335 deletions(-) diff --git a/src/ai/backend/common/api_handlers.py b/src/ai/backend/common/api_handlers.py index def733926da..eaeee109e00 100644 --- a/src/ai/backend/common/api_handlers.py +++ b/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 ( @@ -127,7 +127,7 @@ def from_request(cls, request: web.Request) -> Self: class BaseRequestModel(BaseModel): - pass + model_config = ConfigDict(arbitrary_types_allowed=True) class BaseResponseModel(BaseModel): diff --git a/tests/storage-proxy/vfolder/test_handler.py b/tests/storage-proxy/vfolder/test_handler.py index 68bd6e91326..b1b3eb7f161 100644 --- a/tests/storage-proxy/vfolder/test_handler.py +++ b/tests/storage-proxy/vfolder/test_handler.py @@ -1,376 +1,266 @@ -# import json -# import uuid -# from pathlib import Path -# from unittest.mock import AsyncMock - -# import pytest -# from aiohttp import web - -# from ai.backend.common.types import QuotaConfig, QuotaScopeID, QuotaScopeType, VFolderID -# from ai.backend.storage.api.vfolder.handler import VFolderHandler -# from ai.backend.storage.volumes.types import ( -# NewQuotaScopeCreated, -# NewVFolderCreated, -# QuotaScopeKeyData, -# QuotaScopeMetadata, -# VFolderKeyData, -# VFolderMetadata, -# VolumeKeyData, -# VolumeMetadata, -# VolumeMetadataList, -# ) - -# UUID = uuid.UUID("123e4567-e89b-12d3-a456-426614174000") -# QUOTA_SCOPE_ID = QuotaScopeID( -# scope_type=QuotaScopeType.USER, -# scope_id=UUID, -# ) -# VFOLDER_SCOPE_ID = VFolderID( -# quota_scope_id=QUOTA_SCOPE_ID, -# folder_id=UUID, -# ) - - -# @pytest.fixture -# def mock_vfolder_service(): -# class MockVFolderService: -# async def get_volume(self, volume_data: VolumeKeyData) -> VolumeMetadata: -# return VolumeMetadata( -# volume_id=volume_data.volume_id, -# backend="mock-backend", -# path=Path("/mock/path"), -# fsprefix=None, -# capabilities=["read", "write"], -# ) - -# async def get_volumes(self) -> VolumeMetadataList: -# return VolumeMetadataList( -# volumes=[ -# VolumeMetadata( -# volume_id=UUID, -# backend="mock-backend", -# path=Path("/mock/path"), -# fsprefix=None, -# capabilities=["read", "write"], -# ) -# ] -# ) - -# async def create_quota_scope(self, quota_data: QuotaScopeKeyData) -> NewQuotaScopeCreated: -# return NewQuotaScopeCreated( -# quota_scope_id=QUOTA_SCOPE_ID, -# quota_scope_path=Path("/mock/quota/scope/path"), -# ) - -# async def get_quota_scope(self, quota_data: QuotaScopeKeyData) -> QuotaScopeMetadata: -# return QuotaScopeMetadata( -# used_bytes=1024, -# limit_bytes=2048, -# ) - -# async def update_quota_scope(self, quota_data: QuotaScopeKeyData) -> None: -# pass - -# async def delete_quota_scope(self, quota_data: QuotaScopeKeyData) -> None: -# pass - -# async def create_vfolder(self, vfolder_data: VFolderKeyData) -> NewVFolderCreated: -# return NewVFolderCreated( -# vfolder_id=VFOLDER_SCOPE_ID, -# quota_scope_path=Path("/mock/quota/scope/path"), -# vfolder_path=Path("/mock/vfolder/path"), -# ) - -# async def clone_vfolder(self, vfolder_data: VFolderKeyData) -> NewVFolderCreated: -# return NewVFolderCreated( -# vfolder_id=VFOLDER_SCOPE_ID, -# quota_scope_path=Path("/mock/quota/scope/path"), -# vfolder_path=Path("/mock/vfolder/path"), -# ) - -# async def get_vfolder_info(self, vfolder_data: VFolderKeyData) -> VFolderMetadata: -# return VFolderMetadata( -# mount_path=Path("/mock/mount/path"), -# file_count=100, -# capacity_bytes=1024 * 1024 * 1024, -# used_bytes=1024, -# fs_used_bytes=512000, -# ) - -# async def delete_vfolder(self, vfolder_data: VFolderKeyData) -> None: -# pass - -# return MockVFolderService() - - -# @pytest.mark.asyncio -# async def test_get_volume(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) - -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# request.match_info = {"volume_id": "123e4567-e89b-12d3-a456-426614174000"} -# return request - -# response = await handler.get_volume(await mock_request()) - -# assert isinstance(response, web.Response) -# assert response.status == 200 -# assert response.content_type == "application/json" -# response_volume_id = json.loads(response.text)["volume_id"] -# assert response_volume_id == "123e4567-e89b-12d3-a456-426614174000" - - -# @pytest.mark.asyncio -# async def test_get_volumes(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) - -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# return request - -# response = await handler.get_volumes(await mock_request()) - -# assert isinstance(response, web.Response) -# assert response.status == 200 -# assert response.content_type == "application/json" -# response_volumes = json.loads(response.text)["volumes"] -# assert len(response_volumes) == 1 - - -# @pytest.mark.asyncio -# async def test_create_quota_scope(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) - -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# request.json.return_value = { -# "volume_id": "123e4567-e89b-12d3-a456-426614174000", -# "quota_scope_id": QUOTA_SCOPE_ID, -# } -# return request +import json +import uuid +from pathlib import Path, PurePath +from unittest.mock import AsyncMock + +import pytest +from aiohttp import web + +from ai.backend.common.api_handlers import APIResponse +from ai.backend.common.types import QuotaScopeID, QuotaScopeType, VFolderID +from ai.backend.storage.api.vfolder.handler import VFolderHandler +from ai.backend.storage.volumes.types import QuotaScopeMeta, VFolderMeta, VolumeMeta + +UUID = uuid.UUID("123e4567-e89b-12d3-a456-426614174000") +UUID1 = uuid.UUID("123e4567-e89b-12d3-a456-426614174001") +UUID2 = uuid.UUID("123e4567-e89b-12d3-a456-426614174002") + + +@pytest.fixture +def mock_vfolder_service(): + class MockVFolderService: + async def get_volume(self, volume_id): + return VolumeMeta( + volume_id=volume_id, + backend="vfs", + path=Path("/mnt/test_volume"), + fsprefix=PurePath("vfs-test"), + capabilities=["read", "write"], + ) + + async def get_volumes(self): + volumes = { + UUID1: {"backend": "vfs", "path": "/mnt/volume1", "fsprefix": "vfs-test-1"}, + UUID2: {"backend": "nfs", "path": "/mnt/volume2", "fsprefix": "nfs-test-2"}, + } + return [ + VolumeMeta( + volume_id=volume_id, + backend=info.get("backend", "vfs"), + path=Path(info.get("path", "/mnt/test_volume")), + fsprefix=PurePath(info.get("fsprefix", "vfs-test")), + capabilities=["read", "write"], + ) + for volume_id, info in volumes.items() + ] + + async def create_quota_scope(self, quota_scope_key, options): + pass + + async def get_quota_scope(self, quota_scope_key): + return QuotaScopeMeta.model_validate({ + "used_bytes": 1000, + "limit_bytes": 2000, + }) + + async def update_quota_scope(self, quota_scope_key, options): + pass + + async def delete_quota_scope(self, quota_scope_key): + pass + + async def create_vfolder(self, vfolder_key): + pass + + async def clone_vfolder(self, src_vfolder_key, dst_vfolder_key): + pass -# response = await handler.create_quota_scope(await mock_request()) + async def get_vfolder_info(self, vfolder_key, subpath): + return VFolderMeta.model_validate({ + "mount_path": subpath, + "file_count": 100, + "used_bytes": 1000, + "capacity_bytes": 2000, + "fs_used_bytes": 1000, + }) + + async def delete_vfolder(self, vfolder_key): + pass + + return MockVFolderService() -# assert isinstance(response, web.Response) -# assert response.status == 201 +@pytest.mark.asyncio +async def test_get_volume(mock_vfolder_service): + handler = VFolderHandler(mock_vfolder_service) -# @pytest.mark.asyncio -# async def test_get_quota_scope(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) + mock_request = AsyncMock(web.Request) + mock_request.match_info = {"volume_id": str(UUID)} -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# request.match_info = { -# "volume_id": "123e4567-e89b-12d3-a456-426614174000", -# "quota_scope_id": QUOTA_SCOPE_ID, -# } -# return request + response: APIResponse = await handler.get_volume(mock_request) + + assert isinstance(response, web.Response) + assert response.status == 200 + volume_response = json.loads(response.text)["item"] + assert volume_response["volume_id"] == str(UUID) -# response = await handler.get_quota_scope(await mock_request()) -# assert isinstance(response, web.Response) -# assert response.status == 200 -# assert response.content_type == "application/json" -# response_data = json.loads(response.text)["used_bytes"] -# assert response_data == 1024 +@pytest.mark.asyncio +async def test_get_volumes(mock_vfolder_service): + handler = VFolderHandler(mock_vfolder_service) + mock_request = AsyncMock(web.Request) + response: APIResponse = await handler.get_volumes(mock_request) -# @pytest.mark.asyncio -# async def test_update_quota_scope(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) + assert isinstance(response, web.Response) + assert response.status == 200 + volume_response = json.loads(response.text)["items"] + assert volume_response[0]["volume_id"] == str(UUID1) -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# request.json.return_value = { -# "volume_id": "123e4567-e89b-12d3-a456-426614174000", -# "quota_scope_id": QUOTA_SCOPE_ID, -# "options": QuotaConfig(limit_bytes=2048), # QuotaConfig 객체 사용 -# } -# return request -# response = await handler.update_quota_scope(await mock_request()) +@pytest.mark.asyncio +async def test_create_quota_scope(mock_vfolder_service): + handler = VFolderHandler(mock_vfolder_service) + + mock_request = AsyncMock(web.Request) + mock_request.match_info = { + "volume_id": str(UUID), + "scope_type": "user", + "scope_uuid": str(UUID), + } + mock_request.json.return_value = {"options": None} + + response: APIResponse = await handler.create_quota_scope(mock_request) -# assert isinstance(response, web.Response) -# assert response.status == 204 + assert isinstance(response, web.Response) + assert response.status == 204 -# @pytest.mark.asyncio -# async def test_delete_quota_scope(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) +@pytest.mark.asyncio +async def test_get_quota_scope(mock_vfolder_service): + handler = VFolderHandler(mock_vfolder_service) -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# request.match_info = { -# "volume_id": "123e4567-e89b-12d3-a456-426614174000", -# "quota_scope_id": QUOTA_SCOPE_ID, -# } -# return request + mock_request = AsyncMock(web.Request) + mock_request.match_info = { + "volume_id": str(UUID), + "scope_type": "user", + "scope_uuid": str(UUID), + } -# response = await handler.delete_quota_scope(await mock_request()) + response: APIResponse = await handler.get_quota_scope(mock_request) -# assert isinstance(response, web.Response) -# assert response.status == 204 + assert isinstance(response, web.Response) + assert response.status == 200 + quota_response = json.loads(response.text) + assert quota_response["used_bytes"] == 1000 + assert quota_response["limit_bytes"] == 2000 -# @pytest.mark.asyncio -# async def test_create_vfolder(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) +@pytest.mark.asyncio +async def test_update_quota_scope(mock_vfolder_service): + handler = VFolderHandler(mock_vfolder_service) -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# request.json.return_value = { -# "volume_id": "123e4567-e89b-12d3-a456-426614174000", -# "vfolder_id": VFOLDER_SCOPE_ID, -# } -# return request + mock_request = AsyncMock(web.Request) + mock_request.match_info = { + "volume_id": str(UUID), + "scope_type": "user", + "scope_uuid": str(UUID), + } + mock_request.json.return_value = {"options": None} -# response = await handler.create_vfolder(await mock_request()) + response: APIResponse = await handler.update_quota_scope(mock_request) -# assert isinstance(response, web.Response) -# assert response.status == 201 + assert isinstance(response, web.Response) + assert response.status == 204 -# @pytest.mark.asyncio -# async def test_clone_vfolder(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) +@pytest.mark.asyncio +async def test_delete_quota_scope(mock_vfolder_service): + handler = VFolderHandler(mock_vfolder_service) -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# request.json.return_value = { -# "volume_id": "123e4567-e89b-12d3-a456-426614174000", -# "vfolder_id": VFOLDER_SCOPE_ID, -# "dst_vfolder_id": VFOLDER_SCOPE_ID, -# } -# return request + mock_request = AsyncMock(web.Request) + mock_request.match_info = { + "volume_id": str(UUID), + "scope_type": "user", + "scope_uuid": str(UUID), + } -# response = await handler.clone_vfolder(await mock_request()) + response: APIResponse = await handler.delete_quota_scope(mock_request) -# assert isinstance(response, web.Response) -# assert response.status == 201 + assert isinstance(response, web.Response) + assert response.status == 204 -# @pytest.mark.asyncio -# async def test_get_vfolder_info(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) +@pytest.mark.asyncio +async def test_create_vfolder(mock_vfolder_service): + handler = VFolderHandler(mock_vfolder_service) -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# request.match_info = { -# "volume_id": "123e4567-e89b-12d3-a456-426614174000", -# "vfolder_id": VFOLDER_SCOPE_ID, -# "subpath": "/mock/subpath", -# } -# return request + mock_request = AsyncMock(web.Request) + mock_request.match_info = { + "volume_id": str(UUID), + "scope_type": "user", + "scope_uuid": str(UUID), + "folder_uuid": str(UUID), + } -# response = await handler.get_vfolder_info(await mock_request()) + response: APIResponse = await handler.create_vfolder(mock_request) -# assert isinstance(response, web.Response) -# assert response.status == 200 -# assert response.content_type == "application/json" -# response_data = json.loads(response.text)["mount_path"] -# assert response_data == "/mock/mount/path" + assert isinstance(response, web.Response) + assert response.status == 204 -# @pytest.mark.asyncio -# async def test_get_vfolder_mount(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) +@pytest.mark.asyncio +async def test_clone_vfolder(mock_vfolder_service): + handler = VFolderHandler(mock_vfolder_service) -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# request.match_info = { -# "volume_id": "123e4567-e89b-12d3-a456-426614174000", -# "vfolder_id": VFOLDER_SCOPE_ID, -# "subpath": "/mock/subpath", -# } -# return request + mock_request = AsyncMock(web.Request) + mock_request.match_info = { + "volume_id": str(UUID), + "scope_type": "user", + "scope_uuid": str(UUID), + "folder_uuid": str(UUID), + } + mock_request.json.return_value = { + "dst_vfolder_id": VFolderID( + quota_scope_id=QuotaScopeID(scope_type=QuotaScopeType.USER, scope_id=UUID), + folder_id=UUID, + ) + } -# response = await handler.get_vfolder_mount(await mock_request()) + response: APIResponse = await handler.clone_vfolder(mock_request) -# assert isinstance(response, web.Response) -# assert response.status == 200 -# assert response.content_type == "application/json" -# response_data = json.loads(response.text)["mount_path"] -# assert response_data == "/mock/mount/path" + assert isinstance(response, web.Response) + assert response.status == 204 -# @pytest.mark.asyncio -# async def test_get_vfolder_usage(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) - -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# request.match_info = { -# "volume_id": "123e4567-e89b-12d3-a456-426614174000", -# "vfolder_id": VFOLDER_SCOPE_ID, -# } -# return request +@pytest.mark.asyncio +async def test_get_vfolder_info(mock_vfolder_service): + handler = VFolderHandler(mock_vfolder_service) + + mock_request = AsyncMock(web.Request) + mock_request.match_info = { + "volume_id": str(UUID), + "scope_type": "user", + "scope_uuid": str(UUID), + "folder_uuid": str(UUID), + } + mock_request.json.return_value = {"subpath": "/mnt/test_volume"} -# response = await handler.get_vfolder_usage(await mock_request()) + response: APIResponse = await handler.get_vfolder_info(mock_request) -# assert isinstance(response, web.Response) -# assert response.status == 200 -# assert response.content_type == "application/json" -# response_data = json.loads(response.text)["file_count"] -# assert response_data == 100 - - -# @pytest.mark.asyncio -# async def test_get_vfolder_used_bytes(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) - -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# request.match_info = { -# "volume_id": "123e4567-e89b-12d3-a456-426614174000", -# "vfolder_id": VFOLDER_SCOPE_ID, -# } -# return request + assert isinstance(response, web.Response) + assert response.status == 200 + vfolder_response = json.loads(response.text)["item"] + assert vfolder_response["mount_path"] == "/mnt/test_volume" + assert vfolder_response["file_count"] == 100 + assert vfolder_response["used_bytes"] == 1000 + assert vfolder_response["capacity_bytes"] == 2000 + assert vfolder_response["fs_used_bytes"] == 1000 + + +@pytest.mark.asyncio +async def test_delete_vfolder(mock_vfolder_service): + handler = VFolderHandler(mock_vfolder_service) + + mock_request = AsyncMock(web.Request) + mock_request.match_info = { + "volume_id": str(UUID), + "scope_type": "user", + "scope_uuid": str(UUID), + "folder_uuid": str(UUID), + } -# response = await handler.get_vfolder_used_bytes(await mock_request()) - -# assert isinstance(response, web.Response) -# assert response.status == 200 -# assert response.content_type == "application/json" -# response_data = json.loads(response.text)["used_bytes"] -# assert response_data == 1024 - - -# @pytest.mark.asyncio -# async def test_get_vfolder_fs_usage(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) - -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# request.match_info = { -# "volume_id": "123e4567-e89b-12d3-a456-426614174000", -# "vfolder_id": VFOLDER_SCOPE_ID, -# } -# return request - -# response = await handler.get_vfolder_fs_usage(await mock_request()) - -# assert isinstance(response, web.Response) -# assert response.status == 200 -# assert response.content_type == "application/json" -# response_data = json.loads(response.text)["capacity_bytes"] -# assert response_data == 1024 * 1024 * 1024 - - -# @pytest.mark.asyncio -# async def test_delete_vfolder(mock_vfolder_service): -# handler = VFolderHandler(storage_service=mock_vfolder_service) - -# async def mock_request(): -# request = AsyncMock(spec=web.Request) -# request.match_info = { -# "volume_id": "123e4567-e89b-12d3-a456-426614174000", -# "vfolder_id": VFOLDER_SCOPE_ID, -# } -# return request - -# response = await handler.delete_vfolder(await mock_request()) + response: APIResponse = await handler.delete_vfolder(mock_request) -# assert isinstance(response, web.Response) -# assert response.status == 202 + assert isinstance(response, web.Response) + assert response.status == 204 From 7c5ddafc3002300185b0fdf75deac0ea0bcb4bc3 Mon Sep 17 00:00:00 2001 From: Mincheol Kang Date: Fri, 14 Feb 2025 16:20:02 +0900 Subject: [PATCH 5/7] test: Add test code for the service layer - Added `log_manager_api_entry_new` to correspond with the refactored codes. --- src/ai/backend/storage/services/service.py | 22 +- src/ai/backend/storage/utils.py | 49 ++++ src/ai/backend/storage/volumes/types.py | 6 +- tests/storage-proxy/vfolder/test_service.py | 287 ++++++++++++++++++++ 4 files changed, 350 insertions(+), 14 deletions(-) create mode 100644 tests/storage-proxy/vfolder/test_service.py diff --git a/src/ai/backend/storage/services/service.py b/src/ai/backend/storage/services/service.py index 5b3815f320a..791ef6da483 100644 --- a/src/ai/backend/storage/services/service.py +++ b/src/ai/backend/storage/services/service.py @@ -21,7 +21,7 @@ QuotaScopeNotFoundError, VFolderNotFoundError, ) -from ..utils import log_manager_api_entry +from ..utils import log_manager_api_entry_new from ..volumes.pool import VolumePool from ..volumes.types import ( QuotaScopeKey, @@ -103,7 +103,7 @@ async def _delete_vfolder( ) async def get_volume(self, volume_id: VolumeID) -> VolumeMeta: - await log_manager_api_entry(log, "get_volume", volume_id) + await log_manager_api_entry_new(log, "get_volume", volume_id) volume = self._volume_pool.get_volume_info(volume_id) return VolumeMeta( volume_id=volume_id, @@ -114,7 +114,7 @@ async def get_volume(self, volume_id: VolumeID) -> VolumeMeta: ) async def get_volumes(self) -> list[VolumeMeta]: - await log_manager_api_entry(log, "get_volumes", params=None) + await log_manager_api_entry_new(log, "get_volumes", params=None) volumes = self._volume_pool.list_volumes() return [ VolumeMeta( @@ -131,7 +131,7 @@ async def create_quota_scope( self, quota_scope_key: QuotaScopeKey, options: Optional[QuotaConfig] ) -> None: quota_scope_id = quota_scope_key.quota_scope_id - await log_manager_api_entry(log, "create_quota_scope", quota_scope_key) + await log_manager_api_entry_new(log, "create_quota_scope", quota_scope_key) async with self._volume_pool.get_volume(quota_scope_key.volume_id) as volume: try: async with self._handle_external_errors(): @@ -142,7 +142,7 @@ async def create_quota_scope( raise web.HTTPConflict(reason="Volume already exists with given quota scope.") async def get_quota_scope(self, quota_scope_key: QuotaScopeKey) -> QuotaScopeMeta: - await log_manager_api_entry(log, "get_quota_scope", quota_scope_key) + await log_manager_api_entry_new(log, "get_quota_scope", quota_scope_key) async with self._volume_pool.get_volume(quota_scope_key.volume_id) as volume: async with self._handle_external_errors(): quota_usage = await volume.quota_model.describe_quota_scope( @@ -158,7 +158,7 @@ async def update_quota_scope( self, quota_scope_key: QuotaScopeKey, options: Optional[QuotaConfig] ) -> None: quota_scope_id = quota_scope_key.quota_scope_id - await log_manager_api_entry(log, "update_quota_scope", quota_scope_key) + await log_manager_api_entry_new(log, "update_quota_scope", quota_scope_key) async with self._volume_pool.get_volume(quota_scope_key.volume_id) as volume: async with self._handle_external_errors(): quota_usage = await volume.quota_model.describe_quota_scope(quota_scope_id) @@ -178,7 +178,7 @@ async def update_quota_scope( async def delete_quota_scope(self, quota_scope_key: QuotaScopeKey) -> None: quota_scope_id = quota_scope_key.quota_scope_id - await log_manager_api_entry(log, "delete_quota_scope", quota_scope_key) + await log_manager_api_entry_new(log, "delete_quota_scope", quota_scope_key) async with self._volume_pool.get_volume(quota_scope_key.volume_id) as volume: async with self._handle_external_errors(): quota_usage = await volume.quota_model.describe_quota_scope(quota_scope_id) @@ -190,7 +190,7 @@ async def create_vfolder(self, vfolder_key: VFolderKey) -> None: vfolder_id = vfolder_key.vfolder_id quota_scope_id = vfolder_id.quota_scope_id - await log_manager_api_entry(log, "create_vfolder", vfolder_key) + await log_manager_api_entry_new(log, "create_vfolder", vfolder_key) if quota_scope_id is None: raise InvalidQuotaScopeError("Quota scope ID is not set in the vfolder key.") async with self._volume_pool.get_volume(vfolder_key.volume_id) as volume: @@ -204,13 +204,13 @@ async def create_vfolder(self, vfolder_key: VFolderKey) -> None: raise ExternalError("Failed to create vfolder due to quota scope not found") async def clone_vfolder(self, vfolder_key: VFolderKey, dst_vfolder_id: VFolderID) -> None: - await log_manager_api_entry(log, "clone_vfolder", vfolder_key) + await log_manager_api_entry_new(log, "clone_vfolder", vfolder_key) async with self._volume_pool.get_volume(vfolder_key.volume_id) as volume: await volume.clone_vfolder(vfolder_key.vfolder_id, dst_vfolder_id) async def get_vfolder_info(self, vfolder_key: VFolderKey, subpath: str) -> VFolderMeta: vfolder_id = vfolder_key.vfolder_id - await log_manager_api_entry(log, "get_vfolder_info", vfolder_key) + await log_manager_api_entry_new(log, "get_vfolder_info", vfolder_key) async with self._volume_pool.get_volume(vfolder_key.volume_id) as volume: try: mount_path = await volume.get_vfolder_mount(vfolder_id, subpath) @@ -231,7 +231,7 @@ async def get_vfolder_info(self, vfolder_key: VFolderKey, subpath: str) -> VFold async def delete_vfolder(self, vfolder_key: VFolderKey) -> None: vfolder_id = vfolder_key.vfolder_id - await log_manager_api_entry(log, "delete_vfolder", vfolder_key) + await log_manager_api_entry_new(log, "delete_vfolder", vfolder_key) try: async with self._volume_pool.get_volume(vfolder_key.volume_id) as volume: await volume.get_vfolder_mount(vfolder_id, ".") diff --git a/src/ai/backend/storage/utils.py b/src/ai/backend/storage/utils.py index 4e8e553cd0f..0735d5545e9 100644 --- a/src/ai/backend/storage/utils.py +++ b/src/ai/backend/storage/utils.py @@ -127,3 +127,52 @@ async def log_manager_api_entry( "ManagerAPI::{}()", name.upper(), ) + + +async def log_manager_api_entry_new( + log: Union[logging.Logger, BraceStyleAdapter], + name: str, + params: Any, +) -> None: + if isinstance(params, dict): + if "src_vfid" in params and "dst_vfid" in params: + log.info( + "ManagerAPI::{}(v:{}, f:{} -> dst_v: {}, dst_f:{})", + name.upper(), + params["src_volume"], + params["src_vfid"], + params["dst_volume"], + params["dst_vfid"], + ) + elif "relpaths" in params: + log.info( + "ManagerAPI::{}(v:{}, f:{}, p*:{})", + name.upper(), + params["volume"], + params["vfid"], + str(params["relpaths"][0]) + "...", + ) + elif "relpath" in params: + log.info( + "ManagerAPI::{}(v:{}, f:{}, p:{})", + name.upper(), + params["volume"], + params["vfid"], + params["relpath"], + ) + elif "vfid" in params: + log.info( + "ManagerAPI::{}(v:{}, f:{})", + name.upper(), + params["volume"], + params["vfid"], + ) + elif "volume" in params: + log.info( + "ManagerAPI::{}(v:{})", + name.upper(), + params["volume"], + ) + return + + log.info("ManagerAPI::{}({})", name.upper(), str(params)) diff --git a/src/ai/backend/storage/volumes/types.py b/src/ai/backend/storage/volumes/types.py index a93d8fe586a..6c2d35f6923 100644 --- a/src/ai/backend/storage/volumes/types.py +++ b/src/ai/backend/storage/volumes/types.py @@ -2,7 +2,7 @@ from pathlib import Path, PurePath from typing import Optional, Self -from pydantic import BaseModel, Field +from pydantic import Field from ai.backend.common.dto.storage.field import VFolderMetaField, VolumeMetaField from ai.backend.common.dto.storage.path import QuotaScopeKeyPath, VFolderKeyPath @@ -55,7 +55,7 @@ def to_field(self) -> VolumeMetaField: @dataclass -class VFolderMeta(BaseModel): +class VFolderMeta: mount_path: Path file_count: int used_bytes: int @@ -73,7 +73,7 @@ def to_field(self) -> VFolderMetaField: @dataclass -class QuotaScopeMeta(BaseModel): +class QuotaScopeMeta: used_bytes: Optional[int] = Field(default=0) limit_bytes: Optional[int] = Field(default=0) diff --git a/tests/storage-proxy/vfolder/test_service.py b/tests/storage-proxy/vfolder/test_service.py new file mode 100644 index 00000000000..ac9e5ea41ac --- /dev/null +++ b/tests/storage-proxy/vfolder/test_service.py @@ -0,0 +1,287 @@ +import uuid +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from aiohttp import web + +from ai.backend.common.types import QuotaConfig, QuotaScopeID, QuotaScopeType, VFolderID +from ai.backend.storage.exception import VFolderNotFoundError +from ai.backend.storage.services.service import VolumeService +from ai.backend.storage.services.service import log as service_log +from ai.backend.storage.volumes.pool import VolumePool +from ai.backend.storage.volumes.types import ( + QuotaScopeKey, + QuotaScopeMeta, + VFolderKey, + VolumeMeta, +) + +UUID = uuid.UUID("12345678-1234-5678-1234-567812345678") +UUID1 = uuid.UUID("12345678-1234-5678-1234-567812345679") +UUID2 = uuid.UUID("12345678-1234-5678-1234-567812345680") + + +@pytest.fixture +def mock_volume_pool(): + mock_pool = MagicMock(spec=VolumePool) + return mock_pool + + +@pytest.fixture +def mock_service(mock_volume_pool): + service = VolumeService(volume_pool=mock_volume_pool) + service._get_capabilities = AsyncMock(return_value=["capability1", "capability2"]) + + service.log = service_log + + return service + + +@pytest.mark.asyncio +@patch("ai.backend.storage.services.service.log_manager_api_entry_new", new_callable=AsyncMock) +async def test_get_volume(mock_log, mock_service, mock_volume_pool): + mock_volume_info = MagicMock() + mock_volume_info.backend = "mock-backend" + mock_volume_info.path = "/mock/path" + mock_volume_info.fsprefix = "mock-fsprefix" + + mock_volume_pool.get_volume_info.return_value = mock_volume_info + + volume_id = UUID + result = await mock_service.get_volume(volume_id) + + mock_log.assert_called_once_with(service_log, "get_volume", volume_id) + mock_volume_pool.get_volume_info.assert_called_once_with(volume_id) + mock_service._get_capabilities.assert_called_once_with(volume_id) + + assert isinstance(result, VolumeMeta) + assert result.volume_id == volume_id + assert result.backend == "mock-backend" + assert result.path == "/mock/path" + assert result.fsprefix == "mock-fsprefix" + assert result.capabilities == ["capability1", "capability2"] + + +@pytest.mark.asyncio +@patch("ai.backend.storage.services.service.log_manager_api_entry_new", new_callable=AsyncMock) +async def test_get_volumes(mock_log, mock_service, mock_volume_pool): + mock_volumes = { + str(UUID1): MagicMock(backend="backend1", path="/path1", fsprefix="fsprefix1"), + str(UUID2): MagicMock(backend="backend2", path="/path2", fsprefix="fsprefix2"), + } + mock_volume_pool.list_volumes.return_value = mock_volumes + + mock_service._get_capabilities.side_effect = [ + ["capability1", "capability2"], + ["capability3"], + ] + + result = await mock_service.get_volumes() + + mock_log.assert_called_once_with(service_log, "get_volumes", params=None) + mock_volume_pool.list_volumes.assert_called_once() + + assert len(result) == 2 + assert result[0].volume_id == UUID1 + assert result[0].backend == "backend1" + assert result[0].path == "/path1" + assert result[0].fsprefix == "fsprefix1" + assert result[0].capabilities == ["capability1", "capability2"] + + assert result[1].volume_id == UUID2 + assert result[1].backend == "backend2" + assert result[1].path == "/path2" + assert result[1].fsprefix == "fsprefix2" + assert result[1].capabilities == ["capability3"] + + +@pytest.mark.asyncio +@patch("ai.backend.storage.services.service.log_manager_api_entry_new", new_callable=AsyncMock) +async def test_create_quota_scope(mock_log, mock_service, mock_volume_pool): + mock_volume = MagicMock() + mock_volume.quota_model.create_quota_scope = AsyncMock() + + mock_volume_pool.get_volume.return_value.__aenter__.return_value = mock_volume + + quota_scope_key = QuotaScopeKey( + volume_id=uuid.UUID("12345678-1234-5678-1234-567812345678"), + quota_scope_id=UUID, + ) + options = QuotaConfig(limit_bytes=1024 * 1024 * 1024) + + await mock_service.create_quota_scope(quota_scope_key, options) + + mock_log.assert_called_once_with(service_log, "create_quota_scope", quota_scope_key) + mock_volume.quota_model.create_quota_scope.assert_called_once_with( + quota_scope_id=UUID, options=options, extra_args=None + ) + + +@pytest.mark.asyncio +@patch("ai.backend.storage.services.service.log_manager_api_entry_new", new_callable=AsyncMock) +async def test_get_quota_scope(mock_log, mock_service, mock_volume_pool): + mock_volume = MagicMock() + quota_scope_meta = QuotaScopeMeta(used_bytes=500, limit_bytes=1000) + mock_volume.quota_model.describe_quota_scope = AsyncMock(return_value=quota_scope_meta) + + mock_volume_pool.get_volume.return_value.__aenter__.return_value = mock_volume + + quota_scope_key = QuotaScopeKey(volume_id=UUID, quota_scope_id=UUID) + + result = await mock_service.get_quota_scope(quota_scope_key) + + mock_log.assert_called_once_with(service_log, "get_quota_scope", quota_scope_key) + mock_volume.quota_model.describe_quota_scope.assert_called_once_with(UUID) + + assert result.used_bytes == 500 + assert result.limit_bytes == 1000 + + +@pytest.mark.asyncio +@patch("ai.backend.storage.services.service.log_manager_api_entry_new", new_callable=AsyncMock) +async def test_update_quota_scope(mock_log, mock_service, mock_volume_pool): + mock_volume = MagicMock() + quota_scope_meta = QuotaScopeMeta(used_bytes=500, limit_bytes=1000) + mock_volume.quota_model.describe_quota_scope = AsyncMock(return_value=quota_scope_meta) + mock_volume.quota_model.update_quota_scope = AsyncMock() + + mock_volume_pool.get_volume.return_value.__aenter__.return_value = mock_volume + + quota_scope_key = QuotaScopeKey(volume_id=UUID, quota_scope_id=UUID) + options = QuotaConfig(limit_bytes=2000) + + await mock_service.update_quota_scope(quota_scope_key, options) + + mock_log.assert_called_once_with(service_log, "update_quota_scope", quota_scope_key) + mock_volume.quota_model.describe_quota_scope.assert_called_once_with(UUID) + mock_volume.quota_model.update_quota_scope.assert_called_once_with( + quota_scope_id=UUID, config=options + ) + + +@pytest.mark.asyncio +@patch("ai.backend.storage.services.service.log_manager_api_entry_new", new_callable=AsyncMock) +async def test_delete_quota_scope(mock_log, mock_service, mock_volume_pool): + mock_volume = MagicMock() + mock_volume.quota_model.describe_quota_scope = AsyncMock( + return_value=MagicMock(used_bytes=500, limit_bytes=1000) + ) + mock_volume.quota_model.unset_quota = AsyncMock() + + mock_volume_pool.get_volume.return_value.__aenter__.return_value = mock_volume + + quota_scope_key = QuotaScopeKey(volume_id=UUID, quota_scope_id=UUID) + + await mock_service.delete_quota_scope(quota_scope_key) + + mock_log.assert_called_once_with(service_log, "delete_quota_scope", quota_scope_key) + mock_volume.quota_model.describe_quota_scope.assert_called_once_with(UUID) + mock_volume.quota_model.unset_quota.assert_called_once_with(UUID) + + +@pytest.mark.asyncio +@patch("ai.backend.storage.services.service.log_manager_api_entry_new", new_callable=AsyncMock) +async def test_create_vfolder(mock_log, mock_service, mock_volume_pool): + mock_volume = MagicMock() + mock_volume.create_vfolder = AsyncMock() + mock_volume.quota_model.create_quota_scope = AsyncMock() + + mock_volume_pool.get_volume.return_value.__aenter__.return_value = mock_volume + + vfolder_id = VFolderID( + quota_scope_id=QuotaScopeID(scope_type=QuotaScopeType.USER, scope_id=UUID), + folder_id=UUID, + ) + vfolder_key = VFolderKey(volume_id=UUID, vfolder_id=vfolder_id) + + await mock_service.create_vfolder(vfolder_key) + + mock_log.assert_called_once_with(service_log, "create_vfolder", vfolder_key) + mock_volume.create_vfolder.assert_called_once_with(vfolder_id) + + +@pytest.mark.asyncio +@patch("ai.backend.storage.services.service.log_manager_api_entry_new", new_callable=AsyncMock) +async def test_clone_vfolder(mock_log, mock_service, mock_volume_pool): + mock_volume = MagicMock() + mock_volume.clone_vfolder = AsyncMock() + + mock_volume_pool.get_volume.return_value.__aenter__.return_value = mock_volume + + src_vfolder_id = VFolderID( + quota_scope_id=QuotaScopeID(scope_type=QuotaScopeType.USER, scope_id=UUID), + folder_id=UUID, + ) + dst_vfolder_id = VFolderID( + quota_scope_id=QuotaScopeID(scope_type=QuotaScopeType.USER, scope_id=UUID1), + folder_id=UUID2, + ) + vfolder_key = VFolderKey(volume_id=UUID, vfolder_id=src_vfolder_id) + + await mock_service.clone_vfolder(vfolder_key, dst_vfolder_id) + + mock_log.assert_called_once_with(service_log, "clone_vfolder", vfolder_key) + mock_volume.clone_vfolder.assert_called_once_with(src_vfolder_id, dst_vfolder_id) + + +@pytest.mark.asyncio +@patch("ai.backend.storage.services.service.log_manager_api_entry_new", new_callable=AsyncMock) +async def test_get_vfolder_info(mock_log, mock_service, mock_volume_pool): + mock_volume = MagicMock() + mock_volume.get_vfolder_mount = AsyncMock(return_value=Path("/mock/mount")) + + usage_data = MagicMock(spec=["file_count", "used_bytes"]) + usage_data.file_count = 10 + usage_data.used_bytes = 5000 + + fs_usage_data = MagicMock(spec=["capacity_bytes", "used_bytes"]) + fs_usage_data.capacity_bytes = 100000 + fs_usage_data.used_bytes = 20000 + + mock_volume.get_usage = AsyncMock(return_value=usage_data) + mock_volume.get_fs_usage = AsyncMock(return_value=fs_usage_data) + + mock_volume_pool.get_volume.return_value.__aenter__.return_value = mock_volume + + vfolder_id = VFolderID( + quota_scope_id=QuotaScopeID(scope_type=QuotaScopeType.USER, scope_id=UUID), + folder_id=UUID, + ) + vfolder_key = VFolderKey(volume_id=UUID, vfolder_id=vfolder_id) + subpath = "test_subpath" + + result = await mock_service.get_vfolder_info(vfolder_key, subpath) + + mock_log.assert_called_once_with(service_log, "get_vfolder_info", vfolder_key) + mock_volume.get_vfolder_mount.assert_called_once_with(vfolder_id, subpath) + mock_volume.get_usage.assert_called_once_with(vfolder_id) + mock_volume.get_fs_usage.assert_called_once() + + assert isinstance(result.mount_path, Path) + assert result.mount_path == Path("/mock/mount") + assert result.file_count == 10 + assert result.used_bytes == 5000 + assert result.capacity_bytes == 100000 + assert result.fs_used_bytes == 20000 + + +@pytest.mark.asyncio +@patch("ai.backend.storage.services.service.log_manager_api_entry_new", new_callable=AsyncMock) +async def test_delete_vfolder(mock_log, mock_service, mock_volume_pool): + mock_volume = MagicMock() + mock_volume.get_vfolder_mount = AsyncMock(side_effect=VFolderNotFoundError) + + mock_volume_pool.get_volume.return_value.__aenter__.return_value = mock_volume + + vfolder_id = VFolderID( + quota_scope_id=QuotaScopeID(scope_type=QuotaScopeType.USER, scope_id=UUID), + folder_id=UUID, + ) + vfolder_key = VFolderKey(volume_id=UUID, vfolder_id=vfolder_id) + + with pytest.raises(web.HTTPGone, match="VFolder not found"): + await mock_service.delete_vfolder(vfolder_key) + + mock_log.assert_called_once_with(service_log, "delete_vfolder", vfolder_key) + mock_volume.get_vfolder_mount.assert_called_once_with(vfolder_id, ".") From 233d8eab071ce0b468cdc5604cd609357485d48f Mon Sep 17 00:00:00 2001 From: Mincheol Kang Date: Fri, 14 Feb 2025 16:34:41 +0900 Subject: [PATCH 6/7] fix: Revise unchanged lines (pydantic to dataclass) --- src/ai/backend/storage/volumes/types.py | 6 ++---- tests/storage-proxy/vfolder/test_handler.py | 22 ++++++++++----------- 2 files changed, 13 insertions(+), 15 deletions(-) diff --git a/src/ai/backend/storage/volumes/types.py b/src/ai/backend/storage/volumes/types.py index 6c2d35f6923..50748ecc1de 100644 --- a/src/ai/backend/storage/volumes/types.py +++ b/src/ai/backend/storage/volumes/types.py @@ -2,8 +2,6 @@ from pathlib import Path, PurePath from typing import Optional, Self -from pydantic import Field - from ai.backend.common.dto.storage.field import VFolderMetaField, VolumeMetaField from ai.backend.common.dto.storage.path import QuotaScopeKeyPath, VFolderKeyPath from ai.backend.common.dto.storage.response import QuotaScopeResponse @@ -74,8 +72,8 @@ def to_field(self) -> VFolderMetaField: @dataclass class QuotaScopeMeta: - used_bytes: Optional[int] = Field(default=0) - limit_bytes: Optional[int] = Field(default=0) + used_bytes: Optional[int] = 0 + limit_bytes: Optional[int] = 0 def to_response(self) -> QuotaScopeResponse: return QuotaScopeResponse( diff --git a/tests/storage-proxy/vfolder/test_handler.py b/tests/storage-proxy/vfolder/test_handler.py index b1b3eb7f161..0df015ed6cd 100644 --- a/tests/storage-proxy/vfolder/test_handler.py +++ b/tests/storage-proxy/vfolder/test_handler.py @@ -48,10 +48,10 @@ async def create_quota_scope(self, quota_scope_key, options): pass async def get_quota_scope(self, quota_scope_key): - return QuotaScopeMeta.model_validate({ - "used_bytes": 1000, - "limit_bytes": 2000, - }) + return QuotaScopeMeta( + used_bytes=1000, + limit_bytes=2000, + ) async def update_quota_scope(self, quota_scope_key, options): pass @@ -66,13 +66,13 @@ async def clone_vfolder(self, src_vfolder_key, dst_vfolder_key): pass async def get_vfolder_info(self, vfolder_key, subpath): - return VFolderMeta.model_validate({ - "mount_path": subpath, - "file_count": 100, - "used_bytes": 1000, - "capacity_bytes": 2000, - "fs_used_bytes": 1000, - }) + return VFolderMeta( + mount_path=Path(subpath), + file_count=100, + used_bytes=1000, + capacity_bytes=2000, + fs_used_bytes=1000, + ) async def delete_vfolder(self, vfolder_key): pass From 4b29c56f2222742aa5cfb27e61c7f0ab0f06bba0 Mon Sep 17 00:00:00 2001 From: Mincheol Kang Date: Fri, 14 Feb 2025 16:53:54 +0900 Subject: [PATCH 7/7] fix: Updated field names in log_manager_api_entry_new --- src/ai/backend/storage/utils.py | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/src/ai/backend/storage/utils.py b/src/ai/backend/storage/utils.py index 0735d5545e9..d2bb3724449 100644 --- a/src/ai/backend/storage/utils.py +++ b/src/ai/backend/storage/utils.py @@ -135,43 +135,43 @@ async def log_manager_api_entry_new( params: Any, ) -> None: if isinstance(params, dict): - if "src_vfid" in params and "dst_vfid" in params: + if "vfolder_id" in params and "dst_vfolder_id" in params: log.info( "ManagerAPI::{}(v:{}, f:{} -> dst_v: {}, dst_f:{})", name.upper(), - params["src_volume"], - params["src_vfid"], - params["dst_volume"], - params["dst_vfid"], + params["volume_id"], + params["vfolder_id"], + params["dst_volume_id"], + params["dst_vfolder_id"], ) elif "relpaths" in params: log.info( "ManagerAPI::{}(v:{}, f:{}, p*:{})", name.upper(), - params["volume"], - params["vfid"], + params["volume_id"], + params["vfolder_id"], str(params["relpaths"][0]) + "...", ) elif "relpath" in params: log.info( "ManagerAPI::{}(v:{}, f:{}, p:{})", name.upper(), - params["volume"], - params["vfid"], + params["volume_id"], + params["vfolder_id"], params["relpath"], ) - elif "vfid" in params: + elif "vfolder_id" in params: log.info( "ManagerAPI::{}(v:{}, f:{})", name.upper(), - params["volume"], - params["vfid"], + params["volume_id"], + params["vfolder_id"], ) - elif "volume" in params: + elif "volume_id" in params: log.info( "ManagerAPI::{}(v:{})", name.upper(), - params["volume"], + params["volume_id"], ) return