Files
2026-07-13 12:09:03 +08:00

346 lines
11 KiB
Python

"""Code execution metric: extract, sandbox subprocess, pass-at-k.
Conceptual references:
- ./docs/en.md (this lesson)
- lesson 70 (task spec format) for code_exec task shape
- lesson 71 (classical metrics) for the dispatcher pattern
Stdlib + numpy. Run: python3 code/main.py
"""
from __future__ import annotations
import json
import math
import os
import re
import select
import subprocess
import sys
import tempfile
import time
from dataclasses import dataclass, field
from typing import Iterable
import numpy as np
CODE_BLOCK_RE = re.compile(r"```(?:[a-zA-Z0-9_]*)\n(.*?)```", re.DOTALL)
DEFAULT_TIMEOUT_S = 3.0
MAX_TIMEOUT_S = 30.0
OUTPUT_CAP_BYTES = 256 * 1024
EXIT_PASS = "pass"
EXIT_ASSERTION_FAIL = "assertion_fail"
EXIT_SYNTAX_ERROR = "syntax_error"
EXIT_TIMEOUT = "timeout"
EXIT_ERROR = "error"
EXIT_CODES = (EXIT_PASS, EXIT_ASSERTION_FAIL, EXIT_SYNTAX_ERROR, EXIT_TIMEOUT, EXIT_ERROR)
@dataclass
class ExecResult:
score: float
exit_code: str
passed: int
total: int
detail: str = ""
per_assertion: list[bool] = field(default_factory=list)
def to_dict(self) -> dict:
return {
"score": self.score,
"exit_code": self.exit_code,
"passed": self.passed,
"total": self.total,
"detail": self.detail,
"per_assertion": list(self.per_assertion),
}
RUNNER_TEMPLATE = r"""
import json, sys, traceback
import builtins as _builtins
_DENIED_MODULES = {
'subprocess', 'socket', 'shutil', 'requests', 'urllib',
'urllib.request', 'urllib.error', 'urllib.parse', 'ctypes',
'http.client', 'asyncio.subprocess',
}
_real_import = _builtins.__import__
def _guarded_import(name, *args, **kwargs):
base = name.split('.')[0]
if name in _DENIED_MODULES or base in _DENIED_MODULES:
raise ImportError("denied module: %s" % name)
return _real_import(name, *args, **kwargs)
_builtins.__import__ = _guarded_import
import os as _os
def _denied_system(_cmd):
raise OSError("os.system denied")
_os.system = _denied_system
CODE = __PLACEHOLDER_CODE__
ASSERTIONS = __PLACEHOLDER_ASSERTIONS__
ns = {}
status = {"exit_code": "pass", "results": [], "detail": ""}
try:
compiled = compile(CODE, "<candidate>", "exec")
except SyntaxError as exc:
status["exit_code"] = "syntax_error"
status["detail"] = "SyntaxError: %s" % exc
sys.stdout.write(json.dumps(status))
sys.exit(0)
except Exception as exc:
status["exit_code"] = "error"
status["detail"] = "compile error: %r" % exc
sys.stdout.write(json.dumps(status))
sys.exit(0)
try:
exec(compiled, ns)
except Exception as exc:
status["exit_code"] = "error"
status["detail"] = "exec error: %s" % exc
sys.stdout.write(json.dumps(status))
sys.exit(0)
passed = 0
for assertion in ASSERTIONS:
try:
ok = bool(eval(assertion, ns))
except Exception as exc:
ok = False
status["detail"] = "assertion error: %s" % exc
status["results"].append(ok)
if ok:
passed += 1
if passed < len(ASSERTIONS):
status["exit_code"] = "assertion_fail" if passed > 0 or len(ASSERTIONS) > 0 else status["exit_code"]
if passed == len(ASSERTIONS):
status["exit_code"] = "pass"
elif len(ASSERTIONS) > 0:
status["exit_code"] = "pass"
sys.stdout.write(json.dumps(status))
""".strip()
def extract_code(text: str) -> str | None:
m = CODE_BLOCK_RE.search(text)
if m:
return m.group(1)
if "```" in text:
return None
return text
def _build_runner(code: str, assertions: list[str]) -> str:
script = RUNNER_TEMPLATE
script = script.replace("__PLACEHOLDER_CODE__", repr(code))
script = script.replace("__PLACEHOLDER_ASSERTIONS__", repr(list(assertions)))
return script
def run_candidate(code: str, assertions: list[str], timeout_s: float = DEFAULT_TIMEOUT_S) -> ExecResult:
if not assertions:
return ExecResult(0.0, EXIT_ERROR, 0, 0, detail="no assertions provided")
timeout_s = min(max(0.1, float(timeout_s)), MAX_TIMEOUT_S)
script = _build_runner(code, assertions)
try:
proc = subprocess.Popen(
[sys.executable, "-I", "-c", script],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
except Exception as exc:
return ExecResult(0.0, EXIT_ERROR, 0, len(assertions), detail=f"spawn error: {exc!r}")
with proc:
deadline = time.monotonic() + timeout_s
stdout_chunks: list[bytes] = []
stdout_bytes = 0
overflow = False
timed_out = False
fd = proc.stdout
assert fd is not None
os.set_blocking(fd.fileno(), False)
try:
while True:
remaining = deadline - time.monotonic()
if remaining <= 0:
timed_out = True
break
try:
r, _, _ = select.select([fd], [], [], min(remaining, 0.1))
except (OSError, ValueError):
break
if r:
chunk = fd.read(8192)
if chunk is None:
continue
if chunk == b"":
break
stdout_bytes += len(chunk)
if stdout_bytes > OUTPUT_CAP_BYTES:
overflow = True
break
stdout_chunks.append(chunk)
elif proc.poll() is not None:
tail = fd.read()
if tail:
stdout_bytes += len(tail)
if stdout_bytes > OUTPUT_CAP_BYTES:
overflow = True
else:
stdout_chunks.append(tail)
break
finally:
if proc.poll() is None:
proc.kill()
try:
proc.wait(timeout=1.0)
except subprocess.TimeoutExpired:
pass
try:
stderr_bytes = proc.stderr.read() if proc.stderr else b""
except Exception:
stderr_bytes = b""
if timed_out:
return ExecResult(0.0, EXIT_TIMEOUT, 0, len(assertions),
detail=f"timeout after {timeout_s:.2f}s")
if overflow:
return ExecResult(0.0, EXIT_ERROR, 0, len(assertions), detail="output overflow")
stdout = b"".join(stdout_chunks).decode("utf-8", errors="replace")
if proc.returncode != 0:
stderr_text = stderr_bytes.decode("utf-8", errors="replace") if stderr_bytes else ""
return ExecResult(0.0, EXIT_ERROR, 0, len(assertions),
detail=f"subprocess exit {proc.returncode}: {stderr_text[:200]}")
try:
status = json.loads(stdout) if stdout else {}
except json.JSONDecodeError:
return ExecResult(0.0, EXIT_ERROR, 0, len(assertions), detail="non-json output")
exit_code = status.get("exit_code", EXIT_ERROR)
results = list(status.get("results", []))
detail = status.get("detail", "")
passed = sum(1 for r in results if r)
total = len(assertions)
score = passed / total if total else 0.0
if exit_code not in EXIT_CODES:
exit_code = EXIT_ERROR
if exit_code == EXIT_PASS and passed < total:
exit_code = EXIT_ASSERTION_FAIL
if exit_code == EXIT_ASSERTION_FAIL and passed == total and total > 0:
exit_code = EXIT_PASS
return ExecResult(score, exit_code, passed, total, detail=detail, per_assertion=results)
def score_code_exec(prediction: str, assertions: list[str], timeout_s: float = DEFAULT_TIMEOUT_S) -> float:
code = extract_code(prediction)
if code is None:
return 0.0
return run_candidate(code, assertions, timeout_s=timeout_s).score
def pass_at_k(n: int, c: int, k: int) -> float:
if n < 0 or c < 0 or k <= 0:
raise ValueError("n, c >= 0 and k > 0")
if c > n:
raise ValueError("c cannot exceed n")
if k > n:
return float(c > 0)
if n - c < k:
return 1.0
log_prob = 0.0
for i in range(k):
log_prob += math.log(n - c - i) - math.log(n - i)
return float(1.0 - math.exp(log_prob))
def pass_at_k_estimator(samples: list[list[bool]], ks: list[int]) -> dict[int, float]:
out: dict[int, float] = {}
for k in ks:
per_task = []
for sample in samples:
n = len(sample)
c = sum(1 for s in sample if s)
per_task.append(pass_at_k(n, c, k))
out[k] = float(np.mean(per_task)) if per_task else 0.0
return out
def _demo_examples() -> list[dict]:
return [
{
"name": "passes_all",
"prediction": "```python\ndef add(a, b):\n return a + b\n```",
"assertions": ["add(1, 2) == 3", "add(-5, 5) == 0", "add(100, 1) == 101"],
"expected_exit": EXIT_PASS,
"expected_score": 1.0,
},
{
"name": "one_assertion_fails",
"prediction": "```python\ndef add(a, b):\n return a + b + 1\n```",
"assertions": ["add(1, 2) == 3", "add(-5, 5) == 0"],
"expected_exit": EXIT_ASSERTION_FAIL,
"expected_score": 0.0,
},
{
"name": "syntax_error",
"prediction": "```python\ndef add(a, b)\n return a + b\n```",
"assertions": ["add(1, 2) == 3"],
"expected_exit": EXIT_SYNTAX_ERROR,
"expected_score": 0.0,
},
{
"name": "timeout",
"prediction": "```python\nimport time\ndef add(a, b):\n time.sleep(10)\n return a + b\n```",
"assertions": ["add(1, 2) == 3"],
"expected_exit": EXIT_TIMEOUT,
"expected_score": 0.0,
"timeout_s": 0.6,
},
{
"name": "denylist_subprocess",
"prediction": "```python\nimport subprocess\ndef bad():\n subprocess.run(['ls'])\nbad()\n```",
"assertions": ["bad() is None"],
"expected_exit": EXIT_ERROR,
"expected_score": 0.0,
},
]
def demo() -> int:
failures = 0
for ex in _demo_examples():
result = run_candidate(
extract_code(ex["prediction"]) or "",
ex["assertions"],
timeout_s=ex.get("timeout_s", DEFAULT_TIMEOUT_S),
)
status = "OK" if result.exit_code == ex["expected_exit"] and abs(result.score - ex["expected_score"]) < 1e-6 else "WRONG"
if status == "WRONG":
failures += 1
print(f" [{status}] {ex['name']:24s} exit={result.exit_code:14s} score={result.score:.2f} detail={result.detail[:60]}")
samples_per_task = [
[True, True, False, False, False, False, False, False, False, False],
[True, False, False, False, False, False, False, False, False, False],
[False, False, False, False, False, False, False, False, False, False],
]
p_at_k = pass_at_k_estimator(samples_per_task, [1, 5, 10])
print(f" pass@1={p_at_k[1]:.3f} pass@5={p_at_k[5]:.3f} pass@10={p_at_k[10]:.3f}")
return 0 if failures == 0 else 1
if __name__ == "__main__":
sys.exit(demo())