| """Tiny Lean runner for the 03xx notebooks. Goes ONLY through align.verify.lean. |
| |
| Rules (house + task): <= 2 workers (we use 1), nice 15, MemAvailable >= 45 GB checked |
| before the pool starts, never the heavy-job lock, never touches a GPU. |
| """ |
| from __future__ import annotations |
| import os, sys, re |
| os.environ["CUDA_VISIBLE_DEVICES"] = "" |
| sys.path.insert(0, "/apps/development/rlhf/align") |
| from align.verify.lean import ReplPool, ReplLimits, find_toolchain |
| from align.verify.lean.repl_pool import meminfo_available_gb |
|
|
| HEADER = "import Mathlib\nimport Aesop\nset_option maxHeartbeats 400000\nopen Nat\n" |
|
|
| class Lean: |
| def __init__(self, header: str = HEADER): |
| avail = meminfo_available_gb() |
| if avail < 45.0: |
| raise RuntimeError(f"MemAvailable {avail:.1f} GB < 45 GB: refusing to start Lean") |
| self.tc = find_toolchain() |
| self.pool = ReplPool.from_toolchain(self.tc, n_workers=1, limits=ReplLimits(nice=15)) |
| self.header = header |
| self._cm = self.pool.acquire() |
| self.w = self._cm.__enter__() |
| self.env = self.w.header_env(header) |
|
|
| def run(self, code: str, timeout: float = 120.0, keep: bool = False): |
| """Run `code` on top of the header; return (messages, new_env).""" |
| resp = self.w.command(code, self.env, timeout) |
| if keep and "env" in resp: |
| self.env = resp["env"] |
| return resp |
|
|
| def show(self, code: str, timeout: float = 120.0, keep: bool = False) -> dict: |
| r = self.run(code, timeout, keep) |
| for m in r.get("messages", []): |
| print(f"[{m['severity']}] line {m['pos']['line']}: {m['data']}") |
| for s in r.get("sorries", []): |
| print("[sorries]", s.get("goal", "")[:200]) |
| if not r.get("messages") and not r.get("sorries"): |
| print("[ok] no messages") |
| return r |
|
|
| def close(self): |
| try: self._cm.__exit__(None, None, None) |
| finally: self.pool.close() |
|
|
