35 lines
3 KiB
Python
35 lines
3 KiB
Python
#!/usr/bin/env python3
|
|
import os
|
|
os.environ['OPENBLAS_NUM_THREADS']='1'
|
|
from pathlib import Path
|
|
import subprocess,concurrent.futures,json,time,signal,sys,hashlib,numpy as np
|
|
ROOT=Path('/home/user/run/mnist-altai-native');GROUP=ROOT/(sys.argv[1] if len(sys.argv)>1 else 'gm-shards-v1');GROUP.mkdir(exist_ok=False);START=time.monotonic();N=1000;WORKERS=8
|
|
(GROUP/'SHARDS.json').write_text(json.dumps({'total':N,'shards':WORKERS,'replicas':2,'configuration':'default official GM 4x4'}))
|
|
children=[];logs=[]
|
|
try:
|
|
for i in range(WORKERS):
|
|
ids=np.arange(60000+i*125,60000+(i+1)*125);f=GROUP/f'ids-{i:02}.txt';np.savetxt(f,ids,fmt='%d');name=f'{GROUP.name}/shard-{i:02}';log=(GROUP/f'runner-{i:02}.log').open('w');logs.append(log)
|
|
cmd=[sys.executable,'/home/user/run/altai-official-probe/run_verify.py','gm','125','2',name,str(f)];p=subprocess.Popen(cmd,stdout=log,stderr=subprocess.STDOUT,start_new_session=True);children.append(p)
|
|
(GROUP/'pids.json').write_text(json.dumps([p.pid for p in children]))
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=WORKERS) as pool:
|
|
future={pool.submit(p.wait,1800):p for p in children}
|
|
for f in concurrent.futures.as_completed(future):
|
|
if f.result()!=0:raise RuntimeError('An official GoldenModel shard failed verification')
|
|
hs=[];outs=[];preds=[];summaries=[]
|
|
for i in range(WORKERS):
|
|
r=GROUP/f'shard-{i:02}';summaries.append(json.loads((r/'VERIFIED.json').read_text()));hs.append(np.fromfile(r/'hidden_counts.bin',np.int32).reshape(125,256));outs.append(np.fromfile(r/'output_counts.bin',np.int32).reshape(125,2,8,10));preds.append(np.loadtxt(r/'predictions.csv',delimiter=',',skiprows=1,dtype=int))
|
|
h=np.concatenate(hs);o=np.concatenate(outs);p=np.concatenate(preds);ref=np.load(ROOT/'reference.npz');assert np.array_equal(p[:,1],np.arange(60000,61000));assert np.array_equal(h,ref['hidden_counts'][:N]);assert np.array_equal(o,ref['output_counts'][:N]);assert np.array_equal(p[:,3],ref['predictions'][:N]);assert np.array_equal(p[:,2],ref['labels'][:N])
|
|
h.tofile(GROUP/'hidden_counts.bin');o.tofile(GROUP/'output_counts.bin');p[:,0]=np.arange(N);np.savetxt(GROUP/'predictions.csv',p,fmt='%d',delimiter=',',header='row,image_id,label,prediction',comments='')
|
|
result=dict(total=N,correct=int((p[:,2]==p[:,3]).sum()),accuracy=float((p[:,2]==p[:,3]).mean()),hidden_exact=True,output_exact=True,predictions_exact=True,requested_ids_exact=True,processes=WORKERS,replicas_per_process=2,grid_per_process='default4x4',backend='official1.0.2 AltAIv1_GoldenModel',wall_seconds=time.monotonic()-START,model_sha256=hashlib.sha256((ROOT/'model.bin').read_bytes()).hexdigest(),shards=summaries)
|
|
(GROUP/'VERIFIED.json').write_text(json.dumps(result,indent=2));print(json.dumps(result),flush=True)
|
|
except BaseException as e:
|
|
for p in children:
|
|
if p.poll() is None:
|
|
try:os.killpg(p.pid,signal.SIGTERM)
|
|
except ProcessLookupError:pass
|
|
for p in children:
|
|
try:p.wait(timeout=5)
|
|
except subprocess.TimeoutExpired:os.killpg(p.pid,signal.SIGKILL);p.wait()
|
|
(GROUP/'FAILED.json').write_text(json.dumps({'error':str(e)}));raise
|
|
finally:
|
|
for f in logs:f.close()
|