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_override import forced_mate
8
9PIN = 'e396948acbd80cdcbad3dccfb07e878fba9df8f4'
10ROOT = Path(__file__).resolve().parent
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',choose(model,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 cap*3+1 > model.context_length:
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            before=time.monotonic(); metrics={}
163            if board.turn==color:
164                tactical=forced_mate(board,adjudicate=outcome,max_nodes=mate_nodes,
165                    seconds=min(mate_seconds,remaining) if remaining is not None else mate_seconds) if mate_override else {'status':'disabled'}
166                if tactical['status']=='proven':
167                    move=chess.Move.from_uci(tactical['move'])
168                    metrics={'decision_source':'forced_mate_override','model_query':'not_queried'}
169                else:
170                    remaining=None if walltime is None else walltime-(time.monotonic()-started)
171                    if remaining is not None and remaining <= 0: raise TimeoutError('game walltime limit')
172                    if isinstance(model,ModelWorker) and remaining is not None:
173                        model.timeout=min(model.timeout,remaining)
174                    move,metrics=model.choose(board) if isinstance(model,ModelWorker) else choose(model,board)
175                    metrics.update(decision_source='model_policy',model_query='queried')
176                metrics['mate_search']=tactical
177            else:
178                if engine is None: engine=engine_factory()
179                if remaining is not None: engine.timeout=min(engine.timeout,max(.001,walltime-(time.monotonic()-started)))
180                move=engine.move(board,nodes)
181            if move not in board.legal_moves: raise ValueError('illegal player output: '+move.uci())
182            board.push(move); node=node.add_variation(move)
183            if metrics.get('decision_source')=='forced_mate_override':
184                node.comment=(f"Forced-mate override: mate within {tactical['mate_plies']} plies; "
185                    f"limit 3 plies; {tactical['nodes']} nodes; {tactical['seconds']:.6f}s; "
186                    f"model not queried; full all-defenses proof in game.json at ply {board.ply()}.")
187            rows.append({'ply':board.ply(),'uci':move.uci(),'seconds':time.monotonic()-before,**metrics})
188    except Exception as exc:
189        reason = 'timeout_unfinished' if isinstance(exc,TimeoutError) else 'error_unfinished'
190        error=repr(exc)
191    finally:
192        cleanup_error = None
193        if engine is not None:
194            try: engine.close()
195            except Exception as exc: cleanup_error = repr(exc)
196        mode='model+mate-in-3-ply-override' if mate_override else 'pure-model-policy'
197        record.headers.update(Event=label+'; '+mode, White='Model' if color else 'Sunfish',
198            Black='Sunfish' if color else 'Model',Result=result,Termination=reason)
199        record.headers['ModelPolicy']=mode
200        (out/'game.pgn').write_text(str(record)+'\n')
201        report={'result':result,'termination':reason,'error':error,'cleanup_error':cleanup_error,'plies':board.ply(),
202            'final_fen':board.fen(),'seconds':time.monotonic()-started,'moves':rows,
203            'model_policy':mode,'mate_override':mate_override,'mate_max_plies':3,
204            'mate_node_budget':mate_nodes,'mate_time_budget_seconds':mate_seconds}
205        (out/'game.json').write_text(json.dumps(report,indent=2))
206        if engine is not None:
207            (out/'uci.log').write_text('\n'.join(engine.transcript))
208    return report
209
210def engine_command():
211    return [sys.executable,'-u',str(ROOT/'sunfish'/'sunfish.py')]
212
213def main():
214    p=argparse.ArgumentParser()
215    p.add_argument('--adapter',required=True,help='module:load; see README')
216    p.add_argument('--checkpoint',type=Path,required=True)
217    p.add_argument('--dataset-manifest',type=Path,required=True)
218    p.add_argument('--output',type=Path,required=True)
219    p.add_argument('--verified-checkpoint-sha256',required=True)
220    p.add_argument('--authorize-matches',action='store_true')
221    p.add_argument('--cap',type=int,default=16); p.add_argument('--nodes',type=int,default=1000)
222    p.add_argument('--timeout',type=float,default=10); p.add_argument('--seed',type=int,default=0)
223    p.add_argument('--model-timeout',type=float,default=60)
224    p.add_argument('--opening',nargs='*',default=[])
225    p.add_argument('--mate-override',action=argparse.BooleanOptionalAction,default=True)
226    p.add_argument('--mate-nodes',type=int,default=50000)
227    p.add_argument('--mate-seconds',type=float,default=1.0)
228    args=p.parse_args()
229    if not args.authorize_matches: p.error('matches require explicit authorization flag')
230    if args.cap<1 or args.nodes<1 or args.timeout<=0 or args.model_timeout<=0: p.error('positive budgets required')
231    if args.mate_nodes<1 or not math.isfinite(args.mate_seconds) or args.mate_seconds<=0: p.error('positive finite tactical budgets required')
232    if sha(args.checkpoint)!=args.verified_checkpoint_sha256: p.error('checkpoint hash mismatch')
233    commit=subprocess.check_output(['git','-c','core.autocrlf=true','-c','core.filemode=false','-C',str(ROOT/'sunfish'),'rev-parse','HEAD'],text=True).strip()
234    dirty=subprocess.check_output(['git','-c','core.autocrlf=true','-c','core.filemode=false','-C',str(ROOT/'sunfish'),'status','--porcelain'],text=True)
235    if commit!=PIN or dirty: p.error('Sunfish pin mismatch or modified checkout')
236    random.seed(args.seed)
237    model=ModelWorker(args.adapter,args.checkpoint,args.seed,args.model_timeout)
238    try:
239        args.output.mkdir(parents=True,exist_ok=False)
240        metadata={**{k:str(v) if isinstance(v,Path) else v for k,v in vars(args).items()},
241            'checkpoint_sha256':sha(args.checkpoint),'dataset_manifest_sha256':sha(args.dataset_manifest),
242            'tokenizer_sha256':sha(tokenizer.__file__),'tokenizer_contract':tokenizer.VERSION,
243            'harness_sha256':sha(__file__),'sunfish_commit':commit,'sunfish_license':'GPL-3.0',
244            'mate_search_sha256':sha(ROOT/'mate_override.py'),'mate_max_plies':3,
245            'runtime':sys.version,'executable':sys.executable,'platform':platform.platform(),
246            'python_chess':chess.__version__,'context_length':model.context_length,
247            'draw_policy':'python-chess outcome(claim_draw=True), immediate symmetric claims',
248            'decoding':'global legal full-move joint argmax; unmasked 128-way conditionals',
249            'model_policy':'model+mate-in-3-ply-override' if args.mate_override else 'pure-model-policy'}
250        (args.output/'manifest.json').write_text(json.dumps(metadata,indent=2))
251        for color,name in [(chess.WHITE,'model-white'),(chess.BLACK,'model-black')]:
252            game(model,lambda:UCI(engine_command(),args.timeout),color,args.output/name,args.cap,args.nodes,args.opening,
253                mate_override=args.mate_override,mate_nodes=args.mate_nodes,mate_seconds=args.mate_seconds)
254    finally: model.close()
255
256if __name__=='__main__': main()
257