| import datetime |
| import json |
| import subprocess |
| import sys |
| import time |
| from pathlib import Path |
| ROOT=Path('/home/ubuntu/tsr') |
| OUT=ROOT/'research_round' |
| PLAN=[{'method':'l2sp','epochs':100,'lr':1e-4,'init':'pretrained','patience':30}, |
| {'method':'sam','epochs':20,'lr':1e-5,'init':'validation_winner','patience':0}, |
| {'method':'contrast','epochs':100,'lr':1e-4,'init':'pretrained','patience':30}, |
| {'method':'consistency','epochs':100,'lr':1e-4,'init':'pretrained','patience':30}, |
| {'method':'swa','epochs':20,'lr':1e-5,'init':'validation_winner','patience':0}] |
| (OUT/'plan.json').write_text(json.dumps(PLAN,indent=2)) |
|
|
| def command(trial,smoke=False): |
| method=trial['method'] |
| cmd=[sys.executable,'-u',str(ROOT/'finetune_research.py'),'--pretrain_pth',str(ROOT/'latest.utae.tar'), |
| '--research',method,'--tag',f'32X32_research_{method}', '--out_dir',str(OUT/('smoke' if smoke else 'checkpoints')/method), |
| '--epochs',str(trial['epochs']),'--base_lr',str(trial['lr']),'--encoder_lr_mult','0.5', |
| '--batch','2','--crop_size','32','--fixed_crops','--mask_ce_invalid','--resume','none', |
| '--patience',str(trial['patience']),'--warmup_epochs', '0' if trial['init']=='validation_winner' else '3'] |
| if trial['init']=='validation_winner': |
| selection=OUT/'initialization.json' |
| initial=json.loads(selection.read_text())['checkpoint'] if selection.exists() else str(ROOT/'checkpoints/finetune_v2/h3_mask_ce_invalid/model_best_raw.tar') |
| cmd+=['--init_finetuned',initial] |
| if not smoke and method == 'contrast': |
| cmd+=['--resume','auto'] |
| if smoke: |
| |
| cmd+=['--epochs',str(6 if method=='swa' else 4 if method=='consistency' else 1),'--limit_batches','1','--patience','0'] |
| return cmd |
|
|
| def run(trial,smoke=False): |
| name=trial['method']; key=('smoke_' if smoke else '')+name |
| statusfile=OUT/'queue_status.json' |
| status=json.loads(statusfile.read_text()) if statusfile.exists() else {} |
| status[key]={'status':'running','started':datetime.datetime.now(datetime.timezone.utc).isoformat(),'command':command(trial,smoke)} |
| statusfile.write_text(json.dumps(status,indent=2)) |
| start=time.monotonic() |
| with (OUT/f'{key}.log').open('w') as stream: |
| result=subprocess.run(command(trial,smoke),cwd=ROOT,stdout=stream,stderr=subprocess.STDOUT) |
| status[key].update(status='completed' if result.returncode==0 else 'failed',seconds=time.monotonic()-start,returncode=result.returncode) |
| statusfile.write_text(json.dumps(status,indent=2)) |
| if result.returncode: raise RuntimeError(f'{key} failed; inspect {key}.log') |
| print(key,'completed',status[key]['seconds'],'seconds',flush=True) |
|
|
| def main(): |
| if '--smoke' in sys.argv: |
| for trial in PLAN: run(trial,True) |
| return |
| import hashlib |
| manifest=json.loads((OUT/'protocol_manifest.json').read_text()) |
| for name,digest in manifest.items(): |
| if hashlib.sha256((ROOT/name).read_bytes()).hexdigest()!=digest: |
| raise RuntimeError(f'Protocol or pretrained file changed: {name}') |
| print('Waiting for existing improvement round to finish; no competing GPU training.',flush=True) |
| while 'Improvement round complete' not in (ROOT/'improvement_round/service.log').read_text(): |
| state=subprocess.run(['systemctl','--user','show','tsr-improvement-round.service','--property=ActiveState','--value'],capture_output=True,text=True) |
| if state.returncode==0 and state.stdout.strip() in ['failed','inactive']: |
| raise RuntimeError('Prior improvement round stopped before completion') |
| time.sleep(30) |
| (OUT/'initialization.json').write_text((ROOT/'improvement_round/selected_model.json').read_text()) |
| status=json.loads((OUT/'queue_status.json').read_text()) |
| for trial in PLAN: |
| if status.get('smoke_'+trial['method'],{}).get('status')!='completed': run(trial,True) |
| for trial in PLAN: |
| current=json.loads((OUT/'queue_status.json').read_text()) |
| if current.get(trial['method'],{}).get('status') == 'completed': |
| print('Retaining completed trial:',trial['method'],flush=True) |
| continue |
| run(trial) |
| print('Research queue completed. Test evaluations run in the separate test service.',flush=True) |
|
|
| if __name__=='__main__': main() |
|
|
