srt/mem_cache/storage/backend_factory.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to SGLang project

import importlib
import logging
from typing import TYPE_CHECKING, Any, Dict

from sglang.srt.mem_cache.hicache_storage import HiCacheStorage, HiCacheStorageConfig

if TYPE_CHECKING:
    pass

logger = logging.getLogger(__name__)


class StorageBackendFactory:
    """Factory for creating storage backend instances with support for dynamic loading."""

    _registry: Dict[str, Dict[str, Any]] = {}

    @staticmethod
    def _load_backend_class(
        module_path: str, class_name: str, backend_name: str
    ) -> type[HiCacheStorage]:
        """Load and validate a backend class from module path."""
        try:
            module = importlib.import_module(module_path)
            backend_class = getattr(module, class_name)
            if not issubclass(backend_class, HiCacheStorage):
                raise TypeError(
                    f"Backend class {class_name} must inherit from HiCacheStorage"
                )
            return backend_class
        except ImportError as e:
            raise ImportError(
                f"Failed to import backend '{backend_name}' from '{module_path}': {e}"
            ) from e
        except AttributeError as e:
            raise AttributeError(
                f"Class '{class_name}' not found in module '{module_path}': {e}"
            ) from e

    @classmethod
    def register_backend(cls, name: str, module_path: str, class_name: str) -> None:
        """Register a storage backend with lazy loading.

        Args:
            name: Backend identifier
            module_path: Python module path containing the backend class
            class_name: Name of the backend class
        """
        if name in cls._registry:
            logger.warning(f"Backend '{name}' is already registered, overwriting")

        def loader() -> type[HiCacheStorage]:
            """Lazy loader function to import the backend class."""
            return cls._load_backend_class(module_path, class_name, name)

        cls._registry[name] = {
            "loader": loader,
            "module_path": module_path,
            "class_name": class_name,
        }

    @classmethod
    def create_backend(
        cls,
        backend_name: str,
        storage_config: HiCacheStorageConfig,
        mem_pool_host: Any,
        **kwargs,
    ) -> HiCacheStorage:
        """Create a storage backend instance.
        Args:
            backend_name: Name of the backend to create
            storage_config: Storage configuration
            mem_pool_host: Memory pool host object
            **kwargs: Additional arguments passed to external backends
        Returns:
            Initialized storage backend instance
        Raises:
            ValueError: If backend is not registered and cannot be dynamically loaded
            ImportError: If backend module cannot be imported
            Exception: If backend initialization fails
        """
        # First check if backend is already registered
        if backend_name in cls._registry:
            registry_entry = cls._registry[backend_name]
            backend_class = registry_entry["loader"]()
            logger.info(
                f"Creating storage backend '{backend_name}' "
                f"({registry_entry['module_path']}.{registry_entry['class_name']})"
            )
            return cls._create_builtin_backend(
                backend_name, backend_class, storage_config, mem_pool_host
            )

        # Try to dynamically load backend from extra_config
        if backend_name == "dynamic" and storage_config.extra_config is not None:
            backend_config = storage_config.extra_config
            return cls._create_dynamic_backend(
                backend_config, storage_config, mem_pool_host, **kwargs
            )

        # Backend not found
        available_backends = list(cls._registry.keys())

        raise ValueError(
            f"Unknown storage backend '{backend_name}'. "
            f"Registered backends: {available_backends}. "
        )

    @classmethod
    def _create_dynamic_backend(
        cls,
        backend_config: Dict[str, Any],
        storage_config: HiCacheStorageConfig,
        mem_pool_host: Any,
        **kwargs,
    ) -> HiCacheStorage:
        """Create a backend dynamically from configuration."""
        required_fields = ["backend_name", "module_path", "class_name"]
        for field in required_fields:
            if field not in backend_config:
                raise ValueError(
                    f"Missing required field '{field}' in backend config for 'dynamic' backend"
                )

        backend_name = backend_config["backend_name"]
        module_path = backend_config["module_path"]
        class_name = backend_config["class_name"]

        try:
            # Import the backend class
            backend_class = cls._load_backend_class(
                module_path, class_name, backend_name
            )

            logger.info(
                f"Creating dynamic storage backend '{backend_name}' "
                f"({module_path}.{class_name})"
            )

            # Create the backend instance with storage_config
            return backend_class(storage_config, kwargs)
        except Exception as e:
            logger.error(
                f"Failed to create dynamic storage backend '{backend_name}': {e}"
            )
            raise

    @classmethod
    def _create_builtin_backend(
        cls,
        backend_name: str,
        backend_class: type[HiCacheStorage],
        storage_config: HiCacheStorageConfig,
        mem_pool_host: Any,
    ) -> HiCacheStorage:
        """Create built-in backend with original initialization logic."""
        if backend_name == "file":
            return backend_class(storage_config)
        elif backend_name == "nixl":
            return backend_class(storage_config)
        elif backend_name == "mooncake":
            backend = backend_class(storage_config)
            return backend
        elif backend_name == "aibrix":
            backend = backend_class(storage_config, mem_pool_host)
            return backend
        elif backend_name == "hf3fs":
            # Calculate bytes_per_page based on memory pool layout
            if mem_pool_host.layout in ["page_first", "page_first_direct"]:
                bytes_per_page = (
                    mem_pool_host.get_ksize_per_token() * mem_pool_host.page_size
                )
            elif mem_pool_host.layout == "layer_first":
                bytes_per_page = (
                    mem_pool_host.get_size_per_token() * mem_pool_host.page_size
                )

            dtype = mem_pool_host.dtype
            return backend_class.from_env_config(bytes_per_page, dtype, storage_config)
        elif backend_name == "eic":
            return backend_class(storage_config, mem_pool_host)
        else:
            raise ValueError(f"Unknown built-in backend: {backend_name}")


# Register built-in storage backends
StorageBackendFactory.register_backend(
    "file", "sglang.srt.mem_cache.hicache_storage", "HiCacheFile"
)

StorageBackendFactory.register_backend(
    "nixl",
    "sglang.srt.mem_cache.storage.nixl.hicache_nixl",
    "HiCacheNixl",
)

StorageBackendFactory.register_backend(
    "mooncake",
    "sglang.srt.mem_cache.storage.mooncake_store.mooncake_store",
    "MooncakeStore",
)

StorageBackendFactory.register_backend(
    "hf3fs",
    "sglang.srt.mem_cache.storage.hf3fs.storage_hf3fs",
    "HiCacheHF3FS",
)

StorageBackendFactory.register_backend(
    "aibrix",
    "sglang.srt.mem_cache.storage.aibrix_kvcache.aibrix_kvcache_storage",
    "AibrixKVCacheStorage",
)

StorageBackendFactory.register_backend(
    "eic",
    "sglang.srt.mem_cache.storage.eic.eic_storage",
    "EICStorage",
)