add

Author: mayinghan

eval_protocol/pytest/buffer.py

@@ -5,14 +5,14 @@
 
 from eval_protocol.models import EvaluationRow
 
-class MiniBatchDataBuffer:
+class MicroBatchDataBuffer:
     """
     Buffers evaluation results and writes them to disk in minibatches.
     Waits for all runs of a sample to complete before considering it ready and flush to disk.
     """
-    def __init__(self, num_runs: int, minibatch_size: int, output_path_template: str):
+    def __init__(self, num_runs: int, batch_size: int, output_path_template: str):
         self.num_runs = num_runs
-        self.minibatch_size = minibatch_size
+        self.batch_size = batch_size
         self.output_path_template = output_path_template
         self.pending_samples: Dict[str, List[EvaluationRow]] = defaultdict(list)  # row_id -> list[EvaluationRow]
         self.completed_samples_buffer: List[List[EvaluationRow]] = []  # List[List[EvaluationRow]]
@@ -37,7 +37,7 @@ async def add_result(self, row: EvaluationRow):
                 completed_rows = self.pending_samples.pop(row_id)
                 self.completed_samples_buffer.append(completed_rows)
 
-                if len(self.completed_samples_buffer) >= self.minibatch_size:
+                if len(self.completed_samples_buffer) >= self.batch_size:
                     await self._flush_unsafe()
 
     async def _flush_unsafe(self):

eval_protocol/pytest/evaluation_test.py

@@ -70,7 +70,7 @@
 from eval_protocol.log_utils.init import init_external_logging_from_env
 from eval_protocol.log_utils.rollout_context import rollout_logging_context
 from eval_protocol.utils.browser_utils import is_logs_server_running, open_browser_tab
-from eval_protocol.pytest.buffer import MiniBatchDataBuffer
+from eval_protocol.pytest.buffer import MicroBatchDataBuffer
 from ..common_utils import load_jsonl
 
 
@@ -411,25 +411,24 @@ def _log_eval_error(status: Status, rows: list[EvaluationRow] | None, passed: bo
                     )
 
                     if use_priority_scheduler:
-                        print("Using priority scheduler")
-                        minibatch_output_size = os.environ.get("EP_MINI_BATCH_OUTPUT_SIZE", None)
+                        microbatch_output_size = os.environ.get("EP_MICRO_BATCH_OUTPUT_SIZE", None)
                         output_dir = os.environ.get("EP_OUTPUT_DIR", None)
-                        if minibatch_output_size and output_dir:
-                            output_buffer = MiniBatchDataBuffer(num_runs=num_runs, minibatch_size=int(minibatch_output_size), output_path_template=os.path.join(output_dir, "buffer_{index}.jsonl"))
+                        if microbatch_output_size and output_dir:
+                            output_buffer = MicroBatchDataBuffer(num_runs=num_runs, batch_size=int(microbatch_output_size), output_path_template=os.path.join(output_dir, "buffer_{index}.jsonl"))
                         else:
                             output_buffer = None
+                        
                         priority_results = await execute_priority_rollouts(
                             dataset=data,
                             num_runs=num_runs,
-                            micro_batch_size=int(os.environ.get("EP_MICRO_BATCH_SIZE", "1")),
                             rollout_processor=rollout_processor,
                             config=config,
                             max_concurrent_rollouts=max_concurrent_rollouts,
                             active_logger=active_logger,
                             eval_executor=test_func,
                             max_concurrent_evaluations=max_concurrent_evaluations,
                             mode=mode,
-                            mini_batch_data_buffer=output_buffer,
+                            micro_batch_data_buffer=output_buffer,
                             evaluation_test_kwargs=kwargs.get("evaluation_test_kwargs") or {},
                         )
 

eval_protocol/pytest/priority_scheduler.py

@@ -8,7 +8,7 @@
 from eval_protocol.pytest.types import RolloutProcessorConfig, TestFunction
 from eval_protocol.pytest.rollout_processor import RolloutProcessor
 from eval_protocol.pytest.evaluation_test_utils import rollout_processor_with_retry, add_cost_metrics
