fix: Reflect feedback

Author: jopemachine

src/ai/backend/common/artifact_storage.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Any, Protocol
+from typing import Any
 
 from ai.backend.common.data.storage.registries.types import ModelTarget
 from ai.backend.common.data.storage.types import ArtifactStorageImportStep
@@ -30,25 +30,30 @@ async def get_file_info(self, filepath: str) -> Any:
         raise NotImplementedError
 
 
-class StoragePoolProtocol(Protocol):
-    """Protocol for StoragePool interface"""
+class AbstractStoragePool(ABC):
+    """Abstract base class for storage pool interface"""
 
+    @abstractmethod
     def get_storage(self, name: str) -> AbstractStorage:
         """Get storage by name"""
         ...
 
+    @abstractmethod
     def add_storage(self, name: str, storage: AbstractStorage) -> None:
         """Add a storage to the pool"""
         ...
 
+    @abstractmethod
     def remove_storage(self, name: str) -> None:
         """Remove a storage from the pool"""
         ...
 
+    @abstractmethod
     def list_storages(self) -> list[str]:
         """List all storage names in the pool"""
         ...
 
+    @abstractmethod
     def has_storage(self, name: str) -> bool:
         """Check if storage exists in the pool"""
         ...
@@ -60,6 +65,6 @@ class ImportStepContext:
 
     model: ModelTarget
     registry_name: str
-    storage_pool: StoragePoolProtocol
+    storage_pool: AbstractStoragePool
     storage_step_mappings: dict[ArtifactStorageImportStep, str]
     step_metadata: dict[str, Any]  # For passing data between steps

src/ai/backend/storage/services/artifacts/huggingface.py

@@ -12,7 +12,7 @@
 
 import aiohttp
 
-from ai.backend.common.artifact_storage import StoragePoolProtocol
+from ai.backend.common.artifact_storage import AbstractStoragePool
 from ai.backend.common.bgtask.bgtask import BackgroundTaskManager, ProgressReporter
 from ai.backend.common.data.artifact.types import ArtifactRegistryType
 from ai.backend.common.data.storage.registries.types import (
@@ -795,7 +795,7 @@ async def _download_file_to_storage(
         file_info: FileObjectData,
         model: ModelTarget,
         storage_name: str,
-        storage_pool: StoragePoolProtocol,
+        storage_pool: AbstractStoragePool,
         download_chunk_size: int,
     ) -> str:
         """Download file from HuggingFace to specified storage"""

src/ai/backend/storage/services/artifacts/reservoir.py

@@ -11,7 +11,7 @@
 
 import aiofiles
 
-from ai.backend.common.artifact_storage import AbstractStorage, StoragePoolProtocol
+from ai.backend.common.artifact_storage import AbstractStorage, AbstractStoragePool
 from ai.backend.common.bgtask.bgtask import BackgroundTaskManager, ProgressReporter
 from ai.backend.common.data.artifact.types import ArtifactRegistryType
 from ai.backend.common.data.storage.registries.types import FileObjectData, ModelTarget
@@ -581,7 +581,7 @@ async def _handle_object_storage_download(
         return downloaded_files, bytes_copied
 
     def _get_s3_client(
-        self, storage_pool: StoragePoolProtocol, storage_name: str
+        self, storage_pool: AbstractStoragePool, storage_name: str
     ) -> tuple[S3Client, str]:
         """Get S3 client for the specified storage"""
         # Get storage from pool and verify it's ObjectStorage type
@@ -617,7 +617,7 @@ async def _stream_bucket_to_bucket(
         self,
         source_cfg: ReservoirConfig,
         storage_name: str,
-        storage_pool: StoragePoolProtocol,
+        storage_pool: AbstractStoragePool,
         options: BucketCopyOptions,
         progress_reporter: Optional[ProgressReporter],
         key_prefix: Optional[str] = None,

src/ai/backend/storage/storages/storage_pool.py

@@ -1,7 +1,7 @@
 import logging
 from typing import Optional, Self
 
-from ai.backend.common.artifact_storage import AbstractStorage
+from ai.backend.common.artifact_storage import AbstractStorage, AbstractStoragePool
 from ai.backend.common.data.storage.types import ArtifactStorageType
 from ai.backend.common.exception import GenericNotImplementedError, InvalidConfigError
 from ai.backend.logging.utils import BraceStyleAdapter
@@ -14,7 +14,7 @@
 log = BraceStyleAdapter(logging.getLogger(__spec__.name))
 
 
-class StoragePool:
+class StoragePool(AbstractStoragePool):
     """
     Storage pool that manages different types of storage backends.
     Supports both Object Storage (S3-compatible) and VFS storage.