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/ diff --git a/src/aiperf/accuracy/graders/_codegen_worker.py b/src/aiperf/accuracy/graders/_codegen_worker.py index 352f317cf1..fb885d5deb 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 @@ -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,) @@ -39,38 +46,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 +77,175 @@ 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] = [] + # (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): + 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: + 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, + "ok": False, + "error": f"malformed request: {exc!r}", + } + continue + id_map.append((i, req_id, start, len(all_samples) - start)) + + 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 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( + {j: raw_results[start + j] for j in range(count)}, + 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_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. + + 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 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. + """ + raw = getattr(stdin, "raw", None) + fd = raw.fileno() if raw is not None else -1 + + 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: + # 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 + if b"\n" not in available: + # Only a partial line is buffered; leave it for the next cycle's + # blocking readline() to complete. + break + line = stdin.readline() + line = line.strip() + if line: + lines.append(line) + finally: + _fcntl.fcntl(fd, _fcntl.F_SETFL, flags) # Restore blocking for next cycle + return lines + + 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. + """ + 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}", - } - else: - resp = handle_request(req, codegen_fn) - out.write(orjson.dumps(resp) + b"\n") + batch_raw: list[bytes] = [first] + 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() @@ -200,9 +323,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/src/aiperf/accuracy/graders/_codegen_worker_client.py b/src/aiperf/accuracy/graders/_codegen_worker_client.py index 1ba4074926..da75d3a94c 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,118 @@ 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() + 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 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(), 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 asyncio.CancelledError: + self._pending.pop(req_id, None) + raise + try: + 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) + 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) + 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 + # 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() + # 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,96 +212,91 @@ 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 + 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") + 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: - self._proc.stdin.write(orjson.dumps(req) + b"\n") - await self._proc.stdin.drain() - line = await asyncio.wait_for(self._proc.stdout.readline(), timeout) - 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}" - ) - + 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"): - # 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")) - + fut.set_exception( + CodegenWorkerError(resp.get("error", "unknown grading error")) + ) + self._mark_proven() + return True 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" - ) + 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.""" + if self._proc is None or self._proc.stdout is None: + return + reader = self._proc.stdout + try: + 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) or not self._dispatch_response(resp): + await self._handle_fault() + return + except asyncio.CancelledError: + 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 + # 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: @@ -261,11 +308,24 @@ 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: + 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() + 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 +337,13 @@ 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: + for fut in list(self._pending.values()): + if not fut.done(): + fut.set_exception(CodegenWorkerError("grading worker closed")) + self._pending.clear() + # 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() diff --git a/tests/component_integration/test_lcb_codegen_worker_e2e.py b/tests/component_integration/test_lcb_codegen_worker_e2e.py index 581ed70439..a86071d7d8 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 @@ -32,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: @@ -45,3 +89,32 @@ 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: + """Concurrent grade_codegen() calls with distinct problems all return the + correct per-problem verdict. + + 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() + try: + results = await asyncio.gather( + *[ + worker.grade_codegen(sample, code, timeout=240) + for sample, code, _ in _CONCURRENT_PROBLEMS + ] + ) + 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() diff --git a/tests/unit/accuracy/test_codegen_worker.py b/tests/unit/accuracy/test_codegen_worker.py index 019af1dbb9..d11ac067d7 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 @@ -17,84 +18,334 @@ _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]]: + # 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 -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_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_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] + # 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_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( + 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 + - def test_codegen_exception_becomes_error_response(self) -> None: - req = { - "id": 3, +class TestRunWorkerLoopBatch: + def _run( + self, + payloads: list[dict[str, Any]], + 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. + 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 + + 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} + + 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) - resp = worker.handle_request( - {"id": 9, "evaluation_sample": [{}], "generated_code": [["x"]]}, _nan_inf + @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 + # 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 ) - assert resp["ok"] is True - assert "pass@1" not in resp["metrics"] - assert "extra" not in resp["metrics"] - assert resp["metrics"]["ok"] == 1.0 + 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; + # 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]: 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 +358,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 diff --git a/tests/unit/accuracy/test_codegen_worker_client.py b/tests/unit/accuracy/test_codegen_worker_client.py index 718ac07ae4..78e68d153f 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 @@ -22,39 +23,67 @@ def _write_worker(tmp_path: Path, body: str) -> list[str]: return [sys.executable, str(script)] -# Echoes pass@1=1.0 for every request, correlating id. -_ECHO_OK = """ - import sys, orjson +# 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() """ -# 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 +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 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": 1.0}} 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: @@ -78,19 +107,110 @@ async def test_second_grade_reuses_same_worker(self, tmp_path) -> None: await worker.aclose() -class TestSerialization: - async def test_concurrent_grades_do_not_interleave(self, tmp_path) -> None: - worker = CodegenGradingWorker(worker_cmd=_write_worker(tmp_path, _ECHO_TRACKED)) +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. + result = await w.grade_codegen( + [{"input_output": "{}"}], [["x"]], timeout=10 + ) + assert result == {"pass@1": 1.0} + 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 + + 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( *[ - 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() # The very first grade (client request id==1) hangs forever to trigger a @@ -180,22 +300,39 @@ 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. + # + # 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 request is in flight (worker 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 None # cancellation killed the worker + assert worker._proc is not None # worker stays alive after cancel + # 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 + ) + assert result == {"pass@1": 1.0} finally: await worker.aclose() @@ -269,7 +406,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 +438,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: @@ -403,4 +544,204 @@ 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: + # 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() + 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: + result1 = await w.grade_codegen( + [{"input_output": "{}"}], [["x"]], timeout=10 + ) + assert result1 == {"pass@1": 1.0} + first_proc = w._proc + # Let the worker finish exiting so _ensure_worker sees returncode != None. + if first_proc is not None: + await first_proc.wait() + 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() + + 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)