feat: Add Privileged Storage Worker

Author: fregataa

changes/5779.feature.md

@@ -0,0 +1 @@
+Add Privileged Storage Worker

src/ai/backend/storage/BUILD

@@ -36,6 +36,7 @@ python_distribution(
         "backendai_cli_v10": {
             "storage": "ai.backend.storage.cli.__main__:main",
             "storage.start-server": "ai.backend.storage.server:main",
+            "storage.start-privileged-worker": "ai.backend.storage.privileged.server:main",
         },
     },
     generate_setup=True,

src/ai/backend/storage/privileged/__init__.py

@@ -0,0 +1,17 @@
+from ai.backend.common.bgtask.task.registry import BackgroundTaskHandlerRegistry
+from ai.backend.common.events.dispatcher import EventProducer
+
+from ...bgtask.tasks.delete import VFolderDeleteTaskHandler
+from ...volumes.pool import VolumePool
+
+
+class BgtaskHandlerRegistryCreator:
+    def __init__(self, volume_pool: VolumePool, event_producer: EventProducer) -> None:
+        self._volume_pool = volume_pool
+        self._event_producer = event_producer
+
+    def create(self) -> BackgroundTaskHandlerRegistry:
+        registry = BackgroundTaskHandlerRegistry()
+        registry.register(VFolderDeleteTaskHandler(self._volume_pool, self._event_producer))
+
+        return registry

src/ai/backend/storage/privileged/bgtask/__init__.py

