tqdm progress bar inside the scheduler

Author: mayinghan

eval_protocol/pytest/priority_scheduler.py

@@ -1,9 +1,12 @@
 import asyncio
 import logging
 import os
+import time
 from collections import defaultdict
 from dataclasses import dataclass, field
-from typing import Any, Callable, List, Dict, Optional, Union, Awaitable
+from typing import Any, List, Dict, Optional, Union
+
+from tqdm.asyncio import tqdm as async_tqdm
 
 from eval_protocol.models import EvaluationRow, Status
 from eval_protocol.pytest.types import RolloutProcessorConfig, TestFunction
@@ -14,7 +17,6 @@
 from eval_protocol.human_id import generate_id
 from eval_protocol.log_utils.rollout_context import rollout_logging_context
 from eval_protocol.pytest.execution import execute_pytest_with_exception_handling
-import time
 
 ENABLE_SPECULATION = os.getenv("ENABLE_SPECULATION", "0").strip() == "1"
 
@@ -80,6 +82,10 @@ def __init__(
         self.rollout_n = 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
+        
+        # Progress bars (initialized in run())
+        self.rollout_pbar: Optional[async_tqdm] = None
+        self.eval_pbar: Optional[async_tqdm] = None
 
     async def schedule_dataset(
         self,
@@ -169,9 +175,15 @@ async def _run_eval(rows_to_eval: Union[EvaluationRow, List[EvaluationRow]]):
                 for row in eval_res:
                     row.execution_metadata.eval_duration_seconds = eval_duration
                     self.results.append(row)
+                # Update eval progress bar (groupwise: 1 eval for the group)
+                if self.eval_pbar:
+                    self.eval_pbar.update(1)
             else:
                 eval_res.execution_metadata.eval_duration_seconds = eval_duration
                 self.results.append(eval_res)
+                # Update eval progress bar (pointwise: 1 eval per row)
+                if self.eval_pbar:
+                    self.eval_pbar.update(1)
             return eval_res
 
         # 1. Prepare Config & Row for this micro-batch
@@ -211,10 +223,18 @@ async def _run_eval(rows_to_eval: Union[EvaluationRow, List[EvaluationRow]]):
         batch_results: List[EvaluationRow] = []
         if current_batch_rows:
             for idx, row in current_batch_rows:
+                start_time = time.perf_counter()
                 async for result_row in rollout_processor_with_retry(
                     self.rollout_processor, [row], task.config, idx, disable_tqdm=True
                 ):
+                    rollout_duration = time.perf_counter() - start_time
+                    result_row.execution_metadata.rollout_duration_seconds = rollout_duration
                     batch_results.append(result_row)
+                    
+                    # Update rollout progress bar
+                    if self.rollout_pbar:
+                        self.rollout_pbar.update(1)
+                    
                     # in pointwise, we start evaluation immediately
                     if self.mode == "pointwise":
                         t = asyncio.create_task(_run_eval(result_row))
@@ -300,28 +320,58 @@ def _post_process_result(self, res: EvaluationRow):
     async def run(self, dataset: List[EvaluationRow], num_runs: int, base_config: RolloutProcessorConfig):
         self.num_runs = num_runs
 
-        # 1. Schedule initial tasks
-        await self.schedule_dataset(dataset, base_config)
-        
-        # 2. Start Workers
-        # If we have separate limits, we need enough workers to saturate both stages
-        num_workers = self.max_concurrent_rollouts
-
-        workers = [asyncio.create_task(self.worker()) for _ in range(num_workers)]
-        
-        # 3. Wait for completion
-        await self.queue.join()
+        # Calculate totals for progress bars
+        total_rollouts = len(dataset) * num_runs
+        # In pointwise mode: 1 eval per rollout; in groupwise mode: 1 eval per dataset row
+        total_evals = total_rollouts if self.mode == "pointwise" else len(dataset)
 
-        # Wait for background evaluations to finish
-        if self.background_tasks:
-            await asyncio.gather(*self.background_tasks, return_exceptions=True)
+        # Initialize progress bars
+        self.rollout_pbar = async_tqdm(
+            total=total_rollouts,
+            desc="🚀 Rollouts",
+            unit="row",
+            position=0,
+            leave=True,
+            colour="cyan",
+        )
+        self.eval_pbar = async_tqdm(
+            total=total_evals,
+            desc="📊 Evals",
+            unit="eval",
+            position=1,
+            leave=True,
+            colour="green",
+        )
 
-        # 4. Cleanup
-        for w in workers:
-            w.cancel()
-        
-        if workers:
-            await asyncio.gather(*workers, return_exceptions=True)
+        try:
+            # 1. Schedule initial tasks
+            await self.schedule_dataset(dataset, base_config)
+            
+            # 2. Start Workers
+            # If we have separate limits, we need enough workers to saturate both stages
+            num_workers = self.max_concurrent_rollouts
+
+            workers = [asyncio.create_task(self.worker()) for _ in range(num_workers)]
+            
+            # 3. Wait for completion
+            await self.queue.join()
+            
+            # Wait for background evaluations to finish
+            if self.background_tasks:
+                await asyncio.gather(*self.background_tasks, return_exceptions=True)
+            
+            # 4. Cleanup
+            for w in workers:
+                w.cancel()
+            
+            if workers:
+                await asyncio.gather(*workers, return_exceptions=True)
+        finally:
+            # Close progress bars
+            if self.rollout_pbar:
+                self.rollout_pbar.close()
+            if self.eval_pbar:
+                self.eval_pbar.close()
 
         # Return collected results
         return self.results