update _StatefulRandomSamplerIterator

update state dict if the iterator has finished

add comment about why were updating state dict

run precommit

Author: ramanishsingh

torchdata/stateful_dataloader/sampler.py

@@ -32,7 +32,6 @@ def __next__(self) -> int:
 
             self.yielded = self.next_yielded
             self.next_yielded = None
-
         val = next(self.parent_iterator)
         self.yielded += 1
         return val
@@ -42,6 +41,9 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
         self.sampler.generator.set_state(state_dict[self._GENERATOR])
         self.next_yielded = state_dict[self._YIELDED]
 
+    def update_state_dict(self) -> None:
+        self.generator_state = self.sampler.generator.get_state()
+
     def state_dict(self) -> Dict[str, Any]:
         return {self._GENERATOR: self.generator_state, self._YIELDED: self.yielded}
 
@@ -114,6 +116,10 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
             for _ in range(self.samples_yielded):
                 next(self.sampler_iter)
 
+    def update_state_dict(self) -> None:
+        if isinstance(self.sampler_iter, Stateful) and hasattr(self.sampler_iter, "update_state_dict"):
+            self.sampler_iter.update_state_dict()
+
 
 class BatchSampler(torch.utils.data.sampler.BatchSampler):
     def __init__(self, sampler, batch_size, drop_last):

torchdata/stateful_dataloader/stateful_dataloader.py

@@ -449,6 +449,11 @@ def __next__(self):
         try:
             return super().__next__()
         except StopIteration:
+            # If we are at the end of the iteration, we want to update the state dict of _sampler_iter.
+            # because in __iter__ after self._iterator is set using self._get_iterator() [which makes self.next_iter_state = None],
+            # it is checked if self._iterator._finished is True, and if it is, self._iterator is reset with next_iter_state = None.
+            if hasattr(self._sampler_iter, "update_state_dict"):
+                self._sampler_iter.update_state_dict()
             self._finished = True
             raise
 
@@ -531,11 +536,7 @@ def load_state_dict(self, state_dict):
             self._sampler_iter = iter(self._index_sampler)
             if state_dict[_SAMPLER_ITER_STATE] is not None:
                 self._sampler_iter = try_to_deserialize(self._sampler_iter, state_dict[_SAMPLER_ITER_STATE])
-                if state_dict[_ITERATOR_FINISHED]:
-                    try:
-                        next(self._sampler_iter)
-                    except StopIteration:
-                        pass
+
         else:
             if not isinstance(
                 self._index_sampler,