fix buffer

Author: pan-x-c

tests/buffer/sql_test.py

@@ -75,8 +75,7 @@ async def test_sql_task_buffer_read_write(self) -> None:
         )
         sql_writer = SQLWriter(config.to_storage_config())
         tasks = [
-            {"task_id": i, "raw_task": {"question": f"question_{i}", "answer": f"answer_{i}"}}
-            for i in range(total_samples)
+            {"question": f"question_{i}", "answer": f"answer_{i}"} for i in range(total_samples)
         ]
         self.assertEqual(await sql_writer.acquire(), 1)
         sql_writer.write(tasks)
@@ -89,8 +88,8 @@ async def test_sql_task_buffer_read_write(self) -> None:
         except StopIteration:
             pass
         self.assertEqual(len(read_tasks), total_samples)
-        self.assertIn("question", read_tasks[0]["raw_task"])
-        self.assertIn("answer", read_tasks[0]["raw_task"])
+        self.assertIn("question", read_tasks[0].raw_task)
+        self.assertIn("answer", read_tasks[0].raw_task)
         db_wrapper = ray.get_actor("sql-test_task_buffer")
         self.assertIsNotNone(db_wrapper)
         self.assertEqual(await sql_writer.release(), 0)

trinity/buffer/reader/file_reader.py

@@ -86,9 +86,6 @@ def select_batch(self, indices: List[int]) -> List:
 
 
 class BaseFileReader(BufferReader):
-    def __len__(self):
-        return self.dataset.dataset_size
-
     async def read_async(self, batch_size: Optional[int] = None):
         try:
             return self.read(batch_size)
@@ -123,6 +120,9 @@ def state_dict(self):
     def load_state_dict(self, state_dict):
         return self.reader.load_state_dict(state_dict)
 
+    def __len__(self):
+        return self.reader.__len__()
+
 
 class ExperienceFileReader(BaseFileReader):
     """Reader for SFT / DPO file data."""
@@ -156,6 +156,9 @@ def state_dict(self):
     def load_state_dict(self, state_dict):
         self.dataset.current_offset = state_dict["current_index"]
 
+    def __len__(self):
+        return self.dataset.dataset_size
+
 
 class TaskFileReader(BaseFileReader):
     """A Reader for task file data."""
@@ -205,3 +208,6 @@ def state_dict(self):
 
     def load_state_dict(self, state_dict):
         self.dataset.current_offset = state_dict["current_index"]
+
+    def __len__(self):
+        return self.dataset.dataset_size

trinity/buffer/task_scheduler.py

@@ -121,7 +121,7 @@ def __init__(self, explorer_state: Dict, config: Config):
         self.read_batch_size = config.buffer.batch_size
         taskset_configs = config.buffer.explorer_input.tasksets
 
-        from trinity.buffer.reader.file_reader import TaskFileReader
+        from trinity.buffer.reader.file_reader import FileReader
 
         taskset_states = explorer_state.get(
             "taskset_states", [{"current_index": 0}] * len(taskset_configs)
@@ -131,15 +131,15 @@ def __init__(self, explorer_state: Dict, config: Config):
         for taskset_config, taskset_state in zip(taskset_configs, taskset_states):
             assert not taskset_config.is_eval  # assume drop last
             taskset = get_buffer_reader(taskset_config)
-            if not isinstance(taskset, TaskFileReader):
+            if not isinstance(taskset, FileReader):
                 raise TypeError(
                     f"Taskset '{taskset_config.name}' has an unsupported type '{type(taskset).__name__}'."
-                    f"Currently, only 'TaskFileReader' is supported by TasksetScheduler."
+                    f"Currently, only 'FileReader' is supported by TasksetScheduler."
                 )
 
             # Create selector based on type specified in config (e.g., 'sequential', 'shuffle')
             selector = SELECTORS.get(taskset_config.task_selector.selector_type)(
-                taskset.dataset, taskset_config.task_selector
+                taskset.reader.dataset, taskset_config.task_selector
             )
             selector.load_state_dict(taskset_state)  # Restore any prior state