Digital Drosophila Archive
67 lines
3.3 KiB
Python
67 lines
3.3 KiB
Python
"""One registered AK action per fresh process; no changes to frozen runtimes."""
|
|
import argparse
|
|
import json
|
|
import os
|
|
import urllib.request
|
|
|
|
from context import ROOT, EXP, CFG, check, budget, read
|
|
import runtime_bridge as bridge
|
|
|
|
|
|
def preflight():
|
|
# Establish visibility before the lazy runtime/component Torch import.
|
|
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'
|
|
for port in (8188, 8189, 8190):
|
|
with urllib.request.urlopen(f'http://192.168.10.51:{port}/queue', timeout=5) as response:
|
|
queue = json.load(response)
|
|
bridge.require(not queue['queue_running'] and not queue['queue_pending'], 'ComfyUI busy: defer, never evict')
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument('backend', choices=['gpu', 'oracle', 'component'])
|
|
parser.add_argument('group', choices=bridge.GROUPS)
|
|
parser.add_argument('case')
|
|
parser.add_argument('--repeat', type=int, default=0)
|
|
args = parser.parse_args()
|
|
check(); budget()
|
|
row = bridge.authorize_case(EXP, CFG, args.group, args.case, args.backend, args.repeat)
|
|
case = EXP / 'cases' / args.group / args.case
|
|
spec = bridge.case_spec(case)
|
|
bridge.require(spec['name'] == args.case and spec['budget_root'] == str(EXP) and
|
|
spec['output_budget_bytes'] == CFG['output_budget_bytes'], 'Case identity/new budget differs')
|
|
if args.group == 'controls':
|
|
bridge.control_inputs(case, row['source_case'])
|
|
else:
|
|
bridge.require(spec.get('kind') == 'fullgraph' and spec.get('validation_cohort') == args.group and
|
|
all(spec[key] == row[key] for key in ('mode', 'odor_id', 'odor_index', 'seed')),
|
|
'New fullgraph case differs from its registered cohort cell')
|
|
output = EXP / 'gpu' / args.group / f'{args.case}-r{args.repeat}'
|
|
if args.backend == 'component':
|
|
summary = read(output / 'summary.json')
|
|
bridge.require(summary.get('complete') is True and summary.get('pass_all') is True and
|
|
summary.get('case') == args.case and summary.get('apl_delay_steps') == 18,
|
|
'Component audit needs a complete matching D18 GPU run')
|
|
preflight()
|
|
def execute():
|
|
if args.backend == 'gpu':
|
|
result = bridge.run_gpu(case, output, EXP / 'extension-build')
|
|
elif args.backend == 'oracle':
|
|
result = bridge.run_oracle(case, EXP / 'oracle' / args.group / args.case)
|
|
else:
|
|
result = bridge.component(output, spec['apl_bundle_dir'], EXP / 'component-audit' / args.group / (args.case + '.json'))
|
|
budget()
|
|
bridge.require(result.get('pass_all') is True, 'Runtime/component evidence did not pass its own integrity checks')
|
|
return result
|
|
result = bridge.guarded_action(EXP, args.backend, args.group, args.case, args.repeat, execute)
|
|
print(json.dumps(dict(backend=args.backend, group=args.group, case=args.case, repeat=args.repeat,
|
|
complete=result.get('complete'), pass_all=result.get('pass_all')), allow_nan=False), flush=True)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|