Make Multi Node sampler cycle forever (pytorch#1424)

divyanshk · web-flow · commit daafee4cb118 · 2025-01-21T13:09:32.000-08:00
* multi node sampler cycle forever

* test for test flakiness

* test for test flakiness
diff --git a/test/nodes/test_multi_node_weighted_sampler.py b/test/nodes/test_multi_node_weighted_sampler.py
@@ -30,6 +30,7 @@ def setUp(self) -> None:
             for i in range(self._num_datasets)
         }
         self.weights = {f"ds{i}": self._weights_fn(i) for i in range(self._num_datasets)}
+        self.equal_weights = {f"ds{i}": 1.0 for i in range(self._num_datasets)}
 
     def test_torchdata_nodes_imports(self) -> None:
         try:
@@ -149,6 +150,23 @@ def test_multi_node_weighted_sampler_cycle_until_all_exhausted(self) -> None:
             self.assertEqual(sorted(datasets_in_results), ["ds0", "ds1", "ds2", "ds3"])
             mixer.reset()
 
+    def test_multi_node_weighted_sampler_cycle_forever(self) -> None:
+        """Test MultiNodeWeightedSampler with stop criteria CYCLE_FOREVER"""
+        mixer = MultiNodeWeightedSampler(
+            self.datasets,
+            self.equal_weights,
+            stop_criteria=StopCriteria.CYCLE_FOREVER,
+        )
+
+        num_yielded = 0
+        _it = iter(mixer)
+        while num_yielded < 256:  # total number of samples is 4 * 10 = 40, 256 is an arbitrary larger number
+            next(_it)
+            num_yielded += 1
+
+        mixer_num_yielded = mixer.get_state()[MultiNodeWeightedSampler.NUM_YIELDED_KEY]
+        self.assertEqual(mixer_num_yielded, num_yielded)
+
     @parameterized.expand([(1, 8), (8, 32)])
     def test_multi_node_weighted_batch_sampler_set_rank_world_size(self, rank, world_size):
         """Test MultiNodeWeightedSampler with different rank and world size"""
diff --git a/torchdata/nodes/samplers/multi_node_weighted_sampler.py b/torchdata/nodes/samplers/multi_node_weighted_sampler.py
@@ -91,6 +91,7 @@ def _validate(self) -> None:
             StopCriteria.CYCLE_UNTIL_ALL_DATASETS_EXHAUSTED,
             StopCriteria.ALL_DATASETS_EXHAUSTED,
             StopCriteria.FIRST_DATASET_EXHAUSTED,
+            StopCriteria.CYCLE_FOREVER,
         ]:
             raise ValueError(
                 f"Invalid {self.stop_criteria=}. stop_criteria must be one of: CYCLE_UNTIL_ALL_DATASETS_EXHAUSTED, FIRST_DATASET_EXHAUSTED, ALL_DATASETS_EXHAUSTED"
@@ -144,6 +145,10 @@ def _get_new_weighted_sampler(self, initial_state=None):
         )
 
     def _check_for_stop_iteration(self) -> None:
+        if self.stop_criteria == StopCriteria.CYCLE_FOREVER:
+            # If StopCriteria is CYCLE_FOREVER, we should never raise StopIteration
+            return
+
         if all(self._datasets_exhausted.values()):
             # Raise StopIteration if all datasets are exhausted,
             # this covers for both ALL_DATASETS_EXHAUSTED and CYCLE_UNTIL_ALL_DATASETS_EXHAUSTED
@@ -174,14 +179,14 @@ def next(self) -> T:
                 # Mark the dataset as exhausted
                 self._datasets_exhausted[key] = True
 
-                # Based on updated _check_for_stop_iteration, check if we should raise StopIteration
+                # Based on updated _datasets_exhausted, use _check_for_stop_iteration to check if we should raise StopIteration
                 self._check_for_stop_iteration()
 
                 # If StopCriteria is ALL_DATASETS_EXHAUSTED, move to next key
                 if self.stop_criteria == StopCriteria.ALL_DATASETS_EXHAUSTED:
                     continue
 
-                # If StopCriteria is CYCLE_UNTIL_ALL_DATASETS_EXHAUSTED,
+                # If StopCriteria is CYCLE_UNTIL_ALL_DATASETS_EXHAUSTED or CYCLE_FOREVER,
                 # reset the iterator and try again
                 self.source_nodes[key].reset()
                 item = next(self.source_nodes[key])
diff --git a/torchdata/nodes/samplers/stop_criteria.py b/torchdata/nodes/samplers/stop_criteria.py
@@ -17,8 +17,12 @@ class StopCriteria:
        dataset is seen exactly once. No wraparound or restart will be performed.
 
     3) FIRST_DATASET_EXHAUSTED: Stop when the first dataset is exhausted.
+
+    4) CYCLE_FOREVER: Cycle through the datasets by reinitializing each exhausted source nodes.
+       This is useful when trainer want control over certain number of steps instead of epochs.
     """
 
     CYCLE_UNTIL_ALL_DATASETS_EXHAUSTED = "CYCLE_UNTIL_ALL_DATASETS_EXHAUSTED"
     ALL_DATASETS_EXHAUSTED = "ALL_DATASETS_EXHAUSTED"
     FIRST_DATASET_EXHAUSTED = "FIRST_DATASET_EXHAUSTED"
+    CYCLE_FOREVER = "CYCLE_FOREVER"
