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