[WIP] Accelerator registry follow-up

pytorch_lightning/accelerators/__init__.py

@@ -15,9 +15,4 @@
 from pytorch_lightning.accelerators.gpu import GPUAccelerator  # noqa: F401
 from pytorch_lightning.accelerators.hpu import HPUAccelerator  # noqa: F401
 from pytorch_lightning.accelerators.ipu import IPUAccelerator  # noqa: F401
-from pytorch_lightning.accelerators.registry import AcceleratorRegistry, call_register_accelerators  # noqa: F401
 from pytorch_lightning.accelerators.tpu import TPUAccelerator  # noqa: F401
-
-ACCELERATORS_BASE_MODULE = "pytorch_lightning.accelerators"
-
-call_register_accelerators(ACCELERATORS_BASE_MODULE)

pytorch_lightning/accelerators/accelerator.py

@@ -76,6 +76,8 @@ def auto_device_count() -> int:
     def is_available() -> bool:
         """Detect if the hardware is available."""
 
-    @classmethod
-    def register_accelerators(cls, accelerator_registry: Dict) -> None:
-        pass
+    @staticmethod
+    @abstractmethod
+    def name() -> str:
+        """Name of the Accelerator."""
+        raise NotImplementedError

pytorch_lightning/accelerators/cpu.py

@@ -60,10 +60,7 @@ def is_available() -> bool:
         """CPU is always available for execution."""
         return True
 
-    @classmethod
-    def register_accelerators(cls, accelerator_registry: Dict) -> None:
-        accelerator_registry.register(
-            "cpu",
-            cls,
-            description=f"{cls.__class__.__name__}",
-        )
+    @staticmethod
+    def name() -> str:
+        """Name of the Accelerator."""
+        return "cpu"

pytorch_lightning/accelerators/gpu.py

@@ -90,13 +90,10 @@ def auto_device_count() -> int:
     def is_available() -> bool:
         return torch.cuda.device_count() > 0
 
-    @classmethod
-    def register_accelerators(cls, accelerator_registry: Dict) -> None:
-        accelerator_registry.register(
-            "gpu",
-            cls,
-            description=f"{cls.__class__.__name__}",
-        )
+    @staticmethod
+    def name() -> str:
+        """Name of the Accelerator."""
+        return "gpu"
 
 
 def get_nvidia_gpu_stats(device: _DEVICE) -> Dict[str, float]:  # pragma: no-cover

pytorch_lightning/accelerators/hpu.py

@@ -60,10 +60,7 @@ def auto_device_count() -> int:
     def is_available() -> bool:
         return _HPU_AVAILABLE
 
-    @classmethod
-    def register_accelerators(cls, accelerator_registry: Dict) -> None:
-        accelerator_registry.register(
-            "hpu",
-            cls,
-            description=f"{cls.__class__.__name__}",
-        )
+    @staticmethod
+    def name() -> str:
+        """Name of the Accelerator."""
+        return "hpu"

pytorch_lightning/accelerators/ipu.py

@@ -47,10 +47,7 @@ def auto_device_count() -> int:
     def is_available() -> bool:
         return _IPU_AVAILABLE
 
-    @classmethod
-    def register_accelerators(cls, accelerator_registry: Dict) -> None:
-        accelerator_registry.register(
-            "ipu",
-            cls,
-            description=f"{cls.__class__.__name__}",
-        )
+    @staticmethod
+    def name() -> str:
+        """Name of the Accelerator."""
+        return "ipu"

pytorch_lightning/accelerators/registry.py

@@ -11,20 +11,16 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import importlib
-from inspect import getmembers, isclass
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, List, Optional, Type
 
 from pytorch_lightning.accelerators.accelerator import Accelerator
-from pytorch_lightning.utilities.exceptions import MisconfigurationException
-from pytorch_lightning.utilities.registry import _is_register_method_overridden
 
 
 class _AcceleratorRegistry(dict):
-    """This class is a Registry that stores information about the Accelerators.
+    """This class is a dictionary that stores information about the Accelerators.
 
     The Accelerators are mapped to strings. These strings are names that identify
-    an accelerator, e.g., "gpu". It also returns Optional description and
+    an accelerator, e.g., "gpu". It also includes an optional description and any
     parameters to initialize the Accelerator, which were defined during the
     registration.
 
@@ -34,89 +30,76 @@ class _AcceleratorRegistry(dict):
 
     Example::
 
