feat(BA-3991): Implement artifact storage usage tracking and quota enforcement (WIP)

changes/8390.feature.md

@@ -0,0 +1 @@
+Implement artifact storage usage tracking and quota enforcement

src/ai/backend/manager/errors/vfolder.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from aiohttp import web
+
+from ai.backend.common.exception import (
+    BackendAIError,
+    ErrorCode,
+    ErrorDetail,
+    ErrorDomain,
+    ErrorOperation,
+)
+
+if TYPE_CHECKING:
+    from ai.backend.common.types import QuotaScopeID, VFolderID
+
+
+class VFolderQuotaExceededError(BackendAIError, web.HTTPBadRequest):
+    """
+    Raised when a VFolder quota scope limit would be exceeded by an operation.
+
+    This error is raised during pre-validation before operations like artifact import,
+    preventing partial downloads and wasted resources.
+    """
+
+    error_type = "https://api.backend.ai/probs/vfolder-quota-exceeded"
+    error_title = "VFolder Quota Scope Limit Would Be Exceeded"
+
+    def __init__(
+        self,
+        vfolder_id: VFolderID,
+        quota_scope_id: QuotaScopeID,
+        current_size: int,
+        max_size: int,
+        requested_size: int,
+    ) -> None:
+        self.vfolder_id = vfolder_id
+        self.quota_scope_id = quota_scope_id
+        self.current_size = current_size
+        self.max_size = max_size
+        self.requested_size = requested_size
+        available_bytes = max_size - current_size
+
+        message = (
+            f"VFolder quota scope limit would be exceeded. "
+            f"Current: {current_size} bytes, Max: {max_size} bytes, "
+            f"Requested: {requested_size} bytes, Available: {available_bytes} bytes"
+        )
+        super().__init__(message)
+
+    def error_code(self) -> ErrorCode:
+        return ErrorCode(
+            domain=ErrorDomain.QUOTA_SCOPE,
+            operation=ErrorOperation.CHECK_LIMIT,
+            error_detail=ErrorDetail.BAD_REQUEST,
+        )
+
+    def error_data(self) -> dict[str, Any]:
+        return {
+            "vfolder_id": str(self.vfolder_id),
+            "quota_scope_id": str(self.quota_scope_id),
+            "current_size_bytes": self.current_size,
+            "max_size_bytes": self.max_size,
+            "requested_size_bytes": self.requested_size,
+            "available_bytes": self.max_size - self.current_size,
+        }

src/ai/backend/manager/services/artifact_revision/service.py

@@ -29,7 +29,7 @@
     HuggingFaceImportModelsReq,
     ReservoirImportModelsReq,
 )
-from ai.backend.common.types import VFolderID
+from ai.backend.common.types import QuotaScopeID, QuotaScopeType, VFolderID
 from ai.backend.logging import BraceStyleAdapter
 from ai.backend.manager.clients.artifact_registry.reservoir_client import ReservoirRegistryClient
 from ai.backend.manager.clients.storage_proxy.session_manager import StorageSessionManager
@@ -44,6 +44,7 @@
     ArtifactType,
 )
 from ai.backend.manager.data.artifact_registries.types import ArtifactRegistryData
+from ai.backend.manager.data.vfolder.types import VFolderData
 from ai.backend.manager.dto.request import (
     DelegateImportArtifactsReq,
     ImportArtifactsOptions,
@@ -59,6 +60,7 @@
 )
 from ai.backend.manager.errors.common import ServerMisconfiguredError
 from ai.backend.manager.errors.storage import UnsupportedStorageTypeError
+from ai.backend.manager.errors.vfolder import VFolderQuotaExceededError
 from ai.backend.manager.repositories.artifact.repository import ArtifactRepository
 from ai.backend.manager.repositories.artifact_registry.repository import ArtifactRegistryRepository
 from ai.backend.manager.repositories.huggingface_registry.repository import HuggingFaceRepository
@@ -422,6 +424,10 @@ async def import_revision(
                     vfolder_id = VFolderID(vfolder_data.quota_scope_id, vfolder_data.id)
                     _, volume_name = self._storage_manager.get_proxy_and_volume(vfolder_data.host)
 
+                    # Pre-validate VFolder quota before starting import
+                    if revision_data.size is not None:
+                        await self._check_vfolder_quota(vfolder_data, revision_data.size)
+
             task_id: UUID
             match artifact.registry_type:
                 case ArtifactRegistryType.HUGGINGFACE:
@@ -812,3 +818,77 @@ def _is_latest_commit_hash(
             artifact_revision.status == ArtifactStatus.AVAILABLE
             and artifact_revision.digest == latest_commit_hash
         )
+
+    async def _check_vfolder_quota(
+        self,
+        vfolder_data: VFolderData,
+        additional_size: int,
+    ) -> None:
+        """
+        Check the quota scope limit for a VFolder before importing artifacts.
+
+        This pre-validation prevents failed imports mid-way when storage proxy
+        detects quota exceeded, avoiding partial downloads.
+
+        :param vfolder_data: VFolder data containing quota_scope_id and host
+        :param additional_size: Size of the artifact to be imported (in bytes)
+        :raises VFolderQuotaExceededError: If the quota limit would be exceeded
+        """
+        quota_scope_id = vfolder_data.quota_scope_id
+        if quota_scope_id is None:
+            # No quota scope configured for this VFolder
+            return
+
+        _, volume_name = self._storage_manager.get_proxy_and_volume(vfolder_data.host)
+
+        # 1. Get current quota scope usage from storage proxy
+        storage_client = self._storage_manager.get_manager_facing_client(vfolder_data.host)
+        usage_response = await storage_client.get_quota_scope(volume_name, str(quota_scope_id))
+        used_bytes = int(usage_response.get("used_bytes", 0))
+
+        # 2. Get limit from resource policy
+        max_size = await self._get_quota_scope_limit(quota_scope_id, vfolder_data.domain_name)
+
+        # -1 or 0 means unlimited
+        if max_size <= 0:
+            return
+
+        # 3. Check if limit would be exceeded
+        if used_bytes + additional_size > max_size:
+            vfolder_id = VFolderID(quota_scope_id, vfolder_data.id)
+            raise VFolderQuotaExceededError(
+                vfolder_id=vfolder_id,
+                quota_scope_id=quota_scope_id,
+                current_size=used_bytes,
+                max_size=max_size,
+                requested_size=additional_size,
+            )
+
+    async def _get_quota_scope_limit(self, quota_scope_id: QuotaScopeID, domain_name: str) -> int:
+        """
+        Get the quota scope limit from the appropriate resource policy.
+
+        :param quota_scope_id: The quota scope ID to get the limit for
+        :param domain_name: The domain name for project scope queries
+        :return: The max_quota_scope_size in bytes, or -1 if unlimited
+        """
+        match quota_scope_id.scope_type:
+            case QuotaScopeType.USER:
+                user_result = await self._vfolder_repository.get_user_resource_info(
+                    quota_scope_id.scope_id
+                )
+                if user_result is None:
+                    # User not found, treat as unlimited
+                    return -1
+                _, max_quota_scope_size, _ = user_result
+                return max_quota_scope_size
+
+            case QuotaScopeType.PROJECT:
+                group_result = await self._vfolder_repository.get_group_resource_info(
+                    quota_scope_id.scope_id, domain_name
+                )
+                if group_result is None:
+                    # Group not found, treat as unlimited
+                    return -1
+                _, _, max_quota_scope_size, _ = group_result
+                return max_quota_scope_size