refactor: rewrite humaneval evaluation with functional pipeline design

- fix KeyError race condition in inference cache touch()
- EvalConfig dataclass for centralized configuration
- load->generate->extract->test->score->report pipeline
- two-phase generation+testing for max GPU utilization
- signal-based SIGALRM timeout protection for code exec
- suppress subprocess stdout/stderr pollution

astrai/inference/core/cache.py

@@ -62,7 +62,8 @@ class Allocator:
 
     def touch(self, idx: int):
         with self._lock:
-            self._lru.move_to_end(idx)
+            if idx in self._lru:
+                self._lru.move_to_end(idx)
 
 
 class PrefixCache:

scripts/eval/evaluate_humaneval.py

@@ -1,22 +1,21 @@
-"""HumanEval code generation benchmark.
+"""HumanEval benchmark — functional pipeline design.
 
-Generates n completions per problem, extracts function bodies, executes
+Pipeline:
-against hidden tests, and computes pass@k.
+    load -> generate -> extract -> test -> score -> report
 
-Usage::
+Each stage is a pure function (except GPU/CPU-bound I/O stages).
-
+Config is a single dataclass; side effects are isolated at pipeline boundaries.
-    python scripts/tools/evaluate_humaneval.py --param_path ./params \
-        --data_path HumanEval.jsonl.gz --output results.json \
-        --num_samples 200 --temperature 0.8 --max_tokens 512
 """
 
 import argparse
 import json
 import os
 import re
+import signal
+import sys
+from dataclasses import dataclass
 from math import prod
-from multiprocessing import Process, Queue
+from typing import Dict, Iterator, List, Optional, Sequence, Tuple
-from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 import torch
@@ -26,11 +25,15 @@ from astrai.inference import InferenceEngine
 from astrai.model import AutoModel
 from astrai.tokenize import AutoTokenizer
 
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
 HUMANEVAL_URL = (
     "https://github.com/openai/human-eval/raw/master/data/HumanEval.jsonl.gz"
 )
 
-_STOP_SEQUENCES = [
+STOP_SEQUENCES = [
     "\nclass ",
     "\ndef ",
     "\n# ",
@@ -40,43 +43,82 @@ _STOP_SEQUENCES = [
 ]
 
 
-def _download_humaneval(data_path: str):
+@dataclass
-    if os.path.exists(data_path):
+class EvalConfig:
+    param_path: str = "./params"
+    data_path: str = "./humaneval/HumanEval.jsonl"
+    output: Optional[str] = None
+
+    num_samples: int = 200
+    max_tokens: int = 512
+    temperature: float = 0.8
+    top_p: float = 0.95
+    top_k: int = 50
+    batch_size: int = 32
+    test_timeout: float = 3.0
+    test_workers: int = 8
+    k_values: Tuple[int, ...] = (1, 10, 100)
+    problem_indices: Optional[List[int]] = None
+
+
+def download(url: str, path: str):
+    if os.path.exists(path):
         return
     import gzip
     import urllib.request
 
-    os.makedirs(os.path.dirname(data_path) or ".", exist_ok=True)
+    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
-    print(f"Downloading HumanEval from {HUMANEVAL_URL} ...")
+    print(f"Downloading {url} ...")
-    tmp = data_path + ".tmp"
+    tmp = path + ".tmp"
-    urllib.request.urlretrieve(HUMANEVAL_URL, tmp)
+    urllib.request.urlretrieve(url, tmp)
     with gzip.open(tmp, "rb") as f_in:
-        with open(data_path, "wb") as f_out:
+        with open(path, "wb") as f_out:
             f_out.write(f_in.read())
     os.remove(tmp)
-    print(f"  saved to {data_path}")
+    print(f"  saved to {path}")
 
 
-def _load_problems(data_path: str) -> List[dict]:
+def load_jsonl(path: str) -> List[dict]:
-    problems = []
+    rows = []
-    with open(data_path, "r", encoding="utf-8") as f:
+    with open(path, encoding="utf-8") as f:
         for line in f:
             line = line.strip()
             if line:
-                problems.append(json.loads(line))
+                rows.append(json.loads(line))
-    return problems
+    return rows
 
 
-def _extract_function_body(code: str, entry_point: str) -> Optional[str]:
+def save_json(path: str, data):
-    """Extract the function body from a completion."""
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=2, ensure_ascii=False)
+
+
+def create_engine(param_path: str, batch_size: int) -> InferenceEngine:
+    model = AutoModel.from_pretrained(param_path)
+    tokenizer = AutoTokenizer.from_pretrained(param_path)
+    model.to(device="cuda", dtype=torch.bfloat16)
+    return InferenceEngine(
+        model=model,
+        tokenizer=tokenizer,
+        max_batch_size=batch_size,
+    )
+
+
+def trim_stop(text: str) -> str:
+    for stop in STOP_SEQUENCES:
+        idx = text.find(stop)
+        if idx != -1:
+            text = text[:idx]
+    return text
+
+
+def extract_body(code: str, entry_point: str) -> Optional[str]:
     pattern = rf"def\s+{re.escape(entry_point)}\b[^:]*:"
     match = re.search(pattern, code)
     if not match:
-        # Use the full code as-is if we can't find the function
         return code
 
-    body_start = match.end()
+    lines = code[match.end() :].split("\n")
-    lines = code[body_start:].split("\n")
     body_lines = []
     started = False
 
@@ -94,240 +136,243 @@ def _extract_function_body(code: str, entry_point: str) -> Optional[str]:
         body_lines.append(stripped)
 
     body = "\n".join(body_lines)
-    if not body.strip():
+    return body if body.strip() else None
-        return None
-    return body
 
 
-def _trim_stop_sequences(text: str) -> str:
+def deduplicate(seq: Sequence[str]) -> List[str]:
-    for stop in _STOP_SEQUENCES:
-        idx = text.find(stop)
-        if idx != -1:
-            text = text[:idx]
-    return text
-
-
-def _execute_code(problem: dict, completion: str, timeout: float = 3.0) -> bool:
-    """Run the completion against hidden tests in a subprocess."""
-
-    def _worker(queue, full_code):
-        try:
-            namespace = {}
-            exec(full_code, namespace)
-            check = namespace.get("check")
-            if check is None:
-                queue.put(False)
-                return
-            check(namespace.get(problem["entry_point"]))
-            queue.put(True)
-        except Exception:
-            queue.put(False)
-
-    full_code = problem["prompt"] + completion + "\n" + problem["test"]
-
-    queue: Queue = Queue()
-    proc = Process(target=_worker, args=(queue, full_code))
-    proc.start()
-    proc.join(timeout)
-
-    if proc.is_alive():
-        proc.terminate()
-        proc.join()
-        return False
-
-    try:
-        return queue.get_nowait()
-    except Exception:
-        return False
-
-
-def _pass_at_k(n: int, c: int, k: int) -> float:
-    """Unbiased estimator of pass@k."""
-    if n - c < k:
-        return 1.0
-    return 1.0 - float(prod(1.0 - k / np.arange(n - c + 1, n + 1)))
-
-
-def _deduplicate(completions: List[str]) -> List[str]:
     seen = set()
-    unique = []
+    return [x for x in seq if not (x in seen or seen.add(x))]
-    for c in completions:
-        if c not in seen:
-            seen.add(c)
-            unique.append(c)
-    return unique
 
 
-def _generate(
+def generate_batch(
     engine: InferenceEngine,
     prompt: str,
-    num_samples: int,
+    n: int,
+    batch_size: int,
     max_tokens: int,
     temperature: float,
     top_p: float,
     top_k: int,
-    batch_size: int,
 ) -> List[str]:
-    batches = [prompt] * min(batch_size, num_samples)
     completions = []
-    remaining = num_samples
+    remaining = n
-
     while remaining > 0:
         current = min(batch_size, remaining)
-        batch_prompts = batches[:current]
         outputs = engine.generate(
-            prompt=batch_prompts,
+            prompt=[prompt] * current,
             stream=False,
             max_tokens=max_tokens,
             temperature=temperature,
             top_p=top_p,
             top_k=top_k,
         )
-        if isinstance(outputs, str):
+        completions.extend(outputs if isinstance(outputs, list) else [outputs])
-            outputs = [outputs]
-        completions.extend(outputs)
         remaining -= current
-
+    return deduplicate(completions)
-    return _deduplicate(completions)
 
 
-def evaluate(
+def extract_completions(
+    raw: Sequence[str],
+    entry_point: str,
+) -> List[str]:
+    bodies = []
+    for r in raw:
+        t = trim_stop(r)
+        body = extract_body(t, entry_point)
+        if body:
+            bodies.append(body)
+    return bodies
+
+
+def generate_all(
     engine: InferenceEngine,
-    problems: List[dict],
+    problems: Sequence[dict],
-    num_samples: int,
+    cfg: EvalConfig,
-    max_tokens: int,
+) -> List[dict]:
-    temperature: float,
+    results = []
-    top_p: float,
+    for problem in tqdm.tqdm(problems, desc="Generating", unit="problem"):
-    top_k: int,
+        raw = generate_batch(
-    batch_size: int,
-    k_values: Tuple[int, ...] = (1, 10, 100),
-) -> Dict:
-    results = {}
-    all_pass_at_k = {k: [] for k in k_values}
-
-    for problem in tqdm.tqdm(problems, desc="HumanEval", unit="problem"):
-        task_id = problem["task_id"]
-        prompt = problem["prompt"]
-        entry_point = problem["entry_point"]
-
-        raw_completions = _generate(
             engine,
-            prompt,
+            problem["prompt"],
-            num_samples,
+            cfg.num_samples,
-            max_tokens,
+            cfg.batch_size,
-            temperature,
+            cfg.max_tokens,
-            top_p,
+            cfg.temperature,
-            top_k,
+            cfg.top_p,
-            batch_size,
+            cfg.top_k,
+        )
+        bodies = extract_completions(raw, problem["entry_point"])
+        results.append(
+            dict(
+                task_id=problem["task_id"],
+                entry_point=problem["entry_point"],
+                prompt=problem["prompt"],
+                test=problem["test"],
+                completions=bodies,
+            )
         )
-
-        completions = []
-        for raw in raw_completions:
-            trimmed = _trim_stop_sequences(raw)
-            body = _extract_function_body(trimmed, entry_point)
-            if body:
-                completions.append(body)
-
-        passed = 0
-        for comp in completions:
-            if _execute_code(problem, comp):
-                passed += 1
-
-        n = len(completions)
-        c = passed
-        result = {"task_id": task_id, "n": n, "passed": c}
-        for k in k_values:
-            result[f"pass@{k}"] = round(_pass_at_k(n, c, k), 4)
-            all_pass_at_k[k].append(_pass_at_k(n, c, k))
-        results[task_id] = result
-
-    summary = {}
-    for k in k_values:
-        vals = all_pass_at_k[k]
-        summary[f"pass@{k}"] = round(float(np.mean(vals)), 4)
-    results["_summary"] = summary
-
     return results
 
 
-def main():
+def _timeout_handler(signum, frame):
-    parser = argparse.ArgumentParser(description="HumanEval benchmark")
+    raise TimeoutError("execution timeout")
-    parser.add_argument(
-        "--param_path", type=str, default="./params", help="Model directory"
-    )
-    parser.add_argument(
-        "--data_path",
-        type=str,
-        default="./humaneval/HumanEval.jsonl",
-        help="HumanEval JSONL file (auto-download if missing)",
-    )
-    parser.add_argument("--output", type=str, default=None, help="Output JSON path")
-    parser.add_argument(
-        "--num_samples",
-        type=int,
-        default=200,
-        help="Completions per problem",
-    )
-    parser.add_argument(
-        "--max_tokens", type=int, default=512, help="Max generation tokens"
-    )
-    parser.add_argument(
-        "--temperature", type=float, default=0.8, help="Sampling temperature"
-    )
-    parser.add_argument("--top_p", type=float, default=0.95, help="Top-p sampling")
-    parser.add_argument("--top_k", type=int, default=50, help="Top-k sampling")
-    parser.add_argument(
-        "--batch_size", type=int, default=1, help="Inference batch size"
-    )
-    parser.add_argument(
-        "--problems",
-        type=int,
-        nargs="+",
-        default=None,
-        help="Specific problem indices (0-based)",
-    )
-    args = parser.parse_args()
 
-    _download_humaneval(args.data_path)
-    problems = _load_problems(args.data_path)
-    if args.problems:
-        problems = [problems[i] for i in args.problems if i < len(problems)]
 
-    model = AutoModel.from_pretrained(args.param_path)
+def execute_one(args: tuple) -> bool:
-    tokenizer = AutoTokenizer.from_pretrained(args.param_path)
+    full_code, entry_point, timeout = args
-    model.to(device="cuda", dtype=torch.bfloat16)
+    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:
+            ns = {}
+            exec(full_code, ns)
+            candidate = ns.get(entry_point)
+            check = ns.get("check")
+            if check is None or candidate is None:
+                return False
+            check(candidate)
+            return True
+        except Exception:
+            return False
+        finally:
+            signal.alarm(0)
+            sys.stdout, sys.stderr = old_out, old_err
 
-    engine = InferenceEngine(
-        model=model,
-        tokenizer=tokenizer,
-        max_batch_size=args.batch_size,
-    )
 
-    results = evaluate(
+def test_one(item: dict, cfg: EvalConfig) -> Tuple[str, int, int]:
-        engine=engine,
+    from concurrent.futures import ProcessPoolExecutor
-        problems=problems,
+
+    task_id = item["task_id"]
+    completions = item["completions"]
+    codes = [
+        (
+            item["prompt"] + c + "\n" + item["test"],
+            item["entry_point"],
+            cfg.test_timeout,
+        )
+        for c in completions
+    ]
+    n = len(codes)
+    passed = 0
+    with ProcessPoolExecutor(max_workers=cfg.test_workers) as pool:
+        for ok in pool.map(execute_one, codes):
+            if ok:
+                passed += 1
+    return task_id, n, passed
+
+
+def test_all(
+    items: Sequence[dict],
+    cfg: EvalConfig,
+) -> Iterator[Tuple[str, int, int]]:
+    for item in tqdm.tqdm(items, desc="Testing", unit="problem"):
+        yield test_one(item, cfg)
+
+
+def pass_at_k(n: int, c: int, k: int) -> float:
+    if n - c < k:
+        return 1.0
+    return 1.0 - float(prod(1.0 - k / np.arange(n - c + 1, n + 1)))
+
+
+def score_results(
+    results: Iterator[Tuple[str, int, int]],
+    k_values: Tuple[int, ...],
+) -> Dict:
+    scores = {k: [] for k in k_values}
+    output = {}
+    for task_id, n, passed in results:
+        scores["task_id"] = task_id
+        entry = {"task_id": task_id, "n": n, "passed": passed}
+        for k in k_values:
+            pk = round(pass_at_k(n, passed, k), 4)
+            entry[f"pass@{k}"] = pk
+            scores[k].append(pk)
+        output[task_id] = entry
+
+    summary = {}
+    for k in k_values:
+        vals = scores[k]
+        summary[f"pass@{k}"] = round(float(np.mean(vals)), 4)
+    output["_summary"] = summary
+    return output
+
+
+def run_pipeline(cfg: EvalConfig) -> Dict:
+    download(HUMANEVAL_URL, cfg.data_path)
+
+    problems = load_jsonl(cfg.data_path)
+    if cfg.problem_indices:
+        problems = [problems[i] for i in cfg.problem_indices if i < len(problems)]
+
+    engine = create_engine(cfg.param_path, cfg.batch_size)
+
+    try:
+        generated = generate_all(engine, problems, cfg)
+
+        if cfg.output:
+            mid = cfg.output.replace(".json", "_completions.json")
+            save_json(mid, generated)
+            print(f"Completions saved to {mid}")
+
+        results = test_all(generated, cfg)
+        scored = score_results(results, cfg.k_values)
+    finally:
+        engine.shutdown()
+
+    return scored
+
+
+def parse_args(argv: Optional[List[str]] = None) -> EvalConfig:
+    p = argparse.ArgumentParser(description="HumanEval benchmark")
+    p.add_argument("--param_path", type=str, default="./params")
+    p.add_argument("--data_path", type=str, default="./humaneval/HumanEval.jsonl")
+    p.add_argument("--output", type=str, default=None)
+    p.add_argument("--num_samples", type=int, default=200)
+    p.add_argument("--max_tokens", type=int, default=512)
+    p.add_argument("--temperature", type=float, default=0.8)
+    p.add_argument("--top_p", type=float, default=0.95)
+    p.add_argument("--top_k", type=int, default=50)
+    p.add_argument("--batch_size", type=int, default=32)
+    p.add_argument("--test_workers", type=int, default=8)
+    p.add_argument("--test_timeout", type=float, default=3.0)
+    p.add_argument("--problems", type=int, nargs="+", default=None)
+    args = p.parse_args(argv)
+
+    return EvalConfig(
+        param_path=args.param_path,
+        data_path=args.data_path,
+        output=args.output,
         num_samples=args.num_samples,
         max_tokens=args.max_tokens,
         temperature=args.temperature,
         top_p=args.top_p,
         top_k=args.top_k,
         batch_size=args.batch_size,
-        k_values=(1, 10, 100),
+        test_workers=args.test_workers,
+        test_timeout=args.test_timeout,
+        problem_indices=args.problems,
     )
 
-    summary = results.pop("_summary")
+
+def report(scored: Dict):
+    summary = scored.pop("_summary", {})
     print(f"\n{'=' * 60}")
     for k, v in summary.items():
         print(f"  {k}: {v:.2%}")
     print(f"{'=' * 60}")
+    scored["_summary"] = summary
 
-    if args.output:
-        results["_summary"] = summary
-        with open(args.output, "w", encoding="utf-8") as f:
-            json.dump(results, f, indent=2, ensure_ascii=False)
-        print(f"Results saved to {args.output}")
 
-    engine.shutdown()
+def main():
+    cfg = parse_args()
+    scored = run_pipeline(cfg)
+    report(scored)
+    if cfg.output:
+        save_json(cfg.output, scored)
+        print(f"Results saved to {cfg.output}")
 
 
 if __name__ == "__main__":