| |
| import importlib |
| import importlib.util |
| import inspect |
| import pkgutil |
| import sys |
| import warnings |
| from collections import defaultdict |
| from dataclasses import dataclass, field |
| from timeit import default_timer as timer |
| from typing import Any, Dict, List, Optional, Tuple, Type |
|
|
| from omegaconf import DictConfig |
|
|
| from hydra._internal.sources_registry import SourcesRegistry |
| from hydra.core.singleton import Singleton |
| from hydra.plugins.completion_plugin import CompletionPlugin |
| from hydra.plugins.config_source import ConfigSource |
| from hydra.plugins.launcher import Launcher |
| from hydra.plugins.plugin import Plugin |
| from hydra.plugins.search_path_plugin import SearchPathPlugin |
| from hydra.plugins.sweeper import Sweeper |
| from hydra.types import HydraContext, TaskFunction |
| from hydra.utils import instantiate |
|
|
| PLUGIN_TYPES: List[Type[Plugin]] = [ |
| Plugin, |
| ConfigSource, |
| CompletionPlugin, |
| Launcher, |
| Sweeper, |
| SearchPathPlugin, |
| ] |
|
|
|
|
| @dataclass |
| class ScanStats: |
| total_time: float = 0 |
| total_modules_import_time: float = 0 |
| modules_import_time: Dict[str, float] = field(default_factory=dict) |
|
|
|
|
| class Plugins(metaclass=Singleton): |
| @staticmethod |
| def instance(*args: Any, **kwargs: Any) -> "Plugins": |
| ret = Singleton.instance(Plugins, *args, **kwargs) |
| assert isinstance(ret, Plugins) |
| return ret |
|
|
| def __init__(self) -> None: |
| self.plugin_type_to_subclass_list: Dict[Type[Plugin], List[Type[Plugin]]] = {} |
| self.class_name_to_class: Dict[str, Type[Plugin]] = {} |
| self.stats: Optional[ScanStats] = None |
| self._initialize() |
|
|
| def _initialize(self) -> None: |
| top_level: List[Any] = [] |
| core_plugins = importlib.import_module("hydra._internal.core_plugins") |
| top_level.append(core_plugins) |
|
|
| try: |
| hydra_plugins = importlib.import_module("hydra_plugins") |
| top_level.append(hydra_plugins) |
| except ImportError: |
| |
| pass |
|
|
| self.plugin_type_to_subclass_list = defaultdict(list) |
| self.class_name_to_class = {} |
|
|
| scanned_plugins, self.stats = self._scan_all_plugins(modules=top_level) |
| for clazz in scanned_plugins: |
| self._register(clazz) |
|
|
| def register(self, clazz: Type[Plugin]) -> None: |
| """ |
| Call Plugins.instance().register(MyPlugin) to manually register a plugin class. |
| """ |
| if not _is_concrete_plugin_type(clazz): |
| raise ValueError("Not a valid Hydra Plugin") |
| self._register(clazz) |
|
|
| def _register(self, clazz: Type[Plugin]) -> None: |
| assert _is_concrete_plugin_type(clazz) |
| for plugin_type in PLUGIN_TYPES: |
| if issubclass(clazz, plugin_type): |
| if clazz not in self.plugin_type_to_subclass_list[plugin_type]: |
| self.plugin_type_to_subclass_list[plugin_type].append(clazz) |
| name = f"{clazz.__module__}.{clazz.__name__}" |
| self.class_name_to_class[name] = clazz |
| if issubclass(clazz, ConfigSource): |
| SourcesRegistry.instance().register(clazz) |
|
|
| def _instantiate(self, config: DictConfig) -> Plugin: |
| from hydra._internal import utils as internal_utils |
|
|
| classname = internal_utils._get_cls_name(config, pop=False) |
| try: |
| if classname is None: |
| raise ImportError("class not configured") |
|
|
| if not self.is_in_toplevel_plugins_module(classname): |
| |
| |
| raise RuntimeError( |
| f"Invalid plugin '{classname}': not the hydra_plugins package" |
| ) |
|
|
| if classname not in self.class_name_to_class.keys(): |
| raise RuntimeError(f"Unknown plugin class : '{classname}'") |
| clazz = self.class_name_to_class[classname] |
| plugin = instantiate(config=config, _target_=clazz) |
| assert isinstance(plugin, Plugin) |
|
|
| except ImportError as e: |
| raise ImportError( |
| f"Could not instantiate plugin {classname} : {str(e)}\n\n\tIS THE PLUGIN INSTALLED?\n\n" |
| ) |
|
|
| return plugin |
|
|
| @staticmethod |
| def is_in_toplevel_plugins_module(clazz: str) -> bool: |
| return clazz.startswith("hydra_plugins.") or clazz.startswith( |
| "hydra._internal.core_plugins." |
| ) |
|
|
| def instantiate_sweeper( |
| self, |
| *, |
| hydra_context: HydraContext, |
| task_function: TaskFunction, |
| config: DictConfig, |
| ) -> Sweeper: |
| Plugins.check_usage(self) |
| if config.hydra.sweeper is None: |
| raise RuntimeError("Hydra sweeper is not configured") |
| sweeper = self._instantiate(config.hydra.sweeper) |
| assert isinstance(sweeper, Sweeper) |
| sweeper.setup( |
| hydra_context=hydra_context, task_function=task_function, config=config |
| ) |
| return sweeper |
|
|
| def instantiate_launcher( |
| self, |
| hydra_context: HydraContext, |
| task_function: TaskFunction, |
| config: DictConfig, |
| ) -> Launcher: |
| Plugins.check_usage(self) |
| if config.hydra.launcher is None: |
| raise RuntimeError("Hydra launcher is not configured") |
|
|
| launcher = self._instantiate(config.hydra.launcher) |
| assert isinstance(launcher, Launcher) |
| launcher.setup( |
| hydra_context=hydra_context, task_function=task_function, config=config |
| ) |
| return launcher |
|
|
| @staticmethod |
| def _scan_all_plugins( |
| modules: List[Any], |
| ) -> Tuple[List[Type[Plugin]], ScanStats]: |
| stats = ScanStats() |
| stats.total_time = timer() |
|
|
| scanned_plugins: List[Type[Plugin]] = [] |
|
|
| for mdl in modules: |
| for importer, modname, ispkg in pkgutil.walk_packages( |
| path=mdl.__path__, prefix=mdl.__name__ + ".", onerror=lambda x: None |
| ): |
| try: |
| module_name = modname.rsplit(".", 1)[-1] |
| |
| |
| |
| if module_name.startswith("_") and not module_name.startswith("__"): |
| continue |
| import_time = timer() |
|
|
| with warnings.catch_warnings(record=True) as recorded_warnings: |
| if sys.version_info < (3, 10): |
| m = importer.find_module(modname) |
| assert m is not None |
| loaded_mod = m.load_module(modname) |
| else: |
| spec = importer.find_spec(modname) |
| assert spec is not None |
| if modname in sys.modules: |
| loaded_mod = sys.modules[modname] |
| else: |
| loaded_mod = importlib.util.module_from_spec(spec) |
| if loaded_mod is not None: |
| spec.loader.exec_module(loaded_mod) |
| sys.modules[modname] = loaded_mod |
|
|
| import_time = timer() - import_time |
| if len(recorded_warnings) > 0: |
| sys.stderr.write( |
| f"[Hydra plugins scanner] : warnings from '{modname}'. Please report to plugin author.\n" |
| ) |
| for w in recorded_warnings: |
| warnings.showwarning( |
| message=w.message, |
| category=w.category, |
| filename=w.filename, |
| lineno=w.lineno, |
| file=w.file, |
| line=w.line, |
| ) |
|
|
| stats.total_modules_import_time += import_time |
|
|
| assert modname not in stats.modules_import_time |
| stats.modules_import_time[modname] = import_time |
|
|
| if loaded_mod is not None: |
| for name, obj in inspect.getmembers(loaded_mod): |
| if _is_concrete_plugin_type(obj): |
| scanned_plugins.append(obj) |
| except ImportError as e: |
| warnings.warn( |
| message=f"\n" |
| f"\tError importing '{modname}'.\n" |
| f"\tPlugin is incompatible with this Hydra version or buggy.\n" |
| f"\tRecommended to uninstall or upgrade plugin.\n" |
| f"\t\t{type(e).__name__} : {e}", |
| category=UserWarning, |
| ) |
|
|
| stats.total_time = timer() - stats.total_time |
| return scanned_plugins, stats |
|
|
| def get_stats(self) -> Optional[ScanStats]: |
| return self.stats |
|
|
| def discover( |
| self, plugin_type: Optional[Type[Plugin]] = None |
| ) -> List[Type[Plugin]]: |
| """ |
| :param plugin_type: class of plugin to discover, None for all |
| :return: a list of plugins implementing the plugin type (or all if plugin type is None) |
| """ |
| Plugins.check_usage(self) |
| ret: List[Type[Plugin]] = [] |
| if plugin_type is None: |
| plugin_type = Plugin |
| assert issubclass(plugin_type, Plugin) |
| if plugin_type not in self.plugin_type_to_subclass_list: |
| return [] |
| for clazz in self.plugin_type_to_subclass_list[plugin_type]: |
| ret.append(clazz) |
|
|
| return ret |
|
|
| @staticmethod |
| def check_usage(self_: Any) -> None: |
| if not isinstance(self_, Plugins): |
| raise ValueError( |
| f"Plugins is now a Singleton. usage: Plugins.instance().{inspect.stack()[1][3]}(...)" |
| ) |
|
|
|
|
| def _is_concrete_plugin_type(obj: Any) -> bool: |
| return ( |
| inspect.isclass(obj) and issubclass(obj, Plugin) and not inspect.isabstract(obj) |
| ) |
|
|
