Blame for cpu_queue.py

"""独立プロセスのFIFO。CPUは1件、待機は最大6件。状態APIは待機しない。"""
import json,os,signal,subprocess,sys,threading,time
from collections import deque
from pathlib import Path

class CpuQueue:
 def __init__(self,root,on_progress,on_finish,on_start,max_waiting=6,timeout=2700,launcher=None):
  self.root=Path(root).resolve();self.progress=on_progress;self.finish=on_finish;self.start=on_start
  self.max_waiting=max_waiting;self.timeout=timeout;self.launcher=launcher or self._launch
  self.lock=threading.Lock();self.waiting=deque();self.entries={};self.active=None
 def submit(self,jid,spec):
  with self.lock:
   if jid in self.entries and self.entries[jid]['state'] in ('queued','running'):return True
   if len(self.waiting)>=self.max_waiting:return False
   from cpu_worker import atomic_json
   directory=self.root/jid
   for name in ('cpu-progress.json','cpu-result.json','cpu-error.json'): (directory/name).unlink(missing_ok=True)
   atomic_json(directory/'cpu-request.json',spec)
   atomic_json(directory/'cpu-state.json',{'state':'queued'})
   self.entries[jid]={'state':'queued','cancel':threading.Event()};self.waiting.append(jid)
   if self.active is None:
    self.active='starting';threading.Thread(target=self._drain,daemon=True).start()
   return True
 def info(self,jid):
  with self.lock:
   item=self.entries.get(jid)
   if not item:return None
   return {'execution_device':'cpu','execution_state':item['state'],'queue_position':list(self.waiting).index(jid)+1 if jid in self.waiting else 0,'queue_total':len(self.waiting)+(1 if self.active else 0)}
 def pending(self,jid):
  info=self.info(jid);return bool(info and info['execution_state'] in ('queued','running'))
 def cancel(self,jid):
  queued=False
  with self.lock:
   item=self.entries.get(jid)
   if not item:return
   item['cancel'].set()
   if jid in self.waiting:self.waiting.remove(jid);item['state']='finished';queued=True
  if queued:self.finish(jid,'cancelled',None)
 def _launch(self,jid):
  env=os.environ.copy()
  for key in list(env):
   if any(part in key.upper() for part in ('TOKEN','SECRET','PASSWORD','API_KEY')):env.pop(key,None)
  env.update(CUDA_VISIBLE_DEVICES='',OMP_NUM_THREADS='1',MKL_NUM_THREADS='1',OPENBLAS_NUM_THREADS='1',HF_HUB_DISABLE_IMPLICIT_TOKEN='1')
  return subprocess.Popen([sys.executable,str(Path(__file__).with_name('cpu_worker.py')),str(self.root),jid],env=env,stdin=subprocess.DEVNULL,stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL,start_new_session=os.name!='nt')
 def _stop(self,process):
  if process.poll() is not None:return
  if os.name!='nt':
   try:os.killpg(process.pid,signal.SIGTERM)
   except ProcessLookupError:return
  else:process.terminate()
  try:process.wait(timeout=2)
  except subprocess.TimeoutExpired:
   if os.name!='nt':
    try:os.killpg(process.pid,signal.SIGKILL)
    except ProcessLookupError:pass
   else:process.kill()
   process.wait(timeout=2)
 def _drain(self):
  while True:
   with self.lock:
    if not self.waiting:self.active=None;return
    jid=self.waiting.popleft();self.active=jid;entry=self.entries[jid];entry['state']='running'
   process=None;outcome='processing_failed';result=None
   try:
    import job_control
    with job_control.execution(self.root,jid):
     if entry['cancel'].is_set() or job_control.is_cancelled(self.root,jid):outcome='cancelled'
     else:
      self.start(jid);process=self.launcher(jid);started=time.monotonic();last=None
      while process.poll() is None:
       if entry['cancel'].is_set() or job_control.is_cancelled(self.root,jid):outcome='cancelled';self._stop(process);break
       if time.monotonic()-started>self.timeout:outcome='cpu_timeout';self._stop(process);break
       try:
        data=json.loads((self.root/jid/'cpu-progress.json').read_text(encoding='utf-8'))
        if data!=last:self.progress(jid,data);last=data
       except (OSError,ValueError):pass
       entry['cancel'].wait(.5)
      if entry['cancel'].is_set() or job_control.is_cancelled(self.root,jid):outcome='cancelled'
      elif process.returncode==0 and outcome not in ('cpu_timeout',):
       result=json.loads((self.root/jid/'cpu-result.json').read_text(encoding='utf-8'));outcome='completed'
      elif outcome=='processing_failed':
       try:
        kind=json.loads((self.root/jid/'cpu-error.json').read_text(encoding='utf-8')).get('error_class','')
        if kind in ('ModuleNotFoundError','ImportError','RuntimeError','FileNotFoundError','ValueError','TypeError','MemoryError'):print('[Lunaron CPU] '+kind,flush=True)
       except (OSError,ValueError):pass
   except Exception:pass
   finally:
    if process is not None and process.poll() is None:self._stop(process)
    with self.lock:entry['state']='finished'
    try:self.finish(jid,outcome,result)
    except Exception:pass