-from eval_protocol.pytest.buffer import MiniBatchDataBuffer
+from eval_protocol.pytest.buffer import MicroBatchDataBuffer
 from eval_protocol.dataset_logger.dataset_logger import DatasetLogger
 from eval_protocol.human_id import generate_id
 from eval_protocol.log_utils.rollout_context import rollout_logging_context
@@ -49,10 +49,10 @@ def __init__(
         active_logger: DatasetLogger,
         max_concurrent_evaluations: int,
         eval_executor: TestFunction, # Callback to run evaluation
-        output_buffer: Optional[MiniBatchDataBuffer] = None,
+        output_buffer: Optional[MicroBatchDataBuffer] = None,
         rollout_n: int = 0,
         mode: str = "pointwise",
-        in_group_microbatch_size: int = 0, # for one sample, how many runs to execute at the same time
+        in_group_minibatch_size: int = 0, # for one sample, how many runs to execute at the same time
         evaluation_test_kwargs: Dict[str, Any] = {},
     ):
         self.rollout_processor = rollout_processor
@@ -77,7 +77,7 @@ def __init__(
         self.background_tasks = set() # run evaluations in the background asynchronously
 
         self.rollout_n = rollout_n
-        self.in_group_microbatch_size = in_group_microbatch_size if in_group_microbatch_size > 0 else rollout_n
+        self.in_group_minibatch_size = in_group_minibatch_size if in_group_minibatch_size > 0 else rollout_n
         self.evaluation_test_kwargs = evaluation_test_kwargs
 
     async def schedule_dataset(
@@ -91,7 +91,7 @@ async def schedule_dataset(
         for i, row in enumerate(dataset):
             # Calculate ranges for the first in-group minibatch
             batch_start = 0
-            batch_end = min(self.in_group_microbatch_size, self.rollout_n)
+            batch_end = min(self.in_group_minibatch_size, self.rollout_n)
             run_indices = list(range(batch_start, batch_end))
 
             # Initial priority: Low (1), ordered by dataset index
@@ -243,7 +243,7 @@ async def _run_eval(rows_to_eval: Union[EvaluationRow, List[EvaluationRow]]):
         next_start = last_run_idx + 1
 
         if next_start < self.rollout_n:
-            next_end = min(next_start + self.in_group_microbatch_size, self.rollout_n)
+            next_end = min(next_start + self.in_group_minibatch_size, self.rollout_n)
             next_indices = list(range(next_start, next_end))
             new_history = task.history + current_batch_history_updates
 
@@ -327,27 +327,26 @@ async def run(self, dataset: List[EvaluationRow], num_runs: int, micro_batch_siz
 async def execute_priority_rollouts(
     dataset: List[EvaluationRow],
     num_runs: int,
-    micro_batch_size: int,
     rollout_processor: RolloutProcessor,
     config: RolloutProcessorConfig,
     max_concurrent_rollouts: int,
     active_logger: DatasetLogger,
     eval_executor: TestFunction,
     max_concurrent_evaluations: int = 96,
     mode: str = "pointwise",
-    mini_batch_data_buffer: Optional[MiniBatchDataBuffer] = None,
+    micro_batch_data_buffer: Optional[MicroBatchDataBuffer] = None,
     evaluation_test_kwargs: Dict[str, Any] = {},
 ):
     scheduler = PriorityRolloutScheduler(
         rollout_processor=rollout_processor,
         max_concurrent_rollouts=max_concurrent_rollouts,
         active_logger=active_logger,
         eval_executor=eval_executor,
-        output_buffer=mini_batch_data_buffer,
+        output_buffer=micro_batch_data_buffer,
         max_concurrent_evaluations=max_concurrent_evaluations,
         rollout_n=num_runs,
         mode=mode,
-        in_group_microbatch_size=micro_batch_size,
+        in_group_minibatch_size=(num_runs // 2),
         evaluation_test_kwargs=evaluation_test_kwargs,
     )
     return await scheduler.run(dataset, num_runs, micro_batch_size, config)