From ba937d03b8ac59b36bca60e42230ee9d53d4e688 Mon Sep 17 00:00:00 2001 From: fluffy314 Date: Tue, 21 Jul 2026 19:25:40 +0800 Subject: [PATCH] fix(autoresearch): warm and retry Lean typechecks Prewarm a reduced mathlib environment, classify 45/120-second typecheck retries, kill timed-out process groups, and stop whitespace-only model decode loops. Co-authored-by: Cursor --- KakeyaLeanGate/Prelude.lean | 6 +- autoresearch/prefill/lean_gate.py | 231 +++++++++++++++--- autoresearch/prefill/program.md | 8 + autoresearch/prefill/supervisor.py | 13 + scripts/agent_gan_inference_demo.py | 14 ++ scripts/agent_gan_repl.py | 16 +- .../bench/test_lean_signature_gate.py | 39 +++ .../bridge/test_agent_gan_demo.py | 25 ++ 8 files changed, 319 insertions(+), 33 deletions(-) diff --git a/KakeyaLeanGate/Prelude.lean b/KakeyaLeanGate/Prelude.lean index 2449161..a57e3d3 100644 --- a/KakeyaLeanGate/Prelude.lean +++ b/KakeyaLeanGate/Prelude.lean @@ -1,4 +1,8 @@ -import Mathlib +import Mathlib.Analysis.Complex.Basic +import Mathlib.Analysis.Complex.Hadamard +import Mathlib.Analysis.Complex.JensenFormula +import Mathlib.Analysis.Complex.LocallyUniformLimit +import Mathlib.Analysis.Complex.Order /-! Minimal import target for AutoResearch theorem-signature validation. diff --git a/autoresearch/prefill/lean_gate.py b/autoresearch/prefill/lean_gate.py index 5ac429f..3b01934 100644 --- a/autoresearch/prefill/lean_gate.py +++ b/autoresearch/prefill/lean_gate.py @@ -2,9 +2,12 @@ from __future__ import annotations import hashlib +import os import re +import signal import subprocess import tempfile +import time from dataclasses import dataclass from pathlib import Path @@ -28,7 +31,19 @@ class LeanSignatureResult: source: str signature_hash: str ok: bool + status: str = "FORMALIZED" error: str = "" + attempts: int = 1 + elapsed_s: float = 0.0 + output: str = "" + + +@dataclass(frozen=True) +class _LeanRun: + returncode: int | None + timed_out: bool + elapsed_s: float + output: str def extract_lean_signature_blocks(text: str) -> list[tuple[str, str]]: @@ -43,23 +58,142 @@ def _signature_only(source: str) -> str: return source[:match.start()].strip() if match else source.strip() +def _run_lean( + content: str, + *, + project_root: Path, + timeout_s: float, +) -> _LeanRun: + started = time.monotonic() + with tempfile.NamedTemporaryFile( + mode="w", + suffix=".lean", + encoding="utf-8", + delete=False, + ) as handle: + handle.write(content) + path = Path(handle.name) + process = None + try: + process = subprocess.Popen( + ["lake", "env", "lean", str(path)], + cwd=project_root, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + start_new_session=True, + ) + try: + output, _ = process.communicate(timeout=timeout_s) + return _LeanRun( + process.returncode, + False, + time.monotonic() - started, + output or "", + ) + except subprocess.TimeoutExpired as exc: + partial = ( + exc.stdout.decode(errors="replace") + if isinstance(exc.stdout, bytes) + else (exc.stdout or "") + ) + try: + os.killpg(process.pid, signal.SIGKILL) + except (OSError, ProcessLookupError): + process.kill() + remainder, _ = process.communicate() + return _LeanRun( + None, + True, + time.monotonic() - started, + partial + (remainder or ""), + ) + except OSError as exc: + return _LeanRun( + None, + False, + time.monotonic() - started, + f"{type(exc).__name__}: {exc}", + ) + finally: + if process is not None and process.poll() is None: + process.kill() + process.wait() + path.unlink(missing_ok=True) + + +def warm_lean_environment( + project_root: Path, + *, + timeout_s: float = 120.0, +) -> LeanSignatureResult: + source = "theorem kakeyaLeanWarmup : True := by trivial" + content = ( + "import KakeyaLeanGate\n\n" + "set_option autoImplicit false\n\n" + + source + + "\n" + ) + run = _run_lean( + content, + project_root=project_root, + timeout_s=timeout_s, + ) + if run.timed_out: + return LeanSignatureResult( + source, + "", + False, + status="TYPECHECK_TIMEOUT", + error=f"Lean warmup timed out after {timeout_s:.1f}s", + elapsed_s=run.elapsed_s, + output=run.output, + ) + if run.returncode != 0: + return LeanSignatureResult( + source, + "", + False, + status="ENVIRONMENT_FAILED", + error=f"Lean warmup failed: {run.output[-2000:]}", + elapsed_s=run.elapsed_s, + output=run.output, + ) + return LeanSignatureResult( + source, + "", + True, + status="ENVIRONMENT_READY", + elapsed_s=run.elapsed_s, + output=run.output, + ) + + def validate_lean_signature( source: str, *, project_root: Path, - timeout_s: float = 30.0, + timeout_s: float = 45.0, + retry_timeout_s: float = 120.0, ) -> LeanSignatureResult: source = source.strip() if not source: - return LeanSignatureResult("", "", False, "empty Lean signature") + return LeanSignatureResult( + "", "", False, status="TYPECHECK_FAILED", + error="empty Lean signature", + ) if len(source) > 12_000: - return LeanSignatureResult("", "", False, "Lean signature too large") + return LeanSignatureResult( + "", "", False, status="TYPECHECK_FAILED", + error="Lean signature too large", + ) if _FORBIDDEN.search(source): return LeanSignatureResult( source, "", False, - "forbidden Lean command in generated signature", + status="UNSAFE_REJECTED", + error="forbidden Lean command in generated signature", ) declarations = re.findall(r"^\s*theorem\s+([A-Za-z_][\w']*)", source, re.MULTILINE) if len(declarations) != 1: @@ -67,14 +201,16 @@ def validate_lean_signature( source, "", False, - "expected exactly one theorem declaration", + status="TYPECHECK_FAILED", + error="expected exactly one theorem declaration", ) if not re.search(r"\s*:=\s*by\b", source): return LeanSignatureResult( source, "", False, - "theorem signature must end with `:= by` proof scaffold", + status="TYPECHECK_FAILED", + error="theorem signature must end with `:= by` proof scaffold", ) signature = " ".join(_signature_only(source).split()) signature_hash = hashlib.sha256(signature.encode()).hexdigest() @@ -84,39 +220,72 @@ def validate_lean_signature( + source + "\n" ) - try: - with tempfile.NamedTemporaryFile( - mode="w", - suffix=".lean", - encoding="utf-8", - delete=False, - ) as handle: - handle.write(content) - path = Path(handle.name) - completed = subprocess.run( - ["lake", "env", "lean", str(path)], - cwd=project_root, - capture_output=True, - text=True, - timeout=timeout_s, - check=False, + first = _run_lean( + content, + project_root=project_root, + timeout_s=timeout_s, + ) + attempts = 1 + total_elapsed = first.elapsed_s + output = first.output + run = first + if first.timed_out: + warmup = warm_lean_environment( + project_root, + timeout_s=retry_timeout_s, ) - except (OSError, subprocess.TimeoutExpired) as exc: + total_elapsed += warmup.elapsed_s + output += warmup.output + if not warmup.ok: + return LeanSignatureResult( + source, + signature_hash, + False, + status=warmup.status, + error=warmup.error, + attempts=1, + elapsed_s=total_elapsed, + output=output, + ) + run = _run_lean( + content, + project_root=project_root, + timeout_s=retry_timeout_s, + ) + attempts = 2 + total_elapsed += run.elapsed_s + output += run.output + if run.timed_out: return LeanSignatureResult( source, signature_hash, False, - f"Lean invocation failed: {type(exc).__name__}: {exc}", + status="TYPECHECK_TIMEOUT", + error=( + f"Lean typecheck timed out after {attempts} attempts " + f"({timeout_s:.1f}s/{retry_timeout_s:.1f}s)" + ), + attempts=attempts, + elapsed_s=total_elapsed, + output=output, ) - finally: - if "path" in locals(): - path.unlink(missing_ok=True) - if completed.returncode != 0: - error = (completed.stderr or completed.stdout).strip() + if run.returncode != 0: return LeanSignatureResult( source, signature_hash, False, - f"Lean typecheck failed: {error[-2000:]}", + status="TYPECHECK_FAILED", + error=f"Lean typecheck failed: {run.output[-2000:]}", + attempts=attempts, + elapsed_s=total_elapsed, + output=output, ) - return LeanSignatureResult(source, signature_hash, True) + return LeanSignatureResult( + source, + signature_hash, + True, + status="FORMALIZED", + attempts=attempts, + elapsed_s=total_elapsed, + output=output, + ) diff --git a/autoresearch/prefill/program.md b/autoresearch/prefill/program.md index 321df72..a83da7f 100644 --- a/autoresearch/prefill/program.md +++ b/autoresearch/prefill/program.md @@ -92,6 +92,14 @@ against pinned Lean/mathlib before persistence and records `FORMALIZED` plus a signature hash. Missing, unsafe, ill-typed, or duplicate signatures reject the child. `FORMALIZED` is not `PROVED`: closure still requires a separate proof with no `sorry` and no added axioms. +The supervisor prewarms Lean. Signature checks use a 45-second first attempt; +on timeout the entire Lean process group is killed, the environment is warmed +again, and one 120-second retry is allowed. Distinguish `TYPECHECK_FAILED`, +`TYPECHECK_TIMEOUT`, `UNSAFE_REJECTED`, and `ENVIRONMENT_FAILED`. + +Generator/Critic decode must also make semantic progress. Three consecutive +chunks containing only whitespace or empty decoded text terminate the turn as +`semantic_stall`; never accept an unterminated partial Lean block. Do not optimize output wording, scores, prizes, or other proof-irrelevant content. Prefill performance is a tertiary objective after mathematical diff --git a/autoresearch/prefill/supervisor.py b/autoresearch/prefill/supervisor.py index afd9eb6..b43cc1f 100644 --- a/autoresearch/prefill/supervisor.py +++ b/autoresearch/prefill/supervisor.py @@ -19,6 +19,7 @@ from pathlib import Path from autoresearch.prefill.prepare import _load_candidate, evaluate +from autoresearch.prefill.lean_gate import warm_lean_environment REQUIRED_CANDIDATE_FIELDS = ( @@ -1315,6 +1316,18 @@ def main() -> int: raise SystemExit("strategy-max-prefill-tokens must be > 0") if args.strategy_stagnation_rounds <= 0: raise SystemExit("strategy-stagnation-rounds must be > 0") + lean_warmup = warm_lean_environment( + Path(__file__).resolve().parents[2], + ) + print( + "[autoresearch] phase=lean-warmup " + f"status={lean_warmup.status} " + f"elapsed_s={lean_warmup.elapsed_s:.2f} " + f"error={lean_warmup.error or '(none)'}", + flush=True, + ) + if not lean_warmup.ok: + raise SystemExit(lean_warmup.error) for iteration in range(args.iterations): row = run_iteration(args, iteration) print(json.dumps(row, indent=2, sort_keys=True)) diff --git a/scripts/agent_gan_inference_demo.py b/scripts/agent_gan_inference_demo.py index d370223..7f40d83 100644 --- a/scripts/agent_gan_inference_demo.py +++ b/scripts/agent_gan_inference_demo.py @@ -64,7 +64,11 @@ def _infer( get_stats, on_token=None, max_response_tokens=None, + semantic_progress=None, + max_semantic_stall_chunks: int = 3, ): + if max_semantic_stall_chunks <= 0: + raise ValueError("max_semantic_stall_chunks must be > 0") before = get_stats() started = time.perf_counter() with client.create_session(eos_token_ids=eos_ids, client_label="agent-gan") as s: @@ -79,6 +83,7 @@ def _infer( else int(max_response_tokens) or None ) stop_reason = "unknown" + stalled_chunks = 0 while response_limit is None or len(generated) < response_limit: before_count = len(generated) chunk = ( @@ -92,12 +97,21 @@ def _infer( on_token(generated) if first_at is None: first_at = time.perf_counter() + new_tokens = generated[before_count:] + if semantic_progress is not None and new_tokens: + if semantic_progress(new_tokens): + stalled_chunks = 0 + else: + stalled_chunks += 1 stop_reason = { 1: "max_tokens", 2: "eos", 3: "cancelled", 4: "truncated", }.get(s.last_stop_reason, "unknown") + if stalled_chunks >= max_semantic_stall_chunks: + stop_reason = "semantic_stall" + break if stop_reason != "max_tokens": break if len(generated) == before_count: diff --git a/scripts/agent_gan_repl.py b/scripts/agent_gan_repl.py index 03cf74f..9597cc1 100644 --- a/scripts/agent_gan_repl.py +++ b/scripts/agent_gan_repl.py @@ -1674,6 +1674,12 @@ def get_stats(): get_stats, on_token=generator_printer, max_response_tokens=args.max_response_tokens, + semantic_progress=lambda chunk: bool( + tokenizer.decode( + chunk, + skip_special_tokens=True, + ).strip() + ), ) generator_printer.finish() generator_text = tokenizer.decode( @@ -1768,6 +1774,12 @@ def get_stats(): get_stats, on_token=critic_printer, max_response_tokens=args.max_response_tokens, + semantic_progress=lambda chunk: bool( + tokenizer.decode( + chunk, + skip_special_tokens=True, + ).strip() + ), ) critic_printer.finish() critic_text = tokenizer.decode( @@ -1810,8 +1822,10 @@ def get_stats(): print( "[lean-signature-gate] " f"target={lean_target} " - f"status={'FORMALIZED' if lean_result.ok else 'REJECTED'} " + f"status={lean_result.status} " f"hash={lean_result.signature_hash or '(none)'} " + f"attempts={lean_result.attempts} " + f"elapsed_s={lean_result.elapsed_s:.2f} " f"error={lean_result.error or '(none)'}", flush=True, ) diff --git a/tests/inference_engine/bench/test_lean_signature_gate.py b/tests/inference_engine/bench/test_lean_signature_gate.py index 23f9ae6..ab3566e 100644 --- a/tests/inference_engine/bench/test_lean_signature_gate.py +++ b/tests/inference_engine/bench/test_lean_signature_gate.py @@ -1,5 +1,6 @@ from pathlib import Path +import autoresearch.prefill.lean_gate as lean_gate from autoresearch.prefill.lean_gate import ( extract_lean_signature_blocks, validate_lean_signature, @@ -26,6 +27,7 @@ def test_extract_and_typecheck_minimal_mathlib_signature(): assert target == "RH-C2-leaf" result = validate_lean_signature(source, project_root=ROOT) assert result.ok, result.error + assert result.status == "FORMALIZED" assert len(result.signature_hash) == 64 @@ -35,6 +37,7 @@ def test_lean_gate_rejects_unknown_type(): project_root=ROOT, ) assert not result.ok + assert result.status == "TYPECHECK_FAILED" assert "unknown" in result.error.lower() @@ -47,3 +50,39 @@ def test_lean_gate_rejects_executable_or_axiomatic_commands(): ): result = validate_lean_signature(source, project_root=ROOT) assert not result.ok + assert result.status == "UNSAFE_REJECTED" + + +def test_lean_gate_retries_timeout_after_warmup(monkeypatch): + runs = iter(( + lean_gate._LeanRun(None, True, 45.0, "first partial"), + lean_gate._LeanRun(0, False, 3.0, "warm"), + lean_gate._LeanRun(0, False, 4.0, "retry"), + )) + monkeypatch.setattr(lean_gate, "_run_lean", lambda *args, **kwargs: next(runs)) + result = validate_lean_signature( + "theorem retried : True := by trivial", + project_root=ROOT, + ) + assert result.ok + assert result.status == "FORMALIZED" + assert result.attempts == 2 + assert result.elapsed_s == 52.0 + assert "first partial" in result.output + + +def test_lean_gate_classifies_second_timeout(monkeypatch): + runs = iter(( + lean_gate._LeanRun(None, True, 45.0, "first"), + lean_gate._LeanRun(0, False, 2.0, "warm"), + lean_gate._LeanRun(None, True, 120.0, "second"), + )) + monkeypatch.setattr(lean_gate, "_run_lean", lambda *args, **kwargs: next(runs)) + result = validate_lean_signature( + "theorem timeout : True := by trivial", + project_root=ROOT, + ) + assert not result.ok + assert result.status == "TYPECHECK_TIMEOUT" + assert result.attempts == 2 + assert "firstwarmsecond" in result.output diff --git a/tests/inference_engine/bridge/test_agent_gan_demo.py b/tests/inference_engine/bridge/test_agent_gan_demo.py index dc1fa58..dcbd165 100644 --- a/tests/inference_engine/bridge/test_agent_gan_demo.py +++ b/tests/inference_engine/bridge/test_agent_gan_demo.py @@ -105,6 +105,31 @@ def test_infer_reports_explicit_client_safety_limit(): assert metrics["complete"] is False +def test_infer_stops_repeated_nonsemantic_chunks(): + session = Session([ + ([32, 10], 1), + ([10, 32], 1), + ([32, 32], 1), + ([65], 2), + ]) + tokens, metrics = _infer( + Client(session), + [], + [9], + 2, + lambda: {}, + max_response_tokens=0, + semantic_progress=lambda chunk: bool( + "".join(chr(token) for token in chunk).strip() + ), + max_semantic_stall_chunks=3, + ) + assert tokens == [32, 10, 10, 32, 32, 32] + assert metrics["stop_reason"] == "semantic_stall" + assert metrics["complete"] is False + assert session.calls == 3 + + class CharTokenizer: def encode(self, text, **_kwargs): return [ord(char) for char in text]