File size: 1,953 Bytes
bcc8d35 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | """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()
|