LightningCLI: add top-level argument

[adamjstewart]

I'm playing around with customizing LightningCLI according to https://lightning.ai/docs/pytorch/stable/cli/lightning_cli_expert.html. My code so far:

from lightning.pytorch.cli import LightningCLI

class MyLightningCLI(LightningCLI):
    def add_arguments_to_parser(self, parser):
        parser.add_argument('--version', action='version', version=f'%(prog)s {__version__}')

It took me longer than I would like to admit to realize this, but this adds the flag to myapp fit, not to myapp.

How would I add this argument to the top-level parser?

[adamjstewart]

I got this working using:

class MyLightningCLI(LightningCLI):
    def init_parser(self, **kwargs):
        parser = super().init_parser(**kwargs)
        parser.add_argument('--version', action='version', version=f'%(prog)s {__version__}')
        return parser

I'm just wondering if this is the recommended solution or if there is an easier solution. This also seems to add the argument to all parsers and subparsers, not just the top-level parser.

[mauvilsa]

This does seem to be the best solution to add add top-level arguments. Only for the specific case of version, it should also be possible to do without subclassing LightningCLI, but while instantiating give LightningCLI(..., parser_kwargs={"version": __version__}).

[adamjstewart]

Perfect, that's exactly what I'm looking for! Would the Lightning folks be open to a PR to add -v as a shortcut to --version?

[mauvilsa]

Even with the parser_kwargs the version gets added to the main and subcommand parsers, which is not what you wanted. And currently there isn't a way to do this.

If you want to improve the code, I would prefer if you made possible adding stuff to specific parsers. You want to add only to the main parser. Other people might want to only add to the subcommand parsers, or only to one specific subparser. This could be achieved for example by init_parser having a subcommand: Optional[str] parameter. So for the example above could be like:

class MyLightningCLI(LightningCLI):
    def init_parser(self, subcommand: Optional[str] = None, **kwargs):
        parser = super().init_parser(**kwargs)
        if subcommand is None:  # only add to main parser
            parser.add_argument('-v', '--version', action='version', version=f'%(prog)s {__version__}')
        return parser