kaggle_gpu_server/engine/plugin_registry.py

"""
Plugin Registry — Auto-discovery and management of model plugins.

Scans the `plugins/` directory for Python modules that subclass ModelPlugin.
Plugins are registered by their `capability` and can be retrieved by name or capability.

Usage:
    registry = PluginRegistry(vram_manager)
    registry.discover_plugins()
    
    face_det = registry.get("face_detection")
    face_det.run({"image": pil_img})
    
    all_plugins = registry.list_plugins()
"""

import importlib
import pkgutil
import os
from typing import Optional
from .plugin_base import ModelPlugin, PluginCapability, PluginInfo
from .vram_manager import VRAMManager


class PluginRegistry:
    """
    Central registry for model plugins.
    
    Auto-discovers plugins from the `plugins/` package and manages them
    through the VRAMManager for VRAM-aware loading/unloading.
    """

    def __init__(self, vram_manager: VRAMManager, config: Optional[dict] = None):
        self.vram_manager = vram_manager
        self.config = config or {}
        
        # capability -> list of plugin instances (ordered by priority)
        self._by_capability: dict[PluginCapability, list[ModelPlugin]] = {}
        # name -> plugin instance
        self._by_name: dict[str, ModelPlugin] = {}
        # Disabled plugin names from config
        self._disabled: set[str] = set()
        
        # Parse disabled plugins from config
        plugins_config = self.config.get("plugins", {})
        for cap_name, cap_config in plugins_config.items():
            if isinstance(cap_config, dict) and not cap_config.get("enabled", True):
                self._disabled.add(cap_name)
                
        # Parse engine mode ("all", "generation", "editing")
        engine_config = self.config.get("engine", {})
        self.mode = os.environ.get("ENGINE_MODE", engine_config.get("mode", "all")).lower()
        print(f"⚙️ PluginRegistry initialized in '{self.mode}' mode")


    def discover_plugins(self, package_name: str = "kaggle_gpu_server.plugins") -> int:
        """
        Auto-discover and register all ModelPlugin subclasses from the plugins package.
        
        Returns the number of plugins discovered.
        """
        count = 0
        try:
            package = importlib.import_module(package_name)
        except ImportError as e:
            print(f"❌ Cannot import plugins package '{package_name}': {e}")
            # Try relative import
            try:
                plugins_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "plugins")
                if os.path.isdir(plugins_dir):
                    import sys
                    parent = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
                    if parent not in sys.path:
                        sys.path.insert(0, parent)
                    package = importlib.import_module(package_name)
                else:
                    return 0
            except ImportError:
                return 0

        package_path = os.path.dirname(package.__file__)
        
        for _, module_name, _ in pkgutil.iter_modules([package_path]):
            if module_name.startswith("_"):
                continue
            try:
                full_name = f"{package_name}.{module_name}"
                module = importlib.import_module(full_name)
                
                # Find all ModelPlugin subclasses in the module
                for attr_name in dir(module):
                    attr = getattr(module, attr_name)
                    if (
                        isinstance(attr, type)
                        and issubclass(attr, ModelPlugin)
                        and attr is not ModelPlugin
                        and attr.name  # Must have a name set
                    ):
                        plugin = attr()
                        self.register(plugin)
                        count += 1
                        
            except Exception as e:
                print(f"⚠️ Failed to load plugin module '{module_name}': {e}")
                
        print(f"📦 Discovered {count} plugins from {package_name}")
        return count

    def register(self, plugin: ModelPlugin) -> None:
        """Register a plugin instance."""
        if plugin.name in self._disabled:
            print(f"  Skip disabled plugin: {plugin.name}")
            return
            
        # Filter plugins depending on mode to optimize VRAM
        if self.mode == "generation":
            allowed_caps = {
                PluginCapability.IMAGE_GENERATION,
                PluginCapability.CAPTIONING,
                PluginCapability.NSFW_DETECTION,
                PluginCapability.UPSCALING
            }
            if plugin.capability not in allowed_caps:
                print(f"  Skipping non-generation plugin: {plugin.name}")
                return
        elif self.mode == "editing":
            disallowed_caps = {
                PluginCapability.IMAGE_GENERATION
            }
            if plugin.capability in disallowed_caps:
                print(f"  Skipping generation plugin: {plugin.name}")
                return

        # Register by name
        self._by_name[plugin.name] = plugin
        
        # Register by capability
        cap = plugin.capability
        if cap not in self._by_capability:
            self._by_capability[cap] = []
        self._by_capability[cap].append(plugin)
        
        print(f"  ✅ Registered: {plugin}")


    def get(self, capability: str, preferred_name: Optional[str] = None) -> Optional[ModelPlugin]:
        """
        Get a plugin by capability name.
        
        Args:
            capability: The capability name (e.g., "face_detection")
            preferred_name: Optional specific plugin name to prefer
            
        Returns:
            The plugin instance, or None if no plugin has this capability.
        """
        # Try as PluginCapability enum
        try:
            cap = PluginCapability(capability)
        except ValueError:
            # Try direct name lookup
            return self._by_name.get(capability)
        
        plugins = self._by_capability.get(cap, [])
        if not plugins:
            return None
            
        # If a preferred name is given, try that first
        if preferred_name:
            for p in plugins:
                if p.name == preferred_name:
                    return p
        
        # Check config for default
        plugins_config = self.config.get("plugins", {})
        cap_config = plugins_config.get(capability, {})
        if isinstance(cap_config, dict):
            default_name = cap_config.get("default")
            if default_name:
                for p in plugins:
                    if p.name == default_name:
                        return p
        
        # Return first registered plugin for this capability
        return plugins[0]

    def get_by_name(self, name: str) -> Optional[ModelPlugin]:
        """Get a plugin by its exact name."""
        return self._by_name.get(name)

    def ensure_loaded(self, capability: str, preferred_name: Optional[str] = None) -> ModelPlugin:
        """
        Get a plugin and ensure it's loaded (via VRAMManager).
        Raises RuntimeError if plugin not found or load fails.
        """
        plugin = self.get(capability, preferred_name)
        if plugin is None:
            raise RuntimeError(f"No plugin registered for capability '{capability}'")
        
        if not plugin.is_loaded:
            if not self.vram_manager.request_load(plugin):
                raise RuntimeError(f"Failed to load plugin '{plugin.name}'")
        
        return plugin

    def list_plugins(self) -> list[PluginInfo]:
        """List all registered plugins with their status."""
        return [p.get_info() for p in self._by_name.values()]

    def list_capabilities(self) -> dict[str, list[str]]:
        """List all capabilities and which plugins provide them."""
        return {
            cap.value: [p.name for p in plugins]
            for cap, plugins in self._by_capability.items()
        }

    def get_loaded_plugins(self) -> list[str]:
        """Get names of currently loaded plugins."""
        return [name for name, p in self._by_name.items() if p.is_loaded]

    def __repr__(self):
        total = len(self._by_name)
        loaded = len(self.get_loaded_plugins())
        caps = len(self._by_capability)
        return f"PluginRegistry({total} plugins, {loaded} loaded, {caps} capabilities)"