Refactor use of **kwargs in PL classes for better LightningCLI support

[mauvilsa]

Proposed refactor

All the __init__ parameters from classes in the pytorch-lightning repo that are susceptible to be used via LightningCLI should be resolvable by jsonargparse: parameter name, type and docstring description.

A list of potential classes that would be good to refactor is the following. Probably each class needs to be considered separately and some might have higher priority than others.

• CometLogger
• NeptuneLogger
• TensorBoardLogger
• WandbLogger
• DDPStrategy
• DDPSpawnStrategy
• PyTorchProfiler
• ...

See Possibilities section below for ideas on how to refactor.

Motivation

This has been a recurrent issue, for example:

• How to instantiate plugins with LightningCLI? #8561
• DDPPlugin does not accept find_unused_parameters when used with Lightning CLI #11574
• LightningCLI: Linking arguments from model to checkpoint  #9887 (comment)
• LightningCli: Passing trainer_defaults #8008 (comment)

There is a proposal to allow specifying values for unresolved parameters which would skip validation omni-us/jsonargparse#114. The downside is no support for several nice features (see Pitch). This can be a solution for some cases. But it is not be the ideal solution, particularly for classes in the pytorch-lightning repo which can be more easily considered for modifications.

Pitch

With this refactor all features from the CLIs become available:

• Validation of values based on parameter types and early error messages.
• Description of parameters in the help of CLIs.
• Inclusion in reference config files obtained via --print_config.
• Linking of arguments.

Possibilities

1. Add all parameters to the class. This is how several people have been solving this issue, e.g. LightningCLI: Linking arguments from model to checkpoint  #9887 (comment)

   • Advantages: Very simple solution. All Pitch features available. No need for deprecations or change how the classes are used.
   • Disadvantages: Duplication of code. The parameters to where these **kwargs might change, thus creating problems.

2. Have a specific parameter with dict type. For example in the case of TensorBoardLogger there could be a parameter summary_writer_kwargs: Dict[str, Any]. The existing **kwargs would be deprecated, showing a warning if used.

   • Advantages: Also it is more clear what this parameter is for, unlike a generic **kwargs.
   • Disadvantages: No support for features listed in Pitch. Introduces deprecation. Need to change existing code, e.g. from param1=val, ... to summary_writer_kwargs=dict(param1=val, ...).