-        @AcceleratorRegistry.register("sota", description="Custom sota accelerator", a=1, b=True)
+        @ACCELERATOR_REGISTRY
         class SOTAAccelerator(Accelerator):
-            def __init__(self, a, b):
+            def __init__(self, a):
                 ...
 
-        or
+            @staticmethod
+            def name():
+                return "sota"
 
-        AcceleratorRegistry.register("sota", SOTAAccelerator, description="Custom sota accelerator", a=1, b=True)
+        # or to pass parameters
+        ACCELERATOR_REGISTRY.register(SOTAAccelerator, description="My SoTA accelerator", a=1)
     """
 
+    def __call__(self, *args: Any, **kwargs: Any) -> Type:
+        return self.register(*args, **kwargs)
+
     def register(
         self,
-        name: str,
-        accelerator: Optional[Callable] = None,
-        description: str = "",
+        accelerator: Type[Accelerator],
+        name: Optional[str] = None,
+        description: Optional[str] = None,
         override: bool = False,
-        **init_params: Any,
-    ) -> Callable:
-        """Registers a accelerator mapped to a name and with required metadata.
+        **kwargs: Any,
+    ) -> Type:
+        """Registers an accelerator mapped to a name and with optional metadata.
 
         Args:
-            name : the name that identifies a accelerator, e.g. "gpu"
-            accelerator : accelerator class
-            description : accelerator description
-            override : overrides the registered accelerator, if True
-            init_params: parameters to initialize the accelerator
+            accelerator: The accelerator class.
+            name: The alias for the accelerator, e.g. ``"gpu"``.
+            description: An optional description.
+            override: Whether to override the registered accelerator.
+            **kwargs: parameters to initialize the accelerator.
         """
-        if not (name is None or isinstance(name, str)):
-            raise TypeError(f"`name` must be a str, found {name}")
-
-        if name in self and not override:
-            raise MisconfigurationException(f"'{name}' is already present in the registry. HINT: Use `override=True`.")
-
-        data: Dict[str, Any] = {}
-
-        data["description"] = description
-        data["init_params"] = init_params
-
-        def do_register(name: str, accelerator: Callable) -> Callable:
-            data["accelerator"] = accelerator
-            data["accelerator_name"] = name
-            self[name] = data
-            return accelerator
-
-        if accelerator is not None:
-            return do_register(name, accelerator)
-
-        return do_register
-
-    def get(self, name: str, default: Optional[Any] = None) -> Any:
+        if name is None:
+            name = accelerator.name()
+        if not isinstance(name, str):
+            raise TypeError(f"`name` for {accelerator} must be a str, found {name!r}")
+
+        if name not in self or override:
+            self[name] = {
+                "accelerator": accelerator,
+                "description": description if description is not None else accelerator.__class__.__name__,
+                "kwargs": kwargs,
+            }
+        return accelerator
+
+    def get(self, name: str, default: Optional[Accelerator] = None) -> Accelerator:
         """Calls the registered accelerator with the required parameters and returns the accelerator object.
 
         Args:
-            name (str): the name that identifies a accelerator, e.g. "gpu"
+            name: The name that identifies a accelerator, e.g. "gpu".
+            default: A default value.
+
+        Raises:
+            KeyError: If the key does not exist.
         """
         if name in self:
             data = self[name]
-            return data["accelerator"](**data["init_params"])
-
+            return data["accelerator"](**data["kwargs"])
         if default is not None:
             return default
+        raise KeyError(f"{name!r} not found in registry. {self!s}")
 
-        err_msg = "'{}' not found in registry. Available names: {}"
-        available_names = self.available_accelerators()
-        raise KeyError(err_msg.format(name, available_names))
-
-    def remove(self, name: str) -> None:
-        """Removes the registered accelerator by name."""
-        self.pop(name)
-
-    def available_accelerators(self) -> List[str]:
-        """Returns a list of registered accelerators."""
-        return list(self.keys())
+    @property
+    def names(self) -> List[str]:
+        """Returns the registered names."""
+        return sorted(list(self))
 
     def __str__(self) -> str:
-        return "Registered Accelerators: {}".format(", ".join(self.available_accelerators()))
-
-
-AcceleratorRegistry = _AcceleratorRegistry()
+        return f"Registered Accelerators: {self.names}"
 
 
-def call_register_accelerators(base_module: str) -> None:
-    module = importlib.import_module(base_module)
-    for _, mod in getmembers(module, isclass):
-        if issubclass(mod, Accelerator) and _is_register_method_overridden(mod, Accelerator, "register_accelerators"):
-            mod.register_accelerators(AcceleratorRegistry)
+ACCELERATOR_REGISTRY = _AcceleratorRegistry()

