Improve typing coverage (4/n)

[otaj]

🚀 Typing coverage

Let's improve typing coverage of PyTorch Lightning together!

I'm creating a new issue in order to increase visibility. There are three older issues (#7037, #5023, #4698) which became stale over time.

Plan

Currently, there are 55 files which are excluded from mypy checks so that our CI does not fail. These files vastly differ in difficulty in order to make the typing complete. For this reason, we are introducing difficulty estimate for each file so that community members can choose to work on the files appropriate to their skill level.

Please, comment on this issue in order to reserve a particular file to work on. Once you do so, I will edit this top comment to avoid collisions. Once you think your work is finished, please open a PR referencing this issue which:

• removes the corresponding line from pyproject.toml
• and passes mypy checks with the corresponding line removed. You can test it locally by running mypy from root directory

If you are struggling with pushing it over the finish line, open the PR anyway and someone from our team will help you to get it there. 🚀

Please note, that it can happen that you may need to edit more than just one file. This is fine, but please keep in mind, that the goal of your PR will be to make the check passing for the chosen file. Also, please note that the difficulty is just an educated guess.

For those of you who are not familiar with the process of contributing a PR, we have prepared a simple guide that will walk you through the necessary steps. You can do it! :rocket: :muscle:

List of files and guesstimated difficulty

Difficulty 1 of 3

• [ ] pytorch_lightning/tuner/batch_size_scaling.py @ar90n #13518

Difficulty 3 of 3

• [ ] pytorch_lightning/trainer/trainer.py ~@JustinGoheen #13810~ @BongYang #14204
• [ ] pytorch_lightning/callbacks/progress/rich_progress.py @donlapark

Completed

Difficulty 1 of 3

• [x] pytorch_lightning/core/decorators.py #14044
• [x] pytorch_lightning/profilers/advanced.py @nninept #13792 ~- [ ] pytorch_lightning/profilers/base.py @LeeChanHyuk #13879~
• [x] pytorch_lightning/loggers/base.py @JustinGoheen #13494
• [x] pytorch_lightning/__setup__.py @CyprienRicque #13472 ~- [ ] pytorch_lightning/distributed/dist.py @puhuk #13492~
• [x] pytorch_lightning/strategies/single_device.py @CyprienRicque #13532
• [x] pytorch_lightning/trainer/optimizers.py @gautierdag #13470
• [x] pytorch_lightning/utilities/distributed.py @krishnakalyan3 #13678
• [x] pytorch_lightning/callbacks/finetuning.py @ar90n #13516
• [x] pytorch_lightning/loggers/mlflow.py @JustinGoheen ~~#13690~~ #13691
• [x] pytorch_lightning/tuner/tuning.py @donlapark ~~#13616~~ #13631
• [x] pytorch_lightning/strategies/single_tpu.py @CyprienRicque #13534
• [x] pytorch_lightning/strategies/ddp2.py @CyprienRicque #13535
• [x] pytorch_lightning/strategies/parallel.py @CyprienRicque #13556
• [x] pytorch_lightning/loggers/csv_logs.py @JustinGoheen #13538
• [x] pytorch_lightning/tuner/lr_finder.py @donlapark #13513 #13652
• [x] pytorch_lightning/strategies/dp.py @CyprienRicque #13564
• [x] pytorch_lightning/profilers/simple.py @krishnakalyan3 #14103
• [x] pytorch_lightning/strategies/sharded_spawn.py @krishnakalyan3 #14102
• [x] pytorch_lightning/demos/mnist_datamodule.py @alro923 #13929
• [x] pytorch_lightning/demos/boring_classes.py @krishnakalyan3 #14201

Difficulty 2 of 3