3. Have a specific parameter with a dataclass type. For example in the case of TensorBoardLogger there could be a parameter summary_writer_options: SummaryWriterOptions. The dataclass would be automatically generated from a signature, thus avoiding code duplication or discrepancies in parameters. See below possible function that could create these dataclasses.

   • Advantages: All Pitch features available.
   • Disadvantages: Introduces deprecation. Need to change existing code, e.g. from param1=val, ... to first have an import from ... import SummaryWriterOptions and then when used ... summary_writer_kwargs=SummaryWriterOptions(param1=val, ...)`.

   import inspect
   from typing import Any, Set
   from dataclasses import make_dataclass, field
   
   def dataclass_from_signature(name: str, func, skip: Set = None):
       """Creates a dataclass using as fields parameters from the signature of some callable.
   
       Args:
           name: How the new class should be called, i.e. value to set in ``__name__``.
           func: The callable from which to get the signature.
           skip: Names of parameters to exclude from the signature.
   
       Example:
           ``MyData = dataclass_from_signature('MyData', MyClass.__init__, skip={'self', 'paramX'})``
       """
       fields = []
       params = list(inspect.signature(func).parameters.values())
       for param in params:
           if skip and param.name in skip:
               continue
           field_info = [param.name, Any if param._annotation == inspect._empty else param._annotation]
           if param.default != inspect._empty:
               field_info.append(field(default=param.default))
           fields.append(tuple(field_info))
       new_dataclass = make_dataclass(name, fields)
       # TODO: Add docstring with description of fields.
       return new_dataclass

4. Similar to (3) but instead of a dataclass, it would be a normal class also automatically generated. The __init__ would define the parameters, but it would ignore them. This class would be used as a parent class, so jsonargparse would resolve the parameters via inheritance. For example for TensorBoardLogger the definition could be something like class TensorBoardLogger(SummaryWriterOptions, LightningLoggerBase):.

   • Advantages: All Pitch features available. No need for deprecations or change how the classes are used.
   • Disadvantages: A bit more complex solution and feels a bit more hacky (though might not be).

5. Extend the support of **kwargs parameter discovery in jsonargparse using ast, such that no refactoring is needed.

   • Advantages: All Pitch features available.
   • Disadvantages: Considerably complex and time consuming solution. Might not avoid the need of refactoring of all cases.

6. Please propose more possibilities.

The deprecations mentioned above could be optional. It wouldn't be a problem to keep the new parameter but also the existing **kwargs. This way there is not need for people to change their existing code.

The dataclass_from_signature function very likely will be added to and maintained in jsonargparse irrespective of this.

Additional context

---

If you enjoy Lightning, check out our other projects! ⚡

• Metrics: Machine learning metrics for distributed, scalable PyTorch applications.

• Lite: enables pure PyTorch users to scale their existing code on any kind of device while retaining full control over their own loops and optimization logic.

• Flash: The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, fine-tuning, and solving problems with deep learning.

• Bolts: Pretrained SOTA Deep Learning models, callbacks, and more for research and production with PyTorch Lightning and PyTorch.

• Lightning Transformers: Flexible interface for high-performance research using SOTA Transformers leveraging Pytorch Lightning, Transformers, and Hydra.

cc @Borda @justusschock @awaelchli @rohitgr7 @tchaton @carmocca @mauvilsa @akihironitta

[carmocca]

I don't think (2) or (3) will happen because, from past experience, people are really allergic to having to create these extra objects. And since they would be created dynamically, users could not easily figure out what arguments can be passed and what they do.

My preference is doing (1) for the popular classes used in the CLI to reap the benefits described in your pitch.

[mauvilsa]

I don't think (2) or (3) will happen because, from past experience, people are really allergic to having to create these extra objects. And since they would be created dynamically, users could not easily figure out what arguments can be passed and what they do.

My preference is doing (1) for the popular classes used in the CLI to reap the benefits described in your pitch.

Don't underestimate the downsides of (1). For example for PyTorchProfiler even the docstring says that the parameters depend on the pytorch version. Defining parameters of a class conditional on the installed pytorch version for sure would be ugly.

If (3) is done without deprecation of the current **kwargs, people would not need to use an extra object. The new parameter could be intended only for the CLIs, and described as such in the docstring. The difference would be that these settings in the yaml would be indented in a section, e.g. in summary_writer_kwargs instead of in __unresolved__. This section would be included when doing print_config and in the help, so the usage would be clear I think.

I have just added a fourth option that could also be interesting.

[carmocca]

To me, a killer feature of jsonargparse (and consequently the LightningCLI) is that you don't need to adapt your code to your parser, instead the parser integrates to your code (by using the annotations, docstrings, ...).

To me, doing (2), (3), or (4) is adapting the code to fit the parser needs. If the LightningCLI did not exist, would these changes be valuable by themselves? Probably not.

(1) is somewhat valuable because the user avoids 1 jump to read the argument annotations and docstrings, at the cost of duplication. That's why it's my preference.

Defining parameters of a class conditional on the installed pytorch version for sure would be ugly.

I agree. We could keep the code as-is for these tricky classes and just let go of the validation.

I have just added a fourth option that could also be interesting.

Could this option be done by jsonargparse behind the scenes?

[mauvilsa]

To me, a killer feature of jsonargparse (and consequently the LightningCLI) is that you don't need to adapt your code to your parser, instead the parser integrates to your code (by using the annotations, docstrings, ...).

Indeed this has been the objective of jsonargparse and certainly improving its parameter discovery feature is another possibility to address this LightningCLI recurring issue, now it is (5) above. However, there are endless ways in which people might use **kwargs internally in their code. jsonargparse would need to analyze the source code to trace down the **kwargs. I suppose that with the https://docs.python.org/3/library/ast.html module, analyzing the source code is doable. Though, targeting to support all possible **kwargs uses would be extremely complex and very difficult to achieve in practice. If not all possible **kwargs uses are supported, then unfortunately people need to adapt the code to the parser, if in their case the parameter discovery does not work.

Going this route for certain will require way more effort and time. I have never used ast and also haven't thought it through. It might be going down a rabbit hole and introduce difficult to fix problems. Almost certain that it would not be a support for any **kwargs use. It probably would be a very small number of easy to explain support cases, such that when it does not work, people are able to understand why and modify accordingly.

To me, doing (2), (3), or (4) is adapting the code to fit the parser needs. If the LightningCLI did not exist, would these changes be valuable by themselves? Probably not.

In my opinion it depends on the case and how you look at it. Having (2) or (3) and a parameter called summary_writer_kwargs leads to easier to understand code and could be considered better practice by some people. But anyway, deciding not to adapt the code just because of value under the assumption that LightningCLI did not exist ignores the issue. The fact is that LightningCLI exists and there is a recurring issue that affects many users. An adaptation in PL might be justified to resolve the issue even if it is only temporal.

(1) is somewhat valuable because the user avoids 1 jump to read the argument annotations and docstrings, at the cost of duplication. That's why it's my preference.

That is fine. I am just trying to list out all the possibilities so that better decisions can be made.

[carmocca]

An adaptation in PL might be justified to resolve the issue even if it is only temporal.

I can get on board with this, but this disqualifies (2) and (3) as they would require long deprecation processes.

[mauvilsa]

I will look into ast and later let you what could be done. Regarding (1) it also has the advantage of being simple enough to be implemented by any contributor. And definitely would provide a better user experience than omni-us/jsonargparse#114.

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

[mauvilsa]

An update about this. I have been looking at option (5). As I thought, it is considerably complex. Still, I think it will be eventually implemented in jsonargparse. Though this will take time. Thus, (1) as interim solution could be considered.