| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """MRCR long-context accuracy test. |
| 4 | |
| 5 | Usage: |
| 6 | pytest -s -v tests/evals/mrcr/test_mrcr_correctness.py \ |
| 7 | --config-list-file=configs/models-small.txt |
| 8 | """ |
| 9 | |
| 10 | import shlex |
| 11 | |
| 12 | import yaml |
| 13 | |
| 14 | from tests.utils import RemoteOpenAIServer |
| 15 | |
| 16 | from .mrcr_eval import evaluate_mrcr |
| 17 | |
| 18 | |
| 19 | def _split_host_port(url: str, default_port: int = 8000) -> tuple[str, int]: |
| 20 | if "://" in url: |
| 21 | url = url.split("://", 1)[1] |
| 22 | host_port = url.split("/", 1)[0] |
| 23 | if ":" in host_port: |
| 24 | host, p = host_port.split(":", 1) |
| 25 | return f"http://{host}", int(p) |
| 26 | return f"http://{host_port}", default_port |
| 27 | |
| 28 | |
| 29 | def test_mrcr_correctness(config_filename): |
| 30 | cfg = yaml.safe_load(config_filename.read_text(encoding="utf-8")) |
| 31 | |
| 32 | server_args = shlex.split(cfg.get("server_args", "")) |
| 33 | server_args += ["--trust-remote-code", "--disable-uvicorn-access-log"] |
| 34 | |
| 35 | print( |
| 36 | f"MRCR eval for {cfg['model_name']} (threshold {cfg['match_ratio_threshold']})" |
| 37 | ) |
| 38 | |
| 39 | with RemoteOpenAIServer( |
| 40 | cfg["model_name"], |
| 41 | server_args, |
| 42 | env_dict=cfg.get("env"), |
| 43 | max_wait_seconds=cfg.get("startup_max_wait_seconds", 600), |
| 44 | ) as server: |
| 45 | host, port = _split_host_port(server.url_for("v1")) |
| 46 | results = evaluate_mrcr( |
| 47 | model_name=cfg.get("model_name"), |
| 48 | num_samples=cfg.get("num_samples", 40), |
| 49 | needles=cfg.get("needles", [2, 4, 8]), |
| 50 | max_prompt_tokens=cfg.get("max_prompt_tokens"), |
| 51 | max_tokens=cfg.get("max_tokens", 2048), |
| 52 | host=host, |
| 53 | port=port, |
| 54 | concurrency=cfg.get("concurrency", 8), |
| 55 | extra_body=cfg.get("extra_body"), |
| 56 | ) |
| 57 | |
| 58 | threshold = cfg["match_ratio_threshold"] |
| 59 | tol = cfg.get("tolerance", 0.05) |
| 60 | |
| 61 | print(f" match_ratio: {results['match_ratio']:.4f}") |
| 62 | print(f" prefix_hit_rate: {results['prefix_hit_rate']:.4f}") |
| 63 | for k, v in results["per_needle"].items(): |
| 64 | print(f" {k}: {v:.4f}") |
| 65 | |
| 66 | failures: list[str] = [] |
| 67 | if isinstance(threshold, dict): |
| 68 | for n, expected in threshold.items(): |
| 69 | key = f"match_ratio_n{int(n)}" |
| 70 | measured = results["per_needle"].get(key) |
| 71 | if measured is None: |
| 72 | failures.append(f"{key}: no samples collected") |
| 73 | elif measured < expected - tol: |
| 74 | failures.append(f"{key}: {measured:.4f} < {expected:.4f} - {tol:.4f}") |
| 75 | else: |
| 76 | measured = results["match_ratio"] |
| 77 | if measured < threshold - tol: |
| 78 | failures.append( |
| 79 | f"match_ratio: {measured:.4f} < {threshold:.4f} - {tol:.4f}" |
| 80 | ) |
| 81 | |
| 82 | assert not failures, "MRCR thresholds failed: " + "; ".join(failures) |