inference.py

1"""Single keyed inference owner. Explicit CPU/GPU device; no concurrent Pipe consumers."""
2import multiprocessing as mp,queue,threading,time,uuid,os
3
4def serving_device():
5    device=os.environ.get('LITTLEPLY_DEVICE','cpu')
6    if device not in ('cpu','cuda:0'):raise ValueError('unsupported serving device')
7    return device
8
9def child(conn,checkpoint):
10    import torch
11    import training_core as tc
12    import harness as h
13    from scorer import Scorer
14    from mate_backend import forced_mate
15    try:
16        torch.set_num_threads(4);torch.manual_seed(20261028);torch.use_deterministic_algorithms(True)
17        device=serving_device()
18        net,meta=tc.load_model_checkpoint(checkpoint,device)
19        from chessdata import tokenizer
20        from pathlib import Path
21        import hashlib
22        if hashlib.sha256(Path(tokenizer.__file__).read_bytes()).hexdigest()!=meta['identity']['tokenizer']['sha256']:
23            raise ValueError('checkpoint tokenizer identity mismatch')
24        model=Scorer(net)
25        model.score(h.chess.Board());conn.send(('ready',{}))
26        while True:
27            rid,board,seconds=conn.recv();start=time.monotonic()
28            proof=forced_mate(board,adjudicate=h.outcome,max_nodes=50000,seconds=min(1.,seconds*.2))
29            if proof['status']=='proven':
30                move=h.chess.Move.from_uci(proof['move']);metrics={'decision_source':'forced_mate_override','model_query':'not_queried'}
31            else:
32                move,metrics=model.score(board);metrics.pop('full_legal_scores',None)
33                metrics.update(decision_source='model_policy',model_query='queried')
34            metrics.update(mate_search=proof,decision_seconds=time.monotonic()-start,inference_device=device)
35            conn.send((rid,move.uci(),metrics))
36    except EOFError:pass
37    except Exception as exc:
38        try:conn.send(('error',type(exc).__name__))
39        except Exception:pass
40    finally:conn.close()
41
42class Owner:
43    def __init__(self,checkpoint):
44        self.checkpoint=str(checkpoint);self.jobs=queue.Queue(1);self.results=queue.Queue(2)
45        self.stop=threading.Event();self.ready=False;self.busy=False;self.p=None;self.conn=None;self.poll_floor=0.
46        self.thread=threading.Thread(target=self.run,daemon=True);self.thread.start()
47    def submit(self,key,board,seconds):
48        if self.busy or not self.ready:raise RuntimeError('inference unavailable')
49        self.busy=True;self.jobs.put_nowait((key,board,seconds))
50    def reset(self):
51        self.ready=False
52        if self.p:
53            if self.p.is_alive():self.p.kill()
54            self.p.join(timeout=3)
55        if self.conn:self.conn.close()
56        self.p=None;self.conn=None
57    def run(self):
58        failures=0
59        try:
60            while not self.stop.is_set():
61                if not self.ready:
62                    try:
63                        ctx=mp.get_context('spawn');parent,peer=ctx.Pipe();self.conn=parent
64                        self.p=ctx.Process(target=child,args=(peer,self.checkpoint),daemon=True);self.p.start();peer.close()
65                        if not parent.poll(60) or parent.recv()[0]!='ready':raise RuntimeError('warmup failed')
66                        self.ready=True;failures=0
67                    except Exception:
68                        self.reset();failures+=1
69                        if self.stop.wait(min(60,5*failures)):return
70                        continue
71                try:key,board,seconds=self.jobs.get(timeout=.2)
72                except queue.Empty:continue
73                rid=uuid.uuid4().hex
74                try:
75                    self.conn.send((rid,board,seconds))
76                    if not self.conn.poll(max(seconds,self.poll_floor)):raise TimeoutError('model deadline')
77                    answer=self.conn.recv()
78                    if answer[0]!=rid:raise RuntimeError('wrong inference response identity')
79                    self.results.put((key,'ok',answer[1:]))
80                except Exception as exc:
81                    self.results.put((key,'error',type(exc).__name__));self.reset()
82                finally:self.busy=False
83        finally:self.reset()
84    def close(self):
85        self.stop.set()
86        if self.p and self.p.is_alive():self.p.kill()
87        self.thread.join(timeout=5)
88