scaleup_job.py

1"""Reviewable probe-then-fresh-train job, always inside its existing guardian."""
2import argparse,fcntl,hashlib,json,os,subprocess,sys,time
3from pathlib import Path
4from training_deadline import Deadline,read_deadline
5from scaleup_plan import choose_schedule
6
7def digest(p):return hashlib.sha256(Path(p).read_bytes()).hexdigest()
8def write(p,value):
9    tmp=p.with_suffix(p.suffix+'.tmp')
10    with tmp.open('w') as f:json.dump(value,f,indent=2);f.flush();os.fsync(f.fileno())
11    tmp.replace(p)
12    fd=os.open(p.parent,os.O_DIRECTORY)
13    try:os.fsync(fd)
14    finally:os.close(fd)
15def verify_guard(deadline,command,parent_command,lock_path):
16    if deadline.value['budget_seconds']!=14400:raise ValueError('wrong total job budget')
17    if b'training_deadline.py' not in parent_command or b'--guardian' not in parent_command:raise ValueError('owned guardian parent required')
18    if deadline.value['identity']['command']!=command:raise ValueError('deadline command identity mismatch')
19    with open(lock_path,'a') as lock:
20        try:fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB)
21        except BlockingIOError:pass
22        else:fcntl.flock(lock,fcntl.LOCK_UN);raise ValueError('deadline lock is not held')
23def main():
24    p=argparse.ArgumentParser();p.add_argument('--config',required=True);p.add_argument('--config-sha256',required=True);p.add_argument('--self-sha256',required=True);args=p.parse_args()
25    if digest(__file__)!=args.self_sha256 or digest(args.config)!=args.config_sha256:raise ValueError('reviewed job/config changed')
26    config=json.loads(Path(args.config).read_text());runtime=Path(__file__).resolve().parent
27    for name,sha in config['runtime_sha256'].items():
28        if digest(runtime/name)!=sha:raise ValueError('training runtime changed: '+name)
29    if 'CHESS_JOB_DEADLINE' not in os.environ:raise ValueError('existing guarded total deadline required')
30    deadline=Deadline(read_deadline(os.environ['CHESS_JOB_DEADLINE']))
31    parent=(Path('/proc')/str(os.getppid())/'cmdline').read_bytes()
32    verify_guard(deadline,[sys.executable,str(Path(__file__).resolve()),*sys.argv[1:]],parent,os.environ['CHESS_JOB_DEADLINE']+'.lock')
33    if deadline.soft_stop():raise TimeoutError('original total deadline reached')
34    dataset=Path(config['dataset']);receipt=Path(config['integrity_receipt'])
35    if digest(dataset/'manifest.json')!=config['dataset_manifest_sha256'] or digest(receipt)!=config['integrity_receipt_sha256']:raise ValueError('reviewed dataset changed')
36    root=Path(config['output']);root.mkdir(parents=True,exist_ok=True)
37    if (root/'job-started.json').exists():raise ValueError('job already attempted; no automatic probe/training restart')
38    write(root/'job-started.json',{'utc_unix':time.time(),'deadline':deadline.value,'config_sha256':args.config_sha256,'initialization':'from_scratch','test_evaluated':False})
39    common=[sys.executable,str(runtime/'train_chess.py'),'--launch-ddp','--dataset',str(dataset),'--integrity-receipt',str(receipt),
40        '--epochs','1','--micro',str(config['micro']),'--accum',str(config['accum']),'--bucket',str(config['bucket']),
41        '--seed',str(config['seed']),'--lr',str(config['lr']),'--collective-timeout-seconds',str(config['collective_timeout_seconds']),
42        '--legality-eval-games',str(config['legality_eval_games'])]
43    probe=[*common,'--output',str(root/'probe'),'--mode','probe']
44    write(root/'probe-command.json',{'argv':probe,'deadline_unix':deadline.value['deadline_unix']})
45    with (root/'probe.log').open('w') as log:subprocess.run(probe,stdout=log,stderr=log,check=True)
46    measured=json.loads((root/'probe/probe-results.json').read_text())
47    schedule=choose_schedule(measured,config,deadline.remaining());write(root/'resolved-schedule.json',schedule)
48    train=[*common,'--output',str(root/'candidate'),'--mode','train','--max-updates',str(schedule['max_updates']),
49        '--schedule-steps',str(schedule['schedule_steps']),'--warmup-steps',str(schedule['warmup_steps']),
50        '--validation-every-steps',str(schedule['validation_every_steps'])]
51    write(root/'train-command.json',{'argv':train,'deadline_unix':deadline.value['deadline_unix'],
52        'fresh_initialization':True,'probe_weights_discarded':True,'test_evaluated':False})
53    if deadline.soft_stop():raise TimeoutError('no remaining budget for fresh candidate')
54    with (root/'training.log').open('w') as log:subprocess.run(train,stdout=log,stderr=log,check=True)
55    write(root/'job-complete.json',{'utc_unix':time.time(),'candidate_results_sha256':digest(root/'candidate/results.json'),
56        'remaining_seconds':deadline.remaining(),'test_evaluated':False})
57if __name__=='__main__':main()
58