fix

Author: kevssim

src/twinkle/dataloader/dataloader.py

@@ -1,5 +1,6 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 import copy
+import os
 import warnings
 from functools import partial
 from typing import Callable, Optional, Type, Union
@@ -56,6 +57,7 @@ def __init__(self,
         self._skip_samples = 0
         self._base_batch_sampler = None
         self._base_sampler = None
+        self._retry_sampler_seed = self._resolve_retry_sampler_seed()
         self._set_work_init_fn()
 
     def _set_work_init_fn(self):
@@ -65,6 +67,17 @@ def _set_work_init_fn(self):
             num_workers=num_workers,
             rank=self.device_mesh.data_rank if self.device_mesh else 0)
 
+    @staticmethod
+    def _resolve_retry_sampler_seed() -> int:
+        env_seed = os.environ.get('TWINKLE_SEED')
+        if env_seed is not None:
+            return int(env_seed)
+        try:
+            from twinkle.infra import _seed
+            return int(_seed)
+        except Exception:
+            return 42
+
     @remote_function()
     def __len__(self):
         self._lazy_init_dataloader()
@@ -145,6 +158,7 @@ def _rebuild_sampler_stack(self):
                 self._base_sampler,
                 self.dataset,
                 max_retries=self.max_retries,
+                seed=self._retry_sampler_seed,
             )
             self.dataloader.batch_sampler = DeviceMeshSampler(
                 batch_sampler,
@@ -158,4 +172,5 @@ def _rebuild_sampler_stack(self):
                 self.dataset,
                 max_retries=self.max_retries,
                 skip_samples=self._skip_samples,
+                seed=self._retry_sampler_seed,
             )

src/twinkle/dataloader/retry_sampler.py

@@ -14,11 +14,17 @@ class RetrySampler(Sampler):
         max_retries: The maximum number of retries.
     """
 
-    def __init__(self, original_sampler: Sampler, dataset: Dataset, max_retries=20, skip_samples: int = 0):
+    def __init__(self,
+                 original_sampler: Sampler,
+                 dataset: Dataset,
+                 max_retries=20,
+                 skip_samples: int = 0,
+                 seed: int = 42):
         self.original_sampler = original_sampler
         self.dataset = dataset
         self.max_retries = max_retries
         self.skip_samples = skip_samples
+        self.seed = int(seed)
 
     def __iter__(self):
         emitted = 0
@@ -48,9 +54,9 @@ def __iter__(self):
         if emitted >= target_total:
             return
 
-        for idx in np.random.RandomState().permutation(len(self.dataset)).tolist():
+        for idx in np.random.RandomState(self.seed).permutation(len(self.dataset)).tolist():
             if emitted >= target_total:
-                raise StopIteration
+                return
             for _ in range(self.max_retries):
                 try:
                     # Skip None values and raises