Add Simple-evals supports to maxtext runner

[JamesDeng42]

Description

This PR is an extension of vLLM Eval Framework. It added simple_evals runner to support gpt-oss model family.

• Added a dedicated simple_evals runner.
• Supports tasks: MMLU, GPQA, DROP, MGSM, GSM8K, AIME 2024/2025.
• Added OpenAI-compatible chat sampling with bounded concurrency, transient retries, and optional error continuation.
• Added GPT-OSS Harmony rendering, reasoning-effort support, and final/reasoning output separation.
• Added concurrent chat-request batching for TPU throughput.
• Added warmup, automatic concurrency selection, result reporting, and optional diagnostic reports.
• Added unit coverage for task selection, sampling, batching etc.

Tests

Added dedicated unit tests.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[vlad-karp]

why did you remove this?

[JamesDeng42]

i do not remove this, i move it into the loop in line 107. it is still semaphore bounded but the slot is released before retry backoff.

[github-actions]

🤖 Hi @entrpn, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

🤖 I'm sorry @entrpn, but I was unable to process your request. Please see the logs for more details.

[github-actions]

🤖 Hi @entrpn, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This Pull Request introduces comprehensive support for OpenAI's simple-evals framework to the MaxText runner, enabling grader-free benchmarks like MMLU, GPQA, DROP, MGSM, GSM8K, and AIME. The implementation is highly structured, including concurrent OpenAI-compatible chat sampling, request batching for TPU efficiency, robust server warmup capabilities, and extensive diagnostic/debug reporting.

🔍 General Feedback

• Excellent Architecture: The addition of the in-process _ChatBatchQueue is a standout feature, enabling concurrent requests to safely coalesce into larger TPU generation batches. This solves the serialization bottleneck and maximizes hardware throughput.
• Strong Unit Test Coverage: The PR includes high-quality unit tests covering the task validation, scoring mechanics, server managers, concurrency bounds, and HTTP retry classifiers.
• Robust Warmup Logic: Pre-compiling XLA graphs with representative shapes before running the benchmark is an excellent design choice that avoids latency spikes and compilation timeouts during evaluation.

[github-actions]

🟡 The exception handling block in the `/v1/chat/completions` endpoint catches only `(TypeError, ValueError, KeyError)`. If `tokenizer.apply_chat_template` is missing, is not callable, or raises an error from the underlying Jinja2 template rendering engine (such as `jinja2.exceptions.TemplateError`), the exception will bypass this catch. This will result in an internal HTTP 500 server error instead of a standard HTTP 400 Bad Request. Broadening the catch block to standard `Exception` ensures a more graceful and accurate error propagation to the client.

Suggested change

	sampling_params = SamplingParams(**sp_kwargs)
	except (TypeError, ValueError, KeyError) as exc:
	except Exception as exc: # pylint: disable=broad-except
	raise fastapi.HTTPException(status_code=400, detail=f"Request formatting failed: {exc}") from exc

[github-actions]

🟠 The response parsing block lacks defensive error handling. If a request is successful (HTTP 200) but the server returns a malformed response or an empty `choices` array (e.g., due to filtering or server-side edge cases), accessing `data["choices"][0]` will raise `KeyError` or `IndexError`. Because the `except` block only catches `(aiohttp.ClientError, asyncio.TimeoutError)`, any other exception will propagate out and cause the entire batch execution in `asyncio.gather` to abort. Wrapping this parsing in a nested `try-except` block to record a clean `GenerationResult` error is highly recommended.

Suggested change

	body = await resp.text()
	last_error = f"HTTP {resp.status}: {body[:200]}"
	retryable = _is_retryable_http_status(resp.status)
	else:
	data = await resp.json()
	latency = time.monotonic() - t0
	choice = data["choices"][0]
	usage = data.get("usage", {})
	return GenerationResult(
	else:
	try:
	data = await resp.json()
	latency = time.monotonic() - t0
	choice = data["choices"][0]
	usage = data.get("usage", {})
	return GenerationResult(
	text=choice.get("text", ""),
	prompt_tokens=usage.get("prompt_tokens", 0),
	completion_tokens=usage.get("completion_tokens", 0),
	latency_s=latency,
	)
	except (KeyError, IndexError, ValueError, TypeError) as exc:
	last_error = f"Malformed response: {type(exc).__name__}: {exc}"
	retryable = False

[JamesDeng42]

let's use this fail fast design since abnormal failure case is worth checking.

[github-actions]

🟡 Hardcoding an exact count assertion of `1319` on the canonical gsm8k test dataset split is extremely fragile. If the upstream repository in Hugging Face is updated (e.g. to fix formatting issues, remove broken entries, or add newer test splits), this strict count assertion will raise a fatal `ValueError` and crash the benchmark runner. Consider logging a warning instead of aborting when the count is different, while still ensuring the dataset is not completely empty.

Suggested change

	if len(examples) != 1319:
	raise ValueError(f"Expected 1319 GSM8K main/test examples, found {len(examples)}.")
	if len(examples) != 1319:
	logger.warning("Expected 1319 GSM8K main/test examples, but found %d.", len(examples))
	if not examples:
	raise ValueError("GSM8K dataset split is empty.")

[JamesDeng42]

let's use this fail fast design since this is worth checking if hf modify the dataset.

[github-actions]

🟡 Directly downloading Parquet files from Hugging Face URLs via `pandas.read_parquet(url)` on every execution of `AIMEEval` lacks local disk caching. This leads to redundant network traffic, slower benchmark startups, and complete failures if Hugging Face is temporarily unavailable or if running in an offline/air-gapped environment. For improved robustness, consider using `huggingface_hub.hf_hub_download` to leverage a local cache directory, similar to the pattern implemented in `gsm8k_eval.py`.

Suggested change

	def __init__(self, year: int, num_examples: int | None = None, n_repeats: int = 1):
	if year not in _DATASET_BY_YEAR:
	raise ValueError(f"Unsupported AIME year: {year}. Supported: {sorted(_DATASET_BY_YEAR)}.")
	url, problem_col, answer_col = _DATASET_BY_YEAR[year]
	df = pandas.read_parquet(url)
	def __init__(self, year: int, num_examples: int | None = None, n_repeats: int = 1):
	if year not in _DATASET_BY_YEAR:
	raise ValueError(f"Unsupported AIME year: {year}. Supported: {sorted(_DATASET_BY_YEAR)}.")
	url, problem_col, answer_col = _DATASET_BY_YEAR[year]
	try:
	from huggingface_hub import hf_hub_download
	# Map raw URL to hf_hub_download parameters
	if "AIME_2024" in url:
	parquet_path = hf_hub_download(repo_id="Maxwell-Jia/AIME_2024", repo_type="dataset", filename="default/train/0000.parquet")
	else:
	parquet_path = hf_hub_download(repo_id="math-ai/aime25", repo_type="dataset", filename="default/test/0000.parquet")
	except Exception as exc:
	raise RuntimeError(f"Could not download AIME dataset: {exc}") from exc
	df = pandas.read_parquet(parquet_path)