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