Should Lightning have separate setup and teardown hooks for each stage?

[ananthsub]

Based on #6368

LightningModules and DataModules currently support a setup API which takes an optional stage argument.
#6386 addresses some issues in the setup/teardown lifecycle, so I was wondering if we should take this further.

Pros of making the separate hooks for each stage:

• Clarity in the API that helps forwards compatibility: new stages become opt-in for users. In the current situation, the Lightning trainer can pass an arbitrary value for stage that user code might not handle. As a result, adding new stages becomes opt-in for users, as they have to implement the corresponding hook in their lightning/data module
• This matches the pattern already established for Lightning/data modules which have train/validation/test/predict defined as separate hooks
• On the Lightning internals, we can remove the base datamodule wrapper class, and remove the has_setup_{stage} attributes since it'll be obvious when the hooks are called

Cons:

• This requires a deprecation process and can cause thrash for users
• Users now have to implement more hooks. However, a mitigation is that the refactoring should be straightforward as users can easily share code with a helper function in the lightning/data module.

@PyTorchLightning/core-contributors

Meta-question: should API proposals start out as discussions? Or should this be in an issue instead as an RFC?

[williamFalcon]

+1 for separate hooks. to keep things inline with the rest of the API

[williamFalcon]

but for sure need to be extra helpful for deprecating the API. perhaps wait 3 major releases?

[justusschock]

I also like to have them.

And I feel that we should go over all the hooks we currently have and see if they're still needed and which ones are missing. Especially the latter may be true for some of them after the accelerator refactor.

[ananthsub]

@justusschock i added an RFC here: #6420