BAAI/Brainmu-Spike

1"""Persistent process per GPU; original global sample indices and seeds are preserved."""2import multiprocessing as mp3from multiprocessing.connection import wait4import os,time,traceback,uuid5from pathlib import Path6 7def shard_indices(total,workers):8 return [list(range(rank,total,workers)) for rank in range(workers)]9 10def worker(conn,gpu,config,log_path):11 os.environ['CUDA_VISIBLE_DEVICES']=str(gpu)12 os.environ['OMP_NUM_THREADS']='2'13 with open(log_path,'a',buffering=1) as log:14  os.dup2(log.fileno(),1);os.dup2(log.fileno(),2)15  try:16   from engine import Engine17   engine=Engine();engine.load(**config)18   conn.send({'type':'ready','gpu':gpu})19   while True:20    item=conn.recv()21    if item['type']=='close':break22    dat,gt,index,out,prompt=item['args']23    row=engine.predict(Path(dat),Path(gt),index,Path(out),prompt)24    row['gpu_id']=gpu25    conn.send({'type':'result','row':row})26  except EOFError:pass27  except BaseException as e:28   traceback.print_exc()29   try:conn.send({'type':'error','error':f'GPU {gpu}: {e}'})30   except Exception:pass31  finally:conn.close()32 33class MultiGPUEngine:34 def __init__(self,gpu_ids):35  if not gpu_ids or len(set(gpu_ids))!=len(gpu_ids):raise ValueError('GPU IDs must be nonempty and unique')36  self.gpu_ids=list(gpu_ids);self.procs=[];self.conns=[]37 def load(self,**config):38  ctx=mp.get_context('spawn')39  logs=Path(__file__).resolve().parents[1]/'logs'40  logs.mkdir(parents=True,exist_ok=True)41  tag=uuid.uuid4().hex[:8]42  try:43   for gpu in self.gpu_ids:44    parent,child=ctx.Pipe()45    p=ctx.Process(target=worker,args=(child,gpu,config,str(logs/f'ui_gpu{gpu}_{tag}.log')),daemon=True)46    p.start();child.close();self.procs.append(p);self.conns.append(parent)47   pending=set(self.conns);deadline=time.monotonic()+120048   while pending:49    if time.monotonic()>deadline:raise TimeoutError('GPU model loading exceeded 20 minutes')50    for c in wait(pending,timeout=1):51     msg=c.recv()52     if msg.get('type')!='ready':raise RuntimeError(msg.get('error','Invalid model load response'))53     pending.remove(c)54    self._check()55  except BaseException:self.close();raise56 def _check(self):57  for gpu,p in zip(self.gpu_ids,self.procs):58   if not p.is_alive():raise RuntimeError(f'GPU worker {gpu} exited ({p.exitcode})')59 def predict_many(self,pairs,out,prompt,stop_event):60  if len(self.conns)!=len(self.gpu_ids):raise RuntimeError('GPU pool is not loaded')61  queues=[iter(x) for x in shard_indices(len(pairs),len(self.gpu_ids))]62  active={};seen=set()63  def dispatch(rank):64   if stop_event.is_set():return65   i=next(queues[rank],None)66   if i is None:return67   dat,gt=pairs[i];c=self.conns[rank]68   c.send({'type':'predict','args':(str(dat),str(gt),i,str(out),prompt)})69   active[c]=(rank,i)70  try:71   for rank in range(len(self.conns)):dispatch(rank)72   while active:73    self._check()74    for c in wait(list(active),timeout=1):75     rank,expected=active.pop(c);msg=c.recv()76     if msg.get('type')!='result':raise RuntimeError(msg.get('error','Invalid worker response'))77     row=msg['row']78     if row['index']!=expected or expected in seen:raise RuntimeError('Duplicate or mismatched sample index')79     seen.add(expected)80     yield row81     dispatch(rank)82   if not stop_event.is_set() and seen!=set(range(len(pairs))):raise RuntimeError('Incomplete GPU result coverage')83  except BaseException:self.close();raise84 def close(self):85  for c in self.conns:86   try:c.send({'type':'close'})87   except Exception:pass88  for p in self.procs:89   p.join(timeout=1)90   if p.is_alive():p.terminate();p.join(timeout=5)91  for c in self.conns:92   try:c.close()93   except Exception:pass94  self.conns=[];self.procs=[]95