harness.py

1"""Standard-start, full-history matches. No trained matches run by default."""
2import argparse, hashlib, importlib, json, math, os, platform, queue, random
3import signal, subprocess, sys, threading, time, multiprocessing
4from pathlib import Path
5import chess, chess.pgn
6from chessdata import tokenizer
7from mate_backend import forced_mate
8
9PIN = 'e396948acbd80cdcbad3dccfb07e878fba9df8f4'
10ROOT = Path('/REDACTED')
11
12def sha(path):
13    h = hashlib.sha256()
14    with open(path, 'rb') as f:
15        for block in iter(lambda: f.read(1048576), b''): h.update(block)
16    return h.hexdigest()
17
18def history(board):
19    if board.root().fen() != chess.STARTING_FEN or board.chess960:
20        raise ValueError('only full standard-start history supported')
21    return [tokenizer.BOS] + [t for m in board.move_stack for t in tokenizer.encode_move(m)]
22
23def logprobs(model, prefix):
24    if len(prefix) > model.context_length: raise ValueError('context overflow; never truncate')
25    values = list(map(float, model.next_logits(list(prefix))))
26    if len(values) != 128 or not all(math.isfinite(v) for v in values):
27        raise ValueError('expected 128 finite raw logits')
28    top = max(values)
29    shifted = [v-top for v in values]
30    norm = math.log(sum(math.exp(v) for v in shifted))
31    return [v-norm for v in shifted]
32
33def choose(model, board):
34    prefix = history(board)
35    legal = sorted(board.legal_moves, key=lambda m: m.uci())
36    if not legal: raise ValueError('terminal board')
37    cache = {}
38    def lp(tail):
39        key = tuple(tail)
40        if key not in cache: cache[key] = logprobs(model, prefix + list(tail))
41        return cache[key]
42    scores = []
43    for move in legal:
44        a,b,c = tokenizer.encode_move(move)
45        scores.append(lp([])[a] + lp([a])[b] + lp([a,b])[c])
46    best = max(range(len(legal)), key=lambda i: scores[i])
47    mass = sum(math.exp(v) for v in scores)
48    raw = []
49    for _ in range(3):
50        values = lp(raw)
51        raw.append(max(range(128), key=lambda i: values[i]))
52    triples = [tokenizer.encode_move(m) for m in legal]
53    return legal[best], {'legal_joint_probability': mass,
54        'illegal_or_nonmove_three_token_probability': max(0., 1.-mass),
55        'raw_greedy_tokens': raw, 'raw_greedy_legal': raw in triples,
56        'selected_joint_log_probability': scores[best], 'forward_calls': len(cache)}
57
58class UCI:
59    def __init__(self, command, timeout=10):
60        self.timeout, self.lines, self.transcript = timeout, queue.Queue(), []
61        self.p = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
62            stderr=subprocess.STDOUT, text=True, bufsize=1, start_new_session=True)
63        def reader():
64            for line in self.p.stdout: self.lines.put(line.strip())
65            self.lines.put(None)
66        self.reader = threading.Thread(target=reader, daemon=True)
67        self.reader.start()
68        try:
69            self.send('uci'); self.until('uciok')
70            self.send('isready'); self.until('readyok')
71        except BaseException:
72            self.close(); raise
73    def send(self, line):
74        self.transcript.append('> '+line)
75        self.p.stdin.write(line+'\n'); self.p.stdin.flush()
76    def until(self, word):
77        deadline = time.monotonic()+self.timeout
78        while True:
79            remaining = deadline-time.monotonic()
80            if remaining <= 0: raise TimeoutError('UCI '+word)
81            try: line = self.lines.get(timeout=remaining)
82            except queue.Empty: raise TimeoutError('UCI '+word)
83            if time.monotonic() >= deadline: raise TimeoutError('UCI '+word)
84            if line is None: raise RuntimeError('engine crashed / EOF')
85            self.transcript.append('< '+line)
86            if line.split()[:1] == [word]: return line
87    def move(self, board, nodes):
88        self.send('position startpos moves '+' '.join(m.uci() for m in board.move_stack))
89        self.send('go nodes '+str(nodes))
90        line = self.until('bestmove')
91        if any('aborted' in s for s in self.transcript[-3:]): raise RuntimeError('aborted search')
92        return chess.Move.from_uci(line.split()[1])
93    def close(self):
94        if self.p.poll() is None:
95            if os.name == 'posix': os.killpg(self.p.pid, signal.SIGKILL)
96            else: self.p.kill()
97        self.p.wait(timeout=2)
98        self.reader.join(timeout=2)
99        self.p.stdin.close(); self.p.stdout.close()
100
101def model_worker(conn, adapter, checkpoint, seed):
102    try:
103        mod,fn=adapter.split(':')
104        model=getattr(importlib.import_module(mod),fn)(Path(checkpoint),seed=seed)
105        conn.send(('ok',model.context_length))
106        while True:
107            board=conn.recv()
108            conn.send(('ok',model.choose(board)))
109    except Exception as exc:
110        conn.send(('error',repr(exc)))
111    finally: conn.close()
112
113class ModelWorker:
114    """Spawned process makes model load/decoding enforceably bounded."""
115    def __init__(self,adapter,checkpoint,seed,timeout):
116        self.timeout=timeout
117        ctx=multiprocessing.get_context('spawn')
118        self.conn,child=ctx.Pipe()
119        self.p=ctx.Process(target=model_worker,args=(child,adapter,str(checkpoint),seed),daemon=True)
120        self.p.start(); child.close()
121        try: self.context_length=self.receive()
122        except BaseException: self.close(); raise
123    def receive(self):
124        if not self.conn.poll(self.timeout):
125            self.close(); raise TimeoutError('model load/decoding timeout')
126        status,value=self.conn.recv()
127        if status!='ok': raise RuntimeError(value)
128        return value
129    def choose(self,board):
130        self.conn.send(board)
131        return self.receive()
132    def close(self):
133        if self.p.is_alive(): self.p.kill()
134        self.p.join(timeout=2); self.conn.close()
135
136def outcome(board):
137    # Symmetric immediate claims, including a claim available by declaring a move.
138    result = board.outcome(claim_draw=True)
139    return (result.result(), result.termination.name.lower()) if result else None
140
141def game(model, engine_factory, color, out, cap=16, nodes=1000, opening=(), walltime=None, label='Harness smoke; no Elo claim',
142         mate_override=True, mate_nodes=50000, mate_seconds=1.0):
143    out.mkdir(parents=True, exist_ok=False)
144    board, record, rows, engine = chess.Board(), chess.pgn.Game(), [], None
145    node=record
146    result, reason, error = '*', 'ply_cap_unfinished', None
147    started=time.monotonic()
148    try:
149        if cap < len(opening) or model.context_length != 122:
150            raise ValueError('cap/opening exceeds full-history capacity')
151        for uci in opening:
152            if outcome(board): raise ValueError('opening continues after adjudicated end')
153            move=chess.Move.from_uci(uci)
154            if move not in board.legal_moves: raise ValueError('illegal opening')
155            board.push(move); node=node.add_variation(move)
156        while True:
157            remaining = None if walltime is None else walltime-(time.monotonic()-started)
158            if remaining is not None and remaining <= 0: raise TimeoutError('game walltime limit')
159            end=outcome(board)
160            if end: result,reason=end; break
161            if board.ply()>=cap: break
162            if board.turn==color:
163                from resource_guard import check
164                time.sleep(.25);check()
165                remaining=None if walltime is None else walltime-(time.monotonic()-started)
166                if remaining is not None and remaining<=0:raise TimeoutError('game deadline during throttling')
167            before=time.monotonic(); metrics={}
168            if board.turn==color:
169                tactical=forced_mate(board,adjudicate=outcome,max_nodes=mate_nodes,
170                    seconds=min(mate_seconds,remaining) if remaining is not None else mate_seconds) if mate_override else {'status':'disabled'}
171                if tactical['status']=='proven':
172                    move=chess.Move.from_uci(tactical['move'])
173                    metrics={'decision_source':'forced_mate_override','model_query':'not_queried'}
174                else:
175                    remaining=None if walltime is None else walltime-(time.monotonic()-started)
176                    if remaining is not None and remaining <= 0: raise TimeoutError('game walltime limit')
177                    if isinstance(model,ModelWorker) and remaining is not None:
178                        model.timeout=min(model.timeout,remaining)
179                    move,metrics=model.choose(board) if isinstance(model,ModelWorker) else choose(model,board)
180                    metrics.update(decision_source='model_policy',model_query='queried')
181                metrics['mate_search']=tactical
182            else:
183                if engine is None: engine=engine_factory()
184                if remaining is not None: engine.timeout=min(engine.timeout,max(.001,walltime-(time.monotonic()-started)))
185                move=engine.move(board,nodes)
186            if move not in board.legal_moves: raise ValueError('illegal player output: '+move.uci())
187            board.push(move); node=node.add_variation(move)
188            if metrics.get('decision_source')=='forced_mate_override':
189                node.comment=(f"Forced-mate override: mate within {tactical['mate_plies']} plies; "
190                    f"limit 7 plies; {tactical['nodes']} nodes; {tactical['seconds']:.6f}s; "
191                    f"model not queried; full all-defenses proof in game.json at ply {board.ply()}.")
192            rows.append({'ply':board.ply(),'uci':move.uci(),'seconds':time.monotonic()-before,**metrics})
193    except Exception as exc:
194        reason = 'timeout_unfinished' if isinstance(exc,TimeoutError) else 'error_unfinished'
195        error=repr(exc)
196    finally:
197        cleanup_error = None
198        if engine is not None:
199            try: engine.close()
200            except Exception as exc: cleanup_error = repr(exc)
201        mode='state-v2+rust-iterative-7-ply-override' if mate_override else 'pure-model-policy'
202        record.headers.update(Event=label+'; '+mode, White='Model' if color else 'Sunfish',
203            Black='Sunfish' if color else 'Model',Result=result,Termination=reason)
204        record.headers['ModelPolicy']=mode
205        (out/'game.pgn').write_text(str(record)+'\n')
206        report={'result':result,'termination':reason,'error':error,'cleanup_error':cleanup_error,'plies':board.ply(),
207            'final_fen':board.fen(),'seconds':time.monotonic()-started,'moves':rows,
208            'model_policy':mode,'mate_override':mate_override,'mate_max_plies':7,
209            'mate_node_budget':mate_nodes,'mate_time_budget_seconds':mate_seconds}
210        (out/'game.json').write_text(json.dumps(report,indent=2))
211        if engine is not None:
212            (out/'uci.log').write_text('\n'.join(engine.transcript))
213    return report
214
215def engine_command():
216    return [sys.executable,'-u',str(ROOT/'sunfish'/'sunfish.py')]
217
218def main():
219    p=argparse.ArgumentParser()
220    p.add_argument('--adapter',required=True,help='module:load; see README')
221    p.add_argument('--checkpoint',type=Path,required=True)
222    p.add_argument('--dataset-manifest',type=Path,required=True)
223    p.add_argument('--output',type=Path,required=True)
224    p.add_argument('--verified-checkpoint-sha256',required=True)
225    p.add_argument('--authorize-matches',action='store_true')
226    p.add_argument('--cap',type=int,default=16); p.add_argument('--nodes',type=int,default=1000)
227    p.add_argument('--timeout',type=float,default=10); p.add_argument('--seed',type=int,default=0)
228    p.add_argument('--model-timeout',type=float,default=60)
229    p.add_argument('--opening',nargs='*',default=[])
230    p.add_argument('--mate-override',action=argparse.BooleanOptionalAction,default=True)
231    p.add_argument('--mate-nodes',type=int,default=50000)
232    p.add_argument('--mate-seconds',type=float,default=1.0)
233    args=p.parse_args()
234    if not args.authorize_matches: p.error('matches require explicit authorization flag')
235    if args.cap<1 or args.nodes<1 or args.timeout<=0 or args.model_timeout<=0: p.error('positive budgets required')
236    if args.mate_nodes<1 or not math.isfinite(args.mate_seconds) or args.mate_seconds<=0: p.error('positive finite tactical budgets required')
237    if sha(args.checkpoint)!=args.verified_checkpoint_sha256: p.error('checkpoint hash mismatch')
238    commit=subprocess.check_output(['git','-c','core.autocrlf=true','-c','core.filemode=false','-C',str(ROOT/'sunfish'),'rev-parse','HEAD'],text=True).strip()
239    dirty=subprocess.check_output(['git','-c','core.autocrlf=true','-c','core.filemode=false','-C',str(ROOT/'sunfish'),'status','--porcelain'],text=True)
240    if commit!=PIN or dirty: p.error('Sunfish pin mismatch or modified checkout')
241    random.seed(args.seed)
242    model=ModelWorker(args.adapter,args.checkpoint,args.seed,args.model_timeout)
243    try:
244        args.output.mkdir(parents=True,exist_ok=False)
245        metadata={**{k:str(v) if isinstance(v,Path) else v for k,v in vars(args).items()},
246            'checkpoint_sha256':sha(args.checkpoint),'dataset_manifest_sha256':sha(args.dataset_manifest),
247            'tokenizer_sha256':sha(tokenizer.__file__),'tokenizer_contract':tokenizer.VERSION,
248            'harness_sha256':sha(__file__),'sunfish_commit':commit,'sunfish_license':'GPL-3.0',
249            'mate_search_sha256':sha(ROOT/'mate_override.py'),'mate_max_plies':7,
250            'runtime':sys.version,'executable':sys.executable,'platform':platform.platform(),
251            'python_chess':chess.__version__,'context_length':model.context_length,
252            'draw_policy':'python-chess outcome(claim_draw=True), immediate symmetric claims',
253            'decoding':'global legal full-move joint argmax; unmasked 128-way conditionals',
254            'model_policy':'state-v2+rust-iterative-7-ply-override' if args.mate_override else 'pure-model-policy'}
255        (args.output/'manifest.json').write_text(json.dumps(metadata,indent=2))
256        for color,name in [(chess.WHITE,'model-white'),(chess.BLACK,'model-black')]:
257            game(model,lambda:UCI(engine_command(),args.timeout),color,args.output/name,args.cap,args.nodes,args.opening,
258                mate_override=args.mate_override,mate_nodes=args.mate_nodes,mate_seconds=args.mate_seconds)
259    finally: model.close()
260
261if __name__=='__main__': main()
262