[CLI] Shorthand notation to instantiate callbacks [3/3]

pytorch_lightning/utilities/cli.py

@@ -13,14 +13,26 @@
 # limitations under the License.
 import inspect
 import os
+import re
+import sys
 from argparse import Namespace
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from functools import partial
 from types import MethodType
 from typing import Any, Callable, cast, Dict, List, Optional, Tuple, Type, Union
+from unittest import mock
 
 from torch.optim import Optimizer
 
 from pytorch_lightning import Callback, LightningDataModule, LightningModule, seed_everything, Trainer
 from pytorch_lightning.utilities import _JSONARGPARSE_AVAILABLE, warnings
+from pytorch_lightning.utilities.cli_registries import (
+    CALLBACK_REGISTRIES,
+    OPTIMIZER_REGISTRIES,
+    Registry,
+    SCHEDULER_REGISTRIES,
+)
 from pytorch_lightning.utilities.cloud_io import get_filesystem
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.model_helpers import is_overridden
@@ -34,6 +46,31 @@
     ArgumentParser = object
 
 
+@dataclass
+class ClassInfo:
+    """This class is an helper to easily build the mocked command line"""
+
+    class_arg: str
+    cls: str
+    class_init_args: List[str] = field(default_factory=lambda: [])
+
+    def add_class_init_args(self, args: Dict[str, str]) -> None:
+        if args != self.class_arg:
+            self.class_init_args.append(args)
+
+    @property
+    def class_init(self) -> Dict[str, str]:
+        class_init = {}
+        class_init["class_path"] = self.cls.__module__ + "." + self.cls.__name__
+        init_args = {}
+        for init_arg in self.class_init_args:
+            separator = "=" if "=" in init_arg else " "
+            arg_path, value = init_arg.split(separator)
+            init_args[arg_path.split(".")[-1]] = value
+        class_init["init_args"] = init_args
+        return class_init
+
+
 class LightningArgumentParser(ArgumentParser):
     """Extension of jsonargparse's ArgumentParser for pytorch-lightning"""
 
