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

pytorch_lightning/utilities/cli.py

@@ -12,15 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import inspect
+import json
 import os
+import sys
 from argparse import Namespace
+from contextlib import contextmanager
 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, 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
@@ -332,6 +337,16 @@ 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.keys() if f"--optimizer={optim_name}" in v
+        ):
+            optimizers = tuple(v for v in OPTIMIZER_REGISTRIES.values())
+            self.parser.add_optimizer_args(optimizers)
+
+        if any(True for v in sys.argv for sch_name in SCHEDULER_REGISTRIES.keys() if f"-lr_scheduler={sch_name}" in v):
+            lr_schdulers = tuple(v for v in SCHEDULER_REGISTRIES.values())
+            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 +356,107 @@ 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_optimizer(self):
+        """
+        This context manager is used to simplify optimizer instantiation for Lightning users.
+        """
+        optimizer_args = [v for v in sys.argv if v.startswith("--optimizer")]
+        should_replace = len(optimizer_args) > 0 and not any(v for v in optimizer_args if "class_path" in v)
+        if should_replace:
+            optimizer_arg = {}
+            init_args = {}
+            for v in optimizer_args:
+                if "optimizer." in v:
+                    arg_path, value = v.split("=")
+                    init_args[arg_path.split(".")[-1]] = value
+                else:
+                    class_name = v.split("=")[-1]
+                    optim_cls = OPTIMIZER_REGISTRIES[class_name]
+                    optimizer_arg["class_path"] = optim_cls.__module__ + "." + class_name
+            optimizer_arg["init_args"] = init_args
+            argv = [v for v in sys.argv if not v.startswith("--optimizer")] + [
+                f"--optimizer={json.dumps(optimizer_arg)}"
+            ]
+            with mock.patch("sys.argv", argv):
+                yield
+        else:
+            yield
+
+    @contextmanager
+    def prepare_callbacks(self):
+        """
+        This context manager is used to simplify callbacks instantiation for Lightning users.
+        """
+        all_callbacks_args = [v for v in sys.argv if v.startswith("--trainer.callbacks")]
+        callbacks_args = [v for v in sys.argv if v.startswith("--trainer.callbacks=")]
+        num_callbacks = len(callbacks_args)
+        should_replace = len(all_callbacks_args) > 0 and not any(v for v in all_callbacks_args if "class_path" in v)
+        if should_replace:
+            # FIXME: Add support for combining callbacks.
+            callbacks_argv = {}
+            init_args = {}
+            map_callback_args = {idx: [] for idx in range(num_callbacks)}
+            counter = -1
+            callback_out = []
+            for v in all_callbacks_args:
+                if "--trainer.callbacks=" in v:
+                    counter += 1
+                map_callback_args[counter].append(v)
+            callback_out = []
+            for callback_idx in range(num_callbacks):
+                callback_args = map_callback_args[callback_idx]
+                callbacks_argv = {}
+                init_args = {}
+                for callback_arg in callback_args:
+                    if "--trainer.callbacks=" in callback_arg:
+                        class_name = callback_arg.split("=")[-1]
+                        callback_cls = CALLBACK_REGISTRIES[class_name]
+                        callbacks_argv["class_path"] = callback_cls.__module__ + "." + class_name
+                    else:
+                        arg_path, value = callback_arg.split("=")
+                        init_args[arg_path.split(".")[-1]] = value
+                callbacks_argv["init_args"] = init_args
+                callback_out.append(callbacks_argv)
+            argv = [v for v in sys.argv if not v.startswith("--trainer.callbacks")] + [
+                f"--trainer.callbacks={json.dumps(callback_out)}"
+            ]
+            with mock.patch("sys.argv", argv):
+                yield
+        else:
+            yield
+
+    @contextmanager
+    def prepare_schedulers(self):
+        """
+        This context manager is used to simplify schedulers instantiation for Lightning users.
+        """
+        lr_scheduler_args = [v for v in sys.argv if v.startswith("--lr_scheduler")]
+        should_replace = len(lr_scheduler_args) > 0 and not any(v for v in lr_scheduler_args if "class_path" in v)
+        if should_replace:
+            lr_scheduler_arg = {}
+            init_args = {}
+            for v in lr_scheduler_args:
+                if "lr_scheduler." in v:
+                    arg_path, value = v.split("=")
+                    init_args[arg_path.split(".")[-1]] = value
+                else:
+                    class_name = v.split("=")[-1]
+                    optim_cls = SCHEDULER_REGISTRIES[class_name]
+                    lr_scheduler_arg["class_path"] = optim_cls.__module__ + "." + class_name
+            lr_scheduler_arg["init_args"] = init_args
+            argv = [v for v in sys.argv if not v.startswith("--lr_scheduler")] + [
+                f"--lr_scheduler={json.dumps(lr_scheduler_arg)}"
+            ]
+            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()
+        with self.prepare_optimizer(), self.prepare_callbacks(), self.prepare_schedulers():
+            self.config = parser.parse_args()
 
     def before_instantiate_classes(self) -> None:
         """Implement to run some code before instantiating the classes."""

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)

