aid2e/schedulers/_registry.py

0001 """Registry and factory for scheduler implementations.
0002
0003 Aligns with the optimizer registry pattern: register, get, list_registered, and
0004 is_registered. Uses lazy registration to avoid circular imports with joblib.
0005 """
0006
0007 from typing import Dict, Type, Optional, Callable
0008
0009 from aid2e.schedulers.base import BaseScheduler
0010
0011 # In-memory mapping of scheduler identifiers to their classes
0012 _scheduler_registry: Dict[str, Type[BaseScheduler]] = {}
0013
0014 # Lazy loaders: map of name -> function that returns scheduler class
0015 _scheduler_loaders: Dict[str, Callable[[], Type[BaseScheduler]]] = {
0016     "joblib": lambda: __import__("aid2e.schedulers.JobLib", fromlist=["JobLibScheduler"]).JobLibScheduler,
0017 }
0018
0019
0020 def register(name: str, scheduler_class: Type[BaseScheduler]) -> None:
0021     """Register a scheduler implementation.
0022
0023     Args:
0024         name: Identifier used to retrieve the scheduler (case-insensitive).
0025         scheduler_class: Scheduler class that implements BaseScheduler.
0026
0027     Raises:
0028         ValueError: If name is already registered or scheduler_class is invalid.
0029     """
0030
0031     name_key = name.lower()
0032     if name_key in _scheduler_registry:
0033         raise ValueError(f"Scheduler '{name}' already registered")
0034     if not issubclass(scheduler_class, BaseScheduler):
0035         raise ValueError("Scheduler class must inherit from BaseScheduler")
0036
0037     _scheduler_registry[name_key] = scheduler_class
0038
0039
0040 def get(name: str) -> Type[BaseScheduler]:
0041     """Retrieve a scheduler class by name (with lazy loading).
0042
0043     Args:
0044         name: Identifier that was used during registration.
0045
0046     Returns:
0047         Scheduler class implementing BaseScheduler.
0048
0049     Raises:
0050         KeyError: If the scheduler is not registered.
0051     """
0052
0053     name_key = name.lower()
0054
0055     # First check if already loaded
0056     if name_key in _scheduler_registry:
0057         return _scheduler_registry[name_key]
0058
0059     # Try lazy loader
0060     if name_key in _scheduler_loaders:
0061         scheduler_class = _scheduler_loaders[name_key]()
0062         _scheduler_registry[name_key] = scheduler_class
0063         return scheduler_class
0064
0065     available = list(_scheduler_registry.keys()) + list(_scheduler_loaders.keys())
0066     raise KeyError(f"Scheduler '{name}' not registered. Available: {available}")
0067
0068
0069 def list_registered() -> Dict[str, Type[BaseScheduler]]:
0070     """Return a copy of the registered schedulers mapping.
0071
0072     Loads all lazy-registered schedulers.
0073     """
0074
0075     # Load all lazy schedulers
0076     for name in list(_scheduler_loaders.keys()):
0077         if name not in _scheduler_registry:
0078             try:
0079                 _scheduler_registry[name] = _scheduler_loaders[name]()
0080             except Exception:
0081                 pass  # Skip if fails to load
0082
0083     return _scheduler_registry.copy()
0084
0085
0086 def is_registered(name: str) -> bool:
0087     """Return True if the scheduler name is registered or can be lazy-loaded."""
0088
0089     name_key = name.lower()
0090     return name_key in _scheduler_registry or name_key in _scheduler_loaders
0091
0092
0093 # Backward-compatible aliases
0094 register_scheduler = register
0095 get_scheduler = get
0096 list_registered_schedulers = list_registered
0097 is_scheduler_registered = is_registered