File size: 9,760 Bytes
d9bc75d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | #!/usr/bin/env python3
"""Stop exact old workers, fork verified journals, add SimCT; never launch GPUs."""
import argparse, contextlib, copy, hashlib, json, os, shutil, signal, sys, time
from pathlib import Path
def proc(pid):
p=Path('/proc')/str(pid)
try:
raw=(p/'stat').read_text().rsplit(')',1)[1].split()
return dict(pid=int(pid),ppid=int(raw[1]),group=int(raw[2]),start=raw[19],state=raw[0],
argv=[x.decode() for x in (p/'cmdline').read_bytes().split(b'\0') if x])
except (OSError, UnicodeError): return None
def stop_workers(plan,queue,pids):
handles=[]; roots=[]
try:
for pid in pids:
info=proc(pid)
if info is None: continue
args=info['argv']
if str(queue) not in args or 'worker' not in args or '--plan' not in args or args[args.index('--plan')+1]!=str(plan):
raise ValueError('PID identity mismatch: '+str(pid))
fd=os.pidfd_open(pid); again=proc(pid)
if not again or again['start']!=info['start']: raise ValueError('PID changed')
handles.append(fd); roots.append(info)
# Freeze only validated workers so they cannot start a new job during teardown.
for fd in handles: signal.pidfd_send_signal(fd,signal.SIGSTOP)
table={int(p.name):proc(int(p.name)) for p in Path('/proc').iterdir() if p.name.isdigit()}
descendants={r['pid'] for r in roots}
while True:
more={pid for pid,x in table.items() if x and x['ppid'] in descendants}
if more<=descendants: break
descendants |= more
children=[]
for pid in descendants-{r['pid'] for r in roots}:
before=table[pid]
try: fd=os.pidfd_open(pid)
except ProcessLookupError: continue
after=proc(pid)
if not after or after['start']!=before['start']: os.close(fd);continue
children.append(fd)
try:
for fd in children+handles:
try: signal.pidfd_send_signal(fd,signal.SIGTERM)
except ProcessLookupError: pass
for fd in handles:
try: signal.pidfd_send_signal(fd,signal.SIGCONT)
except ProcessLookupError: pass
limit=time.monotonic()+30
while time.monotonic()<limit:
alive=[pid for pid in descendants if (x:=proc(pid)) and x['state']!='Z' and x['start']==table[pid]['start']]
if not alive: return
time.sleep(.5)
raise RuntimeError('Some owned processes remain; no migration performed: '+str(alive))
finally:
for fd in children: os.close(fd)
finally:
for fd in handles:
try: signal.pidfd_send_signal(fd,signal.SIGCONT)
except ProcessLookupError: pass
os.close(fd)
def main():
a=argparse.ArgumentParser();a.add_argument('--plan',type=Path,required=True);a.add_argument('--source',type=Path,required=True)
a.add_argument('--out',type=Path,required=True);a.add_argument('--worker-pids',nargs='+',type=int,required=True)
a.add_argument('--simct',type=Path,required=True);args=a.parse_args()
old=args.plan.resolve(strict=True); source=args.source.resolve(strict=True); out=args.out.resolve()
if out.exists() or out.is_relative_to(old.parent): raise ValueError('Output must be a new sibling directory')
sys.path.insert(0,str(source/'scripts/evaluation'))
import eval_queue as Q
E,D=Q.E,Q.D
plan=E.read_json(old); oldhash=E.file_hash(old)
if plan['source']!=D.script_hashes(): raise ValueError('Source differs from running plan')
if len(plan['jobs'])!=17 or any(j['mode'] not in ('sft','atomic','fixed') for j in plan['jobs']): raise ValueError('Unexpected old jobs')
for j in plan['jobs']:
if j['mode'] in D.RUNS and D.RUNS[j['mode']] not in Path(j['checkpoint']['path']).parts: raise ValueError('Wrong historical MP run')
summary=E.read_json(args.simct/'run-summary.json')
if summary['kd_algorithm']!='span_ctkd' or summary['status']!='completed' or summary['optimizer_updates']!=312: raise ValueError('SimCT summary invalid')
if summary['student']!=next(j['checkpoint']['path'] for j in plan['jobs'] if j['mode']=='sft'): raise ValueError('Different SFT initialization')
# Hash before stopping to avoid wasting idle GPU time on checkpoint inventory.
extra=[]
for step in (312,156,80,240,40,200,120,280):
print('HASH_SIMCT',step,flush=True)
identity=E.checkpoint_identity(args.simct/f'step{step}')
tier=next(j['tier'] for j in plan['jobs'] if j['mode']=='atomic' and j['step']==step)
extra.append(dict(id=f'simct-{step}',mode='simct',step=step,tier=tier,checkpoint=identity))
for data in plan['data'].values():
if E.file_hash(data['path'])!=data['sha256']: raise ValueError('Data changed')
stop_workers(old,source/'scripts/evaluation/eval_queue.py',args.worker_pids)
with contextlib.ExitStack() as stack:
for gpu in (0,1):
if stack.enter_context(Q.locked(Path(f'/tmp/simct-eval-gpu{gpu}.lock'),blocking=False)) is None: raise ValueError('GPU worker still active')
stack.enter_context(Q.locked(old.parent/'state.lock'))
for j in plan['jobs']:
if stack.enter_context(Q.locked(old.parent/(j['id']+'.lock'),blocking=False)) is None: raise ValueError('Job still active')
state=E.read_json(old.parent/'state.json')
if state['plan_sha256']!=oldhash or E.file_hash(old)!=oldhash: raise ValueError('Plan changed')
if time.time()>=state['admit_until']: raise ValueError('Original admission window expired; no implicit extension')
out.mkdir()
receipt={'from_plan':str(old),'old_plan_sha256':oldhash,'source':str(source),'concurrency':16,'original_files':{},'status':'incomplete'}
E.write_new(out/'migration.json',receipt)
new=copy.deepcopy(plan);new['jobs']+=extra;new['jobs'].sort(key=lambda j:j['tier'])
new['migration']={'from_plan':str(old),'sha256':oldhash,'reason':'Add eight SimCT checkpoints; preserve journals and original clock'}
E.write_new(out/'plan.json',new);newhash=E.file_hash(out/'plan.json')
if (old.parent/'cells').exists(): shutil.copytree(old.parent/'cells',out/'cells')
jobs={j['id']:j for j in plan['jobs']};datasets={b:{x['id']:x for x in E.read_json(d['path'])['items']} for b,d in plan['data'].items()}
totals={'responses':0,'scores':0,'metrics':0}
for f in (old.parent/'cells').rglob('*'):
if f.is_file(): receipt['original_files'][str(f.relative_to(old.parent))]=E.file_hash(f)
for cell in (out/'cells').glob('*/*/*'):
if not cell.is_dir(): continue
jid,b,seed=cell.relative_to(out/'cells').parts;seed=int(seed)
if jid not in jobs or b not in datasets or seed not in plan['seeds']: raise ValueError('Unknown cell')
contract=E.read_json(cell/'contract.json')
expected=Q.cell_contract(oldhash,jobs[jid],plan['data'],b,seed,contract['server'])
if contract!=expected: raise ValueError('Old cell contract mismatch')
metrics=E.read_json(cell/'metrics.json') if (cell/'metrics.json').exists() else None
if metrics:
if metrics['contract']!=contract or metrics['status']!='completed': raise ValueError('Old metrics mismatch')
for name in ('responses','scores'):
if E.file_hash(cell/(name+'.jsonl'))!=metrics[name+'_sha256']: raise ValueError('Completed journal changed')
responses=Q.journal(cell/'responses.jsonl');scores=Q.journal(cell/'scores.jsonl');items=datasets[b]
if not set(scores)<=set(responses)<=set(items): raise ValueError('Unknown result ID')
for key,row in responses.items():
payload=E.generation_payload('eval-gemma',items[key],b,seed)
if row['seed']!=seed or row['request_sha256']!=E.digest(E.encoded(payload)): raise ValueError('Request mismatch')
E.validate_response(row['response'])
for key,row in scores.items():
if type(row.get('passed')) is not bool or row['response_sha256']!=E.digest(E.encoded(responses[key])): raise ValueError('Score mismatch')
if metrics:
if len(scores)!=len(items) or metrics['count']!=len(items) or metrics['score']!=sum(x['passed'] for x in scores.values())/len(items): raise ValueError('Metric count/score mismatch')
metrics['contract']={**contract,'plan_sha256':newhash};Q.atomic_json(cell/'metrics.json',metrics);totals['metrics']+=1
Q.atomic_json(cell/'contract.json',{**contract,'plan_sha256':newhash})
totals['responses']+=len(responses);totals['scores']+=len(scores)
for rel,digest in receipt['original_files'].items():
if E.file_hash(old.parent/rel)!=digest: raise ValueError('Source journal changed during migration')
state['plan_sha256']=newhash
state['jobs']={k:v for k,v in state['jobs'].items() if v['status']=='completed'}
for jid in state['jobs']:
if not all((out/'cells'/jid/b/str(s)/'metrics.json').exists() for b in E.CAPS for s in plan['seeds']): raise ValueError('Completed job lacks cells')
E.write_new(out/'state.json',state)
receipt.update(status='verified',new_plan_sha256=newhash,totals=totals,deadline=state['deadline'])
Q.atomic_json(out/'migration.json',receipt)
print('MIGRATION_PASS',json.dumps(totals),flush=True);print('NEW_PLAN='+str(out/'plan.json'));print('No GPU launched. Original deadline retained.')
if __name__=='__main__': main()
|