Adopt typing_extensions.Override

[carmocca]

Description & Motivation

We use overrides heavily.

Python 3.12 added support for this. See https://peps.python.org/pep-0698/

Mypy should support it: python/mypy#14072

Pitch

Use https://typing-extensions.readthedocs.io/en/latest/#override in our overrides classes.

Using this would require increasing the minimum typing_extensions to 4.4: https://github.com/Lightning-AI/lightning/blob/master/requirements/fabric/base.txt#L8

Alternatives

No response

Additional context

No response

cc @Borda

[seanbethard]

can you clarify what overrides classes are?

[sartq333]

Hey @carmocca, I would like to solve this, can you assign me this issue?

[carmocca]

@seanbethard see https://peps.python.org/pep-0698

@SarthakJain333 Sure! I suggest that you do it for 1 or a few files per PR to avoid having it blow up in size

[sartq333]

Ok, I'll have a look around this framework and will try to solve the issue. Thanks for assigning me !

[jxtngx]

@SarthakJain333 let me know if you need any help!

[seanbethard]

I understand what overrides and decorators are. it's not clear from the issue which classes you're targeting. good luck @SarthakJain333.

[carmocca]

Sorry @seanbethard, I didn't understand what you meant from your previous message.

It would be subclasses of the following base interfaces:

• https://github.com/Lightning-AI/lightning/blob/master/src/lightning/fabric/accelerators/accelerator.py
• https://github.com/Lightning-AI/lightning/blob/master/src/lightning/fabric/loggers/logger.py
• https://github.com/Lightning-AI/lightning/blob/master/src/lightning/fabric/plugins/environments/cluster_environment.py
• https://github.com/Lightning-AI/lightning/blob/master/src/lightning/fabric/plugins/io/checkpoint_io.py
• https://github.com/Lightning-AI/lightning/blob/master/src/lightning/fabric/plugins/precision/precision.py
• https://github.com/Lightning-AI/lightning/blob/master/src/lightning/fabric/strategies/strategy.py
• https://github.com/Lightning-AI/lightning/blob/master/src/lightning/pytorch/accelerators/accelerator.py
• https://github.com/Lightning-AI/lightning/blob/master/src/lightning/pytorch/callbacks/callback.py
• https://github.com/Lightning-AI/lightning/blob/master/src/lightning/pytorch/callbacks/progress/progress_bar.py
• https://github.com/Lightning-AI/lightning/blob/master/src/lightning/pytorch/core/datamodule.py
• https://github.com/Lightning-AI/lightning/blob/master/src/lightning/pytorch/core/module.py
• https://github.com/Lightning-AI/lightning/blob/master/src/lightning/pytorch/loggers/logger.py
• https://github.com/Lightning-AI/lightning/blob/master/src/lightning/pytorch/loops/loop.py
• https://github.com/Lightning-AI/lightning/blob/master/src/lightning/pytorch/plugins/layer_sync.py
• https://github.com/Lightning-AI/lightning/blob/master/src/lightning/pytorch/plugins/precision/precision_plugin.py
• https://github.com/Lightning-AI/lightning/blob/master/src/lightning/pytorch/profilers/profiler.py
• https://github.com/Lightning-AI/lightning/blob/master/src/lightning/pytorch/strategies/strategy.py

As you can see, we make heavy use of overrides in this project

[sbshah97]

Hey if this is not done yet, is this something I can pick up? I am a new contributor looking to help.

[carmocca]

Yes @sbshah97, you can pick a file from the list above

[carmocca]

@seanbethard This is feedback for #18866 but commenting here for everybody else who decides to help.

This addition should only be done to the src/lightning/fabric/* and src/lightning/pytorch/* directories. examples/, docs/, tests/, or any other directories should not have this decorator added.

[sbshah97]

I can volunteer to start some files in the fabric directory to start off. Will try to avoid a large PR.

[seanbethard]

I see. well, src/lightning/fabric/* and src/lightning/pytorch/* are done. I removed overrides from tests following @awaelchli's feedback. happy to remove from the other directories as well.

each commit in #18866 is associated with a single base class. I thought this way of breaking up the work was preferable to adding 17 PRs.

[seanbethard]

removed overrides from all directories except src/lightning/fabric and src/lightning/pytorch. should be all set. please let me know if I missed anything.

[VictorPrins]

@seanbethard You've covered a lot of classes in #18866. I believe https://github.com/Lightning-AI/lightning/blob/master/src/lightning/pytorch/strategies/strategy.py (among others) is still unchanged. Can I do that one?

[seanbethard]

there were a few that I didn't think required any changes because the subclasses only implement abstract methods or define new ones. I believe pytorch Strategy is one such case.

[VictorPrins]

As an example, setup() is defined in Strategy, and overridden in DDPStrategy https://github.com/Lightning-AI/lightning/blob/31b8777350f545fe7b366ea7ae26d63f87e820df/src/lightning/pytorch/strategies/ddp.py#L150

Am I missing something?

[seanbethard]

nice catch! for each of the base classes listed above I only checked for classes that inherit directly from them (children). I didn't check for grandchildren or additional generations.

[VictorPrins]

Alright! I'll cover pytorch/strategies.

[carmocca]

@VictorPrins I might have missed some files/directories on my comment. Feel free to correct me and add the decorator elsewhere if it's needed

[awaelchli]

@VictorPrins do you have an overview of which files are left to do? Or are we done?

[Borda]

@VictorPrins do you have an overview of which files are left to do? Or are we done?

this shall be visible from mypy of who did we validate the changes?

[VictorPrins]

Yes I'll finish this up shortly, thanks for the heads up. The appropriate command is mypy --enable-error-code explicit-override --follow-imports=silent <folder to check>

[carmocca]

@VictorPrins Do you think that setting should be added to our mypy config after all the relevant files have been updated?

[VictorPrins]

Automated checking might be desirable to avoid gradual slipping, especially since the @override decorator is still rare enough in Python codebases that it's easy to forget. That said, there are some issues that must be navigated:

• Some folders must be excluded from the check (e.g. the demos folder)
• Mypy raises some errors for cases where, in my view, the @override decorator isn't necessary. In particular, an error is raised for overrides of methods from builtins.object, such as __str__ and __getstate__. The error can be silenced with # type: ignore[explicit-override], or solved by using the @override decorator for such methods.
• Mypy raises a strange error for overrides of overloaded methods. For instance for line 89 in the code below, I get error: Method "precision_plugin" is not using @override but is overriding a method in class "lightning.pytorch.strategies.strategy.Strategy, even though the @override is clearly used. This only happens for overloaded methods.
  

  pytorch-lightning/src/lightning/pytorch/strategies/xla.py

  Lines 89 to 91 in 628ee0c

  	@property # type: ignore[override]
  	@override
  	def precision_plugin(self) -> XLAPrecision: