refactor(BA-626): Move storage system modules to volumes package (#3563)

src/ai/backend/storage/api/client.py

@@ -40,8 +40,8 @@
 from ..utils import CheckParamSource, check_params
 
 if TYPE_CHECKING:
-    from ..abc import AbstractVolume
     from ..context import RootContext
+    from ..volumes.abc import AbstractVolume
 
 log = BraceStyleAdapter(logging.getLogger(__spec__.name))
 
@@ -446,7 +446,8 @@ async def init_client_app(ctx: RootContext) -> web.Application:
     r.add_route("GET", check_status)
     r = cors.add(app.router.add_resource("/download"))
     r.add_route("GET", download)
-    r = app.router.add_resource("/upload")  # tus handlers handle CORS by themselves
+    # tus handlers handle CORS by themselves
+    r = app.router.add_resource("/upload")
     r.add_route("OPTIONS", tus_options)
     r.add_route("HEAD", tus_check_session)
     r.add_route("PATCH", tus_upload_part)

src/ai/backend/storage/api/manager.py

@@ -60,8 +60,8 @@
 from ..utils import check_params, log_manager_api_entry
 
 if TYPE_CHECKING:
-    from ..abc import AbstractVolume
     from ..context import RootContext
+    from ..volumes.abc import AbstractVolume
 
 log = BraceStyleAdapter(logging.getLogger(__spec__.name))
 
@@ -958,7 +958,8 @@ class Params(TypedDict):
                     t.Key("vfid"): tx.VFolderID(),
                     t.Key("relpath"): tx.PurePath(relative_only=True),
                     t.Key("new_name"): t.String(),
-                    t.Key("is_dir", default=False): t.ToBool,  # ignored since 22.03
+                    # ignored since 22.03
+                    t.Key("is_dir", default=False): t.ToBool,
                 },
             ),
         ),

src/ai/backend/storage/context.py

@@ -24,28 +24,28 @@
 from ai.backend.common.metrics.metric import CommonMetricRegistry
 from ai.backend.logging import BraceStyleAdapter
 
-from .abc import AbstractVolume
 from .api.client import init_client_app
 from .api.manager import init_manager_app
-from .cephfs import CephFSVolume
-from .ddn import EXAScalerFSVolume
-from .dellemc import DellEMCOneFSVolume
 from .exception import InvalidVolumeError
-from .gpfs import GPFSVolume
-from .netapp import NetAppVolume
 from .plugin import (
     BasePluginContext,
     StorageClientWebappPluginContext,
     StorageManagerWebappPluginContext,
     StoragePluginContext,
 )
-from .purestorage import FlashBladeVolume
 from .types import VolumeInfo
-from .vast import VASTVolume
-from .vfs import BaseVolume
+from .volumes.abc import AbstractVolume
+from .volumes.cephfs import CephFSVolume
+from .volumes.ddn import EXAScalerFSVolume
+from .volumes.dellemc import DellEMCOneFSVolume
+from .volumes.gpfs import GPFSVolume
+from .volumes.netapp import NetAppVolume
+from .volumes.purestorage import FlashBladeVolume
+from .volumes.vast import VASTVolume
+from .volumes.vfs import BaseVolume
+from .volumes.weka import WekaVolume
+from .volumes.xfs import XfsVolume
 from .watcher import WatcherClient
-from .weka import WekaVolume
-from .xfs import XfsVolume
 
 log = BraceStyleAdapter(logging.getLogger(__spec__.name))
 

src/ai/backend/storage/migration.py

@@ -26,10 +26,10 @@
 )
 from ai.backend.logging import BraceStyleAdapter, LocalLogger
 
-from .abc import CAP_FAST_SIZE, AbstractVolume
 from .config import load_local_config, load_shared_config
 from .context import EVENT_DISPATCHER_CONSUMER_GROUP, RootContext
 from .types import VFolderID
+from .volumes.abc import CAP_FAST_SIZE, AbstractVolume
 
 log = BraceStyleAdapter(logging.getLogger(__spec__.name))
 

src/ai/backend/storage/plugin.py

@@ -6,8 +6,8 @@
 from aiohttp import web
 
 from ai.backend.common.plugin import AbstractPlugin, BasePluginContext
-from ai.backend.storage.abc import AbstractVolume
 from ai.backend.storage.api.types import CORSOptions, WebMiddleware
+from ai.backend.storage.volumes.abc import AbstractVolume
 
 
 class AbstractStoragePlugin(AbstractPlugin, metaclass=ABCMeta):

src/ai/backend/storage/volumes/BUILD