tests/utilities/test_cli.py

@@ -33,6 +33,7 @@
 from pytorch_lightning.plugins.environments import SLURMEnvironment
 from pytorch_lightning.utilities import _TPU_AVAILABLE
 from pytorch_lightning.utilities.cli import instantiate_class, LightningArgumentParser, LightningCLI, SaveConfigCallback
+from pytorch_lightning.utilities.cli_registries import CALLBACK_REGISTRIES
 from pytorch_lightning.utilities.imports import _TORCHVISION_AVAILABLE
 from tests.helpers import BoringDataModule, BoringModel
 from tests.helpers.runif import RunIf
@@ -687,3 +688,59 @@ def __init__(self, optim1: dict, optim2: dict, scheduler: dict):
     assert isinstance(cli.model.optim1, torch.optim.Adam)
     assert isinstance(cli.model.optim2, torch.optim.SGD)
     assert isinstance(cli.model.scheduler, torch.optim.lr_scheduler.ExponentialLR)
+
+
+@CALLBACK_REGISTRIES
+class CustomCallback(Callback):
+    pass
+
+
+def test_registries(tmpdir):
+
+    assert CALLBACK_REGISTRIES.available_objects() == [
+        "BackboneFinetuning",
+        "BaseFinetuning",
+        "BasePredictionWriter",
+        "Callback",
+        "EarlyStopping",
+        "GPUStatsMonitor",
+        "GradientAccumulationScheduler",
+        "LambdaCallback",
+        "LearningRateMonitor",
+        "ModelCheckpoint",
+        "ModelPruning",
+        "ProgressBar",
+        "ProgressBarBase",
+        "QuantizationAwareTraining",
+        "StochasticWeightAveraging",
+        "Timer",
+        "XLAStatsMonitor",
+        "CustomCallback",
+    ]
+
+    class MyLightningCLI(LightningCLI):
+        def add_arguments_to_parser(self, parser):
+            pass
+
+    class TestModel(BoringModel):
+        def __init__(self):
+            super().__init__()
+
+    cli_args = [
+        f"--trainer.default_root_dir={tmpdir}",
+        "--trainer.max_epochs=1",
+        "--optimizer=Adam",
+        "--optimizer.lr=0.0001",
+        "--trainer.callbacks=LearningRateMonitor",
+        "--trainer.callbacks.logging_interval=epoch",
+        "--trainer.callbacks.log_momentum=True",
+        "--trainer.callbacks=ModelCheckpoint",
+        "--trainer.callbacks.monitor=loss",
+        "--lr_scheduler=StepLR",
+        "--lr_scheduler.step_size=50",
+    ]
+
+    with mock.patch("sys.argv", ["any.py"] + cli_args):
+        cli = MyLightningCLI(TestModel)
+
+    assert isinstance(cli.trainer.optimizers[0], torch.optim.Adam)