@@ -0,0 +1,197 @@
+import asyncio
+from dataclasses import dataclass
+from typing import override
+
+from ai.backend.common.stage.types import (
+    ArgsSpecGenerator,
+    Provisioner,
+    ProvisionStage,
+)
+
+from ...bgtask.tags import ROOT_PRIVILEGED_TAG
+from ..config import StorageProxyPrivilegedWorkerConfig
+from .stage.bgtask import (
+    BgtaskManagerProvisioner,
+    BgtaskManagerResult,
+    BgtaskManagerSpec,
+    BgtaskManagerSpecGenerator,
+    BgtaskManagerStage,
+)
+from .stage.etcd import (
+    EtcdProvisioner,
+    EtcdResult,
+    EtcdSpec,
+    EtcdSpecGenerator,
+    EtcdStage,
+)
+from .stage.event_dispatcher import (
+    EventDispatcherProvisioner,
+    EventDispatcherResult,
+    EventDispatcherSpec,
+    EventDispatcherSpecGenerator,
+    EventDispatcherStage,
+)
+from .stage.logger import (
+    LoggerProvisioner,
+    LoggerResult,
+    LoggerSpec,
+    LoggerSpecGenerator,
+    LoggerStage,
+)
+from .stage.message_queue import (
+    MessageQueueProvisioner,
+    MessageQueueResult,
+    MessageQueueSpec,
+    MessageQueueSpecGenerator,
+    MessageQueueStage,
+)
+from .stage.monitor import (
+    MonitorProvisioner,
+    MonitorResult,
+    MonitorSpec,
+    MonitorSpecGenerator,
+    MonitorStage,
+)
+from .stage.redis_config import (
+    RedisConfigProvisioner,
+    RedisConfigResult,
+    RedisConfigSpec,
+    RedisConfigSpecGenerator,
+    RedisConfigStage,
+)
+from .stage.valkey_client import (
+    ValkeyClientProvisioner,
+    ValkeyClientResult,
+    ValkeyClientSpec,
+    ValkeyClientSpecGenerator,
+    ValkeyClientStage,
+)
+from .stage.volume_pool import (
+    VolumePoolProvisioner,
+    VolumePoolResult,
+    VolumePoolSpec,
+    VolumePoolSpecGenerator,
+    VolumePoolStage,
+)
+
+
+@dataclass
+class BootstrapSpec:
+    loop: asyncio.AbstractEventLoop
+    local_config: StorageProxyPrivilegedWorkerConfig
+    pidx: int
+
+
+class BootstrapSpecGenerator(ArgsSpecGenerator[BootstrapSpec]):
+    pass
+
+
+@dataclass
+class BootstrapResult:
+    logger: LoggerResult
+    monitor: MonitorResult
+    etcd: EtcdResult
+    redis_config: RedisConfigResult
+    message_queue: MessageQueueResult
+    event_dispatcher: EventDispatcherResult
+    valkey_client: ValkeyClientResult
+    volume_pool: VolumePoolResult
+    bgtask_manager: BgtaskManagerResult
+
+
+class BootstrapProvisioner(Provisioner[BootstrapSpec, BootstrapResult]):
+    @property
+    @override
+    def name(self) -> str:
+        return "storage-worker-bootstrap"
+
+    @override
+    async def setup(self, spec: BootstrapSpec) -> BootstrapResult:
+        local_config = spec.local_config
+        sub_logger_stage = LoggerStage(LoggerProvisioner())
+        sub_logger_spec = LoggerSpec(
+            is_master=False,
+            ipc_base_path=local_config.storage_proxy.ipc_base_path,
+            config=spec.local_config.logging,
+        )
+        await sub_logger_stage.setup(LoggerSpecGenerator(sub_logger_spec))
+        logger_result = await sub_logger_stage.wait_for_resource()
+
+        monitor_stage = MonitorStage(MonitorProvisioner())
+        monitor_spec = MonitorSpec(loop=spec.loop, pidx=spec.pidx, local_config=local_config)
+        await monitor_stage.setup(MonitorSpecGenerator(monitor_spec))
+        monitor_result = await monitor_stage.wait_for_resource()
+
+        etcd_stage = EtcdStage(EtcdProvisioner())
+        etcd_spec = EtcdSpec(local_config=local_config)
+        await etcd_stage.setup(EtcdSpecGenerator(etcd_spec))
+        etcd_result = await etcd_stage.wait_for_resource()
+
+        redis_config_stage = RedisConfigStage(RedisConfigProvisioner())
+        redis_config_spec = RedisConfigSpec(etcd=etcd_result.etcd)
+        await redis_config_stage.setup(RedisConfigSpecGenerator(redis_config_spec))
+        redis_config_result = await redis_config_stage.wait_for_resource()
+        redis_profile_target = redis_config_result.redis_config.to_redis_profile_target()
+
+        mq_stage = MessageQueueStage(MessageQueueProvisioner())
+        mq_spec = MessageQueueSpec(
+            local_config=local_config, redis_profile_target=redis_profile_target
+        )
+        await mq_stage.setup(MessageQueueSpecGenerator(mq_spec))
+        mq_result = await mq_stage.wait_for_resource()
+
+        event_dispatcher_stage = EventDispatcherStage(EventDispatcherProvisioner())
+        event_dispatcher_spec = EventDispatcherSpec(
+            message_queue=mq_result.message_queue,
+            log_events=local_config.debug.log_events,
+            event_observer=monitor_result.metric_registry.event,
+            source_id=None,
+        )
+        await event_dispatcher_stage.setup(EventDispatcherSpecGenerator(event_dispatcher_spec))
+        event_dispatcher_result = await event_dispatcher_stage.wait_for_resource()
+
+        valkey_client_stage = ValkeyClientStage(ValkeyClientProvisioner())
+        valkey_client_spec = ValkeyClientSpec(redis_profile_target)
+        await valkey_client_stage.setup(ValkeyClientSpecGenerator(valkey_client_spec))
+        valkey_client_result = await valkey_client_stage.wait_for_resource()
+
+        volume_pool_stage = VolumePoolStage(VolumePoolProvisioner())
+        volume_pool_spec = VolumePoolSpec(
+            local_config=local_config,
+            etcd=etcd_result.etcd,
+            event_dispatcher=event_dispatcher_result.event_dispatcher,
+            event_producer=event_dispatcher_result.event_producer,
+        )
+        await volume_pool_stage.setup(VolumePoolSpecGenerator(volume_pool_spec))
+        volume_pool_result = await volume_pool_stage.wait_for_resource()
+
+        bgtask_manager_stage = BgtaskManagerStage(BgtaskManagerProvisioner())
+        bgtask_manager_spec = BgtaskManagerSpec(
+            volume_pool=volume_pool_result.volume_pool,
+            valkey_client=valkey_client_result.bgtask_client,
+            event_producer=event_dispatcher_result.event_producer,
+            node_id=local_config.storage_proxy.node_id,
+            tags=[ROOT_PRIVILEGED_TAG],
+        )
+        await bgtask_manager_stage.setup(BgtaskManagerSpecGenerator(bgtask_manager_spec))
+        bgtask_manager_result = await bgtask_manager_stage.wait_for_resource()
+
+        return BootstrapResult(
+            logger=logger_result,
+            monitor=monitor_result,
+            etcd=etcd_result,
+            redis_config=redis_config_result,
+            message_queue=mq_result,
+            event_dispatcher=event_dispatcher_result,
+            valkey_client=valkey_client_result,
+            volume_pool=volume_pool_result,
+            bgtask_manager=bgtask_manager_result,
+        )
+
+    @override
+    async def teardown(self, resource: BootstrapResult) -> None:
+        pass
+
+
+class BootstrapStage(ProvisionStage[BootstrapSpec, BootstrapResult]):
+    pass

