Add @override for files in src/lightning/pytorch/utilities (#19315)

Author: VictorPrins

src/lightning/pytorch/utilities/combined_loader.py

@@ -16,7 +16,7 @@
 from typing import Any, Callable, Dict, Iterator, List, Literal, Optional, Tuple, Type, Union
 
 from torch.utils.data.dataloader import _BaseDataLoaderIter, _MultiProcessingDataLoaderIter
-from typing_extensions import Self, TypedDict
+from typing_extensions import Self, TypedDict, override
 
 from lightning.fabric.utilities.data import sized_len
 from lightning.pytorch.utilities._pytree import _map_and_unflatten, _tree_flatten, tree_unflatten
@@ -33,9 +33,11 @@ def __init__(self, iterables: List[Iterable], limits: Optional[List[Union[int, f
         self._idx = 0  # what would be batch_idx
         self.limits = limits
 
+    @override
     def __next__(self) -> _ITERATOR_RETURN:
         raise NotImplementedError
 
+    @override
     def __iter__(self) -> Self:
         self.iterators = [iter(iterable) for iterable in self.iterables]
         self._idx = 0
@@ -66,6 +68,7 @@ def __init__(self, iterables: List[Iterable], limits: Optional[List[Union[int, f
         super().__init__(iterables, limits)
         self._consumed: List[bool] = []
 
+    @override
     def __next__(self) -> _ITERATOR_RETURN:
         n = len(self.iterators)
         out = [None] * n  # values per iterator
@@ -83,29 +86,34 @@ def __next__(self) -> _ITERATOR_RETURN:
         self._idx += 1
         return out, index, 0
 
+    @override
     def __iter__(self) -> Self:
         super().__iter__()
         self._consumed = [False] * len(self.iterables)
         return self
 
+    @override
     def __len__(self) -> int:
         lengths = _get_iterables_lengths(self.iterables)
         if self.limits is not None:
             return max(min(length, limit) for length, limit in zip(lengths, self.limits))  # type: ignore[return-value]
         return max(lengths)  # type: ignore[return-value]
 
+    @override
     def reset(self) -> None:
         super().reset()
         self._consumed = []
 
 
 class _MinSize(_ModeIterator):
+    @override
     def __next__(self) -> _ITERATOR_RETURN:
         out = [next(it) for it in self.iterators]
         index = self._idx
         self._idx += 1
         return out, index, 0
 
+    @override
     def __len__(self) -> int:
         lengths = _get_iterables_lengths(self.iterables)
         return min(lengths + self.limits) if self.limits is not None else min(lengths)  # type: ignore[return-value]
@@ -116,6 +124,7 @@ def __init__(self, iterables: List[Iterable], limits: Optional[List[Union[int, f
         super().__init__(iterables, limits)
         self._iterator_idx = 0  # what would be dataloader_idx
 
+    @override
     def __next__(self) -> _ITERATOR_RETURN:
         n = len(self.iterables)
         if n == 0 or self._iterator_idx >= n:
@@ -138,18 +147,21 @@ def __next__(self) -> _ITERATOR_RETURN:
         self._idx += 1
         return out, index, self._iterator_idx
 
+    @override
     def __iter__(self) -> Self:
         self._iterator_idx = 0
         self._idx = 0
         self._load_current_iterator()
         return self
 
+    @override
     def __len__(self) -> int:
         lengths = _get_iterables_lengths(self.iterables)
         if self.limits is not None:
             return sum(min(length, limit) for length, limit in zip(lengths, self.limits))  # type: ignore[misc]
         return sum(lengths)  # type: ignore[arg-type]
 
+    @override
     def reset(self) -> None:
         super().reset()
         self._iterator_idx = 0
@@ -169,6 +181,7 @@ def _use_next_iterator(self) -> None:
 
 
 class _MaxSize(_ModeIterator):
+    @override
     def __next__(self) -> _ITERATOR_RETURN:
         n = len(self.iterators)
         out = [None] * n
@@ -183,6 +196,7 @@ def __next__(self) -> _ITERATOR_RETURN:
         self._idx += 1
         return out, index, 0
 
+    @override
     def __len__(self) -> int:
         lengths = _get_iterables_lengths(self.iterables)
         if self.limits is not None:
@@ -329,6 +343,7 @@ def __next__(self) -> _ITERATOR_RETURN:
         out, batch_idx, dataloader_idx = out
         return tree_unflatten(out, self._spec), batch_idx, dataloader_idx
 
+    @override
     def __iter__(self) -> Self:
         cls = _SUPPORTED_MODES[self._mode]["iterator"]
         iterator = cls(self.flattened, self._limits)

src/lightning/pytorch/utilities/migration/utils.py

@@ -21,6 +21,7 @@
 from typing import Any, Dict, List, Optional, Tuple, Type
 
 from packaging.version import Version
+from typing_extensions import override
 
 import lightning.pytorch as pl
 from lightning.fabric.utilities.enums import LightningEnum
@@ -188,6 +189,7 @@ class _RedirectingUnpickler(pickle._Unpickler):
 
     """
 
+    @override
     def find_class(self, module: str, name: str) -> Any:
         new_module = _patch_pl_to_mirror_if_necessary(module)
         # this warning won't trigger for standalone as these imports are identical

src/lightning/pytorch/utilities/model_summary/model_summary_deepspeed.py

@@ -19,6 +19,7 @@
 import torch
 from lightning_utilities.core.imports import RequirementCache
 from torch.nn import Parameter
+from typing_extensions import override
 
 from lightning.pytorch.utilities.model_summary.model_summary import (
     NOT_APPLICABLE,
@@ -36,6 +37,7 @@ def deepspeed_param_size(p: torch.nn.Parameter) -> int:
 
 class DeepSpeedLayerSummary(LayerSummary):
     @property
+    @override
     def num_parameters(self) -> int:
         """Returns the number of parameters in this module."""
         return sum(deepspeed_param_size(p) if not _is_lazy_weight_tensor(p) else 0 for p in self._module.parameters())
@@ -51,6 +53,7 @@ def partitioned_size(p: Parameter) -> int:
 
 
 class DeepSpeedSummary(ModelSummary):
+    @override
     def summarize(self) -> Dict[str, DeepSpeedLayerSummary]:  # type: ignore[override]
         summary = OrderedDict((name, DeepSpeedLayerSummary(module)) for name, module in self.named_modules)
         if self._model.example_input_array is not None:
@@ -66,10 +69,12 @@ def summarize(self) -> Dict[str, DeepSpeedLayerSummary]:  # type: ignore[overrid
         return summary
 
     @property
+    @override
     def total_parameters(self) -> int:
         return sum(deepspeed_param_size(p) if not _is_lazy_weight_tensor(p) else 0 for p in self._model.parameters())
 
     @property
+    @override
     def trainable_parameters(self) -> int:
         return sum(
             deepspeed_param_size(p) if not _is_lazy_weight_tensor(p) else 0
@@ -81,6 +86,7 @@ def trainable_parameters(self) -> int:
     def parameters_per_layer(self) -> List[int]:
         return [layer.average_shard_parameters for layer in self._layer_summary.values()]
 
+    @override
     def _get_summary_data(self) -> List[Tuple[str, List[str]]]:
         """Makes a summary listing with:
 
@@ -104,6 +110,7 @@ def _get_summary_data(self) -> List[Tuple[str, List[str]]]:
 
         return arrays
 
+    @override
     def _add_leftover_params_to_summary(self, arrays: List[Tuple[str, List[str]]], total_leftover_params: int) -> None:
         """Add summary of params not associated with module or layer to model summary."""
         super()._add_leftover_params_to_summary(arrays, total_leftover_params)