| """Prepare/run independent seeds on explicitly assigned Student service lanes.""" |
| import argparse |
| from concurrent.futures import ThreadPoolExecutor |
| import copy |
| import json |
| import os |
| import socket |
| import hashlib |
| from pathlib import Path |
| import subprocess |
| import sys |
|
|
| ROOT = Path(__file__).resolve().parent |
| sys.path.insert(0, str(ROOT/'decico_rsi')) |
|
|
|
|
| def inherited_matrix_leases(env): |
| descriptors=[] |
| for key in ('RSI_MATRIX_HOST_LOCK_FD','RSI_MATRIX_RUN_LOCK_FD'): |
| value=env.get(key) |
| if value is None:continue |
| descriptor=int(value) |
| if descriptor<3:raise ValueError('Invalid matrix lease descriptor') |
| os.fstat(descriptor) |
| descriptors.append(descriptor) |
| return {'pass_fds':tuple(descriptors)} if descriptors else {} |
|
|
|
|
| def main(argv=None): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('--output', required=True) |
| parser.add_argument('--seeds', type=int, nargs='+', default=[2026100801,2026100802]) |
| parser.add_argument('--lanes', help='JSON list of {student_api_base, teacher_api_base}; one Student endpoint per concurrent seed') |
| parser.add_argument('--config', default=str(ROOT/'plan/bfcl-only/full.json')) |
| parser.add_argument('--initialization', choices=['baseline','latest_accepted']) |
| parser.add_argument('--prepare-only', action='store_true') |
| parser.add_argument('--resume', action='store_true',help='Resume the same sealed suite without rerunning prepare') |
| from modeling.student import PROFILES,configure as configure_student |
| parser.add_argument('--student-model',choices=list(PROFILES)) |
| parser.add_argument('--models-root') |
| parser.add_argument('--student-model-path') |
| parser.add_argument('--teacher-model-path') |
| parser.add_argument('--student-gdn-backend', choices=['original', 'fla']) |
| parser.add_argument('--teacher-gdn-backend', choices=['original', 'fla']) |
| args = parser.parse_args(argv) |
| from project_paths import initialize, output_path |
| initialize() |
| from common import write_json |
| from controller.seed_service import identity,alive |
| from controller.service_lease import endpoint_key |
| output = output_path(args.output) |
| if output.exists() and not args.resume: |
| raise ValueError('Suite requires a fresh output directory') |
| if args.resume and not output.is_dir(): |
| raise ValueError('Resume requires an existing suite directory') |
| if len(set(args.seeds)) != len(args.seeds): |
| raise ValueError('Seeds must be unique') |
| lanes = json.loads(Path(args.lanes).read_text()) if args.lanes else [] |
| if not args.prepare_only: |
| if len(lanes) < len(args.seeds): |
| raise ValueError('Provide one distinct prestarted Student lane per parallel seed') |
| endpoints = [endpoint_key(lane['student_api_base']) for lane in lanes[:len(args.seeds)]] |
| if len(set(endpoints)) != len(endpoints): |
| raise ValueError('Parallel seeds cannot share a Student service endpoint') |
| config = configure_student(json.loads(Path(args.config).read_text()), |
| name=args.student_model,models_root=args.models_root,model_path=args.student_model_path) |
| from modeling.teacher import configure as configure_teacher |
| config = configure_teacher(config, models_root=args.models_root, model_path=args.teacher_model_path) |
| from modeling.serving_backend import configure as configure_backend |
| config = configure_backend(config, student=args.student_gdn_backend, teacher=args.teacher_gdn_backend) |
| suite_owner = None if args.prepare_only else identity(os.getpid()) |
| suite_identity = {'seeds':args.seeds, 'lanes':lanes, 'student':config['student'], |
| 'teacher':config['teacher'], 'config_sha256':hashlib.sha256( |
| json.dumps(config, sort_keys=True).encode()).hexdigest(), |
| 'initialization':args.initialization or config['training']['initialization'], |
| 'training_serialization':'project-local exclusive lease; Teacher calls take shared read leases', |
| 'cross_seed_pool_or_posterior_sharing':False} |
| if args.resume: |
| previous=json.loads((output/'suite.json').read_text()) |
| if previous != suite_identity: |
| raise ValueError('Suite model, initialization, seeds or lanes changed; use a new run directory') |
| for lease in output.glob('lane-*.json'): |
| prior=json.loads(lease.read_text()) |
| if prior.get('owner_host') != socket.gethostname(): |
| raise ValueError('Resume this started suite on its original server') |
| if alive(prior['owner']): |
| raise ValueError('Suite controller is still alive; refusing concurrent resume') |
| else: |
| output.mkdir(parents=True) |
| write_json(output/'suite.json',suite_identity) |
|
|
| def execute(item): |
| index, seed = item |
| cfg = copy.deepcopy(config) |
| cfg['seed'] = cfg['model_training_seed'] = seed |
| if args.initialization: |
| cfg['training']['initialization'] = args.initialization |
| lane = lanes[index] if index < len(lanes) else {} |
| if lane.get('teacher_api_base'): |
| cfg['data_service']['api_base'] = lane['teacher_api_base'] |
| for role in ('student', 'teacher'): |
| if lane.get(role + '_backend_receipt_dir'): |
| cfg['serving_backend']['receipt_dirs'][role] = lane[role + '_backend_receipt_dir'] |
| cfg_path = output/f'config-{seed}.json' |
| if args.resume and cfg_path.exists(): |
| if json.loads(cfg_path.read_text()) != cfg: |
| raise ValueError('Sealed per-seed configuration changed') |
| else: |
| write_json(cfg_path,cfg) |
| campaign = output/f'seed-{seed}' |
| env = dict(os.environ, PYTHONDONTWRITEBYTECODE='1') |
| env.pop('DECICO_RSI_RESOURCE_HOOK_CHILD',None) |
| env.pop('RSI_SERVICE_LOCK_FD',None) |
| env.pop('RSI_SUITE_LANE_SPEC',None) |
| env['DECICO_RSI_MODEL_NAME'] = cfg['student']['name'] |
| env['DECICO_RSI_BASE_MODEL'] = cfg['student']['path'] |
| env['DECICO_RSI_STUDENT_MODEL'] = cfg['student']['name'] |
| host=socket.gethostname() |
| host_key=hashlib.sha256(host.encode()).hexdigest()[:16] |
| resources=output/'.resources'/host_key |
| env['DECICO_RSI_RUNTIME_ROOT'] = str(output/'.runtime'/f'seed-{seed}') |
| env['DECICO_RSI_RESOURCE_ROOT'] = str(resources) |
| env['DECICO_RSI_ADAPTER_ROOT'] = str(output/'artifacts'/f'seed-{seed}') |
| env['DECICO_RSI_TRAIN_LOCK'] = str(resources/'training.lock') |
| if lane.get('student_api_base'): |
| env['DECICO_RSI_API_BASE'] = lane['student_api_base'] |
| if not args.prepare_only: |
| lease = output/f'lane-{seed}.json' |
| write_json(lease, {'schema_version':2,'owner':suite_owner,'owner_host':host,'seed':seed, |
| 'student_endpoint':env['DECICO_RSI_API_BASE'], |
| 'training_lock':env['DECICO_RSI_TRAIN_LOCK']}) |
| env['RSI_SUITE_LANE_SPEC'] = str(lease) |
| |
| training = cfg['training'] |
| env.setdefault('DECICO_RSI_TRAIN_LOCK_WAIT_SECONDS', str( |
| training.get('max_training_hours',12)*3600 + training.get('gpu_wait_minutes',120)*60 + 600)) |
| write_json(output/f'layout-{seed}.json', { |
| 'execution_root':str(ROOT),'runtime_root':env['DECICO_RSI_RUNTIME_ROOT'], |
| 'resource_root':env['DECICO_RSI_RESOURCE_ROOT'],'adapter_root':env['DECICO_RSI_ADAPTER_ROOT'], |
| 'training_lock':env['DECICO_RSI_TRAIN_LOCK'],'host':host}) |
| process_options=inherited_matrix_leases(env) |
| with (output/f'seed-{seed}.log').open('a' if args.resume else 'w') as log: |
| if args.resume and campaign.exists(): |
| verified=subprocess.run([sys.executable,'-B',str(campaign/'runtime/preflight_local.py'), |
| '--config',str(campaign/'config.json'),'--output',str(campaign/'checks/resume-preflight.json')], |
| cwd=ROOT,env=env,stdout=log,stderr=subprocess.STDOUT,**process_options) |
| code=verified.returncode |
| if code==0: |
| counts=json.loads((campaign/'partition/manifest.json').read_text())['counts'] |
| expected={**cfg,'data':str(campaign/'partition'),'improve_tasks':0, |
| 'improve_bfcl_tasks':counts['improve'],'selection_tasks':counts['selection'], |
| 'audit_tasks':counts['audit']} |
| if json.loads((campaign/'config.json').read_text()) != expected: |
| raise ValueError('Prepared campaign config differs from the sealed suite') |
| else: |
| prepared = subprocess.run([sys.executable,'-B',str(ROOT/'prepare.py'), |
| '--config',str(cfg_path),'--output',str(campaign)], cwd=ROOT,env=env, |
| stdout=log,stderr=subprocess.STDOUT,**process_options) |
| code = prepared.returncode |
| if code == 0 and not args.prepare_only: |
| code = subprocess.run([sys.executable,'-B',str(campaign/'runtime/rsi_loop.py'), |
| '--config',str(campaign/'config.json'),'--output',str(campaign/'run')], |
| cwd=ROOT,env=env,stdout=log,stderr=subprocess.STDOUT,**process_options).returncode |
| return {'seed':seed,'returncode':code,'campaign':str(campaign)} |
|
|
| with ThreadPoolExecutor(max_workers=len(args.seeds)) as pool: |
| results = list(pool.map(execute, enumerate(args.seeds))) |
| write_json(output/'suite-result.json', {'results':results,'prepare_only':args.prepare_only}) |
| return int(any(row['returncode'] for row in results)) |
|
|
|
|
| if __name__ == '__main__': |
| raise SystemExit(main()) |
|
|
