From 18895789edb04c2b4e52d1c0b0d370c9873d3fcb Mon Sep 17 00:00:00 2001 From: Finbarr Timbers Date: Fri, 15 May 2026 13:54:47 -0600 Subject: [PATCH] Report cumulative vLLM generation progress across batches Co-Authored-By: Claude Opus 4.7 --- src/olmo_eval/common/beaker_status.py | 47 +++++++++++++++++++ .../inference/providers/vllm_server.py | 25 +++++++++- src/olmo_eval/runners/asynq/runner.py | 13 +++++ 3 files changed, 84 insertions(+), 1 deletion(-) diff --git a/src/olmo_eval/common/beaker_status.py b/src/olmo_eval/common/beaker_status.py index 2739cace5..1dc68cd20 100644 --- a/src/olmo_eval/common/beaker_status.py +++ b/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 diff --git a/src/olmo_eval/inference/providers/vllm_server.py b/src/olmo_eval/inference/providers/vllm_server.py index fa99bba42..c24083c81 100644 --- a/src/olmo_eval/inference/providers/vllm_server.py +++ b/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] diff --git a/src/olmo_eval/runners/asynq/runner.py b/src/olmo_eval/runners/asynq/runner.py index c821de863..8393ce6e6 100644 --- a/src/olmo_eval/runners/asynq/runner.py +++ b/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)