src/ai/backend/storage/privileged/bgtask/registry.py

@@ -0,0 +1,68 @@
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import override
+
+from ai.backend.common.bgtask.bgtask import BackgroundTaskManager
+from ai.backend.common.bgtask.task.registry import BackgroundTaskHandlerRegistry
+from ai.backend.common.clients.valkey_client.valkey_bgtask.client import ValkeyBgtaskClient
+from ai.backend.common.events.dispatcher import EventProducer
+from ai.backend.common.stage.types import (
+    ArgsSpecGenerator,
+    Provisioner,
+    ProvisionStage,
+)
+
+from ....bgtask.tasks.delete import VFolderDeleteTaskHandler
+from ....volumes.pool import VolumePool
+
+
+@dataclass
+class BgtaskManagerSpec:
+    volume_pool: VolumePool
+    valkey_client: ValkeyBgtaskClient
+    event_producer: EventProducer
+    node_id: str
+    tags: Iterable[str]
+
+
+class BgtaskManagerSpecGenerator(ArgsSpecGenerator[BgtaskManagerSpec]):
+    pass
+
+
+@dataclass
+class BgtaskManagerResult:
+    bgtask_mgr: BackgroundTaskManager
+
+
+class BgtaskManagerProvisioner(Provisioner[BgtaskManagerSpec, BgtaskManagerResult]):
+    @property
+    @override
+    def name(self) -> str:
+        return "storage-worker-bgtask"
+
+    @override
+    async def setup(self, spec: BgtaskManagerSpec) -> BgtaskManagerResult:
+        registry = self._register_bgtask_handlers(spec.volume_pool, spec.event_producer)
+        bgtask_mgr = BackgroundTaskManager(
+            event_producer=spec.event_producer,
+            task_registry=registry,
+            valkey_client=spec.valkey_client,
+            server_id=spec.node_id,
+            tags=spec.tags,
+        )
+        return BgtaskManagerResult(bgtask_mgr)
+
+    def _register_bgtask_handlers(
+        self, volume_pool: VolumePool, event_producer: EventProducer
+    ) -> BackgroundTaskHandlerRegistry:
+        registry = BackgroundTaskHandlerRegistry()
+        registry.register(VFolderDeleteTaskHandler(volume_pool, event_producer))
+        return registry
+
+    @override
+    async def teardown(self, resource: BgtaskManagerResult) -> None:
+        pass
+
+
+class BgtaskManagerStage(ProvisionStage[BgtaskManagerSpec, BgtaskManagerResult]):
+    pass

src/ai/backend/storage/privileged/bootstrap/__init__.py

@@ -0,0 +1,46 @@
+from dataclasses import dataclass
+from typing import override
+
+from ai.backend.common.etcd import AsyncEtcd
+from ai.backend.common.stage.types import (
+    ArgsSpecGenerator,
+    Provisioner,
+    ProvisionStage,
+)
+
+from ....config.loaders import make_etcd
+from ...config import StorageProxyPrivilegedWorkerConfig
+
+
+@dataclass
+class EtcdSpec:
+    local_config: StorageProxyPrivilegedWorkerConfig
+
+
+class EtcdSpecGenerator(ArgsSpecGenerator[EtcdSpec]):
+    pass
+
+
+@dataclass
+class EtcdResult:
+    etcd: AsyncEtcd
+
+
+class EtcdProvisioner(Provisioner[EtcdSpec, EtcdResult]):
+    @property
+    @override
+    def name(self) -> str:
+        return "storage-worker-etcd"
+
+    @override
+    async def setup(self, spec: EtcdSpec) -> EtcdResult:
+        etcd = make_etcd(spec.local_config)
+        return EtcdResult(etcd)
+
+    @override
+    async def teardown(self, resource: EtcdResult) -> None:
+        pass
+
+
+class EtcdStage(ProvisionStage[EtcdSpec, EtcdResult]):
+    pass