fix: reliable test timeout, separate generate/test phases, dynamic pass@k

- Replace SIGALRM+exec() with subprocess.run(timeout=) for test execution
- Add --test_only flag to skip generation and test existing completions
- Add --generate_only flag for generation-only runs
- Derive pass@k values from num_samples (filter k > n)
- Support loading completions from array JSON (not just JSONL)
This commit is contained in:
ViperEkura 2026-07-05 08:47:07 +08:00
parent 17d6eaa2f2
commit 599a51f4f7
1 changed files with 50 additions and 36 deletions

View File

@ -8,10 +8,11 @@ Config is a single dataclass; side effects are isolated at pipeline boundaries.
""" """
import argparse import argparse
import itertools
import json import json
import os import os
import re import re
import signal import subprocess
import sys import sys
from dataclasses import dataclass from dataclasses import dataclass
from math import prod from math import prod
@ -49,6 +50,9 @@ class EvalConfig:
data_path: str = "./humaneval/HumanEval.jsonl" data_path: str = "./humaneval/HumanEval.jsonl"
output: Optional[str] = None output: Optional[str] = None
test_only: Optional[str] = None
generate_only: bool = False
num_samples: int = 200 num_samples: int = 200
max_tokens: int = 512 max_tokens: int = 512
temperature: float = 0.8 temperature: float = 0.8
@ -214,31 +218,19 @@ def generate_all(
return results return results
def _timeout_handler(signum, frame):
raise TimeoutError("execution timeout")
def execute_one(args: tuple) -> bool: def execute_one(args: tuple) -> bool:
full_code, entry_point, timeout = args full_code, entry_point, timeout = args
signal.signal(signal.SIGALRM, _timeout_handler)
signal.alarm(int(timeout))
with open(os.devnull, "w") as devnull:
old_out, old_err = sys.stdout, sys.stderr
sys.stdout, sys.stderr = devnull, devnull
try: try:
ns = {} r = subprocess.run(
exec(full_code, ns) [sys.executable, "-c", full_code],
candidate = ns.get(entry_point) capture_output=True,
check = ns.get("check") timeout=timeout,
if check is None or candidate is None: )
return r.returncode == 0
except subprocess.TimeoutExpired:
return False return False
check(candidate)
return True
except Exception: except Exception:
return False return False
finally:
signal.alarm(0)
sys.stdout, sys.stderr = old_out, old_err
def test_one(item: dict, cfg: EvalConfig) -> Tuple[str, int, int]: def test_one(item: dict, cfg: EvalConfig) -> Tuple[str, int, int]:
@ -281,10 +273,15 @@ def score_results(
results: Iterator[Tuple[str, int, int]], results: Iterator[Tuple[str, int, int]],
k_values: Tuple[int, ...], k_values: Tuple[int, ...],
) -> Dict: ) -> Dict:
# filter to k <= n (peek first result to get n)
first = next(results)
results = itertools.chain([first], results)
n = first[1]
k_values = tuple(k for k in k_values if k <= n)
scores = {k: [] for k in k_values} scores = {k: [] for k in k_values}
output = {} output = {}
for task_id, n, passed in results: for task_id, n, passed in results:
scores["task_id"] = task_id
entry = {"task_id": task_id, "n": n, "passed": passed} entry = {"task_id": task_id, "n": n, "passed": passed}
for k in k_values: for k in k_values:
pk = round(pass_at_k(n, passed, k), 4) pk = round(pass_at_k(n, passed, k), 4)
@ -301,6 +298,10 @@ def score_results(
def run_pipeline(cfg: EvalConfig) -> Dict: def run_pipeline(cfg: EvalConfig) -> Dict:
if cfg.test_only:
with open(cfg.test_only, encoding="utf-8") as f:
generated = json.load(f)
else:
download(HUMANEVAL_URL, cfg.data_path) download(HUMANEVAL_URL, cfg.data_path)
problems = load_jsonl(cfg.data_path) problems = load_jsonl(cfg.data_path)
@ -311,17 +312,19 @@ def run_pipeline(cfg: EvalConfig) -> Dict:
try: try:
generated = generate_all(engine, problems, cfg) generated = generate_all(engine, problems, cfg)
finally:
engine.shutdown()
if cfg.output: if cfg.output:
mid = cfg.output.replace(".json", "_completions.json") mid = cfg.output.replace(".json", "_completions.json")
save_json(mid, generated) save_json(mid, generated)
print(f"Completions saved to {mid}") print(f"Completions saved to {mid}")
if cfg.generate_only:
return {}
results = test_all(generated, cfg) results = test_all(generated, cfg)
scored = score_results(results, cfg.k_values) scored = score_results(results, cfg.k_values)
finally:
engine.shutdown()
return scored return scored
@ -330,6 +333,15 @@ def parse_args(argv: Optional[List[str]] = None) -> EvalConfig:
p.add_argument("--param_path", type=str, default="./params") p.add_argument("--param_path", type=str, default="./params")
p.add_argument("--data_path", type=str, default="./humaneval/HumanEval.jsonl") p.add_argument("--data_path", type=str, default="./humaneval/HumanEval.jsonl")
p.add_argument("--output", type=str, default=None) p.add_argument("--output", type=str, default=None)
p.add_argument(
"--test_only",
type=str,
default=None,
help="Skip generation, test existing completions JSON",
)
p.add_argument(
"--generate_only", action="store_true", help="Only generate, skip testing"
)
p.add_argument("--num_samples", type=int, default=200) p.add_argument("--num_samples", type=int, default=200)
p.add_argument("--max_tokens", type=int, default=512) p.add_argument("--max_tokens", type=int, default=512)
p.add_argument("--temperature", type=float, default=0.8) p.add_argument("--temperature", type=float, default=0.8)
@ -345,6 +357,8 @@ def parse_args(argv: Optional[List[str]] = None) -> EvalConfig:
param_path=args.param_path, param_path=args.param_path,
data_path=args.data_path, data_path=args.data_path,
output=args.output, output=args.output,
test_only=args.test_only,
generate_only=args.generate_only,
num_samples=args.num_samples, num_samples=args.num_samples,
max_tokens=args.max_tokens, max_tokens=args.max_tokens,
temperature=args.temperature, temperature=args.temperature,