Report cumulative vLLM generation progress across batches

src/olmo_eval/common/beaker_status.py

@@ -37,6 +37,7 @@ def __init__(self, min_interval: float = DEFAULT_MIN_INTERVAL) -> None:
         self._git_suffix = _git_suffix()
         self._workload: BeakerWorkload | None = None
         self._last_update: float = float("-inf")
+        self._accumulators: dict[str, dict[str, float]] = {}
         try:
             self._client: Beaker | None = Beaker.from_env()
         except BeakerConfigurationError:
@@ -84,3 +85,49 @@ def _cb(count: int, total: int, *, force: bool = False) -> None:
                 logger.exception("BeakerStatusReporter progress callback failed")
 
         return _cb
+
+    def start_accumulator(self, label: str, total: int, units: str = "gen/sec") -> None:
+        """Register a long-lived progress accumulator for ``label``.
+
+        Subsequent calls to ``accumulating_callback(label, ...)`` add to the
+        shared total instead of reporting per-batch counts.
+        """
+        self._accumulators[label] = {
+            "completed": 0,
+            "total": total,
+            "start": time.monotonic(),
+            "units": units,
+        }
+
+    def accumulating_callback(self, label: str, per_item: int = 1) -> Callable[..., None]:
+        """Return a callback that accumulates into the ``label`` accumulator.
+
+        Each call converts the per-batch ``(count, total)`` into a delta of
+        ``per_item`` units per completed item and adds it to the registered
+        total. Reports cumulative progress.
+        """
+        last_count = [0]
+
+        def _cb(count: int, total: int, *, force: bool = False) -> None:
+            if self._client is None or label not in self._accumulators:
+                return
+            try:
+                state = self._accumulators[label]
+                delta = (count - last_count[0]) * per_item
+                last_count[0] = count
+                state["completed"] += delta
+                completed = int(state["completed"])
+                total_units = int(state["total"])
+                elapsed = max(time.monotonic() - state["start"], 1e-9)
+                rate = completed / elapsed
+                pct = (completed / total_units * 100) if total_units > 0 else 0.0
+                units = state["units"]
+                msg = f"{label} {completed}/{total_units} ({pct:.0f}%) at {rate:.4f} {units}"
+                self.update(msg, force=force or completed >= total_units)
+            except Exception:
+                logger.exception("BeakerStatusReporter accumulating callback failed")
+
+        return _cb
+
+    def has_accumulator(self, label: str) -> bool:
+        return label in self._accumulators

src/olmo_eval/inference/providers/vllm_server.py

@@ -247,6 +247,21 @@ def __init__(
         """
         super().__init__(model_name)
         self._beaker_reporter = BeakerStatusReporter()
+        benchmark = os.environ.get("OLMO_EVAL_BENCHMARK", "").strip()
+        model_basename = os.path.basename(model_name.rstrip("/")) if model_name else ""
+        label_parts = [p for p in (benchmark, model_basename) if p]
+        self._progress_label = " | ".join(label_parts) or "vLLM gen"
+        # If the runner exported a total generation count, configure an
+        # accumulating progress tracker so we report cumulative progress
+        # across batches instead of per-batch request counts.
+        total_gens_env = os.environ.get("OLMO_VLLM_GEN_TOTAL")
+        if total_gens_env:
+            try:
+                total_gens = int(total_gens_env)
+                if total_gens > 0:
+                    self._beaker_reporter.start_accumulator(self._progress_label, total_gens)
+            except ValueError:
+                pass
         self.timeout = timeout
         self.max_concurrency = max_concurrency
         self.max_retries = max_retries
@@ -856,12 +871,20 @@ async def agenerate(
         async def process(req: LMRequest) -> list[LMOutput]:
             return await self._generate_single_async(req, params)
 
+        if self._beaker_reporter.has_accumulator(self._progress_label):
+            on_progress = self._beaker_reporter.accumulating_callback(
+                self._progress_label, per_item=params.num_samples or 1
+            )
+        else:
+            on_progress = self._beaker_reporter.progress_callback(
+                self._progress_label, units="req/sec"
+            )
         results = await dispatch_concurrent(
             requests,
             process,
             max_in_flight=self.max_concurrency,
             max_retries=self.max_retries,
-            on_progress=self._beaker_reporter.progress_callback("vLLM gen", units="req/sec"),
+            on_progress=on_progress,
         )
         return [r if r is not None else [] for r in results]
 

src/olmo_eval/runners/asynq/runner.py

@@ -5,6 +5,7 @@
 import asyncio
 import contextlib
 import multiprocessing as mp
+import os
 import random
 import time
 from collections.abc import Mapping, Sequence
@@ -585,6 +586,18 @@ async def run_async(self) -> dict[str, Any]:
         total_instances = len(items)
         runner_logger.info(f"Total instances: {total_instances}")
 
+        total_generations = 0
+        for tracker in trackers.values():
+            if tracker.task is None:
+                continue
+            tcfg = tracker.task.config
+            ns = tcfg.sampling_params.num_samples if tcfg.sampling_params else 1
+            total_generations += tracker.total_instances * (ns or 1)
+        if total_generations > 0:
+            os.environ["OLMO_VLLM_GEN_TOTAL"] = str(total_generations)
+        if self.experiment_name:
+            os.environ["OLMO_EVAL_BENCHMARK"] = self.experiment_name
+
         # Update harness metrics config with experiment metadata before starting workers
         self._update_metrics_config(experiment_id)