@@ -289,6 +326,14 @@ def __init__(
         self.fit()
         self.after_fit()
 
+    @property
+    def optimizer_registered(self) -> Tuple[Type[Optimizer]]:
+        return tuple(OPTIMIZER_REGISTRIES.values())
+
+    @property
+    def lr_scheduler_registered(self) -> Tuple[LRSchedulerType]:
+        return tuple(SCHEDULER_REGISTRIES.values())
+
     def init_parser(self, **kwargs: Any) -> LightningArgumentParser:
         """Method that instantiates the argument parser."""
         return LightningArgumentParser(**kwargs)
@@ -332,6 +377,25 @@ def add_arguments_to_parser(self, parser: LightningArgumentParser) -> None:
 
     def link_optimizers_and_lr_schedulers(self) -> None:
         """Creates argument links for optimizers and learning rate schedulers that specified a ``link_to``."""
+        if any(
+            True
+            for v in sys.argv
+            for optim_name in OPTIMIZER_REGISTRIES
+            if re.match(fr"^--optimizer[^\S+=]*?{optim_name}?", v)
+        ):
+            if "optimizer" not in self.parser.groups:
+                self.parser.add_optimizer_args(self.optimizer_registered)
+
+        if any(
+            True
+            for v in sys.argv
+            for sch_name in SCHEDULER_REGISTRIES
+            if re.match(fr"^--lr_scheduler[^\S+=]*{sch_name}?", v)
+        ):
+            lr_schdulers = tuple(v for v in SCHEDULER_REGISTRIES.values())
+            if "lr_scheduler" not in self.parser.groups:
+                self.parser.add_lr_scheduler_args(lr_schdulers)
+
         for key, (class_type, link_to) in self.parser.optimizers_and_lr_schedulers.items():
             if link_to == "AUTOMATIC":
                 continue
@@ -341,9 +405,90 @@ def link_optimizers_and_lr_schedulers(self) -> None:
                 add_class_path = _add_class_path_generator(class_type)
                 self.parser.link_arguments(key, link_to, compute_fn=add_class_path)
 
+    @contextmanager
+    def prepare_from_registry(self, registry: Registry):
+        """
+        This context manager is used to simplify unique class instantiation.
+        """
+
+        # find if the users is using shortcut command line.
+        map_user_key_to_info = {}
+        for registered_name, registered_cls in registry.items():
+            for v in sys.argv:
+                separator = "=" if "=" in v else " "
+                if f"{separator}{registered_name}" in v:
+                    key = v.split(separator)[0]
+                    map_user_key_to_info[key] = ClassInfo(class_arg=v, cls=registered_cls)
+
+        if len(map_user_key_to_info) > 0:
+            # for each shortcut command line, add its init arguments and skip them from `sys.argv`.
+            argv = []
+            for v in sys.argv:
+                skip = False
+                for key in map_user_key_to_info:
+                    if key in v:
+                        skip = True
+                        map_user_key_to_info[key].add_class_init_args(v)
+                if not skip:
+                    argv.append(v)
+
+            # re-create the global command line and mock `sys.argv`.
+            argv += [f"{user_key}={info.class_init}" for user_key, info in map_user_key_to_info.items()]
+            with mock.patch("sys.argv", argv):
+                yield
+        else:
+            yield
+
+    @contextmanager
+    def prepare_class_list_from_registry(self, pattern: str, registry: Registry):
+        """
+        This context manager is used to simplify instantiation of a list of class.
+        """
+        argv = [v for v in sys.argv if pattern not in v]
+        all_matched_args = [v for v in sys.argv if pattern in v]
+        all_simplified_args = [v for v in all_matched_args if f"{pattern}" in v and f"{pattern}=[" not in v]
+        all_cls_simplified_args = [v for v in all_simplified_args if f"{pattern}=" in v]
+        all_non_simplified_args = [v for v in all_matched_args if f"{pattern}=" in v and f"{pattern}=[" in v]
+
+        num_simplified_cls = len(all_simplified_args)
+        should_replace = num_simplified_cls > 0 and not all("class_path" in v for v in all_matched_args)
+
+        if should_replace:
+            # verify the user is properly ordering arguments.
+            assert all_cls_simplified_args[0] == all_simplified_args[0]
+            if len(all_non_simplified_args) > 1:
+                raise MisconfigurationException(f"When provided {pattern} as list, please group them under 1 argument.")
+
+            # group arguments per callbacks
+            infos = []
+            for class_arg in all_cls_simplified_args:
+                class_name = class_arg.split("=")[1]
+                registered_cls = registry[class_name]
+                infos.append(ClassInfo(class_arg=class_arg, cls=registered_cls))
+
+            for v in all_simplified_args:
+                if v in all_cls_simplified_args:
+                    current_info = infos[all_cls_simplified_args.index(v)]
+                current_info.add_class_init_args(v)
+
+            class_args = [info.class_init for info in infos]
+            # add other callback arguments.
+            class_args.extend(eval(all_non_simplified_args[0].split("=")[-1]))
+
+            argv += [f"{pattern}={class_args}"]
+            with mock.patch("sys.argv", argv):
+                yield
+        else:
+            yield
+
     def parse_arguments(self, parser: LightningArgumentParser) -> None:
         """Parses command line arguments and stores it in ``self.config``."""
-        self.config = parser.parse_args()
+        # fmt: off
+        with self.prepare_from_registry(OPTIMIZER_REGISTRIES), \
+             self.prepare_from_registry(SCHEDULER_REGISTRIES), \
+             self.prepare_class_list_from_registry("--trainer.callbacks", CALLBACK_REGISTRIES):
+            self.config = parser.parse_args()
+        # fmt: on
 
     def before_instantiate_classes(self) -> None:
         """Implement to run some code before instantiating the classes."""
@@ -375,7 +520,6 @@ def instantiate_trainer(self, config: Dict[str, Any], callbacks: List[Callback])
     def add_configure_optimizers_method_to_model(self) -> None:
         """
         Adds to the model an automatically generated ``configure_optimizers`` method.
-
         If a single optimizer and optionally a scheduler argument groups are added to the parser as 'AUTOMATIC',
         then a `configure_optimizers` method is automatically implemented in the model class.
         """
@@ -421,17 +565,25 @@ def get_automatic(class_type: Union[Type, Tuple[Type, ...]]) -> List[str]:
             if not isinstance(lr_scheduler_class, tuple):
                 lr_scheduler_init = _global_add_class_path(lr_scheduler_class, lr_scheduler_init)
 
-        def configure_optimizers(
-            self: LightningModule,
-        ) -> Union[Optimizer, Tuple[List[Optimizer], List[LRSchedulerType]]]:
-            optimizer = instantiate_class(self.parameters(), optimizer_init)
-            if not lr_scheduler_init:
-                return optimizer
-            lr_scheduler = instantiate_class(optimizer, lr_scheduler_init)
-            return [optimizer], [lr_scheduler]
+        configure_optimizers = partial(
+            self.configure_optimizers, optimizer_init=optimizer_init, lr_scheduler_init=lr_scheduler_init
+        )
+        configure_optimizers.__code__ = self.model.configure_optimizers.__code__
 
         self.model.configure_optimizers = MethodType(configure_optimizers, self.model)
 
+    @staticmethod
+    def configure_optimizers(
+        pl_module: LightningModule,
+        optimizer_init: Union[str, List[str]],
+        lr_scheduler_init: Optional[Union[str, List[str]]] = None,
+    ) -> Union[Optimizer, Tuple[List[Optimizer], List[LRSchedulerType]]]:
+        optimizer = instantiate_class(pl_module.parameters(), optimizer_init)
+        if not lr_scheduler_init:
+            return optimizer
+        lr_scheduler = instantiate_class(optimizer, lr_scheduler_init)
+        return [optimizer], [lr_scheduler]
+
     def prepare_fit_kwargs(self) -> None:
         """Prepares fit_kwargs including datamodule using self.config_init['data'] if given"""
         self.fit_kwargs = {"model": self.model}

pytorch_lightning/utilities/cli_registries.py

@@ -0,0 +1,79 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 inspect
+from collections import UserDict
+from typing import Callable, List, Optional, Type
+
+import torch
+
+import pytorch_lightning as pl
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
+
+
+class Registry(UserDict):
+    def __call__(
+        self,
+        cls: Optional[Type] = None,
+        key: Optional[str] = None,
+        override: bool = False,
+    ) -> Callable:
+        """
+        Registers a plugin mapped to a name and with required metadata.
+
+        Args:
+            key : the name that identifies a plugin, e.g. "deepspeed_stage_3"
+            value : plugin class
+        """
+        if key is None:
+            key = cls.__name__
+        elif not isinstance(key, str):
+            raise TypeError(f"`key` must be a str, found {key}")
+
+        if key in self and not override:
+            raise MisconfigurationException(f"'{key}' is already present in the registry. HINT: Use `override=True`.")
+
+        def do_register(key, cls) -> Callable:
+            self[key] = cls
+            return cls
+
+        do_register(key, cls)
+
+        return do_register
+
+    def register_package(self, module, base_cls: Type) -> None:
+        for obj_name in dir(module):
+            obj_cls = getattr(module, obj_name)
+            if inspect.isclass(obj_cls) and issubclass(obj_cls, base_cls):
+                self(cls=obj_cls)
+
+    def remove(self, name: str) -> None:
+        """Removes the registered plugin by name"""
+        self.pop(name)
+
+    def available_objects(self) -> List:
+        """Returns a list of registered plugins"""
+        return list(self.keys())
+
+    def __str__(self) -> str:
+        return "Registered Plugins: {}".format(", ".join(self.keys()))
+
+
+CALLBACK_REGISTRIES = Registry()
+CALLBACK_REGISTRIES.register_package(pl.callbacks, pl.callbacks.Callback)
+
+OPTIMIZER_REGISTRIES = Registry()
+OPTIMIZER_REGISTRIES.register_package(torch.optim, torch.optim.Optimizer)
+
+SCHEDULER_REGISTRIES = Registry()
+SCHEDULER_REGISTRIES.register_package(torch.optim.lr_scheduler, torch.optim.lr_scheduler._LRScheduler)