feat(BA-2263): Abstract storage-proxy storage types

changes/5750.feature.md

@@ -0,0 +1 @@
+Abstract `storage-proxy` storage types

src/ai/backend/common/dto/storage/request.py

@@ -190,13 +190,6 @@ class HuggingFaceImportModelsReq(BaseRequestModel):
         """,
         examples=["default-minio", "s3-storage", "local-storage"],
     )
-    bucket_name: str = Field(
-        description="""
-        Target bucket name within the storage for all models.
-        The bucket must exist and be writable by the service.
-        """,
-        examples=["models", "huggingface-models", "ai-models"],
-    )
 
 
 class ReservoirImportModelsReq(BaseRequestModel):
@@ -228,13 +221,6 @@ class ReservoirImportModelsReq(BaseRequestModel):
         """,
         examples=["default-minio", "s3-storage", "local-storage"],
     )
-    bucket_name: str = Field(
-        description="""
-        Target bucket name within the storage for all models.
-        The bucket must exist and be writable by the service.
-        """,
-        examples=["models", "huggingface-models", "ai-models"],
-    )
 
 
 class DeleteObjectReq(BaseRequestModel):

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

@@ -185,7 +185,6 @@ async def import_revision(
                         ],
                         registry_name=registry_data.name,
                         storage_name=storage_data.name,
-                        bucket_name=storage_namespace.bucket,
                     )
                 )
             case ArtifactRegistryType.RESERVOIR:
@@ -202,7 +201,6 @@ async def import_revision(
                         ],
                         registry_name=registry_data.name,
                         storage_name=storage_data.name,
-                        bucket_name=storage_namespace.bucket,
                     )
                 )
             case _:

src/ai/backend/storage/api/v1/registries/huggingface.py

@@ -34,7 +34,6 @@
     HuggingFaceService,
     HuggingFaceServiceArgs,
 )
-from ai.backend.storage.services.storages.object_storage import ObjectStorageService
 
 from ....utils import log_client_api_entry
 
