Integrate lightning_utilities get_all_subclasses (#14575)

Author: carmocca

src/pytorch_lightning/CHANGELOG.md

@@ -27,7 +27,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Integrate the `lightning_utilities` package (
   [#14475](https://github.com/Lightning-AI/lightning/issues/14475),
   [#14537](https://github.com/Lightning-AI/lightning/issues/14537),
-  [#14556](https://github.com/Lightning-AI/lightning/issues/14556))
+  [#14556](https://github.com/Lightning-AI/lightning/issues/14556),
+  [#14558](https://github.com/Lightning-AI/lightning/issues/14558),
+  [#14575](https://github.com/Lightning-AI/lightning/issues/14575))
 
 
 ### Changed

src/pytorch_lightning/utilities/cli.py

@@ -18,11 +18,11 @@
 from typing import Any, Generator, List, Optional, Tuple, Type
 
 import torch
+from lightning_utilities.core.inheritance import get_all_subclasses
 from torch.optim import Optimizer
 
 import pytorch_lightning as pl
 import pytorch_lightning.cli as new_cli
-from pytorch_lightning.utilities.meta import _get_all_subclasses
 from pytorch_lightning.utilities.rank_zero import rank_zero_deprecation
 
 _deprecate_registry_message = (
@@ -108,17 +108,17 @@ def _populate_registries(subclasses: bool) -> None:  # Remove in v1.9
     if subclasses:
         rank_zero_deprecation(_deprecate_auto_registry_message)
         # this will register any subclasses from all loaded modules including userland
-        for cls in _get_all_subclasses(torch.optim.Optimizer):
+        for cls in get_all_subclasses(torch.optim.Optimizer):
             OPTIMIZER_REGISTRY(cls, show_deprecation=False)
-        for cls in _get_all_subclasses(torch.optim.lr_scheduler._LRScheduler):
+        for cls in get_all_subclasses(torch.optim.lr_scheduler._LRScheduler):
             LR_SCHEDULER_REGISTRY(cls, show_deprecation=False)
-        for cls in _get_all_subclasses(pl.Callback):
+        for cls in get_all_subclasses(pl.Callback):
             CALLBACK_REGISTRY(cls, show_deprecation=False)
-        for cls in _get_all_subclasses(pl.LightningModule):
+        for cls in get_all_subclasses(pl.LightningModule):
             MODEL_REGISTRY(cls, show_deprecation=False)
-        for cls in _get_all_subclasses(pl.LightningDataModule):
+        for cls in get_all_subclasses(pl.LightningDataModule):
             DATAMODULE_REGISTRY(cls, show_deprecation=False)
-        for cls in _get_all_subclasses(pl.loggers.Logger):
+        for cls in get_all_subclasses(pl.loggers.Logger):
             LOGGER_REGISTRY(cls, show_deprecation=False)
     else:
         # manually register torch's subclasses and our subclasses

src/pytorch_lightning/utilities/data.py

@@ -22,6 +22,7 @@
 
 import torch
 from lightning_utilities.core.apply_func import is_dataclass_instance
+from lightning_utilities.core.inheritance import get_all_subclasses
 from lightning_utilities.core.rank_zero import WarningCache
 from torch import Tensor
 from torch.utils.data import (
@@ -40,7 +41,6 @@
 from pytorch_lightning.utilities.auto_restart import CaptureIterableDataset, CaptureMapDataset, FastForwardSampler
 from pytorch_lightning.utilities.enums import _FaultTolerantMode, LightningEnum
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
-from pytorch_lightning.utilities.meta import _get_all_subclasses
 from pytorch_lightning.utilities.rank_zero import rank_zero_warn
 from pytorch_lightning.utilities.seed import pl_worker_init_function
 
@@ -549,7 +549,7 @@ def _replace_dunder_methods(base_cls: Type, store_explicit_arg: Optional[str] =
 
     It patches the ``__init__``, ``__setattr__`` and ``__delattr__`` methods.
     """
-    classes = _get_all_subclasses(base_cls) | {base_cls}
+    classes = get_all_subclasses(base_cls) | {base_cls}
     for cls in classes:
         # Check that __init__ belongs to the class
         # https://stackoverflow.com/a/5253424

src/pytorch_lightning/utilities/meta.py

@@ -43,21 +43,9 @@ def get_all_subclasses(cls: Type) -> Set[Type]:
         "`pytorch_lightning.utilities.meta.get_all_subclasses` is deprecated in v1.8 and will be removed in v1.9."
         " Please copy its implementation if you have a use for it."
     )
-    return _get_all_subclasses(cls)
+    from lightning_utilities.core.inheritance import get_all_subclasses as new_get_all_subclasses
 
-
-# https://stackoverflow.com/a/63851681/9201239
-def _get_all_subclasses(cls: Type) -> Set[Type]:
-    subclass_list = []
-
-    def recurse(cl: Type) -> None:
-        for subclass in cl.__subclasses__():
-            subclass_list.append(subclass)
-            recurse(subclass)
-
-    recurse(cls)
-
-    return set(subclass_list)
+    return new_get_all_subclasses(cls)
 
 
 def recursively_setattr(root_module: Any, prefix: str, materialized_module: Module) -> None: