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/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..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 ( @@ -126,6 +126,10 @@ def from_request(cls, request: web.Request) -> Self: pass +class BaseRequestModel(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + + class BaseResponseModel(BaseModel): pass 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/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/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 new file mode 100644 index 00000000000..babd5dedae8 --- /dev/null +++ b/src/ai/backend/common/dto/storage/request.py @@ -0,0 +1,26 @@ +from typing import Optional + +from pydantic import AliasChoices, Field + +from ...api_handlers import BaseRequestModel +from ...types import QuotaConfig, VFolderID + + +class QuotaScopeReq(BaseRequestModel): + options: Optional[QuotaConfig] = Field( + default=None, + description="The options for the quota scope.", + ) + + +class GetVFolderMetaReq(BaseRequestModel): + subpath: str = Field( + description="The subpath of the virtual folder.", + ) + + +class CloneVFolderReq(BaseRequestModel): + dst_vfolder_id: VFolderID = Field( + description="The destination virtual folder ID.", + validation_alias=AliasChoices("dst_vfid", "dst_vfolder_id"), + ) 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..fc358a31b41 --- /dev/null +++ b/src/ai/backend/common/dto/storage/response.py @@ -0,0 +1,23 @@ +from typing import Optional + +from pydantic import Field + +from ai.backend.common.api_handlers import BaseResponseModel +from ai.backend.common.dto.storage.field import VFolderMetaField, VolumeMetaField + + +class GetVolumeResponse(BaseResponseModel): + item: VolumeMetaField + + +class GetVolumesResponse(BaseResponseModel): + items: list[VolumeMetaField] + + +class QuotaScopeResponse(BaseResponseModel): + used_bytes: Optional[int] = Field(default=0) + limit_bytes: Optional[int] = Field(default=0) + + +class VFolderMetadataResponse(BaseResponseModel): + item: VFolderMetaField 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 new file mode 100644 index 00000000000..d8e88e241b5 --- /dev/null +++ b/src/ai/backend/storage/api/vfolder/handler.py @@ -0,0 +1,142 @@ +from typing import Optional, Protocol + +from ai.backend.common.api_handlers import APIResponse, BodyParam, PathParam, api_handler +from ai.backend.common.dto.storage.path import QuotaScopeKeyPath, VFolderKeyPath, VolumeIDPath +from ai.backend.common.dto.storage.request import ( + CloneVFolderReq, + GetVFolderMetaReq, + QuotaScopeReq, +) +from ai.backend.common.dto.storage.response import ( + GetVolumeResponse, + GetVolumesResponse, + VFolderMetadataResponse, +) +from ai.backend.common.types import QuotaConfig, VFolderID, VolumeID + +from ...volumes.types import ( + QuotaScopeKey, + QuotaScopeMeta, + VFolderKey, + VFolderMeta, + VolumeMeta, +) + + +class VFolderServiceProtocol(Protocol): + async def get_volume(self, volume_id: VolumeID) -> VolumeMeta: ... + + async def get_volumes(self) -> list[VolumeMeta]: ... + + async def create_quota_scope( + self, quota_scope_key: QuotaScopeKey, options: Optional[QuotaConfig] + ) -> None: ... + + async def get_quota_scope(self, quota_scope_key: QuotaScopeKey) -> QuotaScopeMeta: ... + + async def update_quota_scope( + self, quota_scope_key: QuotaScopeKey, options: Optional[QuotaConfig] + ) -> None: ... + + async def delete_quota_scope(self, quota_scope_key: QuotaScopeKey) -> None: ... + + async def create_vfolder(self, vfolder_key: VFolderKey) -> None: ... + + async def clone_vfolder(self, vfolder_key: VFolderKey, dst_vfolder_id: VFolderID) -> None: ... + + async def get_vfolder_info(self, vfolder_key: VFolderKey, subpath: str) -> VFolderMeta: ... + + async def delete_vfolder(self, vfolder_key: VFolderKey) -> None: ... + + +class VFolderHandler: + _storage_service: VFolderServiceProtocol + + def __init__(self, storage_service: VFolderServiceProtocol) -> None: + self._storage_service = storage_service + + @api_handler + async def get_volume(self, path: PathParam[VolumeIDPath]) -> APIResponse: + volume_meta = await self._storage_service.get_volume(path.parsed.volume_id) + return APIResponse.build( + status_code=200, + response_model=GetVolumeResponse( + item=volume_meta.to_field(), + ), + ) + + @api_handler + async def get_volumes(self) -> APIResponse: + volume_meta_list = await self._storage_service.get_volumes() + return APIResponse.build( + status_code=200, + response_model=GetVolumesResponse( + items=[volume.to_field() for volume in volume_meta_list], + ), + ) + + @api_handler + async def create_quota_scope( + self, path: PathParam[QuotaScopeKeyPath], body: BodyParam[QuotaScopeReq] + ) -> APIResponse: + quota_scope_key = QuotaScopeKey.from_quota_scope_path(path.parsed) + await self._storage_service.create_quota_scope(quota_scope_key, body.parsed.options) + return APIResponse.no_content(status_code=204) + + @api_handler + async def get_quota_scope(self, path: PathParam[QuotaScopeKeyPath]) -> APIResponse: + quota_scope_key = QuotaScopeKey.from_quota_scope_path(path.parsed) + quota_scope = await self._storage_service.get_quota_scope(quota_scope_key) + return APIResponse.build( + status_code=200, + response_model=quota_scope.to_response(), + ) + + @api_handler + async def update_quota_scope( + self, path: PathParam[QuotaScopeKeyPath], body: BodyParam[QuotaScopeReq] + ) -> APIResponse: + quota_scope_key = QuotaScopeKey.from_quota_scope_path(path.parsed) + await self._storage_service.update_quota_scope(quota_scope_key, body.parsed.options) + return APIResponse.no_content(status_code=204) + + @api_handler + async def delete_quota_scope(self, path: PathParam[QuotaScopeKeyPath]) -> APIResponse: + quota_scope_key = QuotaScopeKey.from_quota_scope_path(path.parsed) + await self._storage_service.delete_quota_scope(quota_scope_key) + return APIResponse.no_content(status_code=204) + + @api_handler + async def create_vfolder(self, path: PathParam[VFolderKeyPath]) -> APIResponse: + vfolder_key = VFolderKey.from_vfolder_path(path.parsed) + await self._storage_service.create_vfolder(vfolder_key) + return APIResponse.no_content(status_code=204) + + @api_handler + async def clone_vfolder( + self, path: PathParam[VFolderKeyPath], body: BodyParam[CloneVFolderReq] + ) -> APIResponse: + vfolder_key = VFolderKey.from_vfolder_path(path.parsed) + await self._storage_service.clone_vfolder(vfolder_key, body.parsed.dst_vfolder_id) + return APIResponse.no_content(status_code=204) + + @api_handler + async def get_vfolder_info( + self, path: PathParam[VFolderKeyPath], body: BodyParam[GetVFolderMetaReq] + ) -> APIResponse: + vfolder_key = VFolderKey.from_vfolder_path(path.parsed) + vfolder_meta = await self._storage_service.get_vfolder_info( + vfolder_key, body.parsed.subpath + ) + return APIResponse.build( + status_code=200, + response_model=VFolderMetadataResponse( + item=vfolder_meta.to_field(), + ), + ) + + @api_handler + async def delete_vfolder(self, path: PathParam[VFolderKeyPath]) -> APIResponse: + vfolder_key = VFolderKey.from_vfolder_path(path.parsed) + await self._storage_service.delete_vfolder(vfolder_key) + return APIResponse.no_content(status_code=204) 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/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/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/services/service.py b/src/ai/backend/storage/services/service.py new file mode 100644 index 00000000000..791ef6da483 --- /dev/null +++ b/src/ai/backend/storage/services/service.py @@ -0,0 +1,247 @@ +import asyncio +import json +import logging +import uuid +import weakref +from contextlib import asynccontextmanager as actxmgr +from typing import AsyncIterator, Optional + +from aiohttp import web + +from ai.backend.common.events import VFolderDeletionFailureEvent, VFolderDeletionSuccessEvent +from ai.backend.common.types import QuotaConfig, VFolderID, VolumeID +from ai.backend.logging.utils import BraceStyleAdapter + +from ..exception import ( + ExternalError, + InvalidQuotaConfig, + InvalidQuotaScopeError, + InvalidSubpathError, + QuotaScopeAlreadyExists, + QuotaScopeNotFoundError, + VFolderNotFoundError, +) +from ..utils import log_manager_api_entry_new +from ..volumes.pool import VolumePool +from ..volumes.types import ( + QuotaScopeKey, + QuotaScopeMeta, + VFolderKey, + VFolderMeta, + VolumeMeta, +) + +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_key: VFolderKey, + ) -> None: + volume_id = vfolder_key.volume_id + vfolder_id = vfolder_key.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_id: VolumeID) -> VolumeMeta: + 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, + backend=volume.backend, + path=volume.path, + fsprefix=volume.fsprefix, + capabilities=await self._get_capabilities(volume_id), + ) + + async def get_volumes(self) -> list[VolumeMeta]: + await log_manager_api_entry_new(log, "get_volumes", params=None) + volumes = self._volume_pool.list_volumes() + 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() + ] + + 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_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(): + 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.") + + async def get_quota_scope(self, quota_scope_key: QuotaScopeKey) -> QuotaScopeMeta: + 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( + quota_scope_key.quota_scope_id + ) + if not quota_usage: + raise QuotaScopeNotFoundError + return QuotaScopeMeta( + used_bytes=quota_usage.used_bytes, limit_bytes=quota_usage.limit_bytes + ) + + 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_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) + 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") + + 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_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) + if not quota_usage: + raise QuotaScopeNotFoundError + await volume.quota_model.unset_quota(quota_scope_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_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: + try: + await volume.create_vfolder(vfolder_id) + except QuotaScopeNotFoundError: + 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") + + async def clone_vfolder(self, vfolder_key: VFolderKey, dst_vfolder_id: VFolderID) -> None: + 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_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) + 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 VFolderMeta( + 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_key: VFolderKey) -> None: + vfolder_id = vfolder_key.vfolder_id + 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, ".") + 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_key)) + return None diff --git a/src/ai/backend/storage/utils.py b/src/ai/backend/storage/utils.py index 4e8e553cd0f..d2bb3724449 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 "vfolder_id" in params and "dst_vfolder_id" in params: + log.info( + "ManagerAPI::{}(v:{}, f:{} -> dst_v: {}, dst_f:{})", + name.upper(), + 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_id"], + params["vfolder_id"], + str(params["relpaths"][0]) + "...", + ) + elif "relpath" in params: + log.info( + "ManagerAPI::{}(v:{}, f:{}, p:{})", + name.upper(), + params["volume_id"], + params["vfolder_id"], + params["relpath"], + ) + elif "vfolder_id" in params: + log.info( + "ManagerAPI::{}(v:{}, f:{})", + name.upper(), + params["volume_id"], + params["vfolder_id"], + ) + elif "volume_id" in params: + log.info( + "ManagerAPI::{}(v:{})", + name.upper(), + params["volume_id"], + ) + return + + log.info("ManagerAPI::{}({})", name.upper(), str(params)) 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..149234c4cf4 100644 --- a/src/ai/backend/storage/volumes/pool.py +++ b/src/ai/backend/storage/volumes/pool.py @@ -6,22 +6,23 @@ 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, @@ -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, @@ -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 { @@ -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/types.py b/src/ai/backend/storage/volumes/types.py new file mode 100644 index 00000000000..50748ecc1de --- /dev/null +++ b/src/ai/backend/storage/volumes/types.py @@ -0,0 +1,82 @@ +from dataclasses import dataclass +from pathlib import Path, PurePath +from typing import Optional, Self + +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 + 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) + ) + + +@dataclass +class VFolderKey: + volume_id: VolumeID + vfolder_id: VFolderID + + @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), + ) + + +@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: + 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, + ) + + +@dataclass +class QuotaScopeMeta: + used_bytes: Optional[int] = 0 + limit_bytes: Optional[int] = 0 + + 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 6941bffe066..0df015ed6cd 100644 --- a/tests/storage-proxy/vfolder/test_handler.py +++ b/tests/storage-proxy/vfolder/test_handler.py @@ -1,291 +1,266 @@ -# 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, 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( + 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.get_quota_scope(await mock_request()) + async def get_vfolder_info(self, vfolder_key, subpath): + 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 + + return MockVFolderService() -# assert isinstance(response, QuotaScopeResponseModel) -# assert response.used_bytes == 1024 -# assert response.limit_bytes == 2048 +@pytest.mark.asyncio +async def test_get_volume(mock_vfolder_service): + handler = VFolderHandler(mock_vfolder_service) -# @pytest.mark.asyncio -# async def test_update_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) -# 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 + 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.update_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) +@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) + + assert isinstance(response, web.Response) + assert response.status == 200 + volume_response = json.loads(response.text)["items"] + assert volume_response[0]["volume_id"] == str(UUID1) + + +@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 + + +@pytest.mark.asyncio +async def test_get_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), + } + + response: APIResponse = await handler.get_quota_scope(mock_request) + + 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_update_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.update_quota_scope(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(mock_vfolder_service) + + mock_request = AsyncMock(web.Request) + mock_request.match_info = { + "volume_id": str(UUID), + "scope_type": "user", + "scope_uuid": str(UUID), + } + + response: APIResponse = await handler.delete_quota_scope(mock_request) + + assert isinstance(response, web.Response) + assert response.status == 204 + + +@pytest.mark.asyncio +async def test_create_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: APIResponse = await handler.create_vfolder(mock_request) + + assert isinstance(response, web.Response) + assert response.status == 204 + + +@pytest.mark.asyncio +async def test_clone_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), + } + mock_request.json.return_value = { + "dst_vfolder_id": VFolderID( + quota_scope_id=QuotaScopeID(scope_type=QuotaScopeType.USER, scope_id=UUID), + folder_id=UUID, + ) + } + + response: APIResponse = await handler.clone_vfolder(mock_request) + + assert isinstance(response, web.Response) + assert response.status == 204 + + +@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: APIResponse = await handler.get_vfolder_info(mock_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: APIResponse = await handler.delete_vfolder(mock_request) + + assert isinstance(response, web.Response) + assert response.status == 204 diff --git a/tests/storage-proxy/vfolder/test_pool.py b/tests/storage-proxy/vfolder/test_pool.py new file mode 100644 index 00000000000..f7c95514eaf --- /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.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 + + +@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 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, ".")