@@ -143,7 +142,6 @@ async def import_models(
             registry_name=body.parsed.registry_name,
             models=body.parsed.models,
             storage_name=body.parsed.storage_name,
-            bucket_name=body.parsed.bucket_name,
         )
 
         response = HuggingFaceImportModelsResponse(
@@ -161,8 +159,6 @@ def create_app(ctx: RootContext) -> web.Application:
     app["ctx"] = ctx
     app["prefix"] = "v1/registries/huggingface"
 
-    storage_service = ObjectStorageService(storage_configs=ctx.local_config.storages)
-
     huggingface_registry_configs = dict(
         (r.name, r.config)
         for r in ctx.local_config.registries
@@ -171,7 +167,7 @@ def create_app(ctx: RootContext) -> web.Application:
     huggingface_service = HuggingFaceService(
         HuggingFaceServiceArgs(
             background_task_manager=ctx.background_task_manager,
-            storage_service=storage_service,
+            storage_pool=ctx.storage_pool,
             registry_configs=huggingface_registry_configs,
             event_producer=ctx.event_producer,
         )

src/ai/backend/storage/api/v1/registries/reservoir.py

@@ -52,7 +52,6 @@ async def import_models(
             registry_name=body.parsed.registry_name,
             models=body.parsed.models,
             storage_name=body.parsed.storage_name,
-            bucket_name=body.parsed.bucket_name,
         )
 
         return APIResponse.build(
@@ -74,7 +73,7 @@ def create_app(ctx: RootContext) -> web.Application:
         ReservoirServiceArgs(
             background_task_manager=ctx.background_task_manager,
             event_producer=ctx.event_producer,
-            storage_configs=ctx.local_config.storages,
+            storage_pool=ctx.storage_pool,
             reservoir_registry_configs=reservoir_registry_configs,
         )
     )

src/ai/backend/storage/api/v1/storages/object_storage.py

@@ -26,11 +26,9 @@
     UploadObjectReq,
 )
 from ai.backend.logging import BraceStyleAdapter
-from ai.backend.storage.config.unified import (
-    ObjectStorageConfig,
-)
 
 from ....services.storages.object_storage import ObjectStorageService
+from ....storages.base import StoragePool
 from ....utils import log_client_api_entry
 
 if TYPE_CHECKING:
@@ -42,13 +40,13 @@
 
 
 class ObjectStorageAPIHandler:
-    _storage_configs: list[ObjectStorageConfig]
+    _storage_pool: StoragePool
 
     def __init__(
         self,
-        storage_configs: list[ObjectStorageConfig],
+        storage_pool: StoragePool,
     ) -> None:
-        self._storage_configs = storage_configs
+        self._storage_pool = storage_pool
 
     @api_handler
     async def upload_object(
@@ -69,7 +67,7 @@ async def upload_object(
 
         await log_client_api_entry(log, "upload_object", req)
 
-        storage_service = ObjectStorageService(self._storage_configs)
+        storage_service = ObjectStorageService(self._storage_pool)
 
         file_part = await file_reader.next()
         while file_part and not getattr(file_part, "filename", None):
@@ -115,7 +113,7 @@ async def download_file(
         bucket_name = path.parsed.bucket_name
 
         await log_client_api_entry(log, "download_file", req)
-        storage_service = ObjectStorageService(self._storage_configs)
+        storage_service = ObjectStorageService(self._storage_pool)
         download_stream = storage_service.stream_download(storage_name, bucket_name, filepath)
 
         return APIStreamResponse(
@@ -141,7 +139,7 @@ async def presigned_upload_url(
         bucket_name = path.parsed.bucket_name
 
         await log_client_api_entry(log, "presigned_upload_url", req)
-        storage_service = ObjectStorageService(self._storage_configs)
+        storage_service = ObjectStorageService(self._storage_pool)
         response = await storage_service.generate_presigned_upload_url(
             storage_name, bucket_name, req.key
         )
@@ -167,7 +165,7 @@ async def presigned_download_url(
         bucket_name = path.parsed.bucket_name
 
         await log_client_api_entry(log, "presigned_download_url", req)
-        storage_service = ObjectStorageService(self._storage_configs)
+        storage_service = ObjectStorageService(self._storage_pool)
         response = await storage_service.generate_presigned_download_url(
             storage_name, bucket_name, filepath
         )
@@ -194,7 +192,7 @@ async def get_object_meta(
 
         await log_client_api_entry(log, "get_object_meta", req)
 
-        storage_service = ObjectStorageService(self._storage_configs)
+        storage_service = ObjectStorageService(self._storage_pool)
         response = await storage_service.get_object_info(storage_name, bucket_name, filepath)
 
         return APIResponse.build(
@@ -218,7 +216,7 @@ async def delete_object(
         bucket_name = path.parsed.bucket_name
 
         await log_client_api_entry(log, "delete_object", req)
-        storage_service = ObjectStorageService(self._storage_configs)
+        storage_service = ObjectStorageService(self._storage_pool)
 
         await storage_service.delete_object(storage_name, bucket_name, prefix)
 
@@ -233,7 +231,7 @@ def create_app(ctx: RootContext) -> web.Application:
     app["prefix"] = "v1/storages/s3"
 
     api_handler = ObjectStorageAPIHandler(
-        storage_configs=ctx.local_config.storages,
+        storage_pool=ctx.storage_pool,
     )
     app.router.add_route(
         "GET", "/{storage_name}/buckets/{bucket_name}/object/meta", api_handler.get_object_meta

src/ai/backend/storage/context.py

@@ -36,6 +36,7 @@
     StoragePluginContext,
 )
 from .services.service import VolumeService
+from .storages.base import StoragePool
 from .types import VolumeInfo
 from .volumes.abc import AbstractVolume
 from .volumes.cephfs import CephFSVolume
@@ -117,6 +118,7 @@ class RootContext:
     local_config: StorageProxyUnifiedConfig
     dsn: str | None
     volume_pool: VolumePool
+    storage_pool: StoragePool
     event_producer: EventProducer
     event_dispatcher: EventDispatcher
     watcher: WatcherClient | None
@@ -132,6 +134,7 @@ def __init__(
         etcd: AsyncEtcd,
         *,
         volume_pool: VolumePool,
+        storage_pool: StoragePool,
         background_task_manager: BackgroundTaskManager,
         event_producer: EventProducer,
         event_dispatcher: EventDispatcher,
@@ -158,6 +161,7 @@ def __init__(
         }
         self.metric_registry = metric_registry
         self.volume_pool = volume_pool
+        self.storage_pool = storage_pool
         self.background_task_manager = background_task_manager
 
     async def __aenter__(self) -> None:

src/ai/backend/storage/exception.py

@@ -391,3 +391,16 @@ def error_code(cls) -> ErrorCode:
             operation=ErrorOperation.CREATE,
             error_detail=ErrorDetail.INTERNAL_ERROR,
         )
+
+
+class NotImplementedAPI(BackendAIError, web.HTTPBadRequest):
+    error_type = "https://api.backend.ai/probs/storage/api/not-implemented"
+    error_title = "API Not Implemented"
+
+    @classmethod
+    def error_code(cls) -> ErrorCode:
+        return ErrorCode(
+            domain=ErrorDomain.STORAGE_PROXY,
+            operation=ErrorOperation.GENERIC,
+            error_detail=ErrorDetail.NOT_IMPLEMENTED,
+        )

src/ai/backend/storage/migration.py

@@ -266,6 +266,7 @@ async def check_and_upgrade(
         event_dispatcher=event_dispatcher,
         watcher=None,
         volume_pool=None,  # type: ignore[arg-type]
+        storage_pool=None,  # type: ignore[arg-type]
         background_task_manager=None,  # type: ignore[arg-type]
     )
 

src/ai/backend/storage/server.py

@@ -106,6 +106,8 @@ async def server_main(
     from .bgtask.registry import BgtaskHandlerRegistryCreator
     from .context import RootContext
     from .migration import check_latest
+    from .storages.base import AbstractStorage, StoragePool
+    from .storages.object_storage import ObjectStorage
     from .volumes.pool import VolumePool
     from .watcher import WatcherClient
 
@@ -211,13 +213,21 @@ async def server_main(
             valkey_client=valkey_client,
             server_id=local_config.storage_proxy.node_id,
         )
+
+        # Create StoragePool for object storage
+        storages: dict[str, AbstractStorage] = {}
+        for config in local_config.storages:
+            storages[config.name] = ObjectStorage(config)
+        storage_pool = StoragePool(storages)
+
         ctx = RootContext(
             pid=os.getpid(),
             node_id=local_config.storage_proxy.node_id,
             pidx=pidx,
             local_config=local_config,
             etcd=etcd,
             volume_pool=volume_pool,
+            storage_pool=storage_pool,
             background_task_manager=bgtask_mgr,
             event_producer=event_producer,
             event_dispatcher=event_dispatcher,