Digital Drosophila Archive
42 lines
2.0 KiB
Python
42 lines
2.0 KiB
Python
"""One process per CPU/GPU trial, gate the ordered fixed protocol."""
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
import traceback
|
|
from context import ROOT,MODEL,EXP,CFG,check,save,budget_check
|
|
|
|
def main():
|
|
ap=argparse.ArgumentParser();ap.add_argument('backend',choices=['oracle','gpu']);ap.add_argument('case')
|
|
a=ap.parse_args();check();budget_check()
|
|
names=[x['name'] for x in CFG['trials']];assert a.case in names
|
|
position=names.index(a.case)
|
|
if position:
|
|
previous=json.loads((EXP/'audit'/('completed-'+names[position-1]+'.json')).read_text())
|
|
assert previous['pass_all'], 'Previous fixed trial failed; do not continue'
|
|
case=EXP/'cases'/a.case;out=EXP/a.backend/a.case
|
|
assert not out.exists(), 'Never overwrite an executed trial'
|
|
start=time.perf_counter()
|
|
try:
|
|
if a.backend=='oracle':
|
|
from oracle_full import run_case
|
|
result=run_case(case,out)
|
|
else:
|
|
os.environ['CUDA_VISIBLE_DEVICES']=CFG['gpu_uuid'];os.environ['CUDA_HOME']='/usr/local/cuda'
|
|
os.environ['PATH']=str(ROOT/'models/male-cns-lif/.venv/bin')+':/usr/local/cuda/bin:'+os.environ['PATH']
|
|
os.environ['MAX_JOBS']='2';os.environ['TORCH_CUDA_ARCH_LIST']='8.9'
|
|
sys.path.insert(0,str(ROOT/'models/male-cns-lif/terminal-scope-diagnostic-v1'))
|
|
from diagnostic import queues
|
|
queues()
|
|
from gpu_full import run_gpu
|
|
result=run_gpu(case,out,EXP/'oracle'/a.case,EXP/'extension-build')
|
|
save(EXP/'audit'/(a.backend+'-'+a.case+'.json'),dict(case=a.case,backend=a.backend,
|
|
wall_seconds_including_validation=time.perf_counter()-start,result=result))
|
|
budget_check();print('COMPLETED',a.backend,a.case,json.dumps(result),flush=True)
|
|
except Exception as exc:
|
|
save(EXP/'audit'/('error-'+a.backend+'-'+a.case+'.json'),dict(error=str(exc),type=type(exc).__name__,traceback=traceback.format_exc()))
|
|
raise
|
|
|
|
if __name__=='__main__':main()
|