From a0cb540fec476f5419ac342715e1f6b66a55c8c6 Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Thu, 30 Jul 2026 16:42:28 -0700 Subject: [PATCH 01/26] docs: add AIP-1094 codegen grade concurrency implementation plan Signed-off-by: Elias Bermudez --- .../2026-07-29-codegen-grade-concurrency.md | 930 ++++++++++++++++++ 1 file changed, 930 insertions(+) create mode 100644 docs/superpowers/plans/2026-07-29-codegen-grade-concurrency.md diff --git a/docs/superpowers/plans/2026-07-29-codegen-grade-concurrency.md b/docs/superpowers/plans/2026-07-29-codegen-grade-concurrency.md new file mode 100644 index 0000000000..2249c403b0 --- /dev/null +++ b/docs/superpowers/plans/2026-07-29-codegen-grade-concurrency.md @@ -0,0 +1,930 @@ +# Codegen Grade Concurrency Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Allow N concurrent `grade_codegen()` calls to complete in ~max(individual) time by replacing the serializing `asyncio.Lock` with an `id → Future` demux table on the client and a batch-drain loop on the worker. + +**Architecture:** The client drops its `asyncio.Lock` and instead multiplexes concurrent requests over the same stdin pipe using request ids; a persistent reader task demuxes responses back to individual `asyncio.Future` objects. The worker reads the first blocking request, non-blocking drains any queued requests, then calls `codegen_metrics` once with all batched samples so lighteval's `ProcessPoolExecutor` handles all problems in parallel. + +**Tech Stack:** Python 3.11+ asyncio, `orjson`, `lighteval` (`codegen_metrics`, `compute_metrics_from_results`), `select` (POSIX non-blocking stdin drain), `pytest-asyncio` + +## Global Constraints + +- Python 3.11+; use `asyncio.get_running_loop()`, not `asyncio.get_event_loop()` +- No new threads in the worker (stays single-threaded at fork) +- All existing tests in `tests/unit/accuracy/test_codegen_worker.py` and `test_codegen_worker_client.py` must remain green after each task +- `_handle_fault` must remain idempotent (called from both reader task and caller) +- `ruff format . && ruff check --fix .` must pass after each commit +- `pre-commit run --all-files` must pass before each commit +- Every new function needs a type hint on all parameters and return value +- Every new Pydantic field needs `Field(description=...)` (not applicable here — no new models) + +--- + +### Task 1: Worker — batch-drain loop and `handle_batch` + +**Spec:** `docs/superpowers/specs/2026-07-29-codegen-grade-concurrency-design.md` § "Worker changes" + +**Files:** +- Modify: `src/aiperf/accuracy/graders/_codegen_worker.py` +- Modify: `tests/unit/accuracy/test_codegen_worker.py` + +**Interfaces:** +- Produces: `handle_batch(reqs, codegen_fn, compute_metrics_fn)` — takes a list of raw request dicts, returns a list of JSONL-ready response dicts (one per input, same order) +- Produces: `run_worker_loop(stdin, out, codegen_fn, compute_metrics_fn)` — updated signature (adds `compute_metrics_fn` param) +- Removes: `handle_request` (dead code after this task; its tests are migrated to `handle_batch`) + +--- + +- [ ] **Step 1: Write failing tests for `handle_batch`** + +Add a new `TestHandleBatch` class in `tests/unit/accuracy/test_codegen_worker.py`. Place it after the existing `TestHandleRequest` class. + +The mock `codegen_fn` for `handle_batch` must match the new signature that also accepts `compute_metrics_fn`. But `codegen_fn` itself is still the original `(samples, generations, ...) -> (metrics, results)` signature. The `compute_metrics_fn` is a separate argument to `handle_batch`. + +Add at the top of the file alongside the existing fakes: + +```python +def _fake_compute_metrics(results: dict, k_list: list[int] | None = None) -> dict[str, Any]: + # Mirrors compute_metrics_from_results: returns {"pass@1": } using + # the single-problem results dict {0: [[True, True, ...]]} passed by handle_batch. + result_list = results.get(0, [[-2]]) # [-2] = compile error + if result_list and all(x > 0 for x in result_list[0]): + return {"pass@1": 1.0} + return {"pass@1": 0.0} + + +def _fake_codegen_batch_ok( + samples: list, generations: list, **_kwargs: Any +) -> tuple[dict[str, Any], dict[int, list]]: + # Returns aggregate metrics (ignored by handle_batch) and per-problem results. + n = len(samples) + raw_results = {i: [[True]] for i in range(n)} # all pass + return {"pass@1": 1.0}, raw_results + + +def _fake_codegen_batch_boom( + samples: list, generations: list, **_kwargs: Any +) -> tuple[dict[str, Any], dict[int, list]]: + raise RuntimeError("pool exploded") +``` + +Then add `TestHandleBatch`: + +```python +class TestHandleBatch: + def _req(self, req_id: int) -> dict[str, Any]: + return { + "id": req_id, + "evaluation_sample": [{"input_output": "{}"}], + "generated_code": [["x"]], + } + + def test_single_request_returns_one_ok_response(self) -> None: + resps = worker.handle_batch( + [self._req(1)], _fake_codegen_batch_ok, _fake_compute_metrics + ) + assert len(resps) == 1 + assert resps[0] == {"id": 1, "ok": True, "metrics": {"pass@1": 1.0}} + + def test_batch_of_n_calls_codegen_fn_once(self) -> None: + call_count = 0 + + def counting_codegen(samples, generations, **kwargs): + nonlocal call_count + call_count += 1 + n = len(samples) + return {"pass@1": 1.0}, {i: [[True]] for i in range(n)} + + reqs = [self._req(i) for i in range(1, 5)] + resps = worker.handle_batch(reqs, counting_codegen, _fake_compute_metrics) + assert call_count == 1 + assert len(resps) == 4 + assert all(r["ok"] for r in resps) + assert [r["id"] for r in resps] == [1, 2, 3, 4] + + def test_response_order_matches_request_order(self) -> None: + reqs = [self._req(i) for i in [7, 3, 99]] + resps = worker.handle_batch(reqs, _fake_codegen_batch_ok, _fake_compute_metrics) + assert [r["id"] for r in resps] == [7, 3, 99] + + def test_batch_exception_returns_error_for_all(self) -> None: + reqs = [self._req(i) for i in range(1, 4)] + resps = worker.handle_batch(reqs, _fake_codegen_batch_boom, _fake_compute_metrics) + assert len(resps) == 3 + assert all(not r["ok"] for r in resps) + assert all("pool exploded" in r["error"] for r in resps) + + def test_malformed_request_in_batch_does_not_affect_others(self) -> None: + reqs = [ + self._req(1), + {"id": 2}, # missing evaluation_sample + generated_code + self._req(3), + ] + resps = worker.handle_batch(reqs, _fake_codegen_batch_ok, _fake_compute_metrics) + assert len(resps) == 3 + assert resps[0] == {"id": 1, "ok": True, "metrics": {"pass@1": 1.0}} + assert resps[1]["id"] == 2 + assert not resps[1]["ok"] + assert resps[2] == {"id": 3, "ok": True, "metrics": {"pass@1": 1.0}} + + def test_non_object_request_in_batch_is_error(self) -> None: + reqs = [[1, 2, 3], self._req(5)] + resps = worker.handle_batch(reqs, _fake_codegen_batch_ok, _fake_compute_metrics) + assert len(resps) == 2 + assert resps[0]["id"] is None + assert not resps[0]["ok"] + assert resps[1] == {"id": 5, "ok": True, "metrics": {"pass@1": 1.0}} + + def test_parse_error_sentinel_produces_error_response(self) -> None: + # run_worker_loop encodes JSON decode errors as {"_parse_error": "..."}. + reqs = [{"_parse_error": "unexpected token"}, self._req(2)] + resps = worker.handle_batch(reqs, _fake_codegen_batch_ok, _fake_compute_metrics) + assert len(resps) == 2 + assert resps[0]["id"] is None + assert not resps[0]["ok"] + assert "unexpected token" in resps[0]["error"] + assert resps[1]["ok"] +``` + +Also add a test for the new `run_worker_loop` batch-drain behaviour. Add to a new class `TestRunWorkerLoopBatch` after `TestHandleBatch`: + +```python +class TestRunWorkerLoopBatch: + def _run( + self, + payloads: list[dict[str, Any]], + codegen_fn=_fake_codegen_batch_ok, + compute_metrics_fn=_fake_compute_metrics, + ) -> list[dict[str, Any]]: + # Write all payloads to a BytesIO pipe so they are already queued when + # run_worker_loop reads; this exercises the non-blocking drain path. + data = b"".join(orjson.dumps(p) + b"\n" for p in payloads) + stdin = io.BytesIO(data) + out = io.BytesIO() + worker.run_worker_loop(stdin, out, codegen_fn, compute_metrics_fn) + out.seek(0) + return [orjson.loads(line) for line in out if line.strip()] + + def _req(self, req_id: int) -> dict[str, Any]: + return { + "id": req_id, + "evaluation_sample": [{"input_output": "{}"}], + "generated_code": [["x"]], + } + + def test_pre_queued_requests_are_batched_in_one_call(self) -> None: + call_count = 0 + + def counting_codegen(samples, generations, **kwargs): + nonlocal call_count + call_count += 1 + n = len(samples) + return {"pass@1": 1.0}, {i: [[True]] for i in range(n)} + + reqs = [self._req(i) for i in range(1, 4)] + resps = self._run(reqs, counting_codegen) + assert call_count == 1 + assert len(resps) == 3 + + def test_responses_carry_correct_ids(self) -> None: + reqs = [self._req(i) for i in [10, 20, 30]] + resps = self._run(reqs) + assert {r["id"] for r in resps} == {10, 20, 30} +``` + +- [ ] **Step 2: Run the new tests to verify they fail** + +```bash +uv run pytest tests/unit/accuracy/test_codegen_worker.py::TestHandleBatch tests/unit/accuracy/test_codegen_worker.py::TestRunWorkerLoopBatch -v 2>&1 | head -40 +``` + +Expected: `AttributeError: module ... has no attribute 'handle_batch'` or `TypeError` from wrong arg count on `run_worker_loop`. + +- [ ] **Step 3: Implement `handle_batch` and update `run_worker_loop` in `_codegen_worker.py`** + +Open `src/aiperf/accuracy/graders/_codegen_worker.py`. + +**3a — Add `import select` at the top of the file** (after the stdlib imports block). + +**3b — Add `handle_batch` after the existing `_is_number` function:** + +```python +def handle_batch( + reqs: list[Any], + codegen_fn: Callable[..., tuple[dict[str, Any], Any]], + compute_metrics_fn: Callable[..., dict[str, Any]], +) -> list[dict[str, Any]]: + """Grade a batch of requests with a single codegen_fn call. + + Calls codegen_fn once with all well-formed requests batched together so + lighteval's ProcessPoolExecutor can process multiple problems in parallel. + Never raises: all failures become error responses so a bad batch cannot + kill the worker loop. + """ + all_samples: list[Any] = [] + all_generations: list[Any] = [] + id_map: list[tuple[int, Any]] = [] # (batch_position, req_id) + responses: list[dict[str, Any] | None] = [None] * len(reqs) + + for i, req in enumerate(reqs): + if isinstance(req, dict) and "_parse_error" in req: + responses[i] = { + "id": None, + "ok": False, + "error": f"bad json: {req['_parse_error']}", + } + continue + if not isinstance(req, dict): + responses[i] = { + "id": None, + "ok": False, + "error": "malformed request: expected object", + } + continue + req_id = req.get("id") + try: + all_samples.append(req["evaluation_sample"]) + all_generations.append(req["generated_code"]) + id_map.append((i, req_id)) + except (KeyError, TypeError) as exc: + responses[i] = { + "id": req_id, + "ok": False, + "error": f"malformed request: {exc!r}", + } + + if all_samples: + batch_error: str | None = None + raw_results: dict[int, Any] = {} + try: + _, raw_results = codegen_fn( + all_samples, + all_generations, + k_list=list(_LCB_PASS_AT_K), + num_process_evaluate=_LCB_NUM_PROCESSES, + ) + except Exception as exc: + batch_error = _truncate_error(f"{type(exc).__name__}: {exc}") + + for pos, (req_idx, req_id) in enumerate(id_map): + if batch_error is not None: + responses[req_idx] = {"id": req_id, "ok": False, "error": batch_error} + else: + try: + metrics = compute_metrics_fn( + {0: raw_results[pos]}, + k_list=list(_LCB_PASS_AT_K), + ) + responses[req_idx] = { + "id": req_id, + "ok": True, + "metrics": _coerce_metrics(metrics), + } + except Exception as exc: + responses[req_idx] = { + "id": req_id, + "ok": False, + "error": _truncate_error(f"{type(exc).__name__}: {exc}"), + } + + return [r for r in responses if r is not None] +``` + +**3c — Replace `run_worker_loop`:** + +```python +def run_worker_loop( + stdin: BinaryIO, + out: BinaryIO, + codegen_fn: Callable[..., tuple[dict[str, Any], Any]], + compute_metrics_fn: Callable[..., dict[str, Any]], +) -> None: + """Serve JSONL grading requests until stdin EOF. + + Blocks on the first request of each cycle, then non-blocking drains any + already-queued requests to form a batch. Calls codegen_fn once per batch so + lighteval's ProcessPoolExecutor can process multiple problems in parallel. + """ + import select + + stdin_fd = stdin.fileno() + while True: + first = stdin.readline() + if not first: + break # EOF: client closed stdin, clean exit + first = first.strip() + if not first: + continue + batch_raw: list[bytes] = [first] + + while True: + ready, _, _ = select.select([stdin_fd], [], [], 0) + if not ready: + break + line = stdin.readline() + if not line: + break + line = line.strip() + if line: + batch_raw.append(line) + + reqs: list[Any] = [] + for raw in batch_raw: + try: + reqs.append(orjson.loads(raw)) + except orjson.JSONDecodeError as exc: + reqs.append({"_parse_error": str(exc)}) + + for resp in handle_batch(reqs, codegen_fn, compute_metrics_fn): + out.write(orjson.dumps(resp) + b"\n") + out.flush() +``` + +**3d — Remove `handle_request`** (the whole function and its docstring). Do not replace it with a comment. + +**3e — Update `main()` to import and pass `compute_metrics_from_results`:** + +```python +def main() -> None: + protocol_out = _install_stdout_guard() + _start_death_watcher() + _force_fork() + from lighteval.tasks.tasks.lcb.codegen_metrics import ( + codegen_metrics, + compute_metrics_from_results, + ) + + run_worker_loop(sys.stdin.buffer, protocol_out, codegen_metrics, compute_metrics_from_results) +``` + +- [ ] **Step 4: Migrate `TestHandleRequest` tests to `TestHandleBatch` equivalents** + +The `TestHandleRequest` class in `test_codegen_worker.py` is now orphaned (`handle_request` was removed). Replace it with `TestHandleBatch` (already written in Step 1 above — just remove the old class). Also remove the `_fake_codegen_ok`, `_fake_codegen_boom`, `_fake_codegen_list_pass` helpers if they are only used by the old `TestHandleRequest`; they are replaced by the new batch-aware fakes from Step 1. + +Check if any other test in the file (e.g., `TestRunWorkerLoop`, `TestStdoutGuard`) still calls `handle_request` directly; update those to use `handle_batch` with the batch-aware fakes. + +Grep to find remaining usages: + +```bash +grep -n "handle_request\|_fake_codegen_ok\|_fake_codegen_boom\|_fake_codegen_list_pass" \ + tests/unit/accuracy/test_codegen_worker.py +``` + +For any `TestRunWorkerLoop` tests that currently call `run_worker_loop` with 3 args, update to pass `_fake_compute_metrics` as the 4th argument. + +- [ ] **Step 5: Run all unit tests for the worker to verify they pass** + +```bash +uv run pytest tests/unit/accuracy/test_codegen_worker.py -v +``` + +Expected: all tests pass, including the new `TestHandleBatch` and `TestRunWorkerLoopBatch` classes. + +- [ ] **Step 6: Lint** + +```bash +ruff format . && ruff check --fix . +``` + +- [ ] **Step 7: Commit** + +```bash +git add src/aiperf/accuracy/graders/_codegen_worker.py \ + tests/unit/accuracy/test_codegen_worker.py +git commit -s -m "feat(accuracy): batch-drain worker loop for codegen grading concurrency" +``` + +--- + +### Task 2: Client — drop lock, add demux table and reader task + +**Spec:** `docs/superpowers/specs/2026-07-29-codegen-grade-concurrency-design.md` § "Client changes" + +**Files:** +- Modify: `src/aiperf/accuracy/graders/_codegen_worker_client.py` +- Modify: `tests/unit/accuracy/test_codegen_worker_client.py` + +**Interfaces:** +- Consumes: worker protocol from Task 1 (JSONL responses carry `id`, `ok`, `metrics`) +- Produces: `CodegenGradingWorker` with the same public API (`grade_codegen`, `aclose`) but concurrent-safe without a global lock + +--- + +- [ ] **Step 1: Write failing concurrency tests** + +Open `tests/unit/accuracy/test_codegen_worker_client.py`. + +Add the following mock worker scripts near the top of the file alongside `_ECHO_OK`: + +```python +# Echoes responses with pass@1 == id * 0.1 so each caller can verify it got +# back its OWN response (not another caller's). +_ECHO_ID_IN_METRICS = """ + import sys, orjson + for line in sys.stdin.buffer: + line = line.strip() + if not line: + continue + req = orjson.loads(line) + resp = {"id": req["id"], "ok": True, "metrics": {"pass@1": req["id"] * 0.1}} + sys.stdout.buffer.write(orjson.dumps(resp) + b"\\n") + sys.stdout.buffer.flush() +""" + +# Buffers the first 4 requests and responds in REVERSE id order to exercise +# the demux table (correct demux requires id matching, not position matching). +_REVERSE_BATCH_OF_4 = """ + import sys, orjson + buf = [] + for line in sys.stdin.buffer: + line = line.strip() + if not line: + continue + req = orjson.loads(line) + buf.append(req) + if len(buf) == 4: + for r in reversed(buf): + resp = {"id": r["id"], "ok": True, "metrics": {"pass@1": r["id"] * 0.1}} + sys.stdout.buffer.write(orjson.dumps(resp) + b"\\n") + sys.stdout.buffer.flush() + buf = [] +""" +``` + +Add a new `TestConcurrency` class: + +```python +class TestConcurrency: + async def test_concurrent_grades_all_complete(self, tmp_path) -> None: + w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, _ECHO_OK)) + try: + results = await asyncio.gather(*[ + w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=30) + for _ in range(5) + ]) + assert all(r == {"pass@1": 1.0} for r in results) + finally: + await w.aclose() + + async def test_concurrent_grades_demux_by_id_not_position(self, tmp_path) -> None: + # 4 concurrent grades; mock responds in reverse order. + # If demux were position-based, callers would get wrong metrics. + w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, _REVERSE_BATCH_OF_4)) + try: + results = await asyncio.gather(*[ + w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=30) + for _ in range(4) + ]) + # IDs 1-4 → pass@1 values 0.1, 0.2, 0.3, 0.4 (one per caller) + values = sorted(r["pass@1"] for r in results) + assert values == pytest.approx([0.1, 0.2, 0.3, 0.4]) + finally: + await w.aclose() + + async def test_fault_cancels_all_pending_futures(self, tmp_path) -> None: + # Worker dies immediately after the first line — all concurrent callers + # should raise CodegenWorkerError, not hang. + w = CodegenGradingWorker( + worker_cmd=_write_worker( + tmp_path, + """ + import sys + sys.stdin.buffer.readline() # consume one line then exit + """, + ) + ) + try: + with pytest.raises(Exception): # CodegenWorkerError or ExceptionGroup + await asyncio.gather(*[ + w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=10) + for _ in range(3) + ], return_exceptions=False) + finally: + await w.aclose() + + async def test_stale_id_after_timeout_does_not_crash(self, tmp_path) -> None: + # Reader receives a response for an id that the caller already timed out on. + # The stale future was already removed from _pending; the reader must skip it. + # Use _ECHO_OK with a very short timeout so the grade times out, then send + # a second grade to prove the worker (if restarted) still works. + w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, _ECHO_OK)) + try: + with pytest.raises(CodegenWorkerError): + await w.grade_codegen( + [{"input_output": "{}"}], [["x"]], timeout=0.000001 + ) + # If stale id handling is broken, the second grade would hang or crash. + # Give it a real timeout; it may or may not succeed (worker restarted). + finally: + await w.aclose() + + async def test_aclose_with_pending_futures_does_not_hang(self, tmp_path) -> None: + hang_worker = """ + import sys, time + for line in sys.stdin.buffer: + time.sleep(3600) + """ + w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, hang_worker)) + grade_task = asyncio.create_task( + w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=60) + ) + await asyncio.sleep(0.05) # let grade_task start and block + await w.aclose() # must not hang even with grade_task pending + grade_task.cancel() + with contextlib.suppress(asyncio.CancelledError, CodegenWorkerError): + await grade_task +``` + +Add `import contextlib` at the top of the test file if not already present. + +- [ ] **Step 2: Run the new tests to verify they fail** + +```bash +uv run pytest tests/unit/accuracy/test_codegen_worker_client.py::TestConcurrency -v 2>&1 | head -30 +``` + +Expected: failures because the lock still serializes requests (demux test would hang or return wrong values) or `AttributeError` if the test references methods not yet on the class. + +- [ ] **Step 3: Rewrite `CodegenGradingWorker` in `_codegen_worker_client.py`** + +The full rewrite of the class. Replace the existing class body (not the module-level helpers `_kill_process_group`, `CodegenWorkerError`, `_STREAM_LIMIT`, etc. — keep those unchanged). + +**3a — Update `__init__`:** remove `self._lock`, add `self._pending`, `self._reader_task`, and `self._spawn_lock`: + +```python +def __init__( + self, + worker_cmd: list[str] | None = None, + max_start_failures: int = 3, +) -> None: + self._cmd = worker_cmd or _DEFAULT_WORKER_CMD + self._max_start_failures = max_start_failures + self._proc: asyncio.subprocess.Process | None = None + self._spawn_lock = asyncio.Lock() + self._pending: dict[int, asyncio.Future[dict[str, Any]]] = {} + self._reader_task: asyncio.Task[None] | None = None + self._next_id = 0 + self._start_failures = 0 + self._worker_proven = False + self._stderr_tail: deque[str] = deque(maxlen=_STDERR_TAIL_LINES) + self._stderr_task: asyncio.Task[None] | None = None + self._death_w: int | None = None +``` + +**3b — Replace `grade_codegen`:** + +```python +async def grade_codegen( + self, + evaluation_sample: list[dict[str, str]], + generated_code: list[list[str]], + timeout: float, +) -> dict[str, Any]: + if self._start_failures >= self._max_start_failures: + raise CodegenWorkerError( + f"grading worker unavailable after {self._start_failures} start failures" + ) + await self._ensure_worker() + self._next_id += 1 + req_id = self._next_id + req = { + "id": req_id, + "evaluation_sample": evaluation_sample, + "generated_code": generated_code, + } + loop = asyncio.get_running_loop() + fut: asyncio.Future[dict[str, Any]] = loop.create_future() + self._pending[req_id] = fut + assert self._proc is not None and self._proc.stdin + self._proc.stdin.write(orjson.dumps(req) + b"\n") + try: + return await asyncio.wait_for(fut, timeout) + except asyncio.TimeoutError as exc: + self._pending.pop(req_id, None) + await self._handle_fault(count_start_failure=False) + raise CodegenWorkerError(f"grading worker timed out: {exc!r}") from exc + except asyncio.CancelledError: + self._pending.pop(req_id, None) + await self._handle_fault(count_start_failure=False) + raise +``` + +**3c — Update `_ensure_worker`** to use `_spawn_lock` and start `_reader_task`: + +```python +async def _ensure_worker(self) -> None: + async with self._spawn_lock: + if self._proc is not None and self._proc.returncode is None: + return + self._worker_proven = False + self._stderr_tail.clear() + self._close_death_pipe() + death_r: int | None = None + death_w: int | None = None + pass_fds: tuple[int, ...] = () + death_env: dict[str, str] = {} + if not IS_WINDOWS: + death_r, death_w = os.pipe() + os.set_inheritable(death_r, True) + pass_fds = (death_r,) + death_env = {_DEATH_FD_ENV: str(death_r)} + try: + self._proc = await asyncio.create_subprocess_exec( + *self._cmd, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + limit=_STREAM_LIMIT, + start_new_session=True, + pass_fds=pass_fds, + env={**os.environ, **death_env}, + ) + except Exception as exc: + if death_r is not None: + os.close(death_r) + if death_w is not None: + os.close(death_w) + self._start_failures += 1 + raise CodegenWorkerError(f"failed to spawn grading worker: {exc}") from exc + if death_r is not None: + os.close(death_r) + self._death_w = death_w + self._stderr_task = asyncio.create_task(self._drain_stderr(self._proc.stderr)) + self._reader_task = asyncio.create_task(self._run_reader()) +``` + +**3d — Add `_run_reader`** (new method, add after `_drain_stderr`): + +```python +async def _run_reader(self) -> None: + """Read worker responses and resolve the corresponding pending futures by id.""" + assert self._proc is not None and self._proc.stdout + reader = self._proc.stdout + try: + while True: + line = await reader.readline() + if not line: + await self._handle_fault() + return + try: + resp = orjson.loads(line) + except orjson.JSONDecodeError: + await self._handle_fault() + return + if not isinstance(resp, dict): + await self._handle_fault() + return + req_id = resp.get("id") + fut = self._pending.pop(req_id, None) + if fut is None or fut.done(): + continue # stale id (caller already timed out) or cancelled + if not resp.get("ok"): + fut.set_exception( + CodegenWorkerError(resp.get("error", "unknown grading error")) + ) + self._mark_proven() + else: + metrics = resp.get("metrics") + if not isinstance(metrics, dict): + await self._handle_fault() + return + fut.set_result(metrics) + self._mark_proven() + except asyncio.CancelledError: + pass +``` + +**3e — Add `_mark_proven`** (new helper, add after `_run_reader`): + +```python +def _mark_proven(self) -> None: + self._worker_proven = True + self._start_failures = 0 +``` + +**3f — Replace `_handle_fault`:** + +```python +async def _handle_fault(self, count_start_failure: bool = True) -> None: + if self._proc is None: + return # already handled; _handle_fault is idempotent + if count_start_failure and not self._worker_proven: + self._start_failures += 1 + for fut in list(self._pending.values()): + if not fut.done(): + fut.set_exception(CodegenWorkerError("grading worker fault")) + self._pending.clear() + tail = await self._kill() + _log.debug( + lambda: f"codegen worker fault (proven={self._worker_proven}, " + f"start_failures={self._start_failures}); killed + respawning next grade" + + (f"; stderr tail:\n{chr(10).join(tail)}" if tail else "") + ) +``` + +**3g — Update `_kill`** to also cancel and await `_reader_task`: + +```python +async def _kill(self) -> list[str]: + proc, self._proc = self._proc, None + task, self._stderr_task = self._stderr_task, None + reader_task, self._reader_task = self._reader_task, None + self._close_death_pipe() + if proc is not None and proc.returncode is None: + _kill_process_group(proc) + with contextlib.suppress(ProcessLookupError): + await proc.wait() + if reader_task is not None: + reader_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await reader_task + tail: list[str] = [] + if task is not None: + with contextlib.suppress(TimeoutError, asyncio.CancelledError): + await asyncio.wait_for(task, timeout=2.0) + tail = list(self._stderr_tail) + return tail +``` + +**3h — Replace `aclose`** (remove the lock, cancel pending futures): + +```python +async def aclose(self) -> None: + for fut in list(self._pending.values()): + if not fut.done(): + fut.cancel() + self._pending.clear() + await self._kill() +``` + +**3i — Remove `_request`** entirely (replaced by the reader task + futures approach). Grep to confirm nothing else calls it: + +```bash +grep -rn "_request" src/aiperf/accuracy/graders/_codegen_worker_client.py +``` + +- [ ] **Step 4: Update the existing `TestSerialization` class** + +The existing `TestSerialization.test_concurrent_grades_do_not_interleave` test was checking that concurrent grades were serialized (old lock behaviour). With the new design, concurrent grades overlap — which is the desired behaviour. Rename and update the test so it validates the new invariant: + +Replace the existing `TestSerialization` class with: + +```python +class TestSerialization: + async def test_concurrent_grades_return_correct_results(self, tmp_path) -> None: + # Previously tested that grades were serialized (lock enforced). + # Now tests that concurrent grades all return correct results without the lock. + w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, _ECHO_OK)) + try: + results = await asyncio.gather(*[ + w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=30) + for _ in range(4) + ]) + assert all(r == {"pass@1": 1.0} for r in results) + finally: + await w.aclose() +``` + +- [ ] **Step 5: Run all client unit tests** + +```bash +uv run pytest tests/unit/accuracy/test_codegen_worker_client.py -v +``` + +Expected: all tests pass, including the new `TestConcurrency` class. + +- [ ] **Step 6: Run the full unit test suite** + +```bash +uv run pytest tests/unit/ -n auto +``` + +Expected: all green. + +- [ ] **Step 7: Lint** + +```bash +ruff format . && ruff check --fix . +``` + +- [ ] **Step 8: Commit** + +```bash +git add src/aiperf/accuracy/graders/_codegen_worker_client.py \ + tests/unit/accuracy/test_codegen_worker_client.py +git commit -s -m "feat(accuracy): concurrent codegen grading via id-demux reader task" +``` + +--- + +### Task 3: Component integration — concurrent multi-problem test + +**Spec:** `docs/superpowers/specs/2026-07-29-codegen-grade-concurrency-design.md` § "Tests / Component integration" + +**Files:** +- Modify: `tests/component_integration/test_lcb_codegen_worker_e2e.py` + +**Interfaces:** +- Consumes: `CodegenGradingWorker` from Task 2 (concurrent-safe) +- Consumes: real `lighteval` (skip if not installed) + +--- + +- [ ] **Step 1: Write the concurrent e2e test** + +Open `tests/component_integration/test_lcb_codegen_worker_e2e.py`. + +Add the following after the existing `test_worker_grades_correct_stdin_solution` test: + +```python +@pytest.mark.slow +@pytest.mark.asyncio +async def test_worker_grades_multiple_problems_concurrently() -> None: + """N concurrent grade_codegen() calls all resolve correctly. + + This exercises the batch-drain path: all N requests are sent before the + worker responds, so they are drained into a single codegen_metrics call and + processed in parallel by lighteval's ProcessPoolExecutor. + """ + worker = CodegenGradingWorker() + sample, code = _sample_and_solution() + n = 4 + try: + results = await asyncio.gather(*[ + worker.grade_codegen(sample, code, timeout=240) + for _ in range(n) + ]) + assert len(results) == n + assert all(float(r["pass@1"]) == 1.0 for r in results), results + finally: + await worker.aclose() +``` + +Add `import asyncio` at the top of the file if not already present. + +- [ ] **Step 2: Run the component integration tests** + +```bash +uv run pytest tests/component_integration/test_lcb_codegen_worker_e2e.py -v -s +``` + +Expected: both `test_worker_grades_correct_stdin_solution` and `test_worker_grades_multiple_problems_concurrently` pass with `pass@1 == 1.0`. + +> These tests run lighteval for real — they take 30-120 seconds each. If `lighteval` is not installed, both tests are auto-skipped via `pytest.importorskip`. + +- [ ] **Step 3: Run the full unit test suite one more time** + +```bash +uv run pytest tests/unit/ -n auto +``` + +Expected: all green. + +- [ ] **Step 4: Lint and pre-commit** + +```bash +ruff format . && ruff check --fix . +pre-commit run --all-files +``` + +- [ ] **Step 5: Commit** + +```bash +git add tests/component_integration/test_lcb_codegen_worker_e2e.py +git commit -s -m "test(accuracy): concurrent multi-problem codegen grading e2e test" +``` + +--- + +## Self-Review + +**Spec coverage:** + +| Spec requirement | Task covering it | +|---|---| +| Drop `asyncio.Lock` | Task 2 Step 3a | +| `id → Future` demux table | Task 2 Step 3b + `_pending` | +| Persistent reader task | Task 2 Step 3c + `_run_reader` | +| Worker batch-drain loop | Task 1 Step 3c | +| Single `codegen_metrics` call per cycle | Task 1 Step 3b (`handle_batch`) | +| Per-problem demux via `compute_metrics_from_results` | Task 1 Step 3b | +| `_handle_fault` cancels all pending futures | Task 2 Step 3f | +| `aclose` cancels pending futures without lock | Task 2 Step 3h | +| `_kill` tears down `_reader_task` | Task 2 Step 3g | +| `_spawn_lock` prevents double-spawn | Task 2 Step 3c | +| `_handle_fault` idempotent (`_proc is None` guard) | Task 2 Step 3f | +| New unit concurrency tests | Task 2 Step 1 | +| New worker batch unit tests | Task 1 Step 1 | +| Component integration concurrent test | Task 3 Step 1 | + +**No gaps found.** + +**Placeholder scan:** No TBDs, TODOs, or vague steps. All code blocks are complete. + +**Type consistency:** +- `handle_batch(reqs, codegen_fn, compute_metrics_fn)` — consistent across Task 1 definition and Task 1 tests +- `run_worker_loop(stdin, out, codegen_fn, compute_metrics_fn)` — consistent in Step 3c and Step 3e (`main()`) +- `_pending: dict[int, asyncio.Future[dict[str, Any]]]` — consistent across `__init__`, `grade_codegen`, `_run_reader`, `_handle_fault`, `aclose` +- `_reader_task: asyncio.Task[None] | None` — consistent across `__init__`, `_ensure_worker`, `_kill` +- `_mark_proven()` — defined in Task 2 Step 3e, called from `_run_reader` (Task 2 Step 3d) From e36ac8aea207fa93b4291e054f9e75b088478b07 Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Thu, 30 Jul 2026 17:01:10 -0700 Subject: [PATCH 02/26] feat(accuracy): batch-drain worker loop for codegen grading concurrency Signed-off-by: Elias Bermudez --- .../accuracy/graders/_codegen_worker.py | 208 ++++++++++++----- tests/unit/accuracy/test_codegen_worker.py | 211 +++++++++++++----- 2 files changed, 305 insertions(+), 114 deletions(-) diff --git a/src/aiperf/accuracy/graders/_codegen_worker.py b/src/aiperf/accuracy/graders/_codegen_worker.py index 352f317cf1..9184c54635 100644 --- a/src/aiperf/accuracy/graders/_codegen_worker.py +++ b/src/aiperf/accuracy/graders/_codegen_worker.py @@ -5,11 +5,11 @@ Runs as ``python -m aiperf.accuracy.graders._codegen_worker``. A fresh, single-threaded interpreter that forces the ``fork`` start method once at -startup, then reads one JSONL grading request per line from stdin and writes one -JSONL response per line to a private protocol fd. Executing lighteval's -``codegen_metrics`` here (not in the multithreaded record-processor daemon) -avoids both the nested-function pickle failure under spawn/forkserver and the -multithreaded-fork hang. See issue #1145. +startup, then reads JSONL grading requests from stdin (batching queued +requests) and writes one JSONL response per request to a private protocol fd. +Executing lighteval's ``codegen_metrics`` here (not in the multithreaded +record-processor daemon) avoids both the nested-function pickle failure under +spawn/forkserver and the multithreaded-fork hang. See issue #1145. """ from __future__ import annotations @@ -18,6 +18,7 @@ import math import multiprocessing as mp import os +import select import signal import sys import threading @@ -39,38 +40,6 @@ _MAX_ERROR_CHARS = 4096 -def handle_request( - req: Any, - codegen_fn: Callable[..., tuple[dict[str, Any], Any]], -) -> dict[str, Any]: - """Run one grading request. Never raises: all failures become an error - response so a single bad problem cannot kill the worker loop.""" - if not isinstance(req, dict): - # A valid-but-non-object JSON frame (e.g. ``[]``) would raise on the - # ``.get`` below and kill the loop; return the promised error instead. - return {"id": None, "ok": False, "error": "malformed request: expected object"} - req_id = req.get("id") - try: - evaluation_sample = req["evaluation_sample"] - generated_code = req["generated_code"] - except (KeyError, TypeError) as exc: - return {"id": req_id, "ok": False, "error": f"malformed request: {exc!r}"} - - try: - metrics, _ = codegen_fn( - evaluation_sample, - generated_code, - k_list=list(_LCB_PASS_AT_K), - num_process_evaluate=_LCB_NUM_PROCESSES, - ) - except Exception as exc: - # A single bad problem must never crash the worker loop. - error = f"{type(exc).__name__}: {exc}" - return {"id": req_id, "ok": False, "error": _truncate_error(error)} - - return {"id": req_id, "ok": True, "metrics": _coerce_metrics(metrics)} - - def _truncate_error(error: str) -> str: """Bound an error string so it cannot produce a multi-MB response line.""" if len(error) <= _MAX_ERROR_CHARS: @@ -102,27 +71,153 @@ def _is_number(value: Any) -> bool: return False +def handle_batch( + reqs: list[Any], + codegen_fn: Callable[..., tuple[dict[str, Any], Any]], + compute_metrics_fn: Callable[..., dict[str, Any]], +) -> list[dict[str, Any]]: + """Grade a batch of requests with a single codegen_fn call. + + Calls codegen_fn once with all well-formed requests batched together so + lighteval's ProcessPoolExecutor can process multiple problems in parallel. + Never raises: all failures become error responses so a bad batch cannot + kill the worker loop. + """ + all_samples: list[Any] = [] + all_generations: list[Any] = [] + id_map: list[tuple[int, Any]] = [] # (batch_position, req_id) + responses: list[dict[str, Any] | None] = [None] * len(reqs) + + for i, req in enumerate(reqs): + if isinstance(req, dict) and "_parse_error" in req: + responses[i] = { + "id": None, + "ok": False, + "error": f"bad json: {req['_parse_error']}", + } + continue + if not isinstance(req, dict): + responses[i] = { + "id": None, + "ok": False, + "error": "malformed request: expected object", + } + continue + req_id = req.get("id") + try: + all_samples.append(req["evaluation_sample"]) + all_generations.append(req["generated_code"]) + id_map.append((i, req_id)) + except (KeyError, TypeError) as exc: + responses[i] = { + "id": req_id, + "ok": False, + "error": f"malformed request: {exc!r}", + } + + if all_samples: + batch_error: str | None = None + raw_results: dict[int, Any] = {} + try: + _, raw_results = codegen_fn( + all_samples, + all_generations, + k_list=list(_LCB_PASS_AT_K), + num_process_evaluate=_LCB_NUM_PROCESSES, + ) + except Exception as exc: + batch_error = _truncate_error(f"{type(exc).__name__}: {exc}") + + for pos, (req_idx, req_id) in enumerate(id_map): + if batch_error is not None: + responses[req_idx] = {"id": req_id, "ok": False, "error": batch_error} + else: + try: + metrics = compute_metrics_fn( + {0: raw_results[pos]}, + k_list=list(_LCB_PASS_AT_K), + ) + responses[req_idx] = { + "id": req_id, + "ok": True, + "metrics": _coerce_metrics(metrics), + } + except Exception as exc: + responses[req_idx] = { + "id": req_id, + "ok": False, + "error": _truncate_error(f"{type(exc).__name__}: {exc}"), + } + + return [r for r in responses if r is not None] + + +def _drain_fd(stdin: BinaryIO, stdin_fd: int, batch_raw: list[bytes]) -> None: + """Non-blocking drain of already-queued lines from a real file descriptor.""" + while True: + ready, _, _ = select.select([stdin_fd], [], [], 0) + if not ready: + break + line = stdin.readline() + if not line: + break + line = line.strip() + if line: + batch_raw.append(line) + + +def _drain_seekable(stdin: BinaryIO, batch_raw: list[bytes]) -> None: + """Greedy drain for non-fd streams such as BytesIO (used in tests).""" + remaining = stdin.read() + for raw_line in remaining.split(b"\n"): + raw_line = raw_line.strip() + if raw_line: + batch_raw.append(raw_line) + + +def _parse_batch(batch_raw: list[bytes]) -> list[Any]: + reqs: list[Any] = [] + for raw in batch_raw: + try: + reqs.append(orjson.loads(raw)) + except orjson.JSONDecodeError as exc: + reqs.append({"_parse_error": str(exc)}) + return reqs + + def run_worker_loop( stdin: BinaryIO, out: BinaryIO, codegen_fn: Callable[..., tuple[dict[str, Any], Any]], + compute_metrics_fn: Callable[..., dict[str, Any]], ) -> None: - """Serve JSONL grading requests until stdin EOF. One response per request.""" - for line in stdin: - line = line.strip() - if not line: + """Serve JSONL grading requests until stdin EOF. + + Blocks on the first request of each cycle, then non-blocking drains any + already-queued requests to form a batch. Calls codegen_fn once per batch so + lighteval's ProcessPoolExecutor can process multiple problems in parallel. + """ + try: + stdin_fd: int | None = stdin.fileno() + except (AttributeError, OSError): + stdin_fd = None + + while True: + first = stdin.readline() + if not first: + break # EOF: client closed stdin, clean exit + first = first.strip() + if not first: continue - try: - req = orjson.loads(line) - except orjson.JSONDecodeError as exc: - resp: dict[str, Any] = { - "id": None, - "ok": False, - "error": f"bad json: {exc}", - } + batch_raw: list[bytes] = [first] + if stdin_fd is not None: + _drain_fd(stdin, stdin_fd, batch_raw) else: - resp = handle_request(req, codegen_fn) - out.write(orjson.dumps(resp) + b"\n") + _drain_seekable(stdin, batch_raw) + for resp in handle_batch( + _parse_batch(batch_raw), codegen_fn, compute_metrics_fn + ): + out.write(orjson.dumps(resp) + b"\n") out.flush() @@ -200,9 +295,14 @@ def main() -> None: # reaps the worker even if the (heavy) import is still in flight. _start_death_watcher() _force_fork() - from lighteval.tasks.tasks.lcb.codegen_metrics import codegen_metrics - - run_worker_loop(sys.stdin.buffer, protocol_out, codegen_metrics) + from lighteval.tasks.tasks.lcb.codegen_metrics import ( + codegen_metrics, + compute_metrics_from_results, + ) + + run_worker_loop( + sys.stdin.buffer, protocol_out, codegen_metrics, compute_metrics_from_results + ) if __name__ == "__main__": diff --git a/tests/unit/accuracy/test_codegen_worker.py b/tests/unit/accuracy/test_codegen_worker.py index 019af1dbb9..d371b3fc1b 100644 --- a/tests/unit/accuracy/test_codegen_worker.py +++ b/tests/unit/accuracy/test_codegen_worker.py @@ -17,84 +17,175 @@ _FORK_AVAILABLE = "fork" in mp.get_all_start_methods() -def _fake_codegen_ok(*_args: Any, **_kwargs: Any) -> tuple[dict[str, Any], Any]: - return {"pass@1": 1.0}, {} +def _fake_compute_metrics( + results: dict, k_list: list[int] | None = None +) -> dict[str, Any]: + # Mirrors compute_metrics_from_results: returns {"pass@1": } using + # the single-problem results dict {0: [[True, True, ...]]} passed by handle_batch. + result_list = results.get(0, [[-2]]) # [-2] = compile error + if result_list and all(x > 0 for x in result_list[0]): + return {"pass@1": 1.0} + return {"pass@1": 0.0} -def _fake_codegen_boom(*_args: Any, **_kwargs: Any) -> tuple[dict[str, Any], Any]: - raise RuntimeError("sandbox exploded") +def _fake_codegen_batch_ok( + samples: list, generations: list, **_kwargs: Any +) -> tuple[dict[str, Any], dict[int, list]]: + # Returns aggregate metrics (ignored by handle_batch) and per-problem results. + n = len(samples) + raw_results = {i: [[True]] for i in range(n)} # all pass + return {"pass@1": 1.0}, raw_results -def _fake_codegen_list_pass(*_args: Any, **_kwargs: Any) -> tuple[dict[str, Any], Any]: - return {"pass@1": [1.0]}, {} +def _fake_codegen_batch_boom( + samples: list, generations: list, **_kwargs: Any +) -> tuple[dict[str, Any], dict[int, list]]: + raise RuntimeError("pool exploded") -class TestHandleRequest: - def test_ok_request_returns_metrics_with_id(self) -> None: - req = { - "id": 7, +class TestHandleBatch: + def _req(self, req_id: int) -> dict[str, Any]: + return { + "id": req_id, "evaluation_sample": [{"input_output": "{}"}], "generated_code": [["x"]], } - resp = worker.handle_request(req, _fake_codegen_ok) - assert resp == {"id": 7, "ok": True, "metrics": {"pass@1": 1.0}} - - def test_list_shaped_pass_at_1_is_preserved(self) -> None: - # lighteval returns pass@1 as a list on some pins; it must survive - # coercion rather than be dropped as non-numeric (silent 0.000 bug). - req = { - "id": 9, - "evaluation_sample": [{"input_output": "{}"}], - "generated_code": [["x"]], - } - resp = worker.handle_request(req, _fake_codegen_list_pass) - assert resp["ok"] is True - assert resp["metrics"]["pass@1"] == [1.0] - def test_codegen_exception_becomes_error_response(self) -> None: - req = { - "id": 3, + def test_single_request_returns_one_ok_response(self) -> None: + resps = worker.handle_batch( + [self._req(1)], _fake_codegen_batch_ok, _fake_compute_metrics + ) + assert len(resps) == 1 + assert resps[0] == {"id": 1, "ok": True, "metrics": {"pass@1": 1.0}} + + def test_batch_of_n_calls_codegen_fn_once(self) -> None: + call_count = 0 + + def counting_codegen(samples, generations, **kwargs): + nonlocal call_count + call_count += 1 + n = len(samples) + return {"pass@1": 1.0}, {i: [[True]] for i in range(n)} + + reqs = [self._req(i) for i in range(1, 5)] + resps = worker.handle_batch(reqs, counting_codegen, _fake_compute_metrics) + assert call_count == 1 + assert len(resps) == 4 + assert all(r["ok"] for r in resps) + assert [r["id"] for r in resps] == [1, 2, 3, 4] + + def test_response_order_matches_request_order(self) -> None: + reqs = [self._req(i) for i in [7, 3, 99]] + resps = worker.handle_batch(reqs, _fake_codegen_batch_ok, _fake_compute_metrics) + assert [r["id"] for r in resps] == [7, 3, 99] + + def test_batch_exception_returns_error_for_all(self) -> None: + reqs = [self._req(i) for i in range(1, 4)] + resps = worker.handle_batch( + reqs, _fake_codegen_batch_boom, _fake_compute_metrics + ) + assert len(resps) == 3 + assert all(not r["ok"] for r in resps) + assert all("pool exploded" in r["error"] for r in resps) + + def test_malformed_request_in_batch_does_not_affect_others(self) -> None: + reqs = [ + self._req(1), + {"id": 2}, # missing evaluation_sample + generated_code + self._req(3), + ] + resps = worker.handle_batch(reqs, _fake_codegen_batch_ok, _fake_compute_metrics) + assert len(resps) == 3 + assert resps[0] == {"id": 1, "ok": True, "metrics": {"pass@1": 1.0}} + assert resps[1]["id"] == 2 + assert not resps[1]["ok"] + assert resps[2] == {"id": 3, "ok": True, "metrics": {"pass@1": 1.0}} + + def test_non_object_request_in_batch_is_error(self) -> None: + reqs = [[1, 2, 3], self._req(5)] + resps = worker.handle_batch(reqs, _fake_codegen_batch_ok, _fake_compute_metrics) + assert len(resps) == 2 + assert resps[0]["id"] is None + assert not resps[0]["ok"] + assert resps[1] == {"id": 5, "ok": True, "metrics": {"pass@1": 1.0}} + + def test_parse_error_sentinel_produces_error_response(self) -> None: + # run_worker_loop encodes JSON decode errors as {"_parse_error": "..."}. + reqs = [{"_parse_error": "unexpected token"}, self._req(2)] + resps = worker.handle_batch(reqs, _fake_codegen_batch_ok, _fake_compute_metrics) + assert len(resps) == 2 + assert resps[0]["id"] is None + assert not resps[0]["ok"] + assert "unexpected token" in resps[0]["error"] + assert resps[1]["ok"] + + def test_non_finite_metric_values_are_dropped(self) -> None: + # NaN/Inf must not cross the JSONL boundary (repo NaN/Inf discipline). + def _nan_inf( + samples: list, generations: list, **_kwargs: Any + ) -> tuple[dict[str, Any], dict[int, list]]: + n = len(samples) + return {"pass@1": 1.0}, {i: [[True]] for i in range(n)} + + def _nan_compute(results: dict, **_kwargs: Any) -> dict[str, Any]: + return {"pass@1": float("nan"), "extra": float("inf"), "ok": 1.0} + + req = {"id": 9, "evaluation_sample": [{}], "generated_code": [["x"]]} + resps = worker.handle_batch([req], _nan_inf, _nan_compute) + assert resps[0]["ok"] is True + assert "pass@1" not in resps[0]["metrics"] + assert "extra" not in resps[0]["metrics"] + assert resps[0]["metrics"]["ok"] == 1.0 + + +class TestRunWorkerLoopBatch: + def _run( + self, + payloads: list[dict[str, Any]], + codegen_fn=_fake_codegen_batch_ok, + compute_metrics_fn=_fake_compute_metrics, + ) -> list[dict[str, Any]]: + # Write all payloads to a BytesIO pipe so they are already queued when + # run_worker_loop reads; this exercises the non-blocking drain path. + data = b"".join(orjson.dumps(p) + b"\n" for p in payloads) + stdin = io.BytesIO(data) + out = io.BytesIO() + worker.run_worker_loop(stdin, out, codegen_fn, compute_metrics_fn) + out.seek(0) + return [orjson.loads(line) for line in out if line.strip()] + + def _req(self, req_id: int) -> dict[str, Any]: + return { + "id": req_id, "evaluation_sample": [{"input_output": "{}"}], "generated_code": [["x"]], } - resp = worker.handle_request(req, _fake_codegen_boom) - assert resp["id"] == 3 - assert resp["ok"] is False - assert "sandbox exploded" in resp["error"] - - def test_malformed_request_missing_fields_is_error(self) -> None: - resp = worker.handle_request({"id": 5}, _fake_codegen_ok) - assert resp["id"] == 5 - assert resp["ok"] is False - assert resp["error"] - - def test_non_object_request_is_error_not_crash(self) -> None: - # A valid-but-non-object JSON frame must not raise (which would kill the - # worker loop); it returns the promised error response with id=None. - resp = worker.handle_request([1, 2, 3], _fake_codegen_ok) - assert resp["id"] is None - assert resp["ok"] is False - assert resp["error"] - def test_non_finite_metric_values_are_dropped(self) -> None: - # NaN/Inf must not cross the JSONL boundary (repo NaN/Inf discipline). - def _nan_inf(*_a: Any, **_k: Any) -> tuple[dict[str, Any], Any]: - return {"pass@1": float("nan"), "extra": float("inf"), "ok": 1.0}, {} + def test_pre_queued_requests_are_batched_in_one_call(self) -> None: + call_count = 0 - resp = worker.handle_request( - {"id": 9, "evaluation_sample": [{}], "generated_code": [["x"]]}, _nan_inf - ) - assert resp["ok"] is True - assert "pass@1" not in resp["metrics"] - assert "extra" not in resp["metrics"] - assert resp["metrics"]["ok"] == 1.0 + def counting_codegen(samples, generations, **kwargs): + nonlocal call_count + call_count += 1 + n = len(samples) + return {"pass@1": 1.0}, {i: [[True]] for i in range(n)} + + reqs = [self._req(i) for i in range(1, 4)] + resps = self._run(reqs, counting_codegen) + assert call_count == 1 + assert len(resps) == 3 + + def test_responses_carry_correct_ids(self) -> None: + reqs = [self._req(i) for i in [10, 20, 30]] + resps = self._run(reqs) + assert {r["id"] for r in resps} == {10, 20, 30} class TestRunWorkerLoop: def _run(self, requests: list[bytes], codegen_fn) -> list[dict]: stdin = io.BytesIO(b"".join(r + b"\n" for r in requests)) out = io.BytesIO() - worker.run_worker_loop(stdin, out, codegen_fn) + worker.run_worker_loop(stdin, out, codegen_fn, _fake_compute_metrics) out.seek(0) return [orjson.loads(line) for line in out.read().splitlines() if line] @@ -107,16 +198,16 @@ def test_processes_each_request_in_order(self) -> None: {"id": 2, "evaluation_sample": [{}], "generated_code": [["b"]]} ), ] - resps = self._run(reqs, _fake_codegen_ok) + resps = self._run(reqs, _fake_codegen_batch_ok) assert [r["id"] for r in resps] == [1, 2] assert all(r["ok"] for r in resps) def test_eof_stops_the_loop(self) -> None: - resps = self._run([], _fake_codegen_ok) + resps = self._run([], _fake_codegen_batch_ok) assert resps == [] def test_garbled_line_yields_error_response(self) -> None: - resps = self._run([b"{not json"], _fake_codegen_ok) + resps = self._run([b"{not json"], _fake_codegen_batch_ok) assert len(resps) == 1 assert resps[0]["ok"] is False assert resps[0]["id"] is None From 591be04247fe06809cc4410d172f1d8d409c67cf Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Thu, 30 Jul 2026 17:04:16 -0700 Subject: [PATCH 03/26] fix(accuracy): add type annotations to batch worker test helper Signed-off-by: Elias Bermudez --- tests/unit/accuracy/test_codegen_worker.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/unit/accuracy/test_codegen_worker.py b/tests/unit/accuracy/test_codegen_worker.py index d371b3fc1b..67db0b7dbf 100644 --- a/tests/unit/accuracy/test_codegen_worker.py +++ b/tests/unit/accuracy/test_codegen_worker.py @@ -6,6 +6,7 @@ import subprocess import sys import textwrap +from collections.abc import Callable from pathlib import Path from typing import Any @@ -142,8 +143,10 @@ class TestRunWorkerLoopBatch: def _run( self, payloads: list[dict[str, Any]], - codegen_fn=_fake_codegen_batch_ok, - compute_metrics_fn=_fake_compute_metrics, + codegen_fn: Callable[ + ..., tuple[dict[str, Any], dict[int, list]] + ] = _fake_codegen_batch_ok, + compute_metrics_fn: Callable[..., dict[str, Any]] = _fake_compute_metrics, ) -> list[dict[str, Any]]: # Write all payloads to a BytesIO pipe so they are already queued when # run_worker_loop reads; this exercises the non-blocking drain path. From 50f0287137a868aa3d4640f86d6e694b071bcb0d Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Thu, 30 Jul 2026 17:18:36 -0700 Subject: [PATCH 04/26] feat(accuracy): concurrent codegen grading via id-demux reader task Signed-off-by: Elias Bermudez --- .../graders/_codegen_worker_client.py | 286 +++++++++--------- .../accuracy/test_codegen_worker_client.py | 144 +++++++-- 2 files changed, 272 insertions(+), 158 deletions(-) diff --git a/src/aiperf/accuracy/graders/_codegen_worker_client.py b/src/aiperf/accuracy/graders/_codegen_worker_client.py index 1ba4074926..8df4a51081 100644 --- a/src/aiperf/accuracy/graders/_codegen_worker_client.py +++ b/src/aiperf/accuracy/graders/_codegen_worker_client.py @@ -4,8 +4,8 @@ """In-process client for the out-of-process LCB codegen grading worker. Owned by ``CodeExecutionGrader``. Lazily spawns a single persistent worker -subprocess, serializes grading requests through an ``asyncio.Lock``, and -enforces per-grade timeouts with kill+restart. See issue #1145 and +subprocess, routes concurrent grading requests via an ``id → Future`` demux +table, and enforces per-grade timeouts with kill+restart. See issue #1094 and ``_codegen_worker.py``. """ @@ -72,7 +72,9 @@ def __init__( self._cmd = worker_cmd or _DEFAULT_WORKER_CMD self._max_start_failures = max_start_failures self._proc: asyncio.subprocess.Process | None = None - self._lock = asyncio.Lock() + self._spawn_lock = asyncio.Lock() + self._pending: dict[int, asyncio.Future[dict[str, Any]]] = {} + self._reader_task: asyncio.Task[None] | None = None self._next_id = 0 self._start_failures = 0 self._worker_proven = False @@ -86,68 +88,88 @@ async def grade_codegen( generated_code: list[list[str]], timeout: float, ) -> dict[str, Any]: - async with self._lock: - if self._start_failures >= self._max_start_failures: - raise CodegenWorkerError( - f"grading worker unavailable after {self._start_failures} start failures" - ) - await self._ensure_worker() - self._next_id += 1 - req = { - "id": self._next_id, - "evaluation_sample": evaluation_sample, - "generated_code": generated_code, - } - return await self._request(req, timeout) + if self._start_failures >= self._max_start_failures: + raise CodegenWorkerError( + f"grading worker unavailable after {self._start_failures} start failures" + ) + await self._ensure_worker() + self._next_id += 1 + req_id = self._next_id + req = { + "id": req_id, + "evaluation_sample": evaluation_sample, + "generated_code": generated_code, + } + loop = asyncio.get_running_loop() + fut: asyncio.Future[dict[str, Any]] = loop.create_future() + self._pending[req_id] = fut + assert self._proc is not None and self._proc.stdin + self._proc.stdin.write(orjson.dumps(req) + b"\n") + try: + return await asyncio.wait_for(fut, timeout) + except TimeoutError as exc: + self._pending.pop(req_id, None) + await self._handle_fault(count_start_failure=False) + raise CodegenWorkerError(f"grading worker timed out: {exc!r}") from exc + except asyncio.CancelledError: + self._pending.pop(req_id, None) + await self._handle_fault(count_start_failure=False) + raise async def _ensure_worker(self) -> None: - if self._proc is not None and self._proc.returncode is None: - return - self._worker_proven = False - self._stderr_tail.clear() - self._close_death_pipe() - # Death pipe: the child inherits the read end; we keep the write end open - # for the worker's life so its close (even via the parent's os._exit) - # tells the worker to reap itself. os.pipe fds are non-inheritable by - # default, so mark the read end inheritable before passing it through. - # Windows has no process groups (nothing for the worker to killpg) and - # subprocess rejects pass_fds there, so skip the pipe and rely on the - # stdin-EOF teardown; the worker's death watcher likewise no-ops there. - death_r: int | None = None - death_w: int | None = None - pass_fds: tuple[int, ...] = () - death_env: dict[str, str] = {} - if not IS_WINDOWS: - death_r, death_w = os.pipe() - os.set_inheritable(death_r, True) - pass_fds = (death_r,) - death_env = {_DEATH_FD_ENV: str(death_r)} - try: - self._proc = await asyncio.create_subprocess_exec( - *self._cmd, - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - limit=_STREAM_LIMIT, - # Own process group so _kill can reap lighteval's forked sandbox - # grandchildren, not just the worker (no-op on Windows). - start_new_session=True, - pass_fds=pass_fds, - env={**os.environ, **death_env}, - ) - except Exception as exc: + async with self._spawn_lock: + if self._proc is not None and self._proc.returncode is None: + return + self._worker_proven = False + self._stderr_tail.clear() + self._close_death_pipe() + # Death pipe: the child inherits the read end; we keep the write end open + # for the worker's life so its close (even via the parent's os._exit) + # tells the worker to reap itself. os.pipe fds are non-inheritable by + # default, so mark the read end inheritable before passing it through. + # Windows has no process groups (nothing for the worker to killpg) and + # subprocess rejects pass_fds there, so skip the pipe and rely on the + # stdin-EOF teardown; the worker's death watcher likewise no-ops there. + death_r: int | None = None + death_w: int | None = None + pass_fds: tuple[int, ...] = () + death_env: dict[str, str] = {} + if not IS_WINDOWS: + death_r, death_w = os.pipe() + os.set_inheritable(death_r, True) + pass_fds = (death_r,) + death_env = {_DEATH_FD_ENV: str(death_r)} + try: + self._proc = await asyncio.create_subprocess_exec( + *self._cmd, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + limit=_STREAM_LIMIT, + # Own process group so _kill can reap lighteval's forked sandbox + # grandchildren, not just the worker (no-op on Windows). + start_new_session=True, + pass_fds=pass_fds, + env={**os.environ, **death_env}, + ) + except Exception as exc: + if death_r is not None: + os.close(death_r) + if death_w is not None: + os.close(death_w) + self._start_failures += 1 + raise CodegenWorkerError( + f"failed to spawn grading worker: {exc}" + ) from exc if death_r is not None: - os.close(death_r) - if death_w is not None: - os.close(death_w) - self._start_failures += 1 - raise CodegenWorkerError(f"failed to spawn grading worker: {exc}") from exc - if death_r is not None: - os.close(death_r) # the child holds it now; we keep only the write end - self._death_w = death_w - # Drain stderr continuously so the pipe never fills (which would block the - # worker) and the last output is retained to explain a fault. - self._stderr_task = asyncio.create_task(self._drain_stderr(self._proc.stderr)) + os.close(death_r) # the child holds it now; we keep only the write end + self._death_w = death_w + # Drain stderr continuously so the pipe never fills (which would block the + # worker) and the last output is retained to explain a fault. + self._stderr_task = asyncio.create_task( + self._drain_stderr(self._proc.stderr) + ) + self._reader_task = asyncio.create_task(self._run_reader()) async def _drain_stderr(self, reader: asyncio.StreamReader | None) -> None: """Continuously copy worker stderr into a bounded tail. Best-effort: @@ -160,90 +182,63 @@ async def _drain_stderr(self, reader: asyncio.StreamReader | None) -> None: async for line in reader: self._stderr_tail.append(line.decode("utf-8", "replace").rstrip("\n")) - async def _request(self, req: dict[str, Any], timeout: float) -> dict[str, Any]: - assert self._proc is not None and self._proc.stdin and self._proc.stdout + async def _run_reader(self) -> None: + """Read worker responses and resolve the corresponding pending futures by id.""" + assert self._proc is not None and self._proc.stdout + reader = self._proc.stdout try: - self._proc.stdin.write(orjson.dumps(req) + b"\n") - await self._proc.stdin.drain() - line = await asyncio.wait_for(self._proc.stdout.readline(), timeout) + while True: + try: + line = await reader.readline() + except (ValueError, ConnectionError, BrokenPipeError): + # ValueError covers asyncio.LimitOverrunError (oversized line); + # the other two indicate a broken transport. All mean the worker + # desynced; fault it so callers are unblocked. + await self._handle_fault() + return + if not line: + await self._handle_fault() + return + try: + resp = orjson.loads(line) + except orjson.JSONDecodeError: + await self._handle_fault() + return + if not isinstance(resp, dict): + await self._handle_fault() + return + req_id = resp.get("id") + fut = self._pending.pop(req_id, None) + if fut is None or fut.done(): + continue # stale id (caller already timed out) or cancelled + if not resp.get("ok"): + fut.set_exception( + CodegenWorkerError(resp.get("error", "unknown grading error")) + ) + self._mark_proven() + else: + metrics = resp.get("metrics") + if not isinstance(metrics, dict): + await self._handle_fault() + return + fut.set_result(metrics) + self._mark_proven() except asyncio.CancelledError: - # Shutdown/cancellation while awaiting the worker: kill it (and its - # sandbox group) so a pending request can't desync the next grade or - # leave orphaned children, then propagate. Not a worker failure. - await self._handle_fault(count_start_failure=False) - raise - except TimeoutError as exc: - # The worker is alive but this grade is too slow — a per-grade fault, - # not a worker-startup failure. Kill+respawn, but do NOT count it - # toward the readiness cap, or a few slow problems at the start of a - # run would trip the cap and disable all grading. - await self._handle_fault(count_start_failure=False) - raise CodegenWorkerError(f"grading worker timed out: {exc!r}") from exc - except (ConnectionError, BrokenPipeError, ValueError) as exc: - # Transport broke, or the response overran the StreamReader limit - # (ValueError covers asyncio.LimitOverrunError): the worker died or - # desynced, so this counts toward the startup/readiness cap. - await self._handle_fault() - raise CodegenWorkerError(f"grading worker fault: {exc!r}") from exc - - if not line: # EOF: worker died - await self._handle_fault() - raise CodegenWorkerError("grading worker exited before responding") - - try: - resp = orjson.loads(line) - except orjson.JSONDecodeError as exc: - # Garbage on stdout means the worker desynced; fault it like an EOF - # so it is killed and counted rather than silently reused. - await self._handle_fault() - raise CodegenWorkerError( - f"grading worker emitted non-JSON output: {line!r}" - ) from exc - - if not isinstance(resp, dict): - # Valid JSON that is not an object would make the ok/metrics lookups - # below raise; fault it like the other malformed-response classes. - await self._handle_fault() - raise CodegenWorkerError( - f"grading worker emitted a non-object response: {line!r}" - ) - - if resp.get("id") != req["id"]: - # A mismatched id means the response no longer correlates to the - # request; the stream is desynced, so fault it before trusting - # ok/metrics from a stale or wrong frame. - await self._handle_fault() - raise CodegenWorkerError( - f"grading worker response id mismatch: expected {req['id']}, " - f"got {resp.get('id')!r}" - ) - - if not resp.get("ok"): - # A clean error response is a proven worker; do not restart. - self._worker_proven = True - self._start_failures = 0 - raise CodegenWorkerError(resp.get("error", "unknown grading error")) - - metrics = resp.get("metrics") - if not isinstance(metrics, dict): - # ok:true without a usable metrics dict is a broken worker, not a - # clean result; fault it rather than returning junk to the grader. - await self._handle_fault() - raise CodegenWorkerError( - "grading worker reported success without valid metrics" - ) + pass + def _mark_proven(self) -> None: self._worker_proven = True self._start_failures = 0 - return metrics async def _handle_fault(self, count_start_failure: bool = True) -> None: - # Only failures that mean the worker never became usable (it died or - # emitted garbage before ever succeeding) count toward the startup cap. - # Per-grade timeouts and shutdown cancellation pass count_start_failure= - # False so a slow grade never disables the whole run. + if self._proc is None: + return # already handled; _handle_fault is idempotent if count_start_failure and not self._worker_proven: self._start_failures += 1 + for fut in list(self._pending.values()): + if not fut.done(): + fut.set_exception(CodegenWorkerError("grading worker fault")) + self._pending.clear() tail = await self._kill() _log.debug( lambda: f"codegen worker fault (proven={self._worker_proven}, " @@ -261,11 +256,19 @@ async def _kill(self) -> list[str]: """Kill the worker and return its captured stderr tail (for diagnostics).""" proc, self._proc = self._proc, None task, self._stderr_task = self._stderr_task, None + reader_task, self._reader_task = self._reader_task, None self._close_death_pipe() if proc is not None and proc.returncode is None: _kill_process_group(proc) with contextlib.suppress(ProcessLookupError): await proc.wait() + # Skip cancel/await if _reader_task is calling _kill from within itself to + # avoid self-awaiting deadlock when the reader detects a fault condition. + current = asyncio.current_task() + if reader_task is not None and reader_task is not current: + reader_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await reader_task tail: list[str] = [] if task is not None: # The dead worker's stderr hits EOF, so the drain task finishes; bound @@ -277,7 +280,8 @@ async def _kill(self) -> list[str]: return tail async def aclose(self) -> None: - # Acquire the lock so teardown cannot race with an in-flight _request - # clearing/reading _proc across an await. - async with self._lock: - await self._kill() + for fut in list(self._pending.values()): + if not fut.done(): + fut.cancel() + self._pending.clear() + await self._kill() diff --git a/tests/unit/accuracy/test_codegen_worker_client.py b/tests/unit/accuracy/test_codegen_worker_client.py index 718ac07ae4..af9a1044ee 100644 --- a/tests/unit/accuracy/test_codegen_worker_client.py +++ b/tests/unit/accuracy/test_codegen_worker_client.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import contextlib import os import sys import textwrap @@ -35,26 +36,39 @@ def _write_worker(tmp_path: Path, body: str) -> list[str]: sys.stdout.buffer.flush() """ - -# Worker that records overlap: writes "BUSY" markers around a small delay so a -# second concurrent request would interleave if not serialized. -_ECHO_TRACKED = """ - import sys, orjson, time - inflight = 0 +# Echoes responses with pass@1 == id * 0.1 so each caller can verify it got +# back its OWN response (not another caller's). +_ECHO_ID_IN_METRICS = """ + import sys, orjson for line in sys.stdin.buffer: line = line.strip() if not line: continue req = orjson.loads(line) - inflight += 1 - overlap = inflight > 1 - time.sleep(0.05) - inflight -= 1 - resp = {"id": req["id"], "ok": True, "metrics": {"pass@1": 1.0}, "overlap": overlap} + resp = {"id": req["id"], "ok": True, "metrics": {"pass@1": req["id"] * 0.1}} sys.stdout.buffer.write(orjson.dumps(resp) + b"\\n") sys.stdout.buffer.flush() """ +# Buffers the first 4 requests and responds in REVERSE id order to exercise +# the demux table (correct demux requires id matching, not position matching). +_REVERSE_BATCH_OF_4 = """ + import sys, orjson + buf = [] + for line in sys.stdin.buffer: + line = line.strip() + if not line: + continue + req = orjson.loads(line) + buf.append(req) + if len(buf) == 4: + for r in reversed(buf): + resp = {"id": r["id"], "ok": True, "metrics": {"pass@1": r["id"] * 0.1}} + sys.stdout.buffer.write(orjson.dumps(resp) + b"\\n") + sys.stdout.buffer.flush() + buf = [] +""" + class TestHappyPath: async def test_grade_returns_metrics(self, tmp_path) -> None: @@ -79,18 +93,110 @@ async def test_second_grade_reuses_same_worker(self, tmp_path) -> None: class TestSerialization: - async def test_concurrent_grades_do_not_interleave(self, tmp_path) -> None: - worker = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, _ECHO_TRACKED)) + async def test_concurrent_grades_return_correct_results(self, tmp_path) -> None: + # Previously tested that grades were serialized (lock enforced). + # Now tests that concurrent grades all return correct results without the lock. + w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, _ECHO_OK)) try: results = await asyncio.gather( *[ - worker.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=30) + w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=30) for _ in range(4) ] ) assert all(r == {"pass@1": 1.0} for r in results) finally: - await worker.aclose() + await w.aclose() + + +class TestConcurrency: + async def test_concurrent_grades_all_complete(self, tmp_path) -> None: + w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, _ECHO_OK)) + try: + results = await asyncio.gather( + *[ + w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=30) + for _ in range(5) + ] + ) + assert all(r == {"pass@1": 1.0} for r in results) + finally: + await w.aclose() + + async def test_concurrent_grades_demux_by_id_not_position(self, tmp_path) -> None: + # 4 concurrent grades; mock responds in reverse order. + # If demux were position-based, callers would get wrong metrics. + w = CodegenGradingWorker( + worker_cmd=_write_worker(tmp_path, _REVERSE_BATCH_OF_4) + ) + try: + results = await asyncio.gather( + *[ + w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=30) + for _ in range(4) + ] + ) + # IDs 1-4 → pass@1 values 0.1, 0.2, 0.3, 0.4 (one per caller) + values = sorted(r["pass@1"] for r in results) + assert values == pytest.approx([0.1, 0.2, 0.3, 0.4]) + finally: + await w.aclose() + + async def test_fault_cancels_all_pending_futures(self, tmp_path) -> None: + # Worker dies immediately after the first line — all concurrent callers + # should raise CodegenWorkerError, not hang. + w = CodegenGradingWorker( + worker_cmd=_write_worker( + tmp_path, + """ + import sys + sys.stdin.buffer.readline() # consume one line then exit + """, + ) + ) + try: + with pytest.raises(CodegenWorkerError): + await asyncio.gather( + *[ + w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=10) + for _ in range(3) + ], + return_exceptions=False, + ) + finally: + await w.aclose() + + async def test_stale_id_after_timeout_does_not_crash(self, tmp_path) -> None: + # Reader receives a response for an id that the caller already timed out on. + # The stale future was already removed from _pending; the reader must skip it. + # Use _ECHO_OK with a very short timeout so the grade times out, then send + # a second grade to prove the worker (if restarted) still works. + w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, _ECHO_OK)) + try: + with pytest.raises(CodegenWorkerError): + await w.grade_codegen( + [{"input_output": "{}"}], [["x"]], timeout=0.000001 + ) + # If stale id handling is broken, the second grade would hang or crash. + # Give it a real timeout; it may or may not succeed (worker restarted). + finally: + await w.aclose() + + async def test_aclose_with_pending_futures_does_not_hang(self, tmp_path) -> None: + hang_worker = """ + import sys, time + for line in sys.stdin.buffer: + time.sleep(3600) + """ + w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, hang_worker)) + grade_task = asyncio.create_task( + w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=60) + ) + await asyncio.sleep(0.05) # let grade_task start and block + await w.aclose() # must not hang even with grade_task pending + grade_task.cancel() + with contextlib.suppress(asyncio.CancelledError, CodegenWorkerError): + await grade_task # The very first grade (client request id==1) hangs forever to trigger a @@ -269,7 +375,8 @@ async def test_worker_that_dies_on_start_hits_cap(self, tmp_path) -> None: # Always echoes a fixed WRONG id regardless of the request id, simulating a -# worker whose responses no longer correlate to requests (a desync). +# worker whose responses no longer correlate to requests (a desync). With the +# demux reader, the stale id is silently dropped and the caller times out. _ECHO_WRONG_ID = """ import sys, orjson for line in sys.stdin.buffer: @@ -300,13 +407,16 @@ async def test_oversized_line_faults_and_kills_worker(self, tmp_path) -> None: class TestResponseIdMismatch: async def test_wrong_id_faults_and_kills_worker(self, tmp_path) -> None: + # With the demux reader, a wrong id is treated as a stale response and + # silently dropped. The caller's future is never resolved, so it times out + # and calls _handle_fault, which kills the worker. worker = CodegenGradingWorker( worker_cmd=_write_worker(tmp_path, _ECHO_WRONG_ID) ) try: with pytest.raises(CodegenWorkerError): await worker.grade_codegen( - [{"input_output": "{}"}], [["x"]], timeout=30 + [{"input_output": "{}"}], [["x"]], timeout=0.5 ) assert worker._proc is None finally: From ef17b7f5e954d1130b3b09d23e301162a4f6b79b Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Thu, 30 Jul 2026 17:27:51 -0700 Subject: [PATCH 05/26] test(accuracy): concurrent multi-problem codegen grading e2e test Signed-off-by: Elias Bermudez --- .../test_lcb_codegen_worker_e2e.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tests/component_integration/test_lcb_codegen_worker_e2e.py b/tests/component_integration/test_lcb_codegen_worker_e2e.py index 581ed70439..fde73ce8c1 100644 --- a/tests/component_integration/test_lcb_codegen_worker_e2e.py +++ b/tests/component_integration/test_lcb_codegen_worker_e2e.py @@ -9,6 +9,8 @@ from __future__ import annotations +import asyncio + import orjson import pytest @@ -45,3 +47,25 @@ async def test_worker_grades_correct_stdin_solution() -> None: assert float(metrics["pass@1"]) == 1.0 finally: await worker.aclose() + + +@pytest.mark.slow +@pytest.mark.asyncio +async def test_worker_grades_multiple_problems_concurrently() -> None: + """N concurrent grade_codegen() calls all resolve correctly. + + This exercises the batch-drain path: all N requests are sent before the + worker responds, so they are drained into a single codegen_metrics call and + processed in parallel by lighteval's ProcessPoolExecutor. + """ + worker = CodegenGradingWorker() + sample, code = _sample_and_solution() + n = 4 + try: + results = await asyncio.gather( + *[worker.grade_codegen(sample, code, timeout=240) for _ in range(n)] + ) + assert len(results) == n + assert all(float(r["pass@1"]) == 1.0 for r in results), results + finally: + await worker.aclose() From 5a39efff6bfc78120db8c3575faabd1ffa338983 Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Fri, 31 Jul 2026 09:12:23 -0700 Subject: [PATCH 06/26] fix(accuracy): correct batch drain, aclose exception type, reader robustness - Replace _drain_fd (broken: select() unreliable after BufferedReader pulls kernel data into userspace) and _drain_seekable with _drain_buffered that uses peek() for BufferedReader and a seekable fallback for BytesIO in tests. Remove unused _parse_batch and import select. - aclose() now sets CodegenWorkerError on pending futures instead of calling cancel(), so shutdown does not propagate CancelledError to grader callers that only catch CodegenWorkerError. - Extract _dispatch_response helper so _run_reader complexity stays within the C901 limit; unhashable req_id from a desynced worker now triggers a fault instead of killing the reader task with TypeError. - grade_codegen calls await stdin.drain() after write() to respect backpressure. - Add test_list_shaped_pass_at_1_is_preserved to TestHandleBatch (guards the silent-0.000 bug for list-shaped pass@1 from lighteval). - Remove unused _ECHO_ID_IN_METRICS constant from client tests. Signed-off-by: Elias Bermudez --- .../accuracy/graders/_codegen_worker.py | 80 +++++++++---------- .../graders/_codegen_worker_client.py | 49 +++++++----- tests/unit/accuracy/test_codegen_worker.py | 12 +++ .../accuracy/test_codegen_worker_client.py | 14 ---- 4 files changed, 79 insertions(+), 76 deletions(-) diff --git a/src/aiperf/accuracy/graders/_codegen_worker.py b/src/aiperf/accuracy/graders/_codegen_worker.py index 9184c54635..337a3e7e0a 100644 --- a/src/aiperf/accuracy/graders/_codegen_worker.py +++ b/src/aiperf/accuracy/graders/_codegen_worker.py @@ -18,7 +18,6 @@ import math import multiprocessing as mp import os -import select import signal import sys import threading @@ -152,37 +151,34 @@ def handle_batch( return [r for r in responses if r is not None] -def _drain_fd(stdin: BinaryIO, stdin_fd: int, batch_raw: list[bytes]) -> None: - """Non-blocking drain of already-queued lines from a real file descriptor.""" - while True: - ready, _, _ = select.select([stdin_fd], [], [], 0) - if not ready: - break - line = stdin.readline() - if not line: - break - line = line.strip() - if line: - batch_raw.append(line) - - -def _drain_seekable(stdin: BinaryIO, batch_raw: list[bytes]) -> None: - """Greedy drain for non-fd streams such as BytesIO (used in tests).""" - remaining = stdin.read() - for raw_line in remaining.split(b"\n"): - raw_line = raw_line.strip() - if raw_line: - batch_raw.append(raw_line) - - -def _parse_batch(batch_raw: list[bytes]) -> list[Any]: - reqs: list[Any] = [] - for raw in batch_raw: - try: - reqs.append(orjson.loads(raw)) - except orjson.JSONDecodeError as exc: - reqs.append({"_parse_error": str(exc)}) - return reqs +def _drain_buffered(stdin: BinaryIO) -> list[bytes]: + """Drain all lines already buffered in stdin without blocking. + + For BufferedReader (sys.stdin.buffer in production), uses peek() to check + the userspace buffer; an empty peek means the next readline() would block, + so we stop. This is the correct non-blocking check because readline() pulls + kernel data into userspace first, making select() on the raw fd unreliable. + + For seekable streams without peek() (e.g. BytesIO used in tests), reads all + remaining data at once — safe because BytesIO is already fully in memory. + """ + lines: list[bytes] = [] + if hasattr(stdin, "peek"): + while True: + if not stdin.peek(0): + break + line = stdin.readline() + if not line: + break + line = line.strip() + if line: + lines.append(line) + else: + for raw_line in stdin.read().split(b"\n"): + raw_line = raw_line.strip() + if raw_line: + lines.append(raw_line) + return lines def run_worker_loop( @@ -197,11 +193,6 @@ def run_worker_loop( already-queued requests to form a batch. Calls codegen_fn once per batch so lighteval's ProcessPoolExecutor can process multiple problems in parallel. """ - try: - stdin_fd: int | None = stdin.fileno() - except (AttributeError, OSError): - stdin_fd = None - while True: first = stdin.readline() if not first: @@ -210,13 +201,14 @@ def run_worker_loop( if not first: continue batch_raw: list[bytes] = [first] - if stdin_fd is not None: - _drain_fd(stdin, stdin_fd, batch_raw) - else: - _drain_seekable(stdin, batch_raw) - for resp in handle_batch( - _parse_batch(batch_raw), codegen_fn, compute_metrics_fn - ): + batch_raw.extend(_drain_buffered(stdin)) + reqs: list[Any] = [] + for raw in batch_raw: + try: + reqs.append(orjson.loads(raw)) + except orjson.JSONDecodeError as exc: + reqs.append({"_parse_error": str(exc)}) + for resp in handle_batch(reqs, codegen_fn, compute_metrics_fn): out.write(orjson.dumps(resp) + b"\n") out.flush() diff --git a/src/aiperf/accuracy/graders/_codegen_worker_client.py b/src/aiperf/accuracy/graders/_codegen_worker_client.py index 8df4a51081..8a2574400f 100644 --- a/src/aiperf/accuracy/graders/_codegen_worker_client.py +++ b/src/aiperf/accuracy/graders/_codegen_worker_client.py @@ -105,6 +105,7 @@ async def grade_codegen( self._pending[req_id] = fut assert self._proc is not None and self._proc.stdin self._proc.stdin.write(orjson.dumps(req) + b"\n") + await self._proc.stdin.drain() try: return await asyncio.wait_for(fut, timeout) except TimeoutError as exc: @@ -182,6 +183,34 @@ async def _drain_stderr(self, reader: asyncio.StreamReader | None) -> None: async for line in reader: self._stderr_tail.append(line.decode("utf-8", "replace").rstrip("\n")) + def _dispatch_response(self, resp: dict[str, Any]) -> bool: + """Resolve the pending future for resp['id']. + + Returns True on success (including stale/timed-out ids, which are + silently skipped). Returns False when the worker has desynced and the + caller should fault: unhashable id, or ok=True with no metrics dict. + """ + req_id = resp.get("id") + try: + fut = self._pending.pop(req_id, None) + except TypeError: + # req_id is unhashable (e.g., a list); the worker desynced + return False + if fut is None or fut.done(): + return True # stale id (caller already timed out or cancelled) + if not resp.get("ok"): + fut.set_exception( + CodegenWorkerError(resp.get("error", "unknown grading error")) + ) + self._mark_proven() + return True + metrics = resp.get("metrics") + if not isinstance(metrics, dict): + return False + fut.set_result(metrics) + self._mark_proven() + return True + async def _run_reader(self) -> None: """Read worker responses and resolve the corresponding pending futures by id.""" assert self._proc is not None and self._proc.stdout @@ -204,25 +233,9 @@ async def _run_reader(self) -> None: except orjson.JSONDecodeError: await self._handle_fault() return - if not isinstance(resp, dict): + if not isinstance(resp, dict) or not self._dispatch_response(resp): await self._handle_fault() return - req_id = resp.get("id") - fut = self._pending.pop(req_id, None) - if fut is None or fut.done(): - continue # stale id (caller already timed out) or cancelled - if not resp.get("ok"): - fut.set_exception( - CodegenWorkerError(resp.get("error", "unknown grading error")) - ) - self._mark_proven() - else: - metrics = resp.get("metrics") - if not isinstance(metrics, dict): - await self._handle_fault() - return - fut.set_result(metrics) - self._mark_proven() except asyncio.CancelledError: pass @@ -282,6 +295,6 @@ async def _kill(self) -> list[str]: async def aclose(self) -> None: for fut in list(self._pending.values()): if not fut.done(): - fut.cancel() + fut.set_exception(CodegenWorkerError("grading worker closed")) self._pending.clear() await self._kill() diff --git a/tests/unit/accuracy/test_codegen_worker.py b/tests/unit/accuracy/test_codegen_worker.py index 67db0b7dbf..24291d5138 100644 --- a/tests/unit/accuracy/test_codegen_worker.py +++ b/tests/unit/accuracy/test_codegen_worker.py @@ -120,6 +120,18 @@ def test_parse_error_sentinel_produces_error_response(self) -> None: assert "unexpected token" in resps[0]["error"] assert resps[1]["ok"] + def test_list_shaped_pass_at_1_is_preserved(self) -> None: + # lighteval returns pass@1 as a list on some pins; _coerce_metrics must + # preserve it rather than drop it (silent 0.000 bug if dropped). + def list_metrics(_results: dict, **_kw: Any) -> dict[str, Any]: + return {"pass@1": [1.0]} + + resps = worker.handle_batch( + [self._req(1)], _fake_codegen_batch_ok, list_metrics + ) + assert resps[0]["ok"] is True + assert resps[0]["metrics"]["pass@1"] == [1.0] + def test_non_finite_metric_values_are_dropped(self) -> None: # NaN/Inf must not cross the JSONL boundary (repo NaN/Inf discipline). def _nan_inf( diff --git a/tests/unit/accuracy/test_codegen_worker_client.py b/tests/unit/accuracy/test_codegen_worker_client.py index af9a1044ee..ed249aefce 100644 --- a/tests/unit/accuracy/test_codegen_worker_client.py +++ b/tests/unit/accuracy/test_codegen_worker_client.py @@ -36,20 +36,6 @@ def _write_worker(tmp_path: Path, body: str) -> list[str]: sys.stdout.buffer.flush() """ -# Echoes responses with pass@1 == id * 0.1 so each caller can verify it got -# back its OWN response (not another caller's). -_ECHO_ID_IN_METRICS = """ - import sys, orjson - for line in sys.stdin.buffer: - line = line.strip() - if not line: - continue - req = orjson.loads(line) - resp = {"id": req["id"], "ok": True, "metrics": {"pass@1": req["id"] * 0.1}} - sys.stdout.buffer.write(orjson.dumps(resp) + b"\\n") - sys.stdout.buffer.flush() -""" - # Buffers the first 4 requests and responds in REVERSE id order to exercise # the demux table (correct demux requires id matching, not position matching). _REVERSE_BATCH_OF_4 = """ From ea3b3d0e6be7725bae3e4e5e3dbbcdce37dbca81 Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Fri, 31 Jul 2026 13:05:08 -0700 Subject: [PATCH 07/26] fix(accuracy): non-blocking drain via O_NONBLOCK instead of peek(0); remove plan doc Signed-off-by: Elias Bermudez --- .../2026-07-29-codegen-grade-concurrency.md | 930 ------------------ .../accuracy/graders/_codegen_worker.py | 57 +- 2 files changed, 40 insertions(+), 947 deletions(-) delete mode 100644 docs/superpowers/plans/2026-07-29-codegen-grade-concurrency.md diff --git a/docs/superpowers/plans/2026-07-29-codegen-grade-concurrency.md b/docs/superpowers/plans/2026-07-29-codegen-grade-concurrency.md deleted file mode 100644 index 2249c403b0..0000000000 --- a/docs/superpowers/plans/2026-07-29-codegen-grade-concurrency.md +++ /dev/null @@ -1,930 +0,0 @@ -# Codegen Grade Concurrency Implementation Plan - -> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. - -**Goal:** Allow N concurrent `grade_codegen()` calls to complete in ~max(individual) time by replacing the serializing `asyncio.Lock` with an `id → Future` demux table on the client and a batch-drain loop on the worker. - -**Architecture:** The client drops its `asyncio.Lock` and instead multiplexes concurrent requests over the same stdin pipe using request ids; a persistent reader task demuxes responses back to individual `asyncio.Future` objects. The worker reads the first blocking request, non-blocking drains any queued requests, then calls `codegen_metrics` once with all batched samples so lighteval's `ProcessPoolExecutor` handles all problems in parallel. - -**Tech Stack:** Python 3.11+ asyncio, `orjson`, `lighteval` (`codegen_metrics`, `compute_metrics_from_results`), `select` (POSIX non-blocking stdin drain), `pytest-asyncio` - -## Global Constraints - -- Python 3.11+; use `asyncio.get_running_loop()`, not `asyncio.get_event_loop()` -- No new threads in the worker (stays single-threaded at fork) -- All existing tests in `tests/unit/accuracy/test_codegen_worker.py` and `test_codegen_worker_client.py` must remain green after each task -- `_handle_fault` must remain idempotent (called from both reader task and caller) -- `ruff format . && ruff check --fix .` must pass after each commit -- `pre-commit run --all-files` must pass before each commit -- Every new function needs a type hint on all parameters and return value -- Every new Pydantic field needs `Field(description=...)` (not applicable here — no new models) - ---- - -### Task 1: Worker — batch-drain loop and `handle_batch` - -**Spec:** `docs/superpowers/specs/2026-07-29-codegen-grade-concurrency-design.md` § "Worker changes" - -**Files:** -- Modify: `src/aiperf/accuracy/graders/_codegen_worker.py` -- Modify: `tests/unit/accuracy/test_codegen_worker.py` - -**Interfaces:** -- Produces: `handle_batch(reqs, codegen_fn, compute_metrics_fn)` — takes a list of raw request dicts, returns a list of JSONL-ready response dicts (one per input, same order) -- Produces: `run_worker_loop(stdin, out, codegen_fn, compute_metrics_fn)` — updated signature (adds `compute_metrics_fn` param) -- Removes: `handle_request` (dead code after this task; its tests are migrated to `handle_batch`) - ---- - -- [ ] **Step 1: Write failing tests for `handle_batch`** - -Add a new `TestHandleBatch` class in `tests/unit/accuracy/test_codegen_worker.py`. Place it after the existing `TestHandleRequest` class. - -The mock `codegen_fn` for `handle_batch` must match the new signature that also accepts `compute_metrics_fn`. But `codegen_fn` itself is still the original `(samples, generations, ...) -> (metrics, results)` signature. The `compute_metrics_fn` is a separate argument to `handle_batch`. - -Add at the top of the file alongside the existing fakes: - -```python -def _fake_compute_metrics(results: dict, k_list: list[int] | None = None) -> dict[str, Any]: - # Mirrors compute_metrics_from_results: returns {"pass@1": } using - # the single-problem results dict {0: [[True, True, ...]]} passed by handle_batch. - result_list = results.get(0, [[-2]]) # [-2] = compile error - if result_list and all(x > 0 for x in result_list[0]): - return {"pass@1": 1.0} - return {"pass@1": 0.0} - - -def _fake_codegen_batch_ok( - samples: list, generations: list, **_kwargs: Any -) -> tuple[dict[str, Any], dict[int, list]]: - # Returns aggregate metrics (ignored by handle_batch) and per-problem results. - n = len(samples) - raw_results = {i: [[True]] for i in range(n)} # all pass - return {"pass@1": 1.0}, raw_results - - -def _fake_codegen_batch_boom( - samples: list, generations: list, **_kwargs: Any -) -> tuple[dict[str, Any], dict[int, list]]: - raise RuntimeError("pool exploded") -``` - -Then add `TestHandleBatch`: - -```python -class TestHandleBatch: - def _req(self, req_id: int) -> dict[str, Any]: - return { - "id": req_id, - "evaluation_sample": [{"input_output": "{}"}], - "generated_code": [["x"]], - } - - def test_single_request_returns_one_ok_response(self) -> None: - resps = worker.handle_batch( - [self._req(1)], _fake_codegen_batch_ok, _fake_compute_metrics - ) - assert len(resps) == 1 - assert resps[0] == {"id": 1, "ok": True, "metrics": {"pass@1": 1.0}} - - def test_batch_of_n_calls_codegen_fn_once(self) -> None: - call_count = 0 - - def counting_codegen(samples, generations, **kwargs): - nonlocal call_count - call_count += 1 - n = len(samples) - return {"pass@1": 1.0}, {i: [[True]] for i in range(n)} - - reqs = [self._req(i) for i in range(1, 5)] - resps = worker.handle_batch(reqs, counting_codegen, _fake_compute_metrics) - assert call_count == 1 - assert len(resps) == 4 - assert all(r["ok"] for r in resps) - assert [r["id"] for r in resps] == [1, 2, 3, 4] - - def test_response_order_matches_request_order(self) -> None: - reqs = [self._req(i) for i in [7, 3, 99]] - resps = worker.handle_batch(reqs, _fake_codegen_batch_ok, _fake_compute_metrics) - assert [r["id"] for r in resps] == [7, 3, 99] - - def test_batch_exception_returns_error_for_all(self) -> None: - reqs = [self._req(i) for i in range(1, 4)] - resps = worker.handle_batch(reqs, _fake_codegen_batch_boom, _fake_compute_metrics) - assert len(resps) == 3 - assert all(not r["ok"] for r in resps) - assert all("pool exploded" in r["error"] for r in resps) - - def test_malformed_request_in_batch_does_not_affect_others(self) -> None: - reqs = [ - self._req(1), - {"id": 2}, # missing evaluation_sample + generated_code - self._req(3), - ] - resps = worker.handle_batch(reqs, _fake_codegen_batch_ok, _fake_compute_metrics) - assert len(resps) == 3 - assert resps[0] == {"id": 1, "ok": True, "metrics": {"pass@1": 1.0}} - assert resps[1]["id"] == 2 - assert not resps[1]["ok"] - assert resps[2] == {"id": 3, "ok": True, "metrics": {"pass@1": 1.0}} - - def test_non_object_request_in_batch_is_error(self) -> None: - reqs = [[1, 2, 3], self._req(5)] - resps = worker.handle_batch(reqs, _fake_codegen_batch_ok, _fake_compute_metrics) - assert len(resps) == 2 - assert resps[0]["id"] is None - assert not resps[0]["ok"] - assert resps[1] == {"id": 5, "ok": True, "metrics": {"pass@1": 1.0}} - - def test_parse_error_sentinel_produces_error_response(self) -> None: - # run_worker_loop encodes JSON decode errors as {"_parse_error": "..."}. - reqs = [{"_parse_error": "unexpected token"}, self._req(2)] - resps = worker.handle_batch(reqs, _fake_codegen_batch_ok, _fake_compute_metrics) - assert len(resps) == 2 - assert resps[0]["id"] is None - assert not resps[0]["ok"] - assert "unexpected token" in resps[0]["error"] - assert resps[1]["ok"] -``` - -Also add a test for the new `run_worker_loop` batch-drain behaviour. Add to a new class `TestRunWorkerLoopBatch` after `TestHandleBatch`: - -```python -class TestRunWorkerLoopBatch: - def _run( - self, - payloads: list[dict[str, Any]], - codegen_fn=_fake_codegen_batch_ok, - compute_metrics_fn=_fake_compute_metrics, - ) -> list[dict[str, Any]]: - # Write all payloads to a BytesIO pipe so they are already queued when - # run_worker_loop reads; this exercises the non-blocking drain path. - data = b"".join(orjson.dumps(p) + b"\n" for p in payloads) - stdin = io.BytesIO(data) - out = io.BytesIO() - worker.run_worker_loop(stdin, out, codegen_fn, compute_metrics_fn) - out.seek(0) - return [orjson.loads(line) for line in out if line.strip()] - - def _req(self, req_id: int) -> dict[str, Any]: - return { - "id": req_id, - "evaluation_sample": [{"input_output": "{}"}], - "generated_code": [["x"]], - } - - def test_pre_queued_requests_are_batched_in_one_call(self) -> None: - call_count = 0 - - def counting_codegen(samples, generations, **kwargs): - nonlocal call_count - call_count += 1 - n = len(samples) - return {"pass@1": 1.0}, {i: [[True]] for i in range(n)} - - reqs = [self._req(i) for i in range(1, 4)] - resps = self._run(reqs, counting_codegen) - assert call_count == 1 - assert len(resps) == 3 - - def test_responses_carry_correct_ids(self) -> None: - reqs = [self._req(i) for i in [10, 20, 30]] - resps = self._run(reqs) - assert {r["id"] for r in resps} == {10, 20, 30} -``` - -- [ ] **Step 2: Run the new tests to verify they fail** - -```bash -uv run pytest tests/unit/accuracy/test_codegen_worker.py::TestHandleBatch tests/unit/accuracy/test_codegen_worker.py::TestRunWorkerLoopBatch -v 2>&1 | head -40 -``` - -Expected: `AttributeError: module ... has no attribute 'handle_batch'` or `TypeError` from wrong arg count on `run_worker_loop`. - -- [ ] **Step 3: Implement `handle_batch` and update `run_worker_loop` in `_codegen_worker.py`** - -Open `src/aiperf/accuracy/graders/_codegen_worker.py`. - -**3a — Add `import select` at the top of the file** (after the stdlib imports block). - -**3b — Add `handle_batch` after the existing `_is_number` function:** - -```python -def handle_batch( - reqs: list[Any], - codegen_fn: Callable[..., tuple[dict[str, Any], Any]], - compute_metrics_fn: Callable[..., dict[str, Any]], -) -> list[dict[str, Any]]: - """Grade a batch of requests with a single codegen_fn call. - - Calls codegen_fn once with all well-formed requests batched together so - lighteval's ProcessPoolExecutor can process multiple problems in parallel. - Never raises: all failures become error responses so a bad batch cannot - kill the worker loop. - """ - all_samples: list[Any] = [] - all_generations: list[Any] = [] - id_map: list[tuple[int, Any]] = [] # (batch_position, req_id) - responses: list[dict[str, Any] | None] = [None] * len(reqs) - - for i, req in enumerate(reqs): - if isinstance(req, dict) and "_parse_error" in req: - responses[i] = { - "id": None, - "ok": False, - "error": f"bad json: {req['_parse_error']}", - } - continue - if not isinstance(req, dict): - responses[i] = { - "id": None, - "ok": False, - "error": "malformed request: expected object", - } - continue - req_id = req.get("id") - try: - all_samples.append(req["evaluation_sample"]) - all_generations.append(req["generated_code"]) - id_map.append((i, req_id)) - except (KeyError, TypeError) as exc: - responses[i] = { - "id": req_id, - "ok": False, - "error": f"malformed request: {exc!r}", - } - - if all_samples: - batch_error: str | None = None - raw_results: dict[int, Any] = {} - try: - _, raw_results = codegen_fn( - all_samples, - all_generations, - k_list=list(_LCB_PASS_AT_K), - num_process_evaluate=_LCB_NUM_PROCESSES, - ) - except Exception as exc: - batch_error = _truncate_error(f"{type(exc).__name__}: {exc}") - - for pos, (req_idx, req_id) in enumerate(id_map): - if batch_error is not None: - responses[req_idx] = {"id": req_id, "ok": False, "error": batch_error} - else: - try: - metrics = compute_metrics_fn( - {0: raw_results[pos]}, - k_list=list(_LCB_PASS_AT_K), - ) - responses[req_idx] = { - "id": req_id, - "ok": True, - "metrics": _coerce_metrics(metrics), - } - except Exception as exc: - responses[req_idx] = { - "id": req_id, - "ok": False, - "error": _truncate_error(f"{type(exc).__name__}: {exc}"), - } - - return [r for r in responses if r is not None] -``` - -**3c — Replace `run_worker_loop`:** - -```python -def run_worker_loop( - stdin: BinaryIO, - out: BinaryIO, - codegen_fn: Callable[..., tuple[dict[str, Any], Any]], - compute_metrics_fn: Callable[..., dict[str, Any]], -) -> None: - """Serve JSONL grading requests until stdin EOF. - - Blocks on the first request of each cycle, then non-blocking drains any - already-queued requests to form a batch. Calls codegen_fn once per batch so - lighteval's ProcessPoolExecutor can process multiple problems in parallel. - """ - import select - - stdin_fd = stdin.fileno() - while True: - first = stdin.readline() - if not first: - break # EOF: client closed stdin, clean exit - first = first.strip() - if not first: - continue - batch_raw: list[bytes] = [first] - - while True: - ready, _, _ = select.select([stdin_fd], [], [], 0) - if not ready: - break - line = stdin.readline() - if not line: - break - line = line.strip() - if line: - batch_raw.append(line) - - reqs: list[Any] = [] - for raw in batch_raw: - try: - reqs.append(orjson.loads(raw)) - except orjson.JSONDecodeError as exc: - reqs.append({"_parse_error": str(exc)}) - - for resp in handle_batch(reqs, codegen_fn, compute_metrics_fn): - out.write(orjson.dumps(resp) + b"\n") - out.flush() -``` - -**3d — Remove `handle_request`** (the whole function and its docstring). Do not replace it with a comment. - -**3e — Update `main()` to import and pass `compute_metrics_from_results`:** - -```python -def main() -> None: - protocol_out = _install_stdout_guard() - _start_death_watcher() - _force_fork() - from lighteval.tasks.tasks.lcb.codegen_metrics import ( - codegen_metrics, - compute_metrics_from_results, - ) - - run_worker_loop(sys.stdin.buffer, protocol_out, codegen_metrics, compute_metrics_from_results) -``` - -- [ ] **Step 4: Migrate `TestHandleRequest` tests to `TestHandleBatch` equivalents** - -The `TestHandleRequest` class in `test_codegen_worker.py` is now orphaned (`handle_request` was removed). Replace it with `TestHandleBatch` (already written in Step 1 above — just remove the old class). Also remove the `_fake_codegen_ok`, `_fake_codegen_boom`, `_fake_codegen_list_pass` helpers if they are only used by the old `TestHandleRequest`; they are replaced by the new batch-aware fakes from Step 1. - -Check if any other test in the file (e.g., `TestRunWorkerLoop`, `TestStdoutGuard`) still calls `handle_request` directly; update those to use `handle_batch` with the batch-aware fakes. - -Grep to find remaining usages: - -```bash -grep -n "handle_request\|_fake_codegen_ok\|_fake_codegen_boom\|_fake_codegen_list_pass" \ - tests/unit/accuracy/test_codegen_worker.py -``` - -For any `TestRunWorkerLoop` tests that currently call `run_worker_loop` with 3 args, update to pass `_fake_compute_metrics` as the 4th argument. - -- [ ] **Step 5: Run all unit tests for the worker to verify they pass** - -```bash -uv run pytest tests/unit/accuracy/test_codegen_worker.py -v -``` - -Expected: all tests pass, including the new `TestHandleBatch` and `TestRunWorkerLoopBatch` classes. - -- [ ] **Step 6: Lint** - -```bash -ruff format . && ruff check --fix . -``` - -- [ ] **Step 7: Commit** - -```bash -git add src/aiperf/accuracy/graders/_codegen_worker.py \ - tests/unit/accuracy/test_codegen_worker.py -git commit -s -m "feat(accuracy): batch-drain worker loop for codegen grading concurrency" -``` - ---- - -### Task 2: Client — drop lock, add demux table and reader task - -**Spec:** `docs/superpowers/specs/2026-07-29-codegen-grade-concurrency-design.md` § "Client changes" - -**Files:** -- Modify: `src/aiperf/accuracy/graders/_codegen_worker_client.py` -- Modify: `tests/unit/accuracy/test_codegen_worker_client.py` - -**Interfaces:** -- Consumes: worker protocol from Task 1 (JSONL responses carry `id`, `ok`, `metrics`) -- Produces: `CodegenGradingWorker` with the same public API (`grade_codegen`, `aclose`) but concurrent-safe without a global lock - ---- - -- [ ] **Step 1: Write failing concurrency tests** - -Open `tests/unit/accuracy/test_codegen_worker_client.py`. - -Add the following mock worker scripts near the top of the file alongside `_ECHO_OK`: - -```python -# Echoes responses with pass@1 == id * 0.1 so each caller can verify it got -# back its OWN response (not another caller's). -_ECHO_ID_IN_METRICS = """ - import sys, orjson - for line in sys.stdin.buffer: - line = line.strip() - if not line: - continue - req = orjson.loads(line) - resp = {"id": req["id"], "ok": True, "metrics": {"pass@1": req["id"] * 0.1}} - sys.stdout.buffer.write(orjson.dumps(resp) + b"\\n") - sys.stdout.buffer.flush() -""" - -# Buffers the first 4 requests and responds in REVERSE id order to exercise -# the demux table (correct demux requires id matching, not position matching). -_REVERSE_BATCH_OF_4 = """ - import sys, orjson - buf = [] - for line in sys.stdin.buffer: - line = line.strip() - if not line: - continue - req = orjson.loads(line) - buf.append(req) - if len(buf) == 4: - for r in reversed(buf): - resp = {"id": r["id"], "ok": True, "metrics": {"pass@1": r["id"] * 0.1}} - sys.stdout.buffer.write(orjson.dumps(resp) + b"\\n") - sys.stdout.buffer.flush() - buf = [] -""" -``` - -Add a new `TestConcurrency` class: - -```python -class TestConcurrency: - async def test_concurrent_grades_all_complete(self, tmp_path) -> None: - w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, _ECHO_OK)) - try: - results = await asyncio.gather(*[ - w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=30) - for _ in range(5) - ]) - assert all(r == {"pass@1": 1.0} for r in results) - finally: - await w.aclose() - - async def test_concurrent_grades_demux_by_id_not_position(self, tmp_path) -> None: - # 4 concurrent grades; mock responds in reverse order. - # If demux were position-based, callers would get wrong metrics. - w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, _REVERSE_BATCH_OF_4)) - try: - results = await asyncio.gather(*[ - w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=30) - for _ in range(4) - ]) - # IDs 1-4 → pass@1 values 0.1, 0.2, 0.3, 0.4 (one per caller) - values = sorted(r["pass@1"] for r in results) - assert values == pytest.approx([0.1, 0.2, 0.3, 0.4]) - finally: - await w.aclose() - - async def test_fault_cancels_all_pending_futures(self, tmp_path) -> None: - # Worker dies immediately after the first line — all concurrent callers - # should raise CodegenWorkerError, not hang. - w = CodegenGradingWorker( - worker_cmd=_write_worker( - tmp_path, - """ - import sys - sys.stdin.buffer.readline() # consume one line then exit - """, - ) - ) - try: - with pytest.raises(Exception): # CodegenWorkerError or ExceptionGroup - await asyncio.gather(*[ - w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=10) - for _ in range(3) - ], return_exceptions=False) - finally: - await w.aclose() - - async def test_stale_id_after_timeout_does_not_crash(self, tmp_path) -> None: - # Reader receives a response for an id that the caller already timed out on. - # The stale future was already removed from _pending; the reader must skip it. - # Use _ECHO_OK with a very short timeout so the grade times out, then send - # a second grade to prove the worker (if restarted) still works. - w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, _ECHO_OK)) - try: - with pytest.raises(CodegenWorkerError): - await w.grade_codegen( - [{"input_output": "{}"}], [["x"]], timeout=0.000001 - ) - # If stale id handling is broken, the second grade would hang or crash. - # Give it a real timeout; it may or may not succeed (worker restarted). - finally: - await w.aclose() - - async def test_aclose_with_pending_futures_does_not_hang(self, tmp_path) -> None: - hang_worker = """ - import sys, time - for line in sys.stdin.buffer: - time.sleep(3600) - """ - w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, hang_worker)) - grade_task = asyncio.create_task( - w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=60) - ) - await asyncio.sleep(0.05) # let grade_task start and block - await w.aclose() # must not hang even with grade_task pending - grade_task.cancel() - with contextlib.suppress(asyncio.CancelledError, CodegenWorkerError): - await grade_task -``` - -Add `import contextlib` at the top of the test file if not already present. - -- [ ] **Step 2: Run the new tests to verify they fail** - -```bash -uv run pytest tests/unit/accuracy/test_codegen_worker_client.py::TestConcurrency -v 2>&1 | head -30 -``` - -Expected: failures because the lock still serializes requests (demux test would hang or return wrong values) or `AttributeError` if the test references methods not yet on the class. - -- [ ] **Step 3: Rewrite `CodegenGradingWorker` in `_codegen_worker_client.py`** - -The full rewrite of the class. Replace the existing class body (not the module-level helpers `_kill_process_group`, `CodegenWorkerError`, `_STREAM_LIMIT`, etc. — keep those unchanged). - -**3a — Update `__init__`:** remove `self._lock`, add `self._pending`, `self._reader_task`, and `self._spawn_lock`: - -```python -def __init__( - self, - worker_cmd: list[str] | None = None, - max_start_failures: int = 3, -) -> None: - self._cmd = worker_cmd or _DEFAULT_WORKER_CMD - self._max_start_failures = max_start_failures - self._proc: asyncio.subprocess.Process | None = None - self._spawn_lock = asyncio.Lock() - self._pending: dict[int, asyncio.Future[dict[str, Any]]] = {} - self._reader_task: asyncio.Task[None] | None = None - self._next_id = 0 - self._start_failures = 0 - self._worker_proven = False - self._stderr_tail: deque[str] = deque(maxlen=_STDERR_TAIL_LINES) - self._stderr_task: asyncio.Task[None] | None = None - self._death_w: int | None = None -``` - -**3b — Replace `grade_codegen`:** - -```python -async def grade_codegen( - self, - evaluation_sample: list[dict[str, str]], - generated_code: list[list[str]], - timeout: float, -) -> dict[str, Any]: - if self._start_failures >= self._max_start_failures: - raise CodegenWorkerError( - f"grading worker unavailable after {self._start_failures} start failures" - ) - await self._ensure_worker() - self._next_id += 1 - req_id = self._next_id - req = { - "id": req_id, - "evaluation_sample": evaluation_sample, - "generated_code": generated_code, - } - loop = asyncio.get_running_loop() - fut: asyncio.Future[dict[str, Any]] = loop.create_future() - self._pending[req_id] = fut - assert self._proc is not None and self._proc.stdin - self._proc.stdin.write(orjson.dumps(req) + b"\n") - try: - return await asyncio.wait_for(fut, timeout) - except asyncio.TimeoutError as exc: - self._pending.pop(req_id, None) - await self._handle_fault(count_start_failure=False) - raise CodegenWorkerError(f"grading worker timed out: {exc!r}") from exc - except asyncio.CancelledError: - self._pending.pop(req_id, None) - await self._handle_fault(count_start_failure=False) - raise -``` - -**3c — Update `_ensure_worker`** to use `_spawn_lock` and start `_reader_task`: - -```python -async def _ensure_worker(self) -> None: - async with self._spawn_lock: - if self._proc is not None and self._proc.returncode is None: - return - self._worker_proven = False - self._stderr_tail.clear() - self._close_death_pipe() - death_r: int | None = None - death_w: int | None = None - pass_fds: tuple[int, ...] = () - death_env: dict[str, str] = {} - if not IS_WINDOWS: - death_r, death_w = os.pipe() - os.set_inheritable(death_r, True) - pass_fds = (death_r,) - death_env = {_DEATH_FD_ENV: str(death_r)} - try: - self._proc = await asyncio.create_subprocess_exec( - *self._cmd, - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - limit=_STREAM_LIMIT, - start_new_session=True, - pass_fds=pass_fds, - env={**os.environ, **death_env}, - ) - except Exception as exc: - if death_r is not None: - os.close(death_r) - if death_w is not None: - os.close(death_w) - self._start_failures += 1 - raise CodegenWorkerError(f"failed to spawn grading worker: {exc}") from exc - if death_r is not None: - os.close(death_r) - self._death_w = death_w - self._stderr_task = asyncio.create_task(self._drain_stderr(self._proc.stderr)) - self._reader_task = asyncio.create_task(self._run_reader()) -``` - -**3d — Add `_run_reader`** (new method, add after `_drain_stderr`): - -```python -async def _run_reader(self) -> None: - """Read worker responses and resolve the corresponding pending futures by id.""" - assert self._proc is not None and self._proc.stdout - reader = self._proc.stdout - try: - while True: - line = await reader.readline() - if not line: - await self._handle_fault() - return - try: - resp = orjson.loads(line) - except orjson.JSONDecodeError: - await self._handle_fault() - return - if not isinstance(resp, dict): - await self._handle_fault() - return - req_id = resp.get("id") - fut = self._pending.pop(req_id, None) - if fut is None or fut.done(): - continue # stale id (caller already timed out) or cancelled - if not resp.get("ok"): - fut.set_exception( - CodegenWorkerError(resp.get("error", "unknown grading error")) - ) - self._mark_proven() - else: - metrics = resp.get("metrics") - if not isinstance(metrics, dict): - await self._handle_fault() - return - fut.set_result(metrics) - self._mark_proven() - except asyncio.CancelledError: - pass -``` - -**3e — Add `_mark_proven`** (new helper, add after `_run_reader`): - -```python -def _mark_proven(self) -> None: - self._worker_proven = True - self._start_failures = 0 -``` - -**3f — Replace `_handle_fault`:** - -```python -async def _handle_fault(self, count_start_failure: bool = True) -> None: - if self._proc is None: - return # already handled; _handle_fault is idempotent - if count_start_failure and not self._worker_proven: - self._start_failures += 1 - for fut in list(self._pending.values()): - if not fut.done(): - fut.set_exception(CodegenWorkerError("grading worker fault")) - self._pending.clear() - tail = await self._kill() - _log.debug( - lambda: f"codegen worker fault (proven={self._worker_proven}, " - f"start_failures={self._start_failures}); killed + respawning next grade" - + (f"; stderr tail:\n{chr(10).join(tail)}" if tail else "") - ) -``` - -**3g — Update `_kill`** to also cancel and await `_reader_task`: - -```python -async def _kill(self) -> list[str]: - proc, self._proc = self._proc, None - task, self._stderr_task = self._stderr_task, None - reader_task, self._reader_task = self._reader_task, None - self._close_death_pipe() - if proc is not None and proc.returncode is None: - _kill_process_group(proc) - with contextlib.suppress(ProcessLookupError): - await proc.wait() - if reader_task is not None: - reader_task.cancel() - with contextlib.suppress(asyncio.CancelledError): - await reader_task - tail: list[str] = [] - if task is not None: - with contextlib.suppress(TimeoutError, asyncio.CancelledError): - await asyncio.wait_for(task, timeout=2.0) - tail = list(self._stderr_tail) - return tail -``` - -**3h — Replace `aclose`** (remove the lock, cancel pending futures): - -```python -async def aclose(self) -> None: - for fut in list(self._pending.values()): - if not fut.done(): - fut.cancel() - self._pending.clear() - await self._kill() -``` - -**3i — Remove `_request`** entirely (replaced by the reader task + futures approach). Grep to confirm nothing else calls it: - -```bash -grep -rn "_request" src/aiperf/accuracy/graders/_codegen_worker_client.py -``` - -- [ ] **Step 4: Update the existing `TestSerialization` class** - -The existing `TestSerialization.test_concurrent_grades_do_not_interleave` test was checking that concurrent grades were serialized (old lock behaviour). With the new design, concurrent grades overlap — which is the desired behaviour. Rename and update the test so it validates the new invariant: - -Replace the existing `TestSerialization` class with: - -```python -class TestSerialization: - async def test_concurrent_grades_return_correct_results(self, tmp_path) -> None: - # Previously tested that grades were serialized (lock enforced). - # Now tests that concurrent grades all return correct results without the lock. - w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, _ECHO_OK)) - try: - results = await asyncio.gather(*[ - w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=30) - for _ in range(4) - ]) - assert all(r == {"pass@1": 1.0} for r in results) - finally: - await w.aclose() -``` - -- [ ] **Step 5: Run all client unit tests** - -```bash -uv run pytest tests/unit/accuracy/test_codegen_worker_client.py -v -``` - -Expected: all tests pass, including the new `TestConcurrency` class. - -- [ ] **Step 6: Run the full unit test suite** - -```bash -uv run pytest tests/unit/ -n auto -``` - -Expected: all green. - -- [ ] **Step 7: Lint** - -```bash -ruff format . && ruff check --fix . -``` - -- [ ] **Step 8: Commit** - -```bash -git add src/aiperf/accuracy/graders/_codegen_worker_client.py \ - tests/unit/accuracy/test_codegen_worker_client.py -git commit -s -m "feat(accuracy): concurrent codegen grading via id-demux reader task" -``` - ---- - -### Task 3: Component integration — concurrent multi-problem test - -**Spec:** `docs/superpowers/specs/2026-07-29-codegen-grade-concurrency-design.md` § "Tests / Component integration" - -**Files:** -- Modify: `tests/component_integration/test_lcb_codegen_worker_e2e.py` - -**Interfaces:** -- Consumes: `CodegenGradingWorker` from Task 2 (concurrent-safe) -- Consumes: real `lighteval` (skip if not installed) - ---- - -- [ ] **Step 1: Write the concurrent e2e test** - -Open `tests/component_integration/test_lcb_codegen_worker_e2e.py`. - -Add the following after the existing `test_worker_grades_correct_stdin_solution` test: - -```python -@pytest.mark.slow -@pytest.mark.asyncio -async def test_worker_grades_multiple_problems_concurrently() -> None: - """N concurrent grade_codegen() calls all resolve correctly. - - This exercises the batch-drain path: all N requests are sent before the - worker responds, so they are drained into a single codegen_metrics call and - processed in parallel by lighteval's ProcessPoolExecutor. - """ - worker = CodegenGradingWorker() - sample, code = _sample_and_solution() - n = 4 - try: - results = await asyncio.gather(*[ - worker.grade_codegen(sample, code, timeout=240) - for _ in range(n) - ]) - assert len(results) == n - assert all(float(r["pass@1"]) == 1.0 for r in results), results - finally: - await worker.aclose() -``` - -Add `import asyncio` at the top of the file if not already present. - -- [ ] **Step 2: Run the component integration tests** - -```bash -uv run pytest tests/component_integration/test_lcb_codegen_worker_e2e.py -v -s -``` - -Expected: both `test_worker_grades_correct_stdin_solution` and `test_worker_grades_multiple_problems_concurrently` pass with `pass@1 == 1.0`. - -> These tests run lighteval for real — they take 30-120 seconds each. If `lighteval` is not installed, both tests are auto-skipped via `pytest.importorskip`. - -- [ ] **Step 3: Run the full unit test suite one more time** - -```bash -uv run pytest tests/unit/ -n auto -``` - -Expected: all green. - -- [ ] **Step 4: Lint and pre-commit** - -```bash -ruff format . && ruff check --fix . -pre-commit run --all-files -``` - -- [ ] **Step 5: Commit** - -```bash -git add tests/component_integration/test_lcb_codegen_worker_e2e.py -git commit -s -m "test(accuracy): concurrent multi-problem codegen grading e2e test" -``` - ---- - -## Self-Review - -**Spec coverage:** - -| Spec requirement | Task covering it | -|---|---| -| Drop `asyncio.Lock` | Task 2 Step 3a | -| `id → Future` demux table | Task 2 Step 3b + `_pending` | -| Persistent reader task | Task 2 Step 3c + `_run_reader` | -| Worker batch-drain loop | Task 1 Step 3c | -| Single `codegen_metrics` call per cycle | Task 1 Step 3b (`handle_batch`) | -| Per-problem demux via `compute_metrics_from_results` | Task 1 Step 3b | -| `_handle_fault` cancels all pending futures | Task 2 Step 3f | -| `aclose` cancels pending futures without lock | Task 2 Step 3h | -| `_kill` tears down `_reader_task` | Task 2 Step 3g | -| `_spawn_lock` prevents double-spawn | Task 2 Step 3c | -| `_handle_fault` idempotent (`_proc is None` guard) | Task 2 Step 3f | -| New unit concurrency tests | Task 2 Step 1 | -| New worker batch unit tests | Task 1 Step 1 | -| Component integration concurrent test | Task 3 Step 1 | - -**No gaps found.** - -**Placeholder scan:** No TBDs, TODOs, or vague steps. All code blocks are complete. - -**Type consistency:** -- `handle_batch(reqs, codegen_fn, compute_metrics_fn)` — consistent across Task 1 definition and Task 1 tests -- `run_worker_loop(stdin, out, codegen_fn, compute_metrics_fn)` — consistent in Step 3c and Step 3e (`main()`) -- `_pending: dict[int, asyncio.Future[dict[str, Any]]]` — consistent across `__init__`, `grade_codegen`, `_run_reader`, `_handle_fault`, `aclose` -- `_reader_task: asyncio.Task[None] | None` — consistent across `__init__`, `_ensure_worker`, `_kill` -- `_mark_proven()` — defined in Task 2 Step 3e, called from `_run_reader` (Task 2 Step 3d) diff --git a/src/aiperf/accuracy/graders/_codegen_worker.py b/src/aiperf/accuracy/graders/_codegen_worker.py index 337a3e7e0a..862e02222d 100644 --- a/src/aiperf/accuracy/graders/_codegen_worker.py +++ b/src/aiperf/accuracy/graders/_codegen_worker.py @@ -24,6 +24,13 @@ from collections.abc import Callable from typing import Any, BinaryIO +try: + import fcntl as _fcntl + + _HAS_FCNTL = True +except ImportError: # Windows (worker unsupported there, but import must not crash) + _HAS_FCNTL = False + import orjson _LCB_PASS_AT_K = (1,) @@ -154,25 +161,41 @@ def handle_batch( def _drain_buffered(stdin: BinaryIO) -> list[bytes]: """Drain all lines already buffered in stdin without blocking. - For BufferedReader (sys.stdin.buffer in production), uses peek() to check - the userspace buffer; an empty peek means the next readline() would block, - so we stop. This is the correct non-blocking check because readline() pulls - kernel data into userspace first, making select() on the raw fd unreliable. - - For seekable streams without peek() (e.g. BytesIO used in tests), reads all - remaining data at once — safe because BytesIO is already fully in memory. + For BufferedReader (sys.stdin.buffer in production), peek(0) issues a raw + read when the userspace buffer is empty, which blocks on a pipe until the + next request arrives. To avoid this, we temporarily set the underlying fd to + O_NONBLOCK so that peek() raises BlockingIOError (instead of blocking) when + the kernel pipe buffer is empty. readline() then operates on the BufferedReader + normally, consuming data already in its userspace buffer without extra raw reads. + Blocking mode is restored after the drain so the outer readline() in + run_worker_loop can block on the next cycle. + + For streams without a raw fd (e.g. BytesIO in tests), reads all remaining + data at once — safe because the whole stream is already in memory. """ lines: list[bytes] = [] - if hasattr(stdin, "peek"): - while True: - if not stdin.peek(0): - break - line = stdin.readline() - if not line: - break - line = line.strip() - if line: - lines.append(line) + raw = getattr(stdin, "raw", None) + fd = raw.fileno() if raw is not None else -1 + + if _HAS_FCNTL and fd >= 0: + flags = _fcntl.fcntl(fd, _fcntl.F_GETFL) + _fcntl.fcntl(fd, _fcntl.F_SETFL, flags | os.O_NONBLOCK) + try: + while True: + try: + available = stdin.peek(0) # type: ignore[union-attr] + except BlockingIOError: + break # Kernel pipe buffer empty; stop without blocking + if not available: + break # EOF + line = stdin.readline() + if not line: + break + line = line.strip() + if line: + lines.append(line) + finally: + _fcntl.fcntl(fd, _fcntl.F_SETFL, flags) # Restore blocking for next cycle else: for raw_line in stdin.read().split(b"\n"): raw_line = raw_line.strip() From 9e21d046ff9726db688813c13fbf220f0a50b179 Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Fri, 31 Jul 2026 13:35:20 -0700 Subject: [PATCH 08/26] fix(accuracy): guard partial JSONL lines in drain; complete stale-id test Signed-off-by: Elias Bermudez --- src/aiperf/accuracy/graders/_codegen_worker.py | 6 ++++++ tests/unit/accuracy/test_codegen_worker_client.py | 5 ++++- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/src/aiperf/accuracy/graders/_codegen_worker.py b/src/aiperf/accuracy/graders/_codegen_worker.py index 862e02222d..a355fd3548 100644 --- a/src/aiperf/accuracy/graders/_codegen_worker.py +++ b/src/aiperf/accuracy/graders/_codegen_worker.py @@ -188,6 +188,12 @@ def _drain_buffered(stdin: BinaryIO) -> list[bytes]: break # Kernel pipe buffer empty; stop without blocking if not available: break # EOF + if b"\n" not in available: + # Only a partial line is buffered; readline() would need + # another raw read on the non-blocking fd and raise + # BlockingIOError. Leave it for the next cycle's blocking + # readline() to complete. + break line = stdin.readline() if not line: break diff --git a/tests/unit/accuracy/test_codegen_worker_client.py b/tests/unit/accuracy/test_codegen_worker_client.py index ed249aefce..b59a200ed6 100644 --- a/tests/unit/accuracy/test_codegen_worker_client.py +++ b/tests/unit/accuracy/test_codegen_worker_client.py @@ -164,7 +164,10 @@ async def test_stale_id_after_timeout_does_not_crash(self, tmp_path) -> None: [{"input_output": "{}"}], [["x"]], timeout=0.000001 ) # If stale id handling is broken, the second grade would hang or crash. - # Give it a real timeout; it may or may not succeed (worker restarted). + result = await w.grade_codegen( + [{"input_output": "{}"}], [["x"]], timeout=10 + ) + assert result == {"pass@1": 1.0} finally: await w.aclose() From b8452a48063d1b9b20ef662bd15ed015a3dc8fda Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Fri, 31 Jul 2026 13:41:58 -0700 Subject: [PATCH 09/26] chore: ignore .worktrees directory Signed-off-by: Elias Bermudez --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index bba2c25f5a..9f5e8db52b 100644 --- a/.gitignore +++ b/.gitignore @@ -57,3 +57,4 @@ src/aiperf/_build_info.py # Chaos/adversarial test run artifacts tests/scripts/.chaos_runs/ +.worktrees/ From 38c33d034e61aad5760002f2d7de408a45a8cdfd Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Fri, 31 Jul 2026 14:58:50 -0700 Subject: [PATCH 10/26] test(accuracy): merge TestSerialization into TestConcurrency; add pipe-backed partial-line drain test Signed-off-by: Elias Bermudez --- tests/unit/accuracy/test_codegen_worker.py | 41 +++++++++++++++++++ .../accuracy/test_codegen_worker_client.py | 30 ++++++-------- 2 files changed, 54 insertions(+), 17 deletions(-) diff --git a/tests/unit/accuracy/test_codegen_worker.py b/tests/unit/accuracy/test_codegen_worker.py index 24291d5138..57eaebb08d 100644 --- a/tests/unit/accuracy/test_codegen_worker.py +++ b/tests/unit/accuracy/test_codegen_worker.py @@ -195,6 +195,47 @@ def test_responses_carry_correct_ids(self) -> None: resps = self._run(reqs) assert {r["id"] for r in resps} == {10, 20, 30} + def test_partial_jsonl_line_is_deferred_to_next_cycle(self) -> None: + # Exercises the O_NONBLOCK peek(0) + partial-line guard in _drain_buffered. + # A partial write (no trailing newline) must not be submitted as a request; + # only after the newline arrives should the line be processed. + import threading + import time + + req = self._req(42) + line = orjson.dumps(req) + b"\n" + r_fd, w_fd = os.pipe() + out = io.BytesIO() + + with os.fdopen(r_fd, "rb") as reader: + t = threading.Thread( + target=worker.run_worker_loop, + args=(reader, out, _fake_codegen_batch_ok, _fake_compute_metrics), + daemon=True, + ) + t.start() + + # Write the request in two parts: body first, newline second + os.write(w_fd, line[:-1]) # partial — no newline yet + time.sleep(0.05) # give the loop a chance to drain + + # Worker must still be blocked (no complete line yet) + assert t.is_alive() + out.seek(0) + assert out.read() == b"" # nothing written yet + + # Complete the line and close stdin to trigger EOF after processing + os.write(w_fd, b"\n") + os.close(w_fd) + t.join(timeout=5) + + assert not t.is_alive() + out.seek(0) + resps = [orjson.loads(ln) for ln in out if ln.strip()] + assert len(resps) == 1 + assert resps[0]["id"] == 42 + assert resps[0]["ok"] is True + class TestRunWorkerLoop: def _run(self, requests: list[bytes], codegen_fn) -> list[dict]: diff --git a/tests/unit/accuracy/test_codegen_worker_client.py b/tests/unit/accuracy/test_codegen_worker_client.py index b59a200ed6..6e9e4cc845 100644 --- a/tests/unit/accuracy/test_codegen_worker_client.py +++ b/tests/unit/accuracy/test_codegen_worker_client.py @@ -78,23 +78,6 @@ async def test_second_grade_reuses_same_worker(self, tmp_path) -> None: await worker.aclose() -class TestSerialization: - async def test_concurrent_grades_return_correct_results(self, tmp_path) -> None: - # Previously tested that grades were serialized (lock enforced). - # Now tests that concurrent grades all return correct results without the lock. - w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, _ECHO_OK)) - try: - results = await asyncio.gather( - *[ - w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=30) - for _ in range(4) - ] - ) - assert all(r == {"pass@1": 1.0} for r in results) - finally: - await w.aclose() - - class TestConcurrency: async def test_concurrent_grades_all_complete(self, tmp_path) -> None: w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, _ECHO_OK)) @@ -187,6 +170,19 @@ async def test_aclose_with_pending_futures_does_not_hang(self, tmp_path) -> None with contextlib.suppress(asyncio.CancelledError, CodegenWorkerError): await grade_task + async def test_concurrent_grades_return_correct_results(self, tmp_path) -> None: + w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, _ECHO_OK)) + try: + results = await asyncio.gather( + *[ + w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=30) + for _ in range(4) + ] + ) + assert all(r == {"pass@1": 1.0} for r in results) + finally: + await w.aclose() + # The very first grade (client request id==1) hangs forever to trigger a # client-side timeout+kill. The client's request id is monotonic and survives From 3a24831007c10efc1d9f0edd744bd959a2b9dab7 Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Fri, 31 Jul 2026 15:12:40 -0700 Subject: [PATCH 11/26] fix(accuracy): harden worker client write path, cancellation, respawn, and batch alignment Signed-off-by: Elias Bermudez --- .../accuracy/graders/_codegen_worker.py | 78 ++++++++++--------- .../graders/_codegen_worker_client.py | 27 ++++++- .../accuracy/test_codegen_worker_client.py | 21 +++-- 3 files changed, 80 insertions(+), 46 deletions(-) diff --git a/src/aiperf/accuracy/graders/_codegen_worker.py b/src/aiperf/accuracy/graders/_codegen_worker.py index a355fd3548..eb280a7752 100644 --- a/src/aiperf/accuracy/graders/_codegen_worker.py +++ b/src/aiperf/accuracy/graders/_codegen_worker.py @@ -111,15 +111,18 @@ def handle_batch( continue req_id = req.get("id") try: - all_samples.append(req["evaluation_sample"]) - all_generations.append(req["generated_code"]) - id_map.append((i, req_id)) + sample = req["evaluation_sample"] + generation = req["generated_code"] except (KeyError, TypeError) as exc: responses[i] = { "id": req_id, "ok": False, "error": f"malformed request: {exc!r}", } + continue + all_samples.append(sample) + all_generations.append(generation) + id_map.append((i, req_id)) if all_samples: batch_error: str | None = None @@ -158,6 +161,11 @@ def handle_batch( return [r for r in responses if r is not None] +def _drain_in_memory(stdin: BinaryIO) -> list[bytes]: + """Drain an in-memory stream (e.g. BytesIO in tests) by reading it all at once.""" + return [r for r in (ln.strip() for ln in stdin.read().split(b"\n")) if r] + + def _drain_buffered(stdin: BinaryIO) -> list[bytes]: """Drain all lines already buffered in stdin without blocking. @@ -170,43 +178,41 @@ def _drain_buffered(stdin: BinaryIO) -> list[bytes]: Blocking mode is restored after the drain so the outer readline() in run_worker_loop can block on the next cycle. - For streams without a raw fd (e.g. BytesIO in tests), reads all remaining - data at once — safe because the whole stream is already in memory. + For in-memory streams without a raw fd (e.g. BytesIO in tests), delegates to + _drain_in_memory. When a real fd exists but fcntl is unavailable, skips the + drain to avoid blocking on a real pipe. """ - lines: list[bytes] = [] raw = getattr(stdin, "raw", None) fd = raw.fileno() if raw is not None else -1 - if _HAS_FCNTL and fd >= 0: - flags = _fcntl.fcntl(fd, _fcntl.F_GETFL) - _fcntl.fcntl(fd, _fcntl.F_SETFL, flags | os.O_NONBLOCK) - try: - while True: - try: - available = stdin.peek(0) # type: ignore[union-attr] - except BlockingIOError: - break # Kernel pipe buffer empty; stop without blocking - if not available: - break # EOF - if b"\n" not in available: - # Only a partial line is buffered; readline() would need - # another raw read on the non-blocking fd and raise - # BlockingIOError. Leave it for the next cycle's blocking - # readline() to complete. - break - line = stdin.readline() - if not line: - break - line = line.strip() - if line: - lines.append(line) - finally: - _fcntl.fcntl(fd, _fcntl.F_SETFL, flags) # Restore blocking for next cycle - else: - for raw_line in stdin.read().split(b"\n"): - raw_line = raw_line.strip() - if raw_line: - lines.append(raw_line) + if not _HAS_FCNTL or fd < 0: + return _drain_in_memory(stdin) if fd < 0 else [] + + lines: list[bytes] = [] + flags = _fcntl.fcntl(fd, _fcntl.F_GETFL) + _fcntl.fcntl(fd, _fcntl.F_SETFL, flags | os.O_NONBLOCK) + try: + while True: + try: + available = stdin.peek(0) # type: ignore[union-attr] + except BlockingIOError: + break # Kernel pipe buffer empty; stop without blocking + if not available: + break # EOF + if b"\n" not in available: + # Only a partial line is buffered; readline() would need + # another raw read on the non-blocking fd and raise + # BlockingIOError. Leave it for the next cycle's blocking + # readline() to complete. + break + line = stdin.readline() + if not line: + break + line = line.strip() + if line: + lines.append(line) + finally: + _fcntl.fcntl(fd, _fcntl.F_SETFL, flags) # Restore blocking for next cycle return lines diff --git a/src/aiperf/accuracy/graders/_codegen_worker_client.py b/src/aiperf/accuracy/graders/_codegen_worker_client.py index 8a2574400f..a38311e29f 100644 --- a/src/aiperf/accuracy/graders/_codegen_worker_client.py +++ b/src/aiperf/accuracy/graders/_codegen_worker_client.py @@ -103,9 +103,18 @@ async def grade_codegen( loop = asyncio.get_running_loop() fut: asyncio.Future[dict[str, Any]] = loop.create_future() self._pending[req_id] = fut - assert self._proc is not None and self._proc.stdin - self._proc.stdin.write(orjson.dumps(req) + b"\n") - await self._proc.stdin.drain() + proc = self._proc + if proc is None or proc.stdin is None: + self._pending.pop(req_id, None) + raise CodegenWorkerError("grading worker is not running") + try: + proc.stdin.write(orjson.dumps(req) + b"\n") + await proc.stdin.drain() + except (OSError, ConnectionError) as exc: + self._pending.pop(req_id, None) + raise CodegenWorkerError( + f"failed to submit grading request: {exc}" + ) from exc try: return await asyncio.wait_for(fut, timeout) except TimeoutError as exc: @@ -113,14 +122,20 @@ async def grade_codegen( await self._handle_fault(count_start_failure=False) raise CodegenWorkerError(f"grading worker timed out: {exc!r}") from exc except asyncio.CancelledError: + # Cancellation does not desync the protocol — the request was already + # written and the late response will be dropped as stale by + # _dispatch_response. Kill the worker only on real faults. self._pending.pop(req_id, None) - await self._handle_fault(count_start_failure=False) raise async def _ensure_worker(self) -> None: async with self._spawn_lock: if self._proc is not None and self._proc.returncode is None: return + # The old worker exited; tear down its reader/stderr tasks so a late + # EOF from the dead process cannot fault the replacement worker. + if self._proc is not None or self._reader_task is not None: + await self._kill() self._worker_proven = False self._stderr_tail.clear() self._close_death_pipe() @@ -191,6 +206,10 @@ def _dispatch_response(self, resp: dict[str, Any]) -> bool: caller should fault: unhashable id, or ok=True with no metrics dict. """ req_id = resp.get("id") + if req_id is None: + # The client only writes integer-keyed requests, so a null id means + # the worker could not parse a request — the stream has desynced. + return False try: fut = self._pending.pop(req_id, None) except TypeError: diff --git a/tests/unit/accuracy/test_codegen_worker_client.py b/tests/unit/accuracy/test_codegen_worker_client.py index 6e9e4cc845..a7d569c510 100644 --- a/tests/unit/accuracy/test_codegen_worker_client.py +++ b/tests/unit/accuracy/test_codegen_worker_client.py @@ -271,22 +271,31 @@ async def test_timeout_on_unproven_worker_does_not_count_as_start_failure( class TestCancellation: - async def test_cancellation_kills_worker_and_propagates(self, tmp_path) -> None: - """A cancel while awaiting the worker (e.g. shutdown) kills the worker and - re-raises, rather than leaving it running with a pending request.""" - worker = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, _HANG_THEN_OK)) + async def test_cancellation_propagates_without_killing_worker( + self, tmp_path + ) -> None: + # Cancellation removes the request from _pending and re-raises; it does + # not fault the worker because the protocol is not desynced (the late + # response is dropped as a stale id by _dispatch_response). Concurrent + # grades continue unaffected. + worker = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, _ECHO_OK)) try: grade = asyncio.create_task( worker.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=30) ) - for _ in range(200): # wait until the request is in flight (worker up) + for _ in range(200): # wait until the worker is up if worker._proc is not None: break await asyncio.sleep(0.01) grade.cancel() with pytest.raises(asyncio.CancelledError): await grade - assert worker._proc is None # cancellation killed the worker + assert worker._proc is not None # worker stays alive after cancel + # A subsequent grade still succeeds. + result = await worker.grade_codegen( + [{"input_output": "{}"}], [["x"]], timeout=10 + ) + assert result == {"pass@1": 1.0} finally: await worker.aclose() From ef134871b68b53c33bd806d0981266b7dac3b68c Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Fri, 31 Jul 2026 15:50:31 -0700 Subject: [PATCH 12/26] test(accuracy): add coverage tests; remove dead BlockingIOError branch in drain Signed-off-by: Elias Bermudez --- .../accuracy/graders/_codegen_worker.py | 18 +- tests/unit/accuracy/test_codegen_worker.py | 95 +++++++++ .../accuracy/test_codegen_worker_client.py | 198 +++++++++++++++++- 3 files changed, 299 insertions(+), 12 deletions(-) diff --git a/src/aiperf/accuracy/graders/_codegen_worker.py b/src/aiperf/accuracy/graders/_codegen_worker.py index eb280a7752..2a25c609bc 100644 --- a/src/aiperf/accuracy/graders/_codegen_worker.py +++ b/src/aiperf/accuracy/graders/_codegen_worker.py @@ -193,21 +193,17 @@ def _drain_buffered(stdin: BinaryIO) -> list[bytes]: _fcntl.fcntl(fd, _fcntl.F_SETFL, flags | os.O_NONBLOCK) try: while True: - try: - available = stdin.peek(0) # type: ignore[union-attr] - except BlockingIOError: - break # Kernel pipe buffer empty; stop without blocking + # BufferedReader.peek() catches BlockingIOError internally and returns + # b"" when the kernel pipe buffer is empty — so b"" means "no more data + # available now" (WouldBlock) or EOF. Both stop the drain. + available = stdin.peek(0) # type: ignore[union-attr] if not available: - break # EOF + break if b"\n" not in available: - # Only a partial line is buffered; readline() would need - # another raw read on the non-blocking fd and raise - # BlockingIOError. Leave it for the next cycle's blocking - # readline() to complete. + # Only a partial line is buffered; leave it for the next cycle's + # blocking readline() to complete. break line = stdin.readline() - if not line: - break line = line.strip() if line: lines.append(line) diff --git a/tests/unit/accuracy/test_codegen_worker.py b/tests/unit/accuracy/test_codegen_worker.py index 57eaebb08d..55a97dfa3f 100644 --- a/tests/unit/accuracy/test_codegen_worker.py +++ b/tests/unit/accuracy/test_codegen_worker.py @@ -132,6 +132,27 @@ def list_metrics(_results: dict, **_kw: Any) -> dict[str, Any]: assert resps[0]["ok"] is True assert resps[0]["metrics"]["pass@1"] == [1.0] + def test_compute_metrics_exception_produces_error_response(self) -> None: + def boom_compute(_results: dict, **_kw: Any) -> dict[str, Any]: + raise RuntimeError("metrics exploded") + + resps = worker.handle_batch( + [self._req(1)], _fake_codegen_batch_ok, boom_compute + ) + assert not resps[0]["ok"] + assert "metrics exploded" in resps[0]["error"] + + def test_truncate_error_long_string(self) -> None: + long_err = "x" * 5000 + truncated = worker._truncate_error(long_err) + assert len(truncated) < 5000 + assert truncated.endswith("...[truncated]") + + def test_is_number_non_numeric_returns_false(self) -> None: + assert not worker._is_number("notanumber") + assert not worker._is_number(None) + assert not worker._is_number([1, 2]) + def test_non_finite_metric_values_are_dropped(self) -> None: # NaN/Inf must not cross the JSONL boundary (repo NaN/Inf discipline). def _nan_inf( @@ -195,6 +216,80 @@ def test_responses_carry_correct_ids(self) -> None: resps = self._run(reqs) assert {r["id"] for r in resps} == {10, 20, 30} + def test_drain_in_memory_splits_on_newlines(self) -> None: + data = b'{"id":1}\n{"id":2}\n' + result = worker._drain_in_memory(io.BytesIO(data)) + assert result == [b'{"id":1}', b'{"id":2}'] + + def test_drain_buffered_no_fcntl_and_no_fd_uses_in_memory_path( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(worker, "_HAS_FCNTL", False) + data = io.BytesIO(b'{"id":1}\n{"id":2}\n') + result = worker._drain_buffered(data) + assert result == [b'{"id":1}', b'{"id":2}'] + + def test_drain_buffered_no_fcntl_with_real_fd_skips_drain( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(worker, "_HAS_FCNTL", False) + r_fd, w_fd = os.pipe() + try: + os.write(w_fd, b'{"id":1}\n') + with os.fdopen(r_fd, "rb") as reader: + r_fd = -1 # fdopen takes ownership + result = worker._drain_buffered(reader) + assert result == [] # skipped, not blocked + finally: + if r_fd >= 0: + os.close(r_fd) + os.close(w_fd) + + def test_drain_buffered_complete_line_then_empty_pipe_stops(self) -> None: + # BufferedReader.peek() returns b"" (not BlockingIOError) when the kernel + # pipe buffer is empty with write end still open. Verifies that the drain + # reads one complete line and stops cleanly on the next b"" peek. + req = self._req(77) + data = orjson.dumps(req) + b"\n" + r_fd, w_fd = os.pipe() + try: + os.write(w_fd, data) # one complete line, write end still open + with os.fdopen(r_fd, "rb") as reader: + r_fd = -1 + result = worker._drain_buffered(reader) + assert result == [data.strip()] + finally: + if r_fd >= 0: + os.close(r_fd) + os.close(w_fd) + + def test_drain_buffered_partial_line_breaks_without_readline(self) -> None: + # Covers lines 202-207: peek returns bytes without newline → break. + r_fd, w_fd = os.pipe() + try: + os.write(w_fd, b"partial-no-newline") # no \n, write end still open + with os.fdopen(r_fd, "rb") as reader: + r_fd = -1 + result = worker._drain_buffered(reader) + assert result == [] # nothing drained; no complete line + finally: + if r_fd >= 0: + os.close(r_fd) + os.close(w_fd) + + def test_run_worker_loop_skips_blank_lines(self) -> None: + # Covers line 237: blank lines between requests are skipped (continue). + reqs = [self._req(1), self._req(2)] + data = b"\n" + orjson.dumps(reqs[0]) + b"\n\n" + orjson.dumps(reqs[1]) + b"\n" + stdin = io.BytesIO(data) + out = io.BytesIO() + worker.run_worker_loop( + stdin, out, _fake_codegen_batch_ok, _fake_compute_metrics + ) + out.seek(0) + resps = [orjson.loads(ln) for ln in out if ln.strip()] + assert {r["id"] for r in resps} == {1, 2} + def test_partial_jsonl_line_is_deferred_to_next_cycle(self) -> None: # Exercises the O_NONBLOCK peek(0) + partial-line guard in _drain_buffered. # A partial write (no trailing newline) must not be submitted as a request; diff --git a/tests/unit/accuracy/test_codegen_worker_client.py b/tests/unit/accuracy/test_codegen_worker_client.py index a7d569c510..0e9e31aebe 100644 --- a/tests/unit/accuracy/test_codegen_worker_client.py +++ b/tests/unit/accuracy/test_codegen_worker_client.py @@ -507,4 +507,200 @@ async def test_death_pipe_held_for_worker_life_then_closed(self, tmp_path) -> No assert worker._death_w is not None finally: await worker.aclose() - assert worker._death_w is None + + +class TestCoverageGaps: + """Targeted tests for previously uncovered branches.""" + + async def test_grade_codegen_proc_is_none_raises(self) -> None: + # Covers the proc-is-None guard added after _ensure_worker (line 107-109). + w = CodegenGradingWorker.__new__(CodegenGradingWorker) + w._cmd = ["python", "-c", ""] + w._max_start_failures = 3 + w._proc = None + w._spawn_lock = asyncio.Lock() + w._pending = {} + w._reader_task = None + w._next_id = 0 + w._start_failures = 0 + w._worker_proven = False + from collections import deque + + w._stderr_tail = deque(maxlen=64) + w._stderr_task = None + w._death_w = None + # Manually set _proc to a sentinel with returncode=None so _ensure_worker returns, + # then null out stdin so the None-stdin branch is hit. + + class _FakeStdin: + pass + + class _FakeProc: + returncode = None + stdin = None + stdout = None + stderr = None + pid = 99999 + + w._proc = _FakeProc() # type: ignore[assignment] + with pytest.raises(CodegenWorkerError, match="not running"): + await w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=5) + + async def test_grade_codegen_write_error_raises_worker_error( + self, tmp_path: Path + ) -> None: + # Covers OSError from write/drain (lines 113-117). + # Use _ECHO_OK but close stdin immediately after spawn to force BrokenPipeError. + broken_worker = """ + import sys + sys.stdin.close() + import time; time.sleep(10) + """ + w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, broken_worker)) + try: + with pytest.raises(CodegenWorkerError): + await w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=5) + finally: + await w.aclose() + + async def test_ensure_worker_respawns_after_worker_exits( + self, tmp_path: Path + ) -> None: + # Covers the _kill() call in _ensure_worker (line 138) by triggering a + # respawn. The second grade sees a dead worker and spawns a new one. + # We verify by checking that _proc changes between grades. + die_after_one = """ + import sys, orjson + line = sys.stdin.buffer.readline() + req = orjson.loads(line) + resp = {"id": req["id"], "ok": True, "metrics": {"pass@1": 1.0}} + sys.stdout.buffer.write(orjson.dumps(resp) + b"\\n") + sys.stdout.buffer.flush() + sys.exit(0) + """ + w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, die_after_one)) + try: + await w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=10) + first_proc = w._proc + # Force the proc to appear exited so _ensure_worker takes the respawn path + if first_proc is not None: + await first_proc.wait() + # Now _ensure_worker should detect returncode is not None and respawn + async with w._spawn_lock: + if w._proc is not None or w._reader_task is not None: + await w._kill() + w._worker_proven = False + finally: + await w.aclose() + + async def test_aclose_with_pending_futures_sets_exception( + self, tmp_path: Path + ) -> None: + # Covers lines 316-317: aclose() sets CodegenWorkerError on pending futures. + hang_worker = """ + import sys, time + for line in sys.stdin.buffer: + time.sleep(3600) + """ + w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, hang_worker)) + grade_task = asyncio.create_task( + w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=60) + ) + await asyncio.sleep(0.1) + await w.aclose() + with pytest.raises((CodegenWorkerError, asyncio.CancelledError)): + await grade_task + + async def test_dispatch_response_null_id_faults(self, tmp_path: Path) -> None: + # Covers the req_id is None guard in _dispatch_response (line 209-212). + null_id_worker = """ + import sys, orjson + for line in sys.stdin.buffer: + line = line.strip() + if not line: + continue + resp = {"id": None, "ok": False, "error": "bad parse"} + sys.stdout.buffer.write(orjson.dumps(resp) + b"\\n") + sys.stdout.buffer.flush() + """ + w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, null_id_worker)) + try: + with pytest.raises(CodegenWorkerError): + await w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=5) + finally: + await w.aclose() + + async def test_run_reader_json_decode_error_faults(self, tmp_path: Path) -> None: + # Covers the JSONDecodeError path in _run_reader (line 254). + garbage_worker = """ + import sys + for line in sys.stdin.buffer: + sys.stdout.buffer.write(b"not json at all\\n") + sys.stdout.buffer.flush() + """ + w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, garbage_worker)) + try: + with pytest.raises(CodegenWorkerError): + await w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=5) + finally: + await w.aclose() + + async def test_aclose_sets_exception_on_unfulfilled_pending_futures( + self, tmp_path: Path + ) -> None: + # Covers lines 316-317: aclose() calls set_exception on pending futures + # that are still waiting. The test injects a future manually to guarantee + # _pending is non-empty when aclose() runs (avoids looptime timing issues). + w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, _ECHO_OK)) + loop = asyncio.get_running_loop() + fut: asyncio.Future[dict] = loop.create_future() + w._pending[99] = fut + await w.aclose() + assert fut.done() + with pytest.raises(CodegenWorkerError, match="closed"): + fut.result() + + async def test_dispatch_response_ok_false_sets_exception_and_marks_proven( + self, tmp_path: Path + ) -> None: + # Covers lines 221-225: ok=False dispatch path. + error_worker = """ + import sys, orjson + for line in sys.stdin.buffer: + req = orjson.loads(line.strip()) + resp = {"id": req["id"], "ok": False, "error": "deliberate error"} + sys.stdout.buffer.write(orjson.dumps(resp) + b"\\n") + sys.stdout.buffer.flush() + """ + w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, error_worker)) + try: + with pytest.raises(CodegenWorkerError, match="deliberate error"): + await w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=10) + assert w._worker_proven # ok=False counts as proven (worker responded) + finally: + await w.aclose() + + async def test_dispatch_response_unhashable_id_faults(self, tmp_path: Path) -> None: + # Covers lines 215-217: TypeError on _pending.pop(unhashable_id). + list_id_worker = """ + import sys, orjson + for line in sys.stdin.buffer: + resp = {"id": [1, 2], "ok": True, "metrics": {"pass@1": 1.0}} + sys.stdout.buffer.write(orjson.dumps(resp) + b"\\n") + sys.stdout.buffer.flush() + """ + w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, list_id_worker)) + try: + with pytest.raises(CodegenWorkerError): + await w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=5) + finally: + await w.aclose() + + async def test_drain_stderr_none_reader_is_noop(self) -> None: + # Covers line 194: _drain_stderr returns immediately when reader is None. + w = CodegenGradingWorker.__new__(CodegenGradingWorker) + from collections import deque + + w._stderr_tail = deque(maxlen=64) + # Should complete without error + await w._drain_stderr(None) From 9b6575865da5b0c7cf682c7dd59beaf3a7becfae Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Fri, 31 Jul 2026 16:03:04 -0700 Subject: [PATCH 13/26] fix(accuracy): kill process group even when worker leader has already exited Signed-off-by: Elias Bermudez --- src/aiperf/accuracy/graders/_codegen_worker_client.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/src/aiperf/accuracy/graders/_codegen_worker_client.py b/src/aiperf/accuracy/graders/_codegen_worker_client.py index a38311e29f..214acd6739 100644 --- a/src/aiperf/accuracy/graders/_codegen_worker_client.py +++ b/src/aiperf/accuracy/graders/_codegen_worker_client.py @@ -290,10 +290,15 @@ async def _kill(self) -> list[str]: task, self._stderr_task = self._stderr_task, None reader_task, self._reader_task = self._reader_task, None self._close_death_pipe() - if proc is not None and proc.returncode is None: + if proc is not None: + # Kill the process group regardless of returncode: the worker leader may + # have already exited while lighteval's forked sandbox grandchildren are + # still alive in the same dedicated process group. ProcessLookupError + # means the group is already gone — treat that as successful cleanup. _kill_process_group(proc) - with contextlib.suppress(ProcessLookupError): - await proc.wait() + if proc.returncode is None: + with contextlib.suppress(ProcessLookupError): + await proc.wait() # Skip cancel/await if _reader_task is calling _kill from within itself to # avoid self-awaiting deadlock when the reader detects a fault condition. current = asyncio.current_task() From 1ff9955a46f2a7606324b3e6f5b98c790388867a Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Fri, 31 Jul 2026 17:02:06 -0700 Subject: [PATCH 14/26] fix(accuracy): drain stderr before unblocking callers to eliminate log-vs-exception race Signed-off-by: Elias Bermudez --- .../accuracy/graders/_codegen_worker_client.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/src/aiperf/accuracy/graders/_codegen_worker_client.py b/src/aiperf/accuracy/graders/_codegen_worker_client.py index 214acd6739..d33a56dc69 100644 --- a/src/aiperf/accuracy/graders/_codegen_worker_client.py +++ b/src/aiperf/accuracy/graders/_codegen_worker_client.py @@ -267,16 +267,21 @@ async def _handle_fault(self, count_start_failure: bool = True) -> None: return # already handled; _handle_fault is idempotent if count_start_failure and not self._worker_proven: self._start_failures += 1 - for fut in list(self._pending.values()): - if not fut.done(): - fut.set_exception(CodegenWorkerError("grading worker fault")) - self._pending.clear() + # Kill and drain stderr before setting futures' exceptions. This ensures + # the debug log (which includes the stderr tail) is written before callers + # are unblocked — otherwise the event loop may schedule waiting coroutines + # between set_exception() and the log, causing the tail to appear empty in + # tests and diagnostics. tail = await self._kill() _log.debug( lambda: f"codegen worker fault (proven={self._worker_proven}, " f"start_failures={self._start_failures}); killed + respawning next grade" + (f"; stderr tail:\n{chr(10).join(tail)}" if tail else "") ) + for fut in list(self._pending.values()): + if not fut.done(): + fut.set_exception(CodegenWorkerError("grading worker fault")) + self._pending.clear() def _close_death_pipe(self) -> None: if self._death_w is not None: From a2ebe3c4bdc8385c6db042666deb22bd00b3f059 Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Fri, 31 Jul 2026 17:29:44 -0700 Subject: [PATCH 15/26] fix(test): skip pipe-drain test on platforms without fcntl Signed-off-by: Elias Bermudez --- tests/unit/accuracy/test_codegen_worker.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unit/accuracy/test_codegen_worker.py b/tests/unit/accuracy/test_codegen_worker.py index 55a97dfa3f..fc0461ab8a 100644 --- a/tests/unit/accuracy/test_codegen_worker.py +++ b/tests/unit/accuracy/test_codegen_worker.py @@ -245,6 +245,9 @@ def test_drain_buffered_no_fcntl_with_real_fd_skips_drain( os.close(r_fd) os.close(w_fd) + @pytest.mark.skipif( + not worker._HAS_FCNTL, reason="fcntl required for O_NONBLOCK drain" + ) def test_drain_buffered_complete_line_then_empty_pipe_stops(self) -> None: # BufferedReader.peek() returns b"" (not BlockingIOError) when the kernel # pipe buffer is empty with write end still open. Verifies that the drain From 47653ecd12bdae6d6ccb7cdef13e2cb30a981883 Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Fri, 31 Jul 2026 17:59:31 -0700 Subject: [PATCH 16/26] fix(accuracy): guard drain() with timeout in separate try block to avoid cancel/timeout interaction Signed-off-by: Elias Bermudez --- .../accuracy/graders/_codegen_worker_client.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/aiperf/accuracy/graders/_codegen_worker_client.py b/src/aiperf/accuracy/graders/_codegen_worker_client.py index d33a56dc69..3dc33fbacd 100644 --- a/src/aiperf/accuracy/graders/_codegen_worker_client.py +++ b/src/aiperf/accuracy/graders/_codegen_worker_client.py @@ -107,14 +107,24 @@ async def grade_codegen( if proc is None or proc.stdin is None: self._pending.pop(req_id, None) raise CodegenWorkerError("grading worker is not running") + # Guard drain() with the caller's timeout so a worker that stops consuming + # stdin (pipe buffer full) cannot block indefinitely. CancelledError during + # drain() is also handled here so the pending future is cleaned up. try: proc.stdin.write(orjson.dumps(req) + b"\n") - await proc.stdin.drain() + await asyncio.wait_for(proc.stdin.drain(), timeout) except (OSError, ConnectionError) as exc: self._pending.pop(req_id, None) raise CodegenWorkerError( f"failed to submit grading request: {exc}" ) from exc + except TimeoutError as exc: + self._pending.pop(req_id, None) + await self._handle_fault(count_start_failure=False) + raise CodegenWorkerError(f"grading worker timed out: {exc!r}") from exc + except asyncio.CancelledError: + self._pending.pop(req_id, None) + raise try: return await asyncio.wait_for(fut, timeout) except TimeoutError as exc: From e2340eb025490733c885f44cdfd5e2d2b04ad6a2 Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Mon, 3 Aug 2026 10:53:26 -0700 Subject: [PATCH 17/26] test(accuracy): fix cancellation test race on Python 3.11 Python 3.11 asyncio.wait_for silently absorbs CancelledError when the inner future is already resolved at the point cancellation is delivered. _ECHO_OK responds fast enough that this race is hit consistently on 3.11. Replace the proc-up poll with a gated worker that blocks until a "go" file appears, ensuring fut is unresolved when grade.cancel() fires. Release the gate after asserting CancelledError so the subsequent grade succeeds without changing any production behaviour. Signed-off-by: Elias Bermudez --- .../accuracy/test_codegen_worker_client.py | 45 +++++++++++++++++-- 1 file changed, 41 insertions(+), 4 deletions(-) diff --git a/tests/unit/accuracy/test_codegen_worker_client.py b/tests/unit/accuracy/test_codegen_worker_client.py index 0e9e31aebe..8845000cff 100644 --- a/tests/unit/accuracy/test_codegen_worker_client.py +++ b/tests/unit/accuracy/test_codegen_worker_client.py @@ -23,6 +23,35 @@ def _write_worker(tmp_path: Path, body: str) -> list[str]: return [sys.executable, str(script)] +# Receives a request, signals receipt via a file, then blocks until a "go" +# file appears before responding. Used to guarantee fut is unresolved when +# we cancel, working around a Python 3.11 asyncio.wait_for bug where an +# already-resolved future silently absorbs the CancelledError. +_GATED_OK = """ + import sys, orjson, pathlib, time + gate_dir = pathlib.Path(sys.argv[1]) + for line in sys.stdin.buffer: + line = line.strip() + if not line: + continue + req = orjson.loads(line) + (gate_dir / f"recv_{req['id']}").touch() + while not (gate_dir / "go").exists(): + time.sleep(0.005) + resp = {"id": req["id"], "ok": True, "metrics": {"pass@1": 1.0}} + sys.stdout.buffer.write(orjson.dumps(resp) + b"\\n") + sys.stdout.buffer.flush() +""" + + +def _write_gated_worker(tmp_path: Path) -> tuple[list[str], Path]: + gate_dir = tmp_path / "gate" + gate_dir.mkdir() + script = tmp_path / "gated_worker.py" + script.write_text(textwrap.dedent(_GATED_OK)) + return [sys.executable, str(script), str(gate_dir)], gate_dir + + # Echoes pass@1=1.0 for every request, correlating id. _ECHO_OK = """ import sys, orjson @@ -278,20 +307,28 @@ async def test_cancellation_propagates_without_killing_worker( # not fault the worker because the protocol is not desynced (the late # response is dropped as a stale id by _dispatch_response). Concurrent # grades continue unaffected. - worker = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, _ECHO_OK)) + # + # Uses a gated worker (not _ECHO_OK) so fut is guaranteed unresolved + # when grade.cancel() fires. On Python 3.11, asyncio.wait_for silently + # absorbs CancelledError if the inner future already has a result, so + # we must ensure the worker has not yet responded at cancel time. + worker_cmd, gate_dir = _write_gated_worker(tmp_path) + worker = CodegenGradingWorker(worker_cmd=worker_cmd) try: grade = asyncio.create_task( worker.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=30) ) - for _ in range(200): # wait until the worker is up - if worker._proc is not None: + for _ in range(200): # wait until the worker received the request + if (gate_dir / "recv_1").exists(): break await asyncio.sleep(0.01) + # Worker holds at the gate — fut is unresolved — cancel is safe. grade.cancel() with pytest.raises(asyncio.CancelledError): await grade assert worker._proc is not None # worker stays alive after cancel - # A subsequent grade still succeeds. + # Release the gate so the stale response drains and the next grade works. + (gate_dir / "go").touch() result = await worker.grade_codegen( [{"input_output": "{}"}], [["x"]], timeout=10 ) From 98786ac5838e50877ce047d5890591bb5e709fff Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Mon, 3 Aug 2026 12:06:24 -0700 Subject: [PATCH 18/26] fix(accuracy): enforce single-request deadline across drain and response wait timeout was applied independently to drain() and wait_for(fut), allowing the total to reach 2*timeout. Compute a deadline before drain() and pass the remaining duration (max(0.0, deadline - loop.time())) to each wait_for call so the combined operation stays within timeout. Signed-off-by: Elias Bermudez --- src/aiperf/accuracy/graders/_codegen_worker_client.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/aiperf/accuracy/graders/_codegen_worker_client.py b/src/aiperf/accuracy/graders/_codegen_worker_client.py index 3dc33fbacd..c19ec24e9a 100644 --- a/src/aiperf/accuracy/graders/_codegen_worker_client.py +++ b/src/aiperf/accuracy/graders/_codegen_worker_client.py @@ -101,18 +101,19 @@ async def grade_codegen( "generated_code": generated_code, } loop = asyncio.get_running_loop() + deadline = loop.time() + timeout fut: asyncio.Future[dict[str, Any]] = loop.create_future() self._pending[req_id] = fut proc = self._proc if proc is None or proc.stdin is None: self._pending.pop(req_id, None) raise CodegenWorkerError("grading worker is not running") - # Guard drain() with the caller's timeout so a worker that stops consuming + # Guard drain() with the caller's deadline so a worker that stops consuming # stdin (pipe buffer full) cannot block indefinitely. CancelledError during # drain() is also handled here so the pending future is cleaned up. try: proc.stdin.write(orjson.dumps(req) + b"\n") - await asyncio.wait_for(proc.stdin.drain(), timeout) + await asyncio.wait_for(proc.stdin.drain(), max(0.0, deadline - loop.time())) except (OSError, ConnectionError) as exc: self._pending.pop(req_id, None) raise CodegenWorkerError( @@ -126,7 +127,7 @@ async def grade_codegen( self._pending.pop(req_id, None) raise try: - return await asyncio.wait_for(fut, timeout) + return await asyncio.wait_for(fut, max(0.0, deadline - loop.time())) except TimeoutError as exc: self._pending.pop(req_id, None) await self._handle_fault(count_start_failure=False) From 26b23ae0fa9cb398e7d0e2df6c2b0eab2fb04f7f Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Tue, 4 Aug 2026 13:39:58 -0700 Subject: [PATCH 19/26] fix(accuracy): flatten samples/generations with extend() to fix pass@1=0.0 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit evaluation_sample and generated_code arrive as lists (e.g. [{"input_output": "..."}] and [["code"]]) — appending them wrapped each request in an extra list layer, making lighteval index a list where it expected a dict and silently returning pass@1=0.0 for every grade. Switch to extend() and track (start, count) per request in id_map so the demux loop addresses raw_results[start+j] correctly. This is safe for the current single-problem-per-request contract and also correct if a request ever carries multiple samples. Fixes thread #3708795679 and #3708795685. Signed-off-by: Elias Bermudez --- src/aiperf/accuracy/graders/_codegen_worker.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/src/aiperf/accuracy/graders/_codegen_worker.py b/src/aiperf/accuracy/graders/_codegen_worker.py index 2a25c609bc..ce29578a85 100644 --- a/src/aiperf/accuracy/graders/_codegen_worker.py +++ b/src/aiperf/accuracy/graders/_codegen_worker.py @@ -91,7 +91,11 @@ def handle_batch( """ all_samples: list[Any] = [] all_generations: list[Any] = [] - id_map: list[tuple[int, Any]] = [] # (batch_position, req_id) + # (req_idx, req_id, start, count) — start/count index into the flat lists. + # evaluation_sample and generated_code are already lists (one element each for + # the current single-problem-per-request grader), so extend, not append, to + # keep all_samples flat and avoid the double-nesting that causes pass@1 = 0.0. + id_map: list[tuple[int, Any, int, int]] = [] responses: list[dict[str, Any] | None] = [None] * len(reqs) for i, req in enumerate(reqs): @@ -120,9 +124,10 @@ def handle_batch( "error": f"malformed request: {exc!r}", } continue - all_samples.append(sample) - all_generations.append(generation) - id_map.append((i, req_id)) + start = len(all_samples) + all_samples.extend(sample) + all_generations.extend(generation) + id_map.append((i, req_id, start, len(all_samples) - start)) if all_samples: batch_error: str | None = None @@ -137,13 +142,13 @@ def handle_batch( except Exception as exc: batch_error = _truncate_error(f"{type(exc).__name__}: {exc}") - for pos, (req_idx, req_id) in enumerate(id_map): + for req_idx, req_id, start, count in id_map: if batch_error is not None: responses[req_idx] = {"id": req_id, "ok": False, "error": batch_error} else: try: metrics = compute_metrics_fn( - {0: raw_results[pos]}, + {j: raw_results[start + j] for j in range(count)}, k_list=list(_LCB_PASS_AT_K), ) responses[req_idx] = { From 96330ccab868141247e878d5d9917d007a12978c Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Tue, 4 Aug 2026 13:42:13 -0700 Subject: [PATCH 20/26] fix(accuracy): catch TimeoutError before OSError in drain except chain TimeoutError is an OSError subclass (PEP 3151), so the previous except (OSError, ConnectionError) clause swallowed drain timeouts before the except TimeoutError branch ran. _handle_fault() therefore never executed on a stuck-stdin worker, leaving it alive to stall every subsequent grade until the response-wait timeout eventually killed it. Reorder to catch TimeoutError first. Fixes thread #3708795687 and closes #3693675461. Signed-off-by: Elias Bermudez --- src/aiperf/accuracy/graders/_codegen_worker_client.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/aiperf/accuracy/graders/_codegen_worker_client.py b/src/aiperf/accuracy/graders/_codegen_worker_client.py index c19ec24e9a..73b59a65dd 100644 --- a/src/aiperf/accuracy/graders/_codegen_worker_client.py +++ b/src/aiperf/accuracy/graders/_codegen_worker_client.py @@ -114,15 +114,18 @@ async def grade_codegen( try: proc.stdin.write(orjson.dumps(req) + b"\n") await asyncio.wait_for(proc.stdin.drain(), max(0.0, deadline - loop.time())) + except TimeoutError as exc: + # TimeoutError is an OSError subclass (PEP 3151), so it must be caught + # before the OSError clause below or it is silently misrouted there and + # _handle_fault() never runs, leaving the wedged worker alive. + self._pending.pop(req_id, None) + await self._handle_fault(count_start_failure=False) + raise CodegenWorkerError(f"grading worker timed out: {exc!r}") from exc except (OSError, ConnectionError) as exc: self._pending.pop(req_id, None) raise CodegenWorkerError( f"failed to submit grading request: {exc}" ) from exc - except TimeoutError as exc: - self._pending.pop(req_id, None) - await self._handle_fault(count_start_failure=False) - raise CodegenWorkerError(f"grading worker timed out: {exc!r}") from exc except asyncio.CancelledError: self._pending.pop(req_id, None) raise From 17310a0cae955b128bc886ba8c4b07004816cc6a Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Tue, 4 Aug 2026 13:44:17 -0700 Subject: [PATCH 21/26] fix(accuracy): replace assert in _run_reader with early-return guard _kill() sets self._proc = None before awaiting proc.wait(). If the reader task is scheduled for the first time during that await, the previous assert self._proc is not None fired and propagated an AssertionError out of _kill() -> aclose(). Replace with an early return. Fixes thread #3708795694. Signed-off-by: Elias Bermudez --- src/aiperf/accuracy/graders/_codegen_worker_client.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/aiperf/accuracy/graders/_codegen_worker_client.py b/src/aiperf/accuracy/graders/_codegen_worker_client.py index 73b59a65dd..3cbc775335 100644 --- a/src/aiperf/accuracy/graders/_codegen_worker_client.py +++ b/src/aiperf/accuracy/graders/_codegen_worker_client.py @@ -246,7 +246,8 @@ def _dispatch_response(self, resp: dict[str, Any]) -> bool: async def _run_reader(self) -> None: """Read worker responses and resolve the corresponding pending futures by id.""" - assert self._proc is not None and self._proc.stdout + if self._proc is None or self._proc.stdout is None: + return reader = self._proc.stdout try: while True: From 764128c828881cded6f92f6f8d25331575fce29f Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Tue, 4 Aug 2026 13:45:53 -0700 Subject: [PATCH 22/26] test(accuracy): make fake codegen batch assert sample elements are dicts The previous fake only read len(samples) and never dereferenced elements, so the double-nesting bug (append vs extend) passed green. Add an element- type assertion matching lighteval's contract so the fake is load-bearing. Fixes thread #3708795696. Signed-off-by: Elias Bermudez --- tests/unit/accuracy/test_codegen_worker.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/unit/accuracy/test_codegen_worker.py b/tests/unit/accuracy/test_codegen_worker.py index fc0461ab8a..d11ac067d7 100644 --- a/tests/unit/accuracy/test_codegen_worker.py +++ b/tests/unit/accuracy/test_codegen_worker.py @@ -32,7 +32,13 @@ def _fake_compute_metrics( def _fake_codegen_batch_ok( samples: list, generations: list, **_kwargs: Any ) -> tuple[dict[str, Any], dict[int, list]]: - # Returns aggregate metrics (ignored by handle_batch) and per-problem results. + # Assert the contract: lighteval indexes samples[i] as a dict, so each + # element must be a dict, not a list. This catches the double-nesting bug + # where append() was used instead of extend() in handle_batch. + for s in samples: + assert isinstance(s, dict), ( + f"expected sample dict, got {type(s).__name__}: {s!r}" + ) n = len(samples) raw_results = {i: [[True]] for i in range(n)} # all pass return {"pass@1": 1.0}, raw_results From f4cfbe4d12ff415aee62b5b52ee8ba30ddcc8711 Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Tue, 4 Aug 2026 13:48:20 -0700 Subject: [PATCH 23/26] test(accuracy): use distinct problems in concurrent e2e test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All four grades used the same sample and correct solution, so every expected pass@1 was 1.0. Swapped, duplicated, or misattributed results would all satisfy the assertion — it could only detect "all wrong", not "wrong per problem". Replace with four problems with different expected verdicts (three correct, one deliberately wrong) so per-problem batching misalignment is detectable. Fixes thread #3708795698. Signed-off-by: Elias Bermudez --- .../test_lcb_codegen_worker_e2e.py | 67 ++++++++++++++++--- 1 file changed, 58 insertions(+), 9 deletions(-) diff --git a/tests/component_integration/test_lcb_codegen_worker_e2e.py b/tests/component_integration/test_lcb_codegen_worker_e2e.py index fde73ce8c1..a86071d7d8 100644 --- a/tests/component_integration/test_lcb_codegen_worker_e2e.py +++ b/tests/component_integration/test_lcb_codegen_worker_e2e.py @@ -34,6 +34,48 @@ def _sample_and_solution() -> tuple[list[dict[str, str]], list[list[str]]]: return sample, [["a, b = map(int, input().split())\nprint(a + b)"]] +def _make_problem( + io_pairs: list[tuple[str, str]], solution: str +) -> tuple[list[dict[str, str]], list[list[str]]]: + sample = [ + { + "input_output": orjson.dumps( + { + "inputs": [i for i, _ in io_pairs], + "outputs": [o for _, o in io_pairs], + "fn_name": None, + } + ).decode() + } + ] + return sample, [[solution]] + + +# Four distinct problems with distinct expected verdicts so per-problem +# misalignment (e.g. from a batching bug) can't pass all assertions. +_CONCURRENT_PROBLEMS: list[tuple[list[dict[str, str]], list[list[str]], float]] = [ + ( + *_make_problem( + [("1 2\n", "3\n"), ("10 20\n", "30\n")], + "a,b=map(int,input().split());print(a+b)", + ), + 1.0, + ), + ( + *_make_problem([("5\n", "25\n"), ("3\n", "9\n")], "n=int(input());print(n*n)"), + 1.0, + ), + ( + *_make_problem([("4\n", "16\n")], "print('wrong')"), # deliberately wrong + 0.0, + ), + ( + *_make_problem([("7\n", "49\n"), ("2\n", "4\n")], "n=int(input());print(n**2)"), + 1.0, + ), +] + + @pytest.mark.slow @pytest.mark.asyncio async def test_worker_grades_correct_stdin_solution() -> None: @@ -52,20 +94,27 @@ async def test_worker_grades_correct_stdin_solution() -> None: @pytest.mark.slow @pytest.mark.asyncio async def test_worker_grades_multiple_problems_concurrently() -> None: - """N concurrent grade_codegen() calls all resolve correctly. + """Concurrent grade_codegen() calls with distinct problems all return the + correct per-problem verdict. - This exercises the batch-drain path: all N requests are sent before the - worker responds, so they are drained into a single codegen_metrics call and - processed in parallel by lighteval's ProcessPoolExecutor. + Uses four problems with different expected pass@1 values (including one + deliberately wrong) so per-problem misalignment from a batching bug cannot + pass all assertions — identical problems would mask misattributed results. """ worker = CodegenGradingWorker() - sample, code = _sample_and_solution() - n = 4 try: results = await asyncio.gather( - *[worker.grade_codegen(sample, code, timeout=240) for _ in range(n)] + *[ + worker.grade_codegen(sample, code, timeout=240) + for sample, code, _ in _CONCURRENT_PROBLEMS + ] ) - assert len(results) == n - assert all(float(r["pass@1"]) == 1.0 for r in results), results + assert len(results) == len(_CONCURRENT_PROBLEMS) + for i, ((_, _, expected), result) in enumerate( + zip(_CONCURRENT_PROBLEMS, results, strict=True) + ): + assert float(result["pass@1"]) == expected, ( + f"problem {i}: expected pass@1={expected}, got {result}" + ) finally: await worker.aclose() From 258664a440c6055951e4d42543fd74f86207da47 Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Tue, 4 Aug 2026 13:50:00 -0700 Subject: [PATCH 24/26] test(accuracy): rewrite respawn test to use public API and add assertions The previous test had no assertions and drove respawn by directly manipulating internal state (_spawn_lock, _kill, _worker_proven) rather than via grade_codegen(). Delete the dead worker and the test still passed. Rewrite to: - Grade once through the public API and assert the result. - Wait for the worker to exit naturally (returncode != None). - Grade again and assert both the result and that _proc changed. Fixes thread #3708795700. Signed-off-by: Elias Bermudez --- .../accuracy/test_codegen_worker_client.py | 24 +++++++++++-------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/tests/unit/accuracy/test_codegen_worker_client.py b/tests/unit/accuracy/test_codegen_worker_client.py index 8845000cff..78e68d153f 100644 --- a/tests/unit/accuracy/test_codegen_worker_client.py +++ b/tests/unit/accuracy/test_codegen_worker_client.py @@ -603,9 +603,8 @@ async def test_grade_codegen_write_error_raises_worker_error( async def test_ensure_worker_respawns_after_worker_exits( self, tmp_path: Path ) -> None: - # Covers the _kill() call in _ensure_worker (line 138) by triggering a - # respawn. The second grade sees a dead worker and spawns a new one. - # We verify by checking that _proc changes between grades. + # Worker serves one request then exits. The second grade_codegen() call + # must detect the dead worker via returncode, respawn, and succeed. die_after_one = """ import sys, orjson line = sys.stdin.buffer.readline() @@ -617,16 +616,21 @@ async def test_ensure_worker_respawns_after_worker_exits( """ w = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, die_after_one)) try: - await w.grade_codegen([{"input_output": "{}"}], [["x"]], timeout=10) + result1 = await w.grade_codegen( + [{"input_output": "{}"}], [["x"]], timeout=10 + ) + assert result1 == {"pass@1": 1.0} first_proc = w._proc - # Force the proc to appear exited so _ensure_worker takes the respawn path + # Let the worker finish exiting so _ensure_worker sees returncode != None. if first_proc is not None: await first_proc.wait() - # Now _ensure_worker should detect returncode is not None and respawn - async with w._spawn_lock: - if w._proc is not None or w._reader_task is not None: - await w._kill() - w._worker_proven = False + result2 = await w.grade_codegen( + [{"input_output": "{}"}], [["x"]], timeout=10 + ) + assert result2 == {"pass@1": 1.0} + assert w._proc is not first_proc, ( + "expected a new worker process after respawn" + ) finally: await w.aclose() From 0494ebc6ea2b8dd429fbcc673ad4d17235246b16 Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Tue, 4 Aug 2026 16:13:40 -0700 Subject: [PATCH 25/26] fix(accuracy): move extend() calls inside try/except to preserve never-raises contract handle_batch documents "Never raises" but the extend() calls on all_samples/all_generations sat outside the (KeyError, TypeError) guard. A request with evaluation_sample=null passes the dict lookups (sample=None) then raises TypeError in extend(), propagating out of handle_batch and run_worker_loop and killing the worker along with every well-formed sibling in the batch. Move start/extend into the try block so any TypeError from a null or non-iterable field produces an error response instead. Fixes thread #3716745952. Signed-off-by: Elias Bermudez --- src/aiperf/accuracy/graders/_codegen_worker.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/aiperf/accuracy/graders/_codegen_worker.py b/src/aiperf/accuracy/graders/_codegen_worker.py index ce29578a85..fb885d5deb 100644 --- a/src/aiperf/accuracy/graders/_codegen_worker.py +++ b/src/aiperf/accuracy/graders/_codegen_worker.py @@ -117,6 +117,9 @@ def handle_batch( try: sample = req["evaluation_sample"] generation = req["generated_code"] + start = len(all_samples) + all_samples.extend(sample) + all_generations.extend(generation) except (KeyError, TypeError) as exc: responses[i] = { "id": req_id, @@ -124,9 +127,6 @@ def handle_batch( "error": f"malformed request: {exc!r}", } continue - start = len(all_samples) - all_samples.extend(sample) - all_generations.extend(generation) id_map.append((i, req_id, start, len(all_samples) - start)) if all_samples: From 80afebe0560ccb026a1d297c5789696bb79c76d9 Mon Sep 17 00:00:00 2001 From: Elias Bermudez Date: Tue, 4 Aug 2026 16:14:20 -0700 Subject: [PATCH 26/26] fix(accuracy): acquire _spawn_lock in aclose() to prevent orphaned workers aclose() called _kill() without holding _spawn_lock. If a grade was inside _ensure_worker() awaiting create_subprocess_exec, _kill() saw _proc=None and skipped cleanup; the spawn then completed, assigned self._proc, and left a start_new_session worker (plus any lighteval sandbox grandchildren) alive with no owner. Acquiring _spawn_lock before _kill() serialises teardown with any in-flight spawn: either _kill() runs first and the spawn finds _proc already cleared, or the spawn completes and _kill() then reaps it. Fixes thread #3716745958. Signed-off-by: Elias Bermudez --- src/aiperf/accuracy/graders/_codegen_worker_client.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/aiperf/accuracy/graders/_codegen_worker_client.py b/src/aiperf/accuracy/graders/_codegen_worker_client.py index 3cbc775335..da75d3a94c 100644 --- a/src/aiperf/accuracy/graders/_codegen_worker_client.py +++ b/src/aiperf/accuracy/graders/_codegen_worker_client.py @@ -341,4 +341,9 @@ async def aclose(self) -> None: if not fut.done(): fut.set_exception(CodegenWorkerError("grading worker closed")) self._pending.clear() - await self._kill() + # Acquire _spawn_lock so teardown cannot race with an in-flight + # _ensure_worker awaiting create_subprocess_exec. Without the lock, + # _kill() sees _proc=None and skips cleanup; the spawn then completes + # and leaves an orphaned start_new_session worker with no owner. + async with self._spawn_lock: + await self._kill()