Spaces:
Running on Zero
Running on Zero
| """独立プロセスの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 | |