"""Bounded screening runner. Saves evidence and restores original compose on exit."""
import argparse,json,os,pathlib,subprocess,time,urllib.request
p=argparse.ArgumentParser();p.add_argument('--steps',nargs='+',type=int,default=[2,1,3]);p.add_argument('--batch',type=int,default=2048);p.add_argument('--run-id',required=True);a=p.parse_args()
root=pathlib.Path(__file__).resolve().parents[1];os.chdir(root)
audit=pathlib.Path('/data/flash-next/audit/tuning')/a.run_id;audit.mkdir(parents=True,exist_ok=False)
original=(root/'compose.yaml').read_text();(audit/'original.yaml').write_text(original)
key=(root/'secrets/api-key').read_text().strip()
def cmd(args,**kwargs):return subprocess.run(args,check=True,**kwargs)
def state():return json.loads(subprocess.check_output(['docker','inspect','qwen38-flash-6000d']))[0]
def ready(limit=420):
 start=time.monotonic()
 while time.monotonic()-start<limit:
  try:
   with urllib.request.urlopen('http://127.0.0.1:8000/health',timeout=3) as r:
    if r.status==200:return True
  except Exception:pass
  if state()['State']['Status']=='exited':return False
  time.sleep(5)
 return False
def capture(folder):
 s=subprocess.check_output(['docker','logs','qwen38-flash-6000d'],stderr=subprocess.STDOUT).decode();(folder/'server.log').write_text(s.replace(key,'<REDACTED>'))
 try:
  req=urllib.request.Request('http://127.0.0.1:8000/metrics',headers={'Authorization':'Bearer '+key})
  with urllib.request.urlopen(req,timeout=5) as r:metrics=r.read().decode()
  (folder/'mtp-metrics.txt').write_text('\n'.join(l for l in metrics.splitlines() if 'spec_decode' in l))
 except Exception:pass
monlog=(audit/'resources.jsonl').open('x');monitor=subprocess.Popen(['python3','scripts/monitor.py'],stdout=monlog,stderr=subprocess.STDOUT)
(audit/'monitor.pid').write_text(str(monitor.pid))
try:
 for n in a.steps:
  folder=audit/f'mtp{n}-b{a.batch}';folder.mkdir()
  config=original.replace('num_speculative_tokens":2','num_speculative_tokens":'+str(n)).replace('"cudagraph_capture_sizes":[1,3]','"cudagraph_capture_sizes":[1,'+str(n+1)+']').replace('--max-num-batched-tokens 2048','--max-num-batched-tokens '+str(a.batch)).replace('restart: unless-stopped','restart: "no"')
  (folder/'compose.yaml').write_text(config)
  # Run current MTP2 as-is for fresh paired control; all later candidates restart.
  same=n==2 and a.batch==2048 and a.steps.index(n)==0
  if not same:
   (root/'compose.yaml').write_text(config);cmd(['docker','compose','--progress','plain','up','-d','--force-recreate'],timeout=120)
  if not ready():
   capture(folder);(folder/'status.json').write_text(json.dumps({'startup':False}));print('STARTUP_FAILED',folder,flush=True);continue
  for script,name,extra in [('qualification.py','screen.jsonl',['--screen']),('quality_tasks.py','quality-medium.jsonl',['--effort','medium'])]:
   with (folder/(name+'.log')).open('x') as log:
    result=subprocess.run(['python3','scripts/'+script,'--output',str(folder/name)]+extra,stdout=log,stderr=subprocess.STDOUT,timeout=1200)
   if result.returncode!=0:print('SCRIPT_FAILED',folder,script,result.returncode,flush=True)
  capture(folder);rows=[]
  for name in ['screen.jsonl','quality-medium.jsonl']:
   f=folder/name
   if f.exists():rows += [json.loads(l) for l in f.read_text().splitlines()]
  status={'startup':True,'records':len(rows),'failures':[r['name'] for r in rows if not r['passed']]}
  (folder/'status.json').write_text(json.dumps(status,indent=2));print('SCREEN_COMPLETE',folder,status,flush=True)
finally:
 (root/'compose.yaml').write_text(original)
 cmd(['docker','compose','--progress','plain','up','-d','--force-recreate'],timeout=120)
 restored=ready();(audit/'restored.json').write_text(json.dumps({'healthy':restored,'time':time.time()}))
 monitor.terminate();monitor.wait(timeout=15);monlog.close()
 print('RESTORED',restored,flush=True)