Fix resuming after ds.set_epoch(new_epoch) (#7451)

* fix resuming with new epoch

* more readable states

* add test

* make style

Author: lhoestq

src/datasets/commands/convert.py

@@ -167,7 +167,7 @@ def run(self):
                 output_file = os.path.join(output_dir, f_name)
                 os.makedirs(output_dir, exist_ok=True)
                 self._logger.info(f"Adding directory {output_dir}")
-                imports_to_builder_map.update({imp: output_dir for imp in tfds_imports})
+                imports_to_builder_map.update(dict.fromkeys(tfds_imports, output_dir))
             else:
                 # Utilities will be moved at the end
                 utils_files.append(output_file)

src/datasets/dataset_dict.py

@@ -931,7 +931,7 @@ def map(
         """
         self._check_values_type()
         if cache_file_names is None:
-            cache_file_names = {k: None for k in self}
+            cache_file_names = dict.fromkeys(self)
 
         dataset_dict = {}
         for split, dataset in self.items():
@@ -1051,7 +1051,7 @@ def filter(
         """
         self._check_values_type()
         if cache_file_names is None:
-            cache_file_names = {k: None for k in self}
+            cache_file_names = dict.fromkeys(self)
         return DatasetDict(
             {
                 k: dataset.filter(
@@ -1109,7 +1109,7 @@ def flatten_indices(
         """
         self._check_values_type()
         if cache_file_names is None:
-            cache_file_names = {k: None for k in self}
+            cache_file_names = dict.fromkeys(self)
         return DatasetDict(
             {
                 k: dataset.flatten_indices(
@@ -1176,7 +1176,7 @@ def sort(
         """
         self._check_values_type()
         if indices_cache_file_names is None:
-            indices_cache_file_names = {k: None for k in self}
+            indices_cache_file_names = dict.fromkeys(self)
         return DatasetDict(
             {
                 k: dataset.sort(
@@ -1254,13 +1254,13 @@ def shuffle(
             raise ValueError("Please specify seed or seeds, but not both")
         seeds = seed if seed is not None else seeds
         if seeds is None:
-            seeds = {k: None for k in self}
+            seeds = dict.fromkeys(self)
         elif not isinstance(seeds, dict):
-            seeds = {k: seeds for k in self}
+            seeds = dict.fromkeys(self, seeds)
         if generators is None:
-            generators = {k: None for k in self}
+            generators = dict.fromkeys(self)
         if indices_cache_file_names is None:
-            indices_cache_file_names = {k: None for k in self}
+            indices_cache_file_names = dict.fromkeys(self)
         return DatasetDict(
             {
                 k: dataset.shuffle(
@@ -1326,7 +1326,7 @@ def save_to_disk(
         fs, _ = url_to_fs(dataset_dict_path, **(storage_options or {}))
 
         if num_shards is None:
-            num_shards = {k: None for k in self}
+            num_shards = dict.fromkeys(self)
         elif not isinstance(num_shards, dict):
             raise ValueError(
                 "Please provide one `num_shards` per dataset in the dataset dictionary, e.g. {{'train': 128, 'test': 4}}"
@@ -1696,7 +1696,7 @@ def push_to_hub(
         ```
         """
         if num_shards is None:
-            num_shards = {k: None for k in self}
+            num_shards = dict.fromkeys(self)
         elif not isinstance(num_shards, dict):
             raise ValueError(
                 "Please provide one `num_shards` per dataset in the dataset dictionary, e.g. {{'train': 128, 'test': 4}}"

src/datasets/formatting/formatting.py

@@ -270,7 +270,7 @@ def __init__(self, pa_table: pa.Table, formatter: "Formatter"):
         self.pa_table = pa_table
         self.formatter = formatter
 
-        self.data = {key: None for key in pa_table.column_names}
+        self.data = dict.fromkeys(pa_table.column_names)
         self.keys_to_format = set(self.data.keys())
 
     def __len__(self):

src/datasets/iterable_dataset.py

@@ -212,7 +212,7 @@ def __init__(self, generate_examples_fn: Callable[..., tuple[Key, dict]], kwargs
         self.kwargs = kwargs
 
     def _init_state_dict(self) -> dict:
-        self._state_dict = {"shard_idx": 0, "shard_example_idx": 0}
+        self._state_dict = {"shard_idx": 0, "shard_example_idx": 0, "type": self.__class__.__name__}
         return self._state_dict
 
     def __iter__(self):
@@ -250,7 +250,7 @@ def __init__(
         self.generator = deepcopy(generator)
 
     def _init_state_dict(self) -> dict:
-        self._state_dict = {"shard_idx": 0, "shard_example_idx": 0}
+        self._state_dict = {"shard_idx": 0, "shard_example_idx": 0, "type": self.__class__.__name__}
         return self._state_dict
 
     def __iter__(self):
@@ -290,7 +290,7 @@ def iter_arrow(self):
         return self._iter_arrow
 
     def _init_state_dict(self) -> dict:
-        self._state_dict = {"shard_idx": 0, "shard_example_idx": 0}
+        self._state_dict = {"shard_idx": 0, "shard_example_idx": 0, "type": self.__class__.__name__}
         return self._state_dict
 
     def __iter__(self):
@@ -357,7 +357,7 @@ def __init__(
         self.generator = deepcopy(generator)
 
     def _init_state_dict(self) -> dict:
-        self._state_dict = {"shard_idx": 0, "shard_example_idx": 0}
+        self._state_dict = {"shard_idx": 0, "shard_example_idx": 0, "type": self.__class__.__name__}
         return self._state_dict
 
     def __iter__(self):
@@ -437,11 +437,12 @@ def features(self):
 
     def _init_state_dict(self) -> dict:
         self._state_dict = {
-            "ex_iterable": self.ex_iterable._init_state_dict(),
+            "examples_iterable": self.ex_iterable._init_state_dict(),
             "previous_state": None,
             "batch_idx": 0,
             "num_chunks_since_previous_state": 0,
             "cropped_chunk_length": 0,
+            "type": self.__class__.__name__,
         }
         return self._state_dict
 
@@ -680,6 +681,7 @@ def _init_state_dict(self) -> dict:
             "ex_iterables": [ex_iterable._init_state_dict() for ex_iterable in self.ex_iterables],
             "previous_states": [None] * len(self.ex_iterables),
             "is_exhausted": [False] * len(self.ex_iterables),
+            "type": self.__class__.__name__,
         }
         return self._state_dict
 
@@ -778,6 +780,7 @@ def _init_state_dict(self) -> dict:
         self._state_dict = {
             "ex_iterable_idx": 0,
             "ex_iterables": [ex_iterable._init_state_dict() for ex_iterable in self.ex_iterables],
+            "type": self.__class__.__name__,
         }
         return self._state_dict
 
@@ -858,7 +861,10 @@ def features(self):
         return self.ex_iterables[0].features
 
     def _init_state_dict(self) -> dict:
-        self._state_dict = {"ex_iterables": [ex_iterable._init_state_dict() for ex_iterable in self.ex_iterables]}
+        self._state_dict = {
+            "ex_iterables": [ex_iterable._init_state_dict() for ex_iterable in self.ex_iterables],
+            "type": self.__class__.__name__,
+        }
         return self._state_dict
 
     def __iter__(self):
@@ -960,6 +966,7 @@ def _init_state_dict(self) -> dict:
             "ex_iterables": [ex_iterable._init_state_dict() for ex_iterable in self.ex_iterables],
             "previous_states": [None] * len(self.ex_iterables),
             "is_exhausted": [False] * len(self.ex_iterables),
+            "type": self.__class__.__name__,
         }
         return self._state_dict
 
@@ -1060,10 +1067,11 @@ def features(self):
 
     def _init_state_dict(self) -> dict:
         self._state_dict = {
-            "ex_iterable": self.ex_iterable._init_state_dict(),
+            "examples_iterable": self.ex_iterable._init_state_dict(),
             "previous_state": None,
             "num_examples_since_previous_state": 0,
             "previous_state_example_idx": 0,
+            "type": self.__class__.__name__,
         }
         return self._state_dict
 
@@ -1578,7 +1586,11 @@ def features(self):
         return self.ex_iterable.features
 
     def _init_state_dict(self) -> dict:
-        self._state_dict = {"skipped": False, "ex_iterable": self.ex_iterable._init_state_dict()}
+        self._state_dict = {
+            "skipped": False,
+            "examples_iterable": self.ex_iterable._init_state_dict(),
+            "type": self.__class__.__name__,
+        }
         return self._state_dict
 
     def __iter__(self):
@@ -1642,7 +1654,8 @@ def __init__(
     def _init_state_dict(self) -> dict:
         self._state_dict = {
             "repeat_index": 0,
-            "ex_iterable": self.ex_iterable._init_state_dict(),
+            "examples_iterable": self.ex_iterable._init_state_dict(),
+            "type": self.__class__.__name__,
         }
         return self._state_dict
 
@@ -1655,7 +1668,7 @@ def __iter__(self):
             repeat_index += 1
             if self._state_dict:
                 self._state_dict["repeat_index"] = repeat_index
-                self._state_dict["ex_iterable"] = self.ex_iterable._init_state_dict()
+                self._state_dict["examples_iterable"] = self.ex_iterable._init_state_dict()
 
     def shuffle_data_sources(self, generator: np.random.Generator) -> "RepeatExamplesIterable":
         """Shuffle the underlying iterable, then repeat."""
@@ -1697,7 +1710,11 @@ def features(self):
         return self.ex_iterable.features
 
     def _init_state_dict(self) -> dict:
-        self._state_dict = {"num_taken": 0, "ex_iterable": self.ex_iterable._init_state_dict()}
+        self._state_dict = {
+            "num_taken": 0,
+            "examples_iterable": self.ex_iterable._init_state_dict(),
+            "type": self.__class__.__name__,
+        }
         return self._state_dict
 
     def __iter__(self):
@@ -1956,9 +1973,8 @@ def __init__(
         self._token_per_repo_id: dict[str, Union[str, bool, None]] = token_per_repo_id or {}
         self._epoch: Union[int, "torch.Tensor"] = _maybe_share_with_torch_persistent_workers(0)
         self._starting_state_dict: Optional[dict] = None
-        self._prepared_ex_iterable = self._prepare_ex_iterable_for_iteration()
-        self._state_dict = self._prepared_ex_iterable._init_state_dict()
-        _maybe_add_torch_iterable_dataset_parent_class(self.__class__)
+        self._prepare_ex_iterable_for_iteration()  # set state_dict
+        _maybe_add_torch_iterable_dataset_parent_class(self.__class__)  # subclass of torch IterableDataset
 
     def state_dict(self) -> dict:
         """Get the current state_dict of the dataset.
@@ -2061,7 +2077,6 @@ def load_state_dict(self, state_dict: dict) -> None:
         >>> dataloader.load_state_dict(state_dict)  # uses ds.load_state_dict() under the hood
         ```
         """
-        self._prepared_ex_iterable.load_state_dict(state_dict)
         self._starting_state_dict = state_dict
 
     def __repr__(self):
@@ -2136,9 +2151,12 @@ def _iter_pytorch(self):
             ex_iterable = ex_iterable.shard_data_sources(
                 num_shards=worker_info.num_workers, index=worker_info.id, contiguous=False
             )
-            self._state_dict = ex_iterable._init_state_dict()
-            if self._starting_state_dict:
-                ex_iterable.load_state_dict(self._starting_state_dict)
+            self._state_dict = {
+                "examples_iterable": ex_iterable._init_state_dict(),
+                "epoch": self.epoch,
+            }
+            if self._starting_state_dict and self.epoch == self._starting_state_dict["epoch"]:
+                ex_iterable.load_state_dict(self._starting_state_dict["examples_iterable"])
 
             if self._formatting and (ex_iterable.iter_arrow or self._formatting.is_table):
                 formatter = get_formatter(self._formatting.format_type, features=self.features)
@@ -2216,9 +2234,12 @@ def _prepare_ex_iterable_for_iteration(
                 token_per_repo_id=self._token_per_repo_id,
             )
 
-        self._state_dict = ex_iterable._init_state_dict()
-        if self._starting_state_dict:
-            ex_iterable.load_state_dict(self._starting_state_dict)
+        self._state_dict = {
+            "examples_iterable": ex_iterable._init_state_dict(),
+            "epoch": self.epoch,
+        }
+        if self._starting_state_dict and self.epoch == self._starting_state_dict["epoch"]:
+            ex_iterable.load_state_dict(self._starting_state_dict["examples_iterable"])
         return ex_iterable
 
     def __iter__(self):

tests/test_iterable_dataset.py

@@ -1581,6 +1581,17 @@ def test_iterable_dataset_set_epoch(dataset: IterableDataset):
     assert dataset._epoch == 42
 
 
+def test_iterable_dataset_set_epoch_resuming(dataset: IterableDataset):
+    dataset_length = len(list(dataset))
+    assert len(list(dataset)) == dataset_length > 0
+    dataset.load_state_dict(dataset.state_dict())
+    assert len(list(dataset)) == 0
+    dataset.set_epoch(1)
+    assert len(list(dataset)) == dataset_length > 0
+    dataset.load_state_dict(dataset.state_dict())
+    assert len(list(dataset)) == 0
+
+
 @pytest.mark.parametrize("seed", [None, 42, 1337])
 @pytest.mark.parametrize("epoch", [None, 0, 1, 10])
 def test_iterable_dataset_set_epoch_of_shuffled_dataset(dataset: IterableDataset, seed, epoch):