• [x] pytorch_lightning/loops/epoch/training_epoch_loop.py @himkt #13555
• [x] pytorch_lightning/core/mixins/device_dtype_mixin.py @krishnakalyan3 #13704
• [x] pytorch_lightning/loggers/comet.py @JustinGoheen #13689
• [x] pytorch_lightning/loggers/tensorboard.py @JustinGoheen #13688
• [x] pytorch_lightning/strategies/horovod.py @CyprienRicque #13570
• [x] pytorch_lightning/callbacks/model_checkpoint.py @BongYang #13617
• [x] pytorch_lightning/strategies/fully_sharded.py @BongYang #13941
• [x] pytorch_lightning/loggers/neptune.py @JustinGoheen #13692
• [x] pytorch_lightning/utilities/meta.py @nninept #13763 #13868
• [x] pytorch_lightning/strategies/tpu_spawn.py @BongYang #13813
• [x] pytorch_lightning/loggers/logger.py @JustinGoheen #13541
• [x] pytorch_lightning/loggers/wandb.py @gautierdag #13483
• [x] pytorch_lightning/callbacks/stochastic_weight_avg.py @donlapark #13685 #13860
• [x] pytorch_lightning/strategies/strategy.py @CyprienRicque #13519
• [x] pytorch_lightning/strategies/deepspeed.py @donlapark #13832
• [x] pytorch_lightning/strategies/ddp_spawn.py @donlapark #13865
• [x] pytorch_lightning/strategies/ipu.py @HalestormAI #13786
• [x] pytorch_lightning/trainer/connectors/callback_connector.py @krishnakalyan3 #13750
• [x] pytorch_lightning/strategies/ddp.py @lijm1358 #13885
• [x] pytorch_lightning/core/saving.py @JustinGoheen #13932
• [x] pytorch_lightning/callbacks/quantization.py @krishnakalyan3 #13782
• [x] pytorch_lightning/strategies/sharded.py @lijm1358 #14184
• [x] pytorch_lightning/core/datamodule.py @JustinGoheen #13693

Difficulty 3 of 3

~- [ ] pytorch_lightning/trainer/callback_hook.py @JustinGoheen #13807 ~

• [x] pytorch_lightning/core/module.py @JustinGoheen #13603
• [x] pytorch_lightning/trainer/connectors/data_connector.py @JustinGoheen #13806
• [x] pytorch_lightning/utilities/auto_restart.py @donlapark #13904
• [x] pytorch_lightning/trainer/supporters.py @donlapark #14633
• [x] pytorch_lightning/profilers/pytorch.py @krishnakalyan3 #14405
• [x] pytorch_lightning/utilities/data.py @nandwalritik #13901

cc @borda @justusschock @awaelchli @rohitgr7 @Borda @tchaton @aniketmaurya @kingjuno @alat-rights @carmocca @akihironitta @stancld as you were all involved in previous issues

[kingjuno]

@otaj I will start working on these 3 for now.

• pytorch_lightning/core/decorators.py
• pytorch_lightning/profilers/advanced.py
• pytorch_lightning/profilers/base.py

[jxtngx]

@otaj I can take:

• pytorch_lightning/loggers/base.py
• pytorch_lightning/loggers/csv_logs.py
• pytorch_lightning/loggers/logger.py

[carmocca]

A recommendation for contributors: Unless a file requires <10 lines of fixes, I strongly suggest you work on separate PRs per file if possible. Sometimes mypy requires you to update the typing annotations for a different file, but it does not mean you need to completely annotate all modified files.

[otaj]

@JustinGoheen @kingjuno Thank you very much, you were added to the list. However, I have to agree with @carmocca, keep one PR for fixing one file please.

[barufa]

I would like to make my first contribution to lightning. I can take pytorch_lightning/demos/boring_classes.py and pytorch_lightning/demos/mnist_datamodule.py.

[otaj]

@barufa Thank you very much, added you to the list!

[donlapark]

I would like to contribute as well. I can take pytorch_lightning/tuner/lr_finder.py.

[otaj]

@donlapark thank you as well. Added to the list:)

[puhuk]

I will take pytorch_lightning/distributed/dist.py

[otaj]

@puhuk Thank you very much! You were added to the list!

[CyprienRicque]

Hello! I can work on:

• pytorch_lightning/__setup__.py https://github.com/Lightning-AI/lightning/pull/13472
• pytorch_lightning/strategies/single_device.py https://github.com/Lightning-AI/lightning/pull/13532
• pytorch_lightning/strategies/strategy.py https://github.com/Lightning-AI/lightning/pull/13519

[otaj]

@Cyprien-Ricque Awesome, thanks! Added your reservations to the master comment.

[gautierdag]

Hi! Just did a quick PR for pytorch_lightning/trainer/optimizers.py here: https://github.com/Lightning-AI/lightning/pull/13470

