feat(eval): add MRCR mode for long-context acceptance rate benchmarks #1064
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,254 @@ | ||
| """Long-context acceptance benchmarking via OpenAI's MRCR dataset. | ||
|
|
||
| Loads the ``openai/mrcr`` 2-needle dataset and buckets samples by | ||
| prompt+answer token count, using the same bin edges and ``o200k_base`` | ||
| tiktoken encoding OpenAI uses for its own leaderboard. Each bucket becomes an | ||
| Inspect AI ``Task`` that is run against the target server to drive realistic | ||
| long-context requests and measure acceptance rates across context lengths. | ||
|
|
||
| Correctness is intentionally not graded: MRCR is used here purely as a source | ||
| of long, multi-turn prompts to see how spec-decode acceptance holds up as | ||
| context length grows, not to measure retrieval accuracy. | ||
| """ | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import json | ||
| import logging | ||
| import sys | ||
| from pathlib import Path | ||
| from urllib.error import URLError | ||
| from urllib.request import Request, urlopen | ||
|
|
||
| import pandas as pd | ||
| import tiktoken | ||
| from huggingface_hub import hf_hub_download | ||
|
|
||
| from inspect_ai import Task, eval as inspect_eval | ||
| from inspect_ai.dataset import MemoryDataset, Sample | ||
| from inspect_ai.model import ChatMessageAssistant, ChatMessageSystem, ChatMessageUser | ||
| from inspect_ai.solver import generate | ||
|
|
||
| from perf_utils import CsvWriter, acceptance_csv_columns, extract_spec_decode_metrics, print_acceptance_report | ||
|
|
||
| logger = logging.getLogger("evaluate") | ||
|
|
||
| MRCR_REPO = "openai/mrcr" | ||
| MRCR_ENCODING = "o200k_base" | ||
| MRCR_SHARDS = ("2needle/2needle_0.parquet", "2needle/2needle_1.parquet") | ||
|
|
||
| # OpenAI's own bin edges, in tokens: first bin is closed on both ends, | ||
| # remaining bins are left-open/right-closed. | ||
| MRCR_BIN_EDGES = (4096, 8192, 16384, 32768, 65536, 131072, 262144, 524288, 1048576) | ||
|
|
||
| _ROLE_TO_MESSAGE = { | ||
| "system": ChatMessageSystem, | ||
| "user": ChatMessageUser, | ||
| "assistant": ChatMessageAssistant, | ||
| } | ||
|
|
||
|
|
||
| def _bucket_label(n_tokens: int) -> str | None: | ||
| lo0, hi0 = MRCR_BIN_EDGES[0], MRCR_BIN_EDGES[1] | ||
| if lo0 <= n_tokens <= hi0: | ||
| return f"{lo0}-{hi0}" | ||
| for lo, hi in zip(MRCR_BIN_EDGES[1:], MRCR_BIN_EDGES[2:]): | ||
| if lo < n_tokens <= hi: | ||
| return f"{lo}-{hi}" | ||
| return None | ||
|
|
||
|
|
||
| def _count_tokens(enc: tiktoken.Encoding, messages: list[dict], answer: str) -> int: | ||
| total = sum(len(enc.encode(m["content"])) for m in messages) | ||
| return total + len(enc.encode(answer)) | ||
|
|
||
|
|
||
| def _iter_records(): | ||
| for shard in MRCR_SHARDS: | ||
| path = hf_hub_download(MRCR_REPO, shard, repo_type="dataset") | ||
| yield from pd.read_parquet(path).to_dict("records") | ||
|
|
||
|
|
||
| def load_mrcr_buckets( | ||
| max_context: int, | ||
| max_samples_per_bucket: int | None = None, | ||
| ) -> dict[str, list[Sample]]: | ||
| """Download the MRCR 2-needle dataset and group samples into context-length buckets. | ||
|
|
||
| Buckets whose lower edge exceeds *max_context* are dropped entirely, since | ||
| those prompts would exceed a server started with a smaller max context. | ||
| """ | ||
| enc = tiktoken.get_encoding(MRCR_ENCODING) | ||
| buckets: dict[str, list[Sample]] = {} | ||
|
|
||
| for record in _iter_records(): | ||
| messages = json.loads(record["prompt"]) | ||
| n_tokens = _count_tokens(enc, messages, record["answer"]) | ||
| if n_tokens > max_context: | ||
| continue | ||
| label = _bucket_label(n_tokens) | ||
| if label is None: | ||
| continue | ||
|
|
||
| bucket = buckets.setdefault(label, []) | ||
| if max_samples_per_bucket is not None and len(bucket) >= max_samples_per_bucket: | ||
| continue | ||
|
|
||
| bucket.append( | ||
| Sample( | ||
| input=[_ROLE_TO_MESSAGE[m["role"]](content=m["content"]) for m in messages], | ||
| target="", | ||
| metadata={ | ||
| "bucket": label, | ||
| "n_tokens": n_tokens, | ||
| "random_string_to_prepend": record["random_string_to_prepend"], | ||
| }, | ||
| ) | ||
| ) | ||
|
|
||
| for label, samples in buckets.items(): | ||
| logger.info(" mrcr/2needle/%s: %d samples", label, len(samples)) | ||
|
|
||
| return dict(sorted(buckets.items(), key=lambda kv: int(kv[0].split("-")[0]))) | ||
|
|
||
|
|
||
| def build_task(samples: list[Sample]) -> Task: | ||
| return Task(dataset=MemoryDataset(samples), solver=[generate()]) | ||
|
|
||
|
|
||
| def real_token_count(root_url: str, model: str, sample: Sample) -> int: | ||
| """Tokenize *sample* through the target server's own ``/tokenize`` endpoint. | ||
|
|
||
| Unlike the o200k_base counts used for bucketing, this reflects the actual | ||
| tokenizer and chat template the server will use, and (unlike the | ||
| ``/v1/chat/completions/render`` endpoint) doesn't reject oversized prompts — | ||
| it just reports how long they really are. | ||
| """ | ||
| messages = [{"role": m.role, "content": m.content} for m in sample.input] | ||
| body = json.dumps({"model": model, "messages": messages}).encode() | ||
| req = Request( | ||
| f"{root_url.rstrip('/')}/tokenize", | ||
| data=body, | ||
| headers={"Content-Type": "application/json"}, | ||
| method="POST", | ||
| ) | ||
| with urlopen(req, timeout=180) as resp: # noqa: S310 | ||
| return json.loads(resp.read())["count"] | ||
|
|
||
|
|
||
| def warn_if_oversized( | ||
| root_url: str, | ||
| model: str, | ||
| subset: str, | ||
| samples: list[Sample], | ||
| max_model_len: int, | ||
| ) -> None: | ||
| """Log a warning for any sample whose real tokenized length exceeds *max_model_len*. | ||
|
|
||
| MRCR's own bucket labels are based on OpenAI's o200k_base tokenizer, which | ||
| can undercount relative to the model actually being served (different | ||
| vocab, plus chat-template overhead) — this checks against the real count. | ||
| """ | ||
| oversized = [] | ||
| for sample in samples: | ||
| try: | ||
| n_tokens = real_token_count(root_url, model, sample) | ||
| except (URLError, OSError, KeyError) as e: | ||
| logger.warning("[%s] Could not tokenize sample for length check: %s", subset, e) | ||
| continue | ||
| if n_tokens > max_model_len: | ||
| oversized.append(n_tokens) | ||
|
|
||
| if oversized: | ||
| logger.warning( | ||
| "[%s] %d/%d samples exceed the server's max_model_len=%d once " | ||
| "tokenized by the real model (largest: %d tokens, bucket label " | ||
| "was computed with o200k_base and may undercount)", | ||
| subset, | ||
| len(oversized), | ||
| len(samples), | ||
| max_model_len, | ||
| max(oversized), | ||
| ) | ||
|
|
||
|
|
||
| def run_mrcr( | ||
| target: str, | ||
| model_info: dict | None, | ||
| metrics_url: str, | ||
| output_dir: Path, | ||
| max_concurrency: int, | ||
| max_context: int, | ||
| max_samples_per_bucket: int, | ||
| require_metrics, | ||
| ) -> None: | ||
| """Run MRCR long-context evaluation. | ||
|
|
||
| Drives MRCR prompts through the target server bucket-by-bucket, diffing | ||
| spec-decode Prometheus counters around each bucket's run. | ||
| """ | ||
| if model_info is None: | ||
| logger.error("Could not determine served model info from %s/models", target) | ||
| sys.exit(1) | ||
| model_name = model_info["id"] | ||
| max_model_len = model_info.get("max_model_len") | ||
| logger.info( | ||
| "Server reports model=%s max_model_len=%s", model_name, max_model_len | ||
| ) | ||
| logger.info( | ||
| "Loading MRCR 2-needle dataset (max_context=%d, max_samples_per_bucket=%s)...", | ||
| max_context, | ||
| max_samples_per_bucket, | ||
| ) | ||
| buckets = load_mrcr_buckets( | ||
| max_context=max_context, | ||
| max_samples_per_bucket=max_samples_per_bucket, | ||
| ) | ||
| if not buckets: | ||
| logger.error( | ||
| "No MRCR buckets available at or below --mrcr-max-context=%d", | ||
| max_context, | ||
| ) | ||
| sys.exit(1) | ||
|
|
||
| root_url = target.rstrip("/").removesuffix("/v1") | ||
| log_dir = output_dir / "artifacts" / "mrcr_logs" | ||
| acceptance_csv: CsvWriter | None = None | ||
|
|
||
| for label, samples in buckets.items(): | ||
| subset = f"mrcr/2needle/{label}" | ||
| logger.info("[%s] Starting (%d samples)", subset, len(samples)) | ||
|
|
||
| if max_model_len is not None: | ||
| warn_if_oversized(root_url, model_name, subset, samples, max_model_len) | ||
|
Contributor
There was a problem hiding this comment. 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win 🔎 Supported by static analysis🏁 Script executed: printf '%s\n' '--- applicable repository conventions ---'
for f in /tmp/coderabbit-repo-knowledge/vllm-project-speculators-14d6fbcd/*/*.md; do
case "$f" in
*learnings*|*architecture*) continue ;;
esac
printf '\n### %s\n' "$f"
head -80 "$f"
done
printf '%s\n' '--- target structure ---'
ast-grep outline scripts/evaluate/mrcr_bench.py
printf '%s\n' '--- relevant source ---'
sed -n '1,270p' scripts/evaluate/mrcr_bench.py
printf '%s\n' '--- vLLM deployment references ---'
rg -n -S --glob '!*.lock' --glob '!*.jsonl' 'max-model-len|max_model_len|mrcr_bench|/tokenize|fail_on_error|inspect_eval' . | head -240Repository: vllm-project/speculators Length of output: 20854 🏁 Script executed: printf '%s\n' '--- evaluation server contract and invocation ---'
sed -n '60,105p' scripts/evaluate/evaluate.py
sed -n '285,320p' scripts/evaluate/evaluate.py
sed -n '460,500p' scripts/evaluate/evaluate.py
printf '%s\n' '--- declared dependency versions ---'
rg -n -S --glob 'pyproject.toml' --glob 'requirements*.txt' --glob 'uv.lock' --glob 'poetry.lock' --glob 'setup.cfg' --glob 'setup.py' \
'vllm|inspect-ai|inspect_ai'Repository: vllm-project/speculators Length of output: 4303 🌐 Web query:
💡 Result: The error message indicating that a request exceeds the maximum context length in vLLM occurs because vLLM treats the total of the input prompt tokens and the requested max_tokens (output) as a hard constraint against the model's configured maximum context window [1][2]. Unlike some implementations that automatically clamp the output to fit the remaining window, vLLM validates that the combined length does not exceed --max-model-len [1][2]. To resolve or address this error: 1. Understand the Constraint: The limit is defined by --max-model-len [3][4]. If this is not set, vLLM derives it from the model config (e.g., config.json), but the effective limit is often further constrained by the amount of KV cache memory actually available on your GPU [2][5]. 2. Adjust Server Configuration: - If the error occurs during server startup (e.g., 'max seq len is larger than the maximum number of tokens that can be stored in KV cache'), you must decrease --max-model-len to match your GPU's actual capacity [2][5]. - Alternatively, increase available memory by raising --gpu-memory-utilization, using --kv-cache-dtype fp8, or reducing --max-num-seqs [2]. - You can set --max-model-len to -1 or 'auto' to allow vLLM to automatically profile and select the maximum length that fits in the available GPU memory [3][6]. 3. Manage Client Requests: If the server is correctly configured but you are still receiving this error for specific requests, you must reduce the length of your input prompt or lower the max_tokens value in your API call [7][2]. Note that vLLM requires the total (prompt + requested output) to be within the server's effective --max-model-len [1][2]. Recent versions of vLLM have improved error reporting to provide a more specific code, 'context_length_exceeded' (often with an HTTP 400 or 413 status), to help distinguish these events from other types of bad requests [8][9]. If you observe inconsistent or confusing error messages, ensure you are running a recent version of vLLM, as there have been multiple fixes related to error surfacing and stale config values after auto-fitting [7][10]. Citations:
🏁 Script executed: printf '%s\n' '--- metric extraction contract ---'
fd -i 'perf_utils' .
rg -n -S 'def extract_spec_decode_metrics|num_drafts|num_accepted|accepted|draft' scripts/evaluate
printf '%s\n' '--- Inspect AI error handling references ---'
rg -n -S --glob '*.py' --glob '*.toml' --glob '*.txt' 'fail_on_error|SampleError|sample.*error|eval\(' scripts tests pyproject.toml requirements*.txt 2>/dev/null | head -220Repository: vllm-project/speculators Length of output: 3273 🏁 Script executed: sed -n '285,355p' scripts/evaluate/perf_utils.py
sed -n '235,275p' scripts/evaluate/evaluate.py
printf '%s\n' '--- local Inspect AI dependency metadata ---'
find . -maxdepth 3 -type f \( -name 'pyproject.toml' -o -name 'requirements*.txt' -o -name 'uv.lock' \) -printRepository: vllm-project/speculators Length of output: 4155 🏁 Script executed: cat scripts/evaluate/requirements.txt
rg -n -S 'inspect-ai|inspect_ai|vllm' pyproject.toml scripts/evaluate/requirements.txtRepository: vllm-project/speculators Length of output: 1347 🌐 Web query:
💡 Result: In Inspect AI, when using fail_on_error=False, the evaluation process is configured to suppress failure of the entire evaluation task due to individual sample errors [1][2]. By default, Inspect AI will terminate an evaluation on the first sample error [1][3]. Setting fail_on_error=False changes this behavior such that individual sample errors do not trigger an overall evaluation failure, allowing the process to continue running other samples [1][2][4]. Regarding error handling during the generation phase and solver execution: 1. Handling of Sample Errors: When a sample fails (for instance, during solver execution or model generation), the error is caught by the evaluation runner [5][2]. 2. Impact of continue_on_fail: If you have configured fail_on_error thresholds (e.g., a specific count or proportion of allowed errors), the continue_on_fail parameter determines when the termination occurs [1][2]. Setting continue_on_fail=True allows the evaluation to continue through all samples and only potentially fail at the very end if the failure threshold is ultimately met [1][6]. 3. Scoring Failed Samples: If you wish to obtain scores for samples that encountered errors rather than simply marking them as failures, you can use score_on_error=True [2][6]. Note that even if a sample is scored, it may still be counted toward the fail_on_error threshold used to mark the final evaluation log status as 'error' [2][3]. 4. Retries: You can specify --retry-on-error (or the corresponding parameter in code) to automatically retry samples a set number of times before they are considered permanent errors [4][3]. For managing solver-specific errors, developers are encouraged to use robust custom tool definitions and implement structured error handling within their custom solvers to manage runtime exceptions gracefully [7][8]. Citations:
🌐 Web query:
💡 Result: In the vLLM project, speculative decoding metrics such as vllm:spec_decode_num_drafts are cumulative counters [1]. These metrics are designed to be incremented by the scheduler based on performance statistics aggregated across requests [2]. There is no direct mechanism where a context length rejection or a "request rejected" event triggers a specific decrement or "reject" increment for these metrics; rather, the counters track successful drafting and acceptance events [2][3]. If a draft is rejected or invalidated (e.g., due to grammar constraints or asynchronous scheduling limits), it is typically handled by either excluding those tokens from the accepted token counts or by treating them as invalid in ways that are subtracted from or never added to the relevant success metrics (e.g., num_draft_tokens or num_accepted_tokens) [4]. Because these are Prometheus-style counters, they are strictly monotonically increasing [1][5]. If a request is rejected due to context length, it simply stops being processed by the scheduler, and therefore does not contribute to further increments of speculative decoding metrics for that request [2][6]. Key technical points: - The metrics vllm:spec_decode_num_drafts, vllm:spec_decode_num_draft_tokens, and vllm:spec_decode_num_accepted_tokens are cumulative [1][3]. - Logic to calculate statistics (like acceptance rate) typically handles invalid or skipped drafts by checking the state of the scheduler or the output runner, rather than by modifying these counters directly [4]. - Prometheus implementations of these counters require caution to avoid "sum-of-sums" issues; loggers often compute deltas from these cumulative values to report per-interval performance [1][3]. Citations:
Exclude samples that exceed When 🤖 Prompt for AI Agents |
||
|
|
||
| baseline = require_metrics(metrics_url) | ||
| inspect_eval( | ||
| build_task(samples), | ||
| model=f"vllm/{model_name}", | ||
| model_base_url=target, | ||
| max_connections=max_concurrency, | ||
| log_dir=str(log_dir / label.replace("-", "_")), | ||
| display="none", | ||
| fail_on_error=False, | ||
| ) | ||
| current = require_metrics(metrics_url) | ||
|
|
||
| spec = extract_spec_decode_metrics(current, baseline_metrics=baseline) | ||
| if not spec or spec.get("num_drafts", 0) <= 0: | ||
| logger.warning("[%s] No speculative decoding metrics found", subset) | ||
| continue | ||
|
|
||
| spec["subset"] = subset | ||
| print_acceptance_report(spec) | ||
| if acceptance_csv is None: | ||
| acceptance_csv = CsvWriter( | ||
| output_dir / "acceptance.csv", | ||
| ["subset"] + acceptance_csv_columns(spec), | ||
| ) | ||
| acceptance_csv.append(spec) | ||
| logger.info("[%s] Complete", subset) | ||
|
|
||
| if acceptance_csv is None: | ||
| logger.error("No acceptance metrics collected from any MRCR bucket") | ||
| sys.exit(1) | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Reject non-positive MRCR limits during argument parsing.
A value of
0or a negative value makesload_mrcr_buckets()discard every sample. The script then downloads and scans the dataset before it exits with an empty-bucket error.Use a shared positive-integer
typevalidator for both options.Proposed fix
As per path instructions, scripts must handle argument parsing robustly.
Also applies to: 489-490
🤖 Prompt for AI Agents
Source: Path instructions