Make DataSource implement __len__ to standardize the data source contract (#1518)

Co-authored-by: Nan Jiang <59716405+nanjiangwill@users.noreply.github.com>

Author: TSunny007

docs/en/get_started/customization.md

@@ -350,6 +350,9 @@ class CustomDataSource(DataSource):
 
     def load(self, rollout_id=None):
         """Load state from checkpoint"""
+
+    def __len__(self):
+        """Length of the data source. May change when samples are added/fetched."""
 ```
 
 ---

docs/zh/get_started/customization.md

@@ -350,6 +350,11 @@ class CustomDataSource(DataSource):
 
     def load(self, rollout_id=None):
         """从 ckpt 加载状态"""
+    
+    def __len__(self) -> int:
+    """
+        返回当前数据源中可用样本的数量。该数量可能会随着样本的获取或添加而变化。
+    """
 ```
 
 ---

slime/ray/rollout.py

@@ -125,7 +125,7 @@ def get_rollout_engines_and_lock(self):
 
     def get_num_rollout_per_epoch(self):
         assert self.args.rollout_global_dataset
-        return len(self.data_source.dataset) // self.args.rollout_batch_size
+        return len(self.data_source) // self.args.rollout_batch_size
 
     def generate(self, rollout_id):
         start_time = time.time()

slime/rollout/data_source.py

@@ -39,6 +39,12 @@ def load(self, rollout_id=None):
         Load the state of the data source
         """
 
+    @abc.abstractmethod
+    def __len__(self) -> int:
+        """
+        Length of the data source. May change when samples are added/fetched.
+        """
+
 
 # TODO may further refactor data-loading part later
 class RolloutDataSource(DataSource):
@@ -153,6 +159,9 @@ def load(self, rollout_id=None):
         if self.args.rollout_global_dataset and self.args.rollout_shuffle:
             self.dataset.shuffle(self.epoch_id)
 
+    def __len__(self) -> int:
+        return len(self.dataset)
+
 
 class RolloutDataSourceWithBuffer(RolloutDataSource):
     def __init__(self, args):