[ar90n]

Hi! I want to take the followings.

• pytorch_lightning/callbacks/finetuning.py
• pytorch_lightning/tuner/batch_size_scaling.py

[otaj]

@gautierdag , thank you! I saw it was a fairly small PR, so it doesn't matter much, but we would still prefer in the future if you reserve it as to avoid collisions with other people who might already start working on that. But anyway, thank you :)

@ar90n, you were also added to the list, thank you :)

[gautierdag]

@otaj Cheers, I'll try to take on pytorch_lightning/loggers/wandb.py today :)

update: #13483

[otaj]

@gautierdag, perfect, thank you, added to the list!

[jxtngx]

@otaj @carmocca

two errors in loggers/base.py are attributed to None types being set as the return type for deprecation warnings; should these be left alone or set as the return type for the future

[carmocca]

@JustinGoheen I would fix their return annotation. Since they are simple wrappers for a deprecation, you can just mark the inputs and outptus as Any

-def rank_zero_experiment(*args, **kwargs) -> None:  # type: ignore[no-untyped-def]
+def rank_zero_experiment(*args: Any, **kwargs: Any) -> Any:
     rank_zero_deprecation(
         "The `pytorch_lightning.loggers.base.rank_zero_experiment` is deprecated in v1.7"
         " and will be removed in v1.9. Please use `pytorch_lightning.loggers.logger.rank_zero_experiment` instead."
@@ -77,7 +78,7 @@ class DummyLogger(logger.DummyLogger):
         super().__init__(*args, **kwargs)
 
 
-def merge_dicts(*args, **kwargs) -> None:  # type: ignore[no-untyped-def]
+def merge_dicts(*args: Any, **kwargs: Any) -> Any:
     rank_zero_deprecation(
         "The `pytorch_lightning.loggers.base.merge_dicts` is deprecated in v1.7"
         " and will be removed in v1.9. Please use `pytorch_lightning.loggers.logger.merge_dicts` instead."

[CyprienRicque]

Hello! I am working on:

• pytorch_lightning/strategies/single_tpu.py https://github.com/Lightning-AI/lightning/pull/13534
• pytorch_lightning/strategies/ddp2.py https://github.com/Lightning-AI/lightning/pull/13535
• pytorch_lightning/strategies/parallel.py https://github.com/Lightning-AI/lightning/pull/13556
• pytorch_lightning/strategies/sharded_spawn.py (waiting for pytorch_lightning/strategies/parallel.py to be merged to reuse code)
• pytorch_lightning/strategies/dp.py https://github.com/Lightning-AI/lightning/pull/13564

[jiny419]

Hi! I'll try to take on pytorch_lightning/profilers/simple.py

[jxtngx]

@awaelchli I'd like to be assigned to the following:

• pytorch_lightning/loggers/tensorboard.py
• pytorch_lightning/core/module.py

[awaelchli]

Great! Updated

[jxtngx]

@awaelchli I can take core.datamodule too, can I also be assigned pytorch_lightning/core/module.py.

I'll work on core.module first, then the tensorboard logger and datamodule.

[himkt]

I want to try pytorch_lightning/trainer/trainer.py. 🚋

[awaelchli]

@himkt That's nice! Thanks for the help. I predict that the trainer will be very difficult right now, almost impossible due to it's entanglement with all other files related to it. It probably should be left for later. May I suggest one of the easier ones to get started, e.g. pytorch_lightning/loops/epoch/training_epoch_loop.py? I could be wrong, feel free to try it if you want.

[himkt]

Thank you for the advice! OK, let me try training_epoch_loop.py!

[donlapark]

Hi all! I'd like to take pytorch_lightning/tuner/tuning.py next.

[krishnakalyan3]

I would like to work on pytorch_lightning/utilities/distributed.py.

[CyprienRicque]

Hello @awaelchli, can I please be assigned to this list and these new ones if possible :blush::

• pytorch_lightning/strategies/sharded.py
• pytorch_lightning/strategies/horovod.py https://github.com/Lightning-AI/lightning/pull/13570
• pytorch_lightning/strategies/fully_sharded.py
• pytorch_lightning/strategies/tpu_spawn.py

Also this one is not assigned.