pytorch_lightning/accelerators/tpu.py

@@ -65,10 +65,7 @@ def auto_device_count() -> int:
     def is_available() -> bool:
         return _TPU_AVAILABLE
 
-    @classmethod
-    def register_accelerators(cls, accelerator_registry: Dict) -> None:
-        accelerator_registry.register(
-            "tpu",
-            cls,
-            description=f"{cls.__class__.__name__}",
-        )
+    @staticmethod
+    def name() -> str:
+        """Name of the Accelerator."""
+        return "tpu"

pytorch_lightning/strategies/strategy_registry.py

@@ -17,7 +17,7 @@
 
 from pytorch_lightning.strategies.strategy import Strategy
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
-from pytorch_lightning.utilities.registry import _is_register_method_overridden
+from pytorch_lightning.utilities.model_helpers import is_overridden
 
 
 class _StrategyRegistry(dict):
@@ -118,5 +118,5 @@ def __str__(self) -> str:
 def call_register_strategies(base_module: str) -> None:
     module = importlib.import_module(base_module)
     for _, mod in getmembers(module, isclass):
-        if issubclass(mod, Strategy) and _is_register_method_overridden(mod, Strategy, "register_strategies"):
+        if issubclass(mod, Strategy) and is_overridden("register_strategies", mod, Strategy):
             mod.register_strategies(StrategyRegistry)

pytorch_lightning/trainer/connectors/accelerator_connector.py

@@ -24,7 +24,7 @@
 from pytorch_lightning.accelerators.gpu import GPUAccelerator
 from pytorch_lightning.accelerators.hpu import HPUAccelerator
 from pytorch_lightning.accelerators.ipu import IPUAccelerator
-from pytorch_lightning.accelerators.registry import AcceleratorRegistry
+from pytorch_lightning.accelerators.registry import ACCELERATOR_REGISTRY
 from pytorch_lightning.accelerators.tpu import TPUAccelerator
 from pytorch_lightning.plugins import (
     ApexMixedPrecisionPlugin,
@@ -80,6 +80,7 @@
 )
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.imports import _HOROVOD_AVAILABLE, _HPU_AVAILABLE, _IPU_AVAILABLE, _TPU_AVAILABLE
+from pytorch_lightning.utilities.meta import get_all_subclasses
 
 log = logging.getLogger(__name__)
 
@@ -156,7 +157,8 @@ def __init__(
         # 1. Parsing flags
         # Get registered strategies, built-in accelerators and precision plugins
         self._registered_strategies = StrategyRegistry.available_strategies()
-        self._accelerator_types = AcceleratorRegistry.available_accelerators()
+        _populate_registries()
+        self._accelerator_types = ACCELERATOR_REGISTRY.names
         self._precision_types = ("16", "32", "64", "bf16", "mixed")
 
         # Raise an exception if there are conflicts between flags
@@ -484,16 +486,16 @@ def _set_parallel_devices_and_init_accelerator(self) -> None:
         else:
             assert self._accelerator_flag is not None
             self._accelerator_flag = self._accelerator_flag.lower()
-            if self._accelerator_flag not in AcceleratorRegistry:
+            if self._accelerator_flag not in ACCELERATOR_REGISTRY:
                 raise MisconfigurationException(
                     "When passing string value for the `accelerator` argument of `Trainer`,"
                     f" it can only be one of {self._accelerator_types}."
                 )
-            self.accelerator = AcceleratorRegistry.get(self._accelerator_flag)
+            self.accelerator = ACCELERATOR_REGISTRY.get(self._accelerator_flag)
 
         if not self.accelerator.is_available():
             available_accelerator = [
-                acc_str for acc_str in self._accelerator_types if AcceleratorRegistry.get(acc_str).is_available()
+                acc_str for acc_str in self._accelerator_types if ACCELERATOR_REGISTRY.get(acc_str).is_available()
             ]
             raise MisconfigurationException(
                 f"{self.accelerator.__class__.__qualname__} can not run on your system"
@@ -820,3 +822,9 @@ def is_distributed(self) -> bool:
         if isinstance(self.accelerator, TPUAccelerator):
             is_distributed |= self.strategy.is_distributed
         return is_distributed
+
+
+def _populate_registries() -> None:
+    # automatically register accelerators
+    for cls in get_all_subclasses(Accelerator):
+        ACCELERATOR_REGISTRY(cls)