79 lines
3.9 KiB
Python
Executable File
79 lines
3.9 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Exclusive operational handoff: ASR on GPU1, or GPU1 released for future tests.
|
|
Never starts/stops ComfyUI, never kills unknown GPU processes, never prints API keys.
|
|
"""
|
|
import argparse,fcntl,json,pathlib,subprocess,time,urllib.request
|
|
ROOT=pathlib.Path(__file__).resolve().parent
|
|
CONTAINER='followup-audio-asr'
|
|
|
|
def command(args):
|
|
p=subprocess.run(args,capture_output=True,text=True,timeout=700)
|
|
if p.returncode: raise RuntimeError('COMMAND_FAILED: '+args[0])
|
|
return p.stdout
|
|
|
|
def gpu_pids():
|
|
text=command(['nvidia-smi','--id=1','--query-compute-apps=pid','--format=csv,noheader,nounits'])
|
|
return {int(line.strip()) for line in text.splitlines() if line.strip()}
|
|
|
|
def container_state():
|
|
p=subprocess.run(['docker','inspect','--format','{{json .State}}',CONTAINER],capture_output=True,text=True)
|
|
if p.returncode:
|
|
if 'no such object' in p.stderr.lower() or 'no such container' in p.stderr.lower():return {'Running':False,'Status':'absent'}
|
|
raise RuntimeError('CONTAINER_STATE_UNKNOWN')
|
|
return json.loads(p.stdout)
|
|
|
|
def owned_pids(state):
|
|
if not state.get('Running'):return set()
|
|
lines=command(['docker','top',CONTAINER,'-eo','pid']).splitlines()[1:]
|
|
return {int(line.strip()) for line in lines if line.strip()}
|
|
|
|
def ensure_exclusive(active,owned):
|
|
if active-owned: raise RuntimeError('GPU1_HAS_OTHER_WORKLOAD: refusing to stop or replace it')
|
|
|
|
def idle_values(body):
|
|
values=[]
|
|
for name in ['vllm:num_requests_running','vllm:num_requests_waiting']:
|
|
rows=[float(l.rsplit(' ',1)[1]) for l in body.splitlines() if l.startswith(name+'{')]
|
|
if not rows:raise RuntimeError('ASR_DRAIN_METRICS_MISSING')
|
|
values.append(sum(rows))
|
|
return values
|
|
|
|
def drain():
|
|
zeros=0;deadline=time.monotonic()+600
|
|
while time.monotonic()<deadline:
|
|
with urllib.request.urlopen('http://127.0.0.1:8001/metrics',timeout=5) as r: values=idle_values(r.read().decode())
|
|
zeros=zeros+1 if values==[0,0] else 0
|
|
if zeros>=3:return
|
|
time.sleep(5)
|
|
raise RuntimeError('ASR_BUSY: unchanged; wait for requests to finish')
|
|
|
|
def main():
|
|
p=argparse.ArgumentParser();p.add_argument('mode',choices=['status','asr','test']);args=p.parse_args()
|
|
lock=open(ROOT/'.gpu1-mode.lock','a');fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB)
|
|
state=container_state();active=gpu_pids();owned=owned_pids(state)
|
|
if args.mode=='status':
|
|
print(json.dumps({'asr_status':state['Status'],'asr_running':state['Running'],'gpu1_processes':sorted(active),'other_gpu1_processes':sorted(active-owned),'test_service_started':False}));return
|
|
ensure_exclusive(active,owned)
|
|
compose=['docker','compose','--project-directory',str(ROOT),'--env-file',str(ROOT/'.env'),'-f',str(ROOT/'compose.yaml')]
|
|
if args.mode=='test':
|
|
if state['Running']:
|
|
drain();ensure_exclusive(gpu_pids(),owned_pids(container_state()))
|
|
command(compose+['stop','--timeout','660','asr'])
|
|
if gpu_pids():raise RuntimeError('GPU1_NOT_FREE')
|
|
result={'mode':'test','asr_stopped':True,'gpu1_free':True,'comfyui_test_started':False}
|
|
else:
|
|
if not state['Running']:
|
|
if active:raise RuntimeError('GPU1_NOT_FREE')
|
|
command(compose+['up','-d','--no-build','--pull','never','asr'])
|
|
deadline=time.monotonic()+480
|
|
while time.monotonic()<deadline:
|
|
state=container_state()
|
|
if state.get('Health',{}).get('Status')=='healthy':break
|
|
if state['Status'] in ['exited','dead','restarting']:raise RuntimeError('ASR_START_FAILED')
|
|
time.sleep(5)
|
|
else:raise RuntimeError('ASR_READINESS_TIMEOUT')
|
|
result={'mode':'asr','asr_healthy':True,'gpu1_processes':sorted(gpu_pids()),'comfyui_test_started':False}
|
|
target=ROOT/'.gpu1-mode.json';stage=target.with_suffix('.tmp');stage.write_text(json.dumps(result)+'\n');stage.chmod(0o600);stage.replace(target)
|
|
print(json.dumps(result))
|
|
if __name__=='__main__':main()
|