Drop the DataLoader iterator when pickling (#17130)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: carmocca

src/lightning/pytorch/CHANGELOG.md

@@ -12,6 +12,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed issue where pickling the module instance would fail with a DataLoader error ([#17130](https://github.com/Lightning-AI/lightning/pull/17130))
 - Fixed WandbLogger not showing "best" aliases for model checkpoints when `ModelCheckpoint(save_top_k>0)` is used ([#17121](https://github.com/Lightning-AI/lightning/pull/17121))
 
 

src/lightning/pytorch/utilities/combined_loader.py

@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from collections.abc import Iterable
-from typing import Any, Callable, Iterator, List, Literal, Optional, Tuple, Type, TypeVar, Union
+from typing import Any, Callable, Dict, Iterator, List, Literal, Optional, Tuple, Type, TypeVar, Union
 
-from torch.utils.data.dataloader import _MultiProcessingDataLoaderIter
+from torch.utils.data.dataloader import _BaseDataLoaderIter, _MultiProcessingDataLoaderIter
 from typing_extensions import Self, TypedDict
 
 from lightning.fabric.utilities.data import sized_len
@@ -38,6 +38,18 @@ def __iter__(self) -> Self:
     def reset(self) -> None:
         self.iterators = []
 
+    def __getstate__(self) -> Dict[str, Any]:
+        state = self.__dict__.copy()
+
+        # workaround an inconvenient `NotImplementedError`:
+        # https://github.com/pytorch/pytorch/blob/v2.0.0/torch/utils/data/dataloader.py#L652-L658
+        state["iterators"] = [
+            None if isinstance(iterator, _BaseDataLoaderIter) else iterator_state
+            for iterator, iterator_state in zip(self.iterators, state["iterators"])
+        ]
+
+        return state
+
 
 class _MaxSizeCycle(_ModeIterator[List]):
     def __init__(self, iterables: List[Iterable]) -> None:

tests/tests_pytorch/utilities/test_combined_loader.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import math
+import pickle
 from typing import Any, get_args, NamedTuple, Sequence
 
 import pytest
@@ -536,3 +537,22 @@ def test_combined_dataloader_for_training_with_ddp(use_distributed_sampler, mode
 
 def test_supported_modes():
     assert set(_SUPPORTED_MODES) == set(get_args(_LITERAL_SUPPORTED_MODES))
+
+
+def test_combined_loader_can_be_pickled():
+    dataloader = DataLoader([0, 1, 2, 3])
+
+    # sanity check that and error would be raised. if this ever changes, `_ModeIterator.__getstate__` should be updated
+    iterator = iter(dataloader)
+    with pytest.raises(NotImplementedError, match="cannot be pickled"):
+        pickle.dumps(iterator)
+
+    numbers = list(range(10))
+    cl = CombinedLoader([dataloader, numbers])
+    iter(cl)
+
+    iterator = cl._iterator
+    assert iterator.__getstate__() == {"iterables": [dataloader, numbers], "iterators": [None, iterator.iterators[1]]}
+
+    # no error
+    pickle.dumps(cl)