feat(BA-939): Add dependency verification for storage proxy (#6760)

Co-authored-by: Claude <noreply@anthropic.com>

Author: HyeockJinKim

changes/6760.feature.md

@@ -0,0 +1 @@
+Add dependency verification for storage proxy

src/ai/backend/common/message_queue/redis_queue/subscriber.py

@@ -116,7 +116,7 @@ async def _read_broadcast_messages_loop(self) -> None:
         """
         Background task to read broadcast messages from subscribed channels.
         """
-        log.info("Starting read broadcast messages loop for channels {}", self._channels)
+        log.debug("Starting read broadcast messages loop for channels {}", self._channels)
 
         while not self._closed:
             try:

src/ai/backend/storage/cli/__main__.py

@@ -55,3 +55,8 @@ def main(
 @main.group(cls=LazyGroup, import_name="ai.backend.storage.cli.config:cli")
 def config() -> None:
     """Command set for configuration management."""
+
+
+@main.group(cls=LazyGroup, import_name="ai.backend.storage.cli.dependencies:cli")
+def dependencies() -> None:
+    """Command set for dependency verification."""

src/ai/backend/storage/cli/dependencies.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import asyncio
+import sys
+from pathlib import Path
+
+import click
+
+from ai.backend.common.dependencies.stacks.visualizing import VisualizingDependencyStack
+
+from ..dependencies.composer import StorageDependencyComposer
+from .context import CLIContext
+
+
+@click.group()
+def cli() -> None:
+    """Dependency verification commands for storage proxy setup."""
+    pass
+
+
+@cli.command(name="verify")
+@click.option(
+    "--no-timestamps",
+    is_flag=True,
+    help="Hide timestamps in output",
+)
+@click.option(
+    "--log-level",
+    type=click.Choice(["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], case_sensitive=False),
+    default="WARNING",
+    help="Set the logging level (default: WARNING)",
+)
+@click.pass_obj
+def verify(
+    cli_ctx: CLIContext,
+    no_timestamps: bool,
+    log_level: str,
+) -> None:
+    """
+    Verify all storage proxy dependencies can be initialized successfully.
+
+    This command attempts to initialize all storage proxy dependencies (infrastructure,
+    components) and reports any failures during the setup process.
+
+    The output shows dependency initialization with status indicators:
+    - ▶ Starting (for composers)
+    - ✓ Completed (for dependencies)
+    - ✗ Failed (for dependencies)
+
+    Examples:
+
+        # Verify all dependencies with default WARNING log level
+        $ backend.ai storage dependencies verify
+
+        # Hide timestamps in output
+        $ backend.ai storage dependencies verify --no-timestamps
+
+        # Set log level to DEBUG to see detailed initialization logs
+        $ backend.ai storage dependencies verify --log-level DEBUG
+    """
+    import logging as std_logging
+
+    # Set logging level BEFORE any dependency initialization
+    std_logging.basicConfig(
+        level=getattr(std_logging, log_level.upper()),
+        format="%(levelname)s:%(name)s:%(message)s",
+        force=True,
+    )
+
+    config_path = cli_ctx.config_path
+
+    async def _verify_dependencies() -> bool:
+        """Verify all storage proxy dependencies can be initialized."""
+        print("\n" + "=" * 60)
+        print("Storage Proxy Dependency Verification")
+        print("=" * 60)
+        print(f"Config: {config_path or 'default locations'}")
+        print(f"Log Level: {log_level.upper()}")
+        print("=" * 60 + "\n")
+
+        # Create visualizing stack
+        stack = VisualizingDependencyStack(
+            output=sys.stdout,
+            show_timestamps=not no_timestamps,
+        )
+
+        try:
+            async with stack:
+                from ..dependencies.composer import DependencyInput
+
+                dependency_input = DependencyInput(
+                    config_path=config_path or Path("storage-proxy.toml"),
+                )
+                composer = StorageDependencyComposer()
+
+                # Initialize all dependencies through the composer
+                await stack.enter_composer(composer, dependency_input)
+
+        except Exception as e:
+            print(f"\n❌ Dependency initialization failed: {e}\n")
+
+        # Print summary
+        print()
+        stack.print_summary()
+
+        return not stack.has_failures()
+
+    # Run the verification
+    success = asyncio.run(_verify_dependencies())
+
+    if success:
+        print("All dependencies verified successfully!\n")
+        sys.exit(0)
+    else:
+        print("Some dependencies failed. Check the output above for details.\n")
+        sys.exit(1)

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

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+from .composer import DependencyResources, StorageDependencyComposer
+
+__all__ = [
+    "DependencyResources",
+    "StorageDependencyComposer",
+]

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

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from .composer import BootstrapComposer, BootstrapInput, BootstrapResources
+from .config import ConfigProvider, ConfigProviderInput
+
+__all__ = [
+    "BootstrapComposer",
+    "BootstrapInput",
+    "BootstrapResources",
+    "ConfigProvider",
+    "ConfigProviderInput",
+]

src/ai/backend/storage/dependencies/bootstrap/composer.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from pathlib import Path
+
+from ai.backend.common.dependencies import DependencyComposer, DependencyStack
+from ai.backend.storage.config.unified import StorageProxyUnifiedConfig
+
+from .config import ConfigProvider, ConfigProviderInput
+
+
+@dataclass
+class BootstrapInput:
+    """Input required for bootstrap stage."""
+
+    config_path: Path
+
+
+@dataclass
+class BootstrapResources:
+    """Container for bootstrap stage resources."""
+
+    config: StorageProxyUnifiedConfig
+
+
+class BootstrapComposer(DependencyComposer[BootstrapInput, BootstrapResources]):
+    """Composes bootstrap dependencies."""
+
+    @property
+    def stage_name(self) -> str:
+        return "bootstrap"
+
+    @asynccontextmanager
+    async def compose(
+        self,
+        stack: DependencyStack,
+        setup_input: BootstrapInput,
+    ) -> AsyncIterator[BootstrapResources]:
+        """Compose bootstrap dependencies."""
+        # Load config
+        config_provider = ConfigProvider()
+        config = await stack.enter_dependency(
+            config_provider,
+            ConfigProviderInput(config_path=setup_input.config_path),
+        )
+
+        yield BootstrapResources(config=config)

src/ai/backend/storage/dependencies/bootstrap/config.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from pathlib import Path
+
+from ai.backend.common.dependencies import DependencyProvider
+from ai.backend.storage.config.loaders import load_local_config
+from ai.backend.storage.config.unified import StorageProxyUnifiedConfig
+
+
+@dataclass
+class ConfigProviderInput:
+    """Input for config provider."""
+
+    config_path: Path
+
+
+class ConfigProvider(DependencyProvider[ConfigProviderInput, StorageProxyUnifiedConfig]):
+    """Provider for storage proxy configuration."""
+
+    @property
+    def stage_name(self) -> str:
+        return "config"
+
+    @asynccontextmanager
+    async def provide(
+        self, setup_input: ConfigProviderInput
+    ) -> AsyncIterator[StorageProxyUnifiedConfig]:
+        """Load and provide storage proxy configuration."""
+        config = load_local_config(setup_input.config_path)
+        yield config

src/ai/backend/storage/dependencies/composer.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from pathlib import Path
+
+from ai.backend.common.dependencies import DependencyComposer, DependencyStack
+
+from .bootstrap import BootstrapComposer, BootstrapInput, BootstrapResources
+from .infrastructure.composer import (
+    InfrastructureComposer,
+    InfrastructureComposerInput,
+    InfrastructureResources,
+)
+from .storage.composer import StorageComposer, StorageComposerInput, StorageResources
+
+
+@dataclass
+class DependencyInput:
+    """Input required for complete dependency setup."""
+
+    config_path: Path
+
+
+@dataclass
+class DependencyResources:
+    """All dependency resources for storage proxy."""
+
+    bootstrap: BootstrapResources
+    infrastructure: InfrastructureResources
+    storage: StorageResources
+
+
+class StorageDependencyComposer(DependencyComposer[DependencyInput, DependencyResources]):
+    """Main composer for all storage proxy dependencies."""
+
+    @property
+    def stage_name(self) -> str:
+        return "storage-proxy"
+
+    @asynccontextmanager
+    async def compose(
+        self,
+        stack: DependencyStack,
+        setup_input: DependencyInput,
+    ) -> AsyncIterator[DependencyResources]:
+        """Compose all storage proxy dependencies."""
+        # Stage 1: Bootstrap (config loading)
+        bootstrap = await stack.enter_composer(
+            BootstrapComposer(),
+            BootstrapInput(config_path=setup_input.config_path),
+        )
+
+        # Stage 2: Infrastructure (etcd, redis)
+        infrastructure = await stack.enter_composer(
+            InfrastructureComposer(),
+            InfrastructureComposerInput(local_config=bootstrap.config),
+        )
+
+        # Stage 3: Storage (storage-pool)
+        storage = await stack.enter_composer(
+            StorageComposer(),
+            StorageComposerInput(local_config=bootstrap.config),
+        )
+
+        yield DependencyResources(
+            bootstrap=bootstrap,
+            infrastructure=infrastructure,
+            storage=storage,
+        )

src/ai/backend/storage/dependencies/infrastructure/__init__.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from .composer import InfrastructureComposer, InfrastructureComposerInput, InfrastructureResources
+from .etcd import EtcdProvider
+from .redis import RedisProvider, StorageProxyValkeyClients
+
+__all__ = [
+    "EtcdProvider",
+    "InfrastructureComposer",
+    "InfrastructureComposerInput",
+    "InfrastructureResources",
+    "RedisProvider",
+    "StorageProxyValkeyClients",
+]