@@ -0,0 +1 @@
+python_sources(name="src")

src/ai/backend/storage/abc.py → src/ai/backend/storage/volumes/abc.py

@@ -20,8 +20,8 @@
 from ai.backend.common.types import BinarySize, HardwareMetadata, QuotaScopeID
 from ai.backend.logging import BraceStyleAdapter
 
-from .exception import InvalidSubpathError, VFolderNotFoundError
-from .types import (
+from ..exception import InvalidSubpathError, VFolderNotFoundError
+from ..types import (
     CapacityUsage,
     DirEntry,
     FSPerfMetric,
@@ -35,9 +35,12 @@
 CAP_VFOLDER: Final = "vfolder"  # ability to create vfolder
 CAP_METRIC: Final = "metric"  # ability to report disk related metrics
 CAP_QUOTA: Final = "quota"  # ability to manage quota limits
-CAP_FAST_FS_SIZE: Final = "fast-fs-size"  # ability to scan filesystem size fast (e.g. by API)
-CAP_FAST_SCAN: Final = "fast-scan"  # ability to scan number of files in vFolder fast (e.g. by API)
-CAP_FAST_SIZE: Final = "fast-size"  # ability to scan vFolder size fast (e.g. by API)
+# ability to scan filesystem size fast (e.g. by API)
+CAP_FAST_FS_SIZE: Final = "fast-fs-size"
+# ability to scan number of files in vFolder fast (e.g. by API)
+CAP_FAST_SCAN: Final = "fast-scan"
+# ability to scan vFolder size fast (e.g. by API)
+CAP_FAST_SIZE: Final = "fast-size"
 
 log = BraceStyleAdapter(logging.getLogger(__spec__.name))
 

src/ai/backend/storage/cephfs/__init__.py → src/ai/backend/storage/volumes/cephfs/__init__.py

@@ -8,10 +8,16 @@
 
 from ai.backend.common.types import BinarySize, QuotaScopeID
 from ai.backend.storage.exception import QuotaScopeNotFoundError
-
-from ..abc import CAP_FAST_SIZE, CAP_QUOTA, CAP_VFOLDER, AbstractFSOpModel, AbstractQuotaModel
-from ..subproc import run
-from ..types import CapacityUsage, Optional, QuotaConfig, QuotaUsage, TreeUsage
+from ai.backend.storage.subproc import run
+from ai.backend.storage.types import CapacityUsage, Optional, QuotaConfig, QuotaUsage, TreeUsage
+
+from ..abc import (
+    CAP_FAST_SIZE,
+    CAP_QUOTA,
+    CAP_VFOLDER,
+    AbstractFSOpModel,
+    AbstractQuotaModel,
+)
 from ..vfs import BaseFSOpModel, BaseQuotaModel, BaseVolume
 
 
@@ -34,9 +40,11 @@ async def describe_quota_scope(self, quota_scope_id: QuotaScopeID) -> Optional[Q
         loop = asyncio.get_running_loop()
 
         def read_attrs() -> tuple[int, int]:
-            used_bytes = int(os.getxattr(qspath, "ceph.dir.rbytes").decode())  # type: ignore[attr-defined]
+
+            used_bytes = int(os.getxattr(qspath, "ceph.dir.rbytes").decode()) # type: ignore[attr-defined]
             try:
-                limit_bytes = int(os.getxattr(qspath, "ceph.quota.max_bytes").decode())  # type: ignore[attr-defined]
+
+                limit_bytes = int(os.getxattr(qspath, "ceph.quota.max_bytes").decode()) # type: ignore[attr-defined]
             except OSError as e:
                 match e.errno:
                     case 61:
@@ -82,7 +90,8 @@ async def unset_quota(self, quota_scope_id: QuotaScopeID) -> None:
             None,
             # without type: ignore mypy will raise error when trying to run on macOS
             # because os.setxattr() exists only for linux
-            lambda: os.setxattr(qspath, "ceph.quota.max_bytes", b"0"),  # type: ignore[attr-defined]
+
+            lambda: os.setxattr(qspath, "ceph.quota.max_bytes", b"0"), # type: ignore[attr-defined]
         )
 
 
@@ -92,8 +101,8 @@ async def scan_tree_usage(self, path: Path) -> TreeUsage:
         raw_reports = await loop.run_in_executor(
             None,
             lambda: (
-                os.getxattr(path, "ceph.dir.rentries"),  # type: ignore[attr-defined]
-                os.getxattr(path, "ceph.dir.rbytes"),  # type: ignore[attr-defined]
+                os.getxattr(path, "ceph.dir.rentries"), # type: ignore[attr-defined]
+                os.getxattr(path, "ceph.dir.rbytes"), # type: ignore[attr-defined]
             ),
         )
         file_count = int(raw_reports[0].strip().decode())
@@ -104,7 +113,7 @@ async def scan_tree_size(self, path: Path) -> BinarySize:
         loop = asyncio.get_running_loop()
         raw_report = await loop.run_in_executor(
             None,
-            lambda: os.getxattr(path, "ceph.dir.rbytes"),  # type: ignore[attr-defined]
+            lambda: os.getxattr(path, "ceph.dir.rbytes"), # type: ignore[attr-defined]
         )
         return BinarySize(raw_report.strip().decode())
 

src/ai/backend/storage/ddn/__init__.py → src/ai/backend/storage/volumes/ddn/__init__.py

@@ -9,10 +9,10 @@
 from ai.backend.common.etcd import AsyncEtcd
 from ai.backend.common.types import QuotaScopeID
 from ai.backend.storage.exception import QuotaScopeAlreadyExists, QuotaScopeNotFoundError
+from ai.backend.storage.subproc import run
+from ai.backend.storage.types import Optional, QuotaConfig, QuotaUsage
 
 from ..abc import CAP_QUOTA, CAP_VFOLDER, AbstractQuotaModel
-from ..subproc import run
-from ..types import Optional, QuotaConfig, QuotaUsage
 from ..vfs import BaseQuotaModel, BaseVolume
 
 FIRST_PROJECT_ID: Final = 100
@@ -62,7 +62,8 @@ async def _read_main_project_id(self) -> int:
         return val
 
     async def _set_quota_by_project(self, pid: int, path: Path, options: QuotaConfig) -> None:
-        quota_limit = _byte_to_kilobyte(options.limit_bytes)  # default unit for DDN quota is KB
+        # default unit for DDN quota is KB
+        quota_limit = _byte_to_kilobyte(options.limit_bytes)
         try:
             await run([
                 b"sudo",

src/ai/backend/storage/dellemc/__init__.py → src/ai/backend/storage/volumes/dellemc/__init__.py

@@ -9,10 +9,10 @@
 from ai.backend.common.etcd import AsyncEtcd
 from ai.backend.common.events import EventDispatcher, EventProducer
 from ai.backend.common.types import HardwareMetadata, QuotaScopeID
+from ai.backend.storage.exception import NotEmptyError
+from ai.backend.storage.types import CapacityUsage, FSPerfMetric, QuotaConfig, QuotaUsage
 
 from ..abc import CAP_FAST_FS_SIZE, CAP_METRIC, CAP_QUOTA, CAP_VFOLDER, AbstractQuotaModel
-from ..exception import NotEmptyError
-from ..types import CapacityUsage, FSPerfMetric, QuotaConfig, QuotaUsage
 from ..vfs import BaseQuotaModel, BaseVolume
 from .config import config_iv
 from .exceptions import DellNoMetricError
@@ -197,7 +197,8 @@ async def get_performance_metric(self) -> FSPerfMetric:
         workload = await self.get_workload_stats()
         return FSPerfMetric(
             iops_read=protocol_stats["disk"]["iops"] or 0,
-            iops_write=0,  # Dell does not support IOPS Read/Write, They support only IOPS.
+            # Dell does not support IOPS Read/Write, They support only IOPS.
+            iops_write=0,
             io_bytes_read=protocol_stats["onefs"]["out"] or 0,
             io_bytes_write=protocol_stats["onefs"]["in"] or 0,
             io_usec_read=workload["latency_write"] or 0,

src/ai/backend/storage/gpfs/__init__.py → src/ai/backend/storage/volumes/gpfs/__init__.py

@@ -7,6 +7,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.types import CapacityUsage, FSPerfMetric
 
 from ..abc import (
     CAP_FAST_FS_SIZE,
@@ -18,7 +19,6 @@
     QuotaConfig,
     QuotaUsage,
 )
-from ..types import CapacityUsage, FSPerfMetric
 from ..vfs import BaseFSOpModel, BaseQuotaModel, BaseVolume
 from .exceptions import GPFSNoMetricError
 from .gpfs_client import GPFSAPIClient

src/ai/backend/storage/gpfs/gpfs_client.py → src/ai/backend/storage/volumes/gpfs/gpfs_client.py

@@ -18,8 +18,8 @@
 
 from ai.backend.common.types import BinarySize
 from ai.backend.logging import BraceStyleAdapter
+from ai.backend.storage.exception import ExternalError
 
-from ..exception import ExternalError
 from .exceptions import (
     GPFSAPIError,
     GPFSInvalidBodyError,

src/ai/backend/storage/netapp/__init__.py → src/ai/backend/storage/volumes/netapp/__init__.py

@@ -29,24 +29,14 @@
 
 from ai.backend.common.types import BinarySize, HardwareMetadata, QuotaScopeID
 from ai.backend.logging import BraceStyleAdapter
-
-from ..abc import (
-    CAP_FAST_FS_SIZE,
-    CAP_FAST_SIZE,
-    CAP_METRIC,
-    CAP_QUOTA,
-    CAP_VFOLDER,
-    AbstractFSOpModel,
-    AbstractQuotaModel,
-)
-from ..exception import (
+from ai.backend.storage.exception import (
     ExecutionError,
     InvalidQuotaScopeError,
     NotEmptyError,
     QuotaScopeNotFoundError,
 )
-from ..subproc import spawn_and_watch
-from ..types import (
+from ai.backend.storage.subproc import spawn_and_watch
+from ai.backend.storage.types import (
     SENTINEL,
     CapacityUsage,
     DirEntry,
@@ -58,7 +48,17 @@
     Stat,
     TreeUsage,
 )
-from ..utils import fstime2datetime
+from ai.backend.storage.utils import fstime2datetime
+
+from ..abc import (
+    CAP_FAST_FS_SIZE,
+    CAP_FAST_SIZE,
+    CAP_METRIC,
+    CAP_QUOTA,
+    CAP_VFOLDER,
+    AbstractFSOpModel,
+    AbstractQuotaModel,
+)
 from ..vfs import BaseFSOpModel, BaseQuotaModel, BaseVolume
 from .netappclient import JobResponseCode, NetAppClient, StorageID, VolumeID
 

src/ai/backend/storage/netapp/netappclient.py → src/ai/backend/storage/volumes/netapp/netappclient.py

@@ -75,8 +75,8 @@
 
 import aiohttp
 
-from ..exception import ExternalError
-from ..types import QuotaConfig, QuotaUsage
+from ai.backend.storage.exception import ExternalError
+from ai.backend.storage.types import QuotaConfig, QuotaUsage
 
 StorageID: TypeAlias = uuid.UUID
 VolumeID: TypeAlias = uuid.UUID

src/ai/backend/storage/purestorage/__init__.py → src/ai/backend/storage/volumes/purestorage/__init__.py

@@ -8,9 +8,15 @@
 
 from ai.backend.common.types import HardwareMetadata
 from ai.backend.logging import BraceStyleAdapter
-
-from ..abc import CAP_FAST_FS_SIZE, CAP_FAST_SCAN, CAP_METRIC, CAP_VFOLDER, AbstractFSOpModel
-from ..types import CapacityUsage, FSPerfMetric
+from ai.backend.storage.types import CapacityUsage, FSPerfMetric
+
+from ..abc import (
+    CAP_FAST_FS_SIZE,
+    CAP_FAST_SCAN,
+    CAP_METRIC,
+    CAP_VFOLDER,
+    AbstractFSOpModel,
+)
 from ..vfs import BaseVolume
 from .purity import PurityClient
 from .rapidfiles import RapidFileToolsFSOpModel

src/ai/backend/storage/purestorage/rapidfiles.py → src/ai/backend/storage/volumes/purestorage/rapidfiles.py

@@ -6,10 +6,10 @@
 from typing import AsyncIterator
 
 from ai.backend.common.types import BinarySize
+from ai.backend.storage.subproc import run
+from ai.backend.storage.types import DirEntry, DirEntryType, Stat, TreeUsage
+from ai.backend.storage.utils import fstime2datetime
 
-from ..subproc import run
-from ..types import DirEntry, DirEntryType, Stat, TreeUsage
-from ..utils import fstime2datetime
 from ..vfs import BaseFSOpModel
 
 

src/ai/backend/storage/purestorage/rapidfiles_v2.py → src/ai/backend/storage/volumes/purestorage/rapidfiles_v2.py

@@ -5,10 +5,10 @@
 from subprocess import CalledProcessError
 from typing import AsyncIterator
 
+from ai.backend.storage.subproc import run
+from ai.backend.storage.types import DirEntry, DirEntryType, Stat, TreeUsage
 from ai.backend.storage.utils import fstime2datetime
 
-from ..subproc import run
-from ..types import DirEntry, DirEntryType, Stat, TreeUsage
 from .rapidfiles import RapidFileToolsFSOpModel