56 lines
3.8 KiB
Python
56 lines
3.8 KiB
Python
"""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)
|