"""Detach a bounded training campaign, optionally calibrating/evaluating/deploying. No service is unloaded; all child processes are individually tracked. The smoke lock prevents overlapping training campaigns. The final deadline includes eval. Resume restarts interrupted training from its durable checkpoint into a new run. """ import argparse from datetime import datetime, timezone, timedelta import fcntl import hashlib import json import math import os from pathlib import Path import signal import subprocess import sys import time ROOT = Path(__file__).resolve().parents[1] if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) from selection import SELECTION_METRIC RUN_ROOT = Path.home() / "ai/opensysone/runs" STOP = False CHILD = None def write_json(path, value): tmp = path.with_suffix(".tmp") tmp.write_text(json.dumps(value, indent=2) + "\n") tmp.replace(path) def checksum(path): return hashlib.sha256(path.read_bytes()).hexdigest() def evaluation_budget(parent): """Reserve both base/tuned passes, using the pilot's full validation timing.""" predictions = parent/'initial_validation_predictions.json' if not predictions.exists(): predictions = parent/'resumed_initial_predictions.json' checks = parent/'correctness_initial.json' audit = parent/'data_filter.json' if not all(p.exists() for p in (predictions,checks,audit)): return 10800, None n = len(json.loads(predictions.read_text())) seconds = predictions.stat().st_mtime-checks.stat().st_mtime if n <= 0 or seconds <= 0: return 10800, None counts = json.loads(audit.read_text()) evaluation_rows = 2*sum(counts[k]['retained'] for k in ('calibration','test','holdout')) estimate = seconds/n*evaluation_rows # 30% margin plus ten minutes for load, parity, calibration and API startup. reserve = max(7200,math.ceil((estimate*1.3+600)/300)*300) return reserve, {'pilot_validation_decisions':n,'pilot_validation_seconds':seconds, 'evaluation_decisions':evaluation_rows,'estimated_prediction_seconds':estimate} def training_cutoff(final_deadline, reserve_seconds, requested=None, now=None): latest = final_deadline - reserve_seconds cutoff = latest if requested is None else requested if cutoff > latest: raise ValueError('Training deadline must retain the measured evaluation reserve') if cutoff <= (time.time() if now is None else now): raise ValueError('Training window has already expired') return cutoff def stop(*args): global STOP STOP = True if CHILD is not None and CHILD.poll() is None: CHILD.terminate() def child_stage(campaign, state, stage, command, deadline): global CHILD log = campaign / f"{stage}.log" with log.open("a") as handle: CHILD = subprocess.Popen(command, cwd=ROOT, stdin=subprocess.DEVNULL, stdout=handle, stderr=subprocess.STDOUT) state.update(stage=stage, child_pid=CHILD.pid, child_command=command, stage_started_utc=datetime.now(timezone.utc).isoformat()) write_json(campaign / "state.json", state) termination_at = None while CHILD.poll() is None: if (STOP or time.time() >= deadline) and termination_at is None: CHILD.terminate() termination_at = time.monotonic() state["termination_reason"] = "requested_stop" if STOP else "stage_deadline" if termination_at is not None and time.monotonic() - termination_at >= 30: CHILD.kill() state["heartbeat_utc"] = datetime.now(timezone.utc).isoformat() write_json(campaign / "state.json", state) time.sleep(5) code = CHILD.returncode (campaign / f"{stage}_exit_code").write_text(str(code) + "\n") state[stage + "_exit_code"] = code state["child_pid"] = None write_json(campaign / "state.json", state) CHILD = None return code def trainer_command(config, output, resume, deadline): command = [sys.executable, str(ROOT / "experiment.py"), "train", "--output", str(output), "--resume", str(resume), "--deadline", datetime.fromtimestamp(deadline,timezone.utc).isoformat()] for key in ("model","dataset","epochs","rank","alpha","lr","head_lr","seed","effective_batch", "branch_batch_size","max_tokens","save_steps","save_seconds","eval_steps", "validation_per_family","patience","schedule_steps","selection_metric"): value = config.get(key) if value is not None: command.extend(["--" + key.replace("_","-"), str(value)]) if not config["adapters"]: command.append("--head-only") if config.get('two_pass'): command.append('--two-pass') return command def run_campaign(campaign): campaign = Path(campaign).resolve() plan = json.loads((campaign / "plan.json").read_text()) signal.signal(signal.SIGTERM,stop) signal.signal(signal.SIGINT,stop) Path("/proc/self/oom_score_adj").write_text("0") state = {"supervisor_pid":os.getpid(),"campaign":str(campaign),"stage":"starting", "status":"running", "source_commit":subprocess.check_output(['git','rev-parse','HEAD'],cwd=ROOT,text=True).strip()} write_json(campaign / "state.json",state) RUN_ROOT.mkdir(parents=True,exist_ok=True) with (RUN_ROOT / ".smoke.lock").open("w") as lock: try: fcntl.flock(lock,fcntl.LOCK_EX | fcntl.LOCK_NB) except BlockingIOError: state.update(status="failed",error="another project run holds the lock") write_json(campaign / "state.json",state) return 1 try: final_deadline = datetime.fromisoformat(plan["final_deadline"]).timestamp() requested_deadline = (datetime.fromisoformat(plan['training_deadline']).timestamp() if plan.get('training_deadline') else None) training_deadline = training_cutoff(final_deadline,plan['eval_reserve_seconds'],requested_deadline) pilot = Path(plan["resume_checkpoint"]) if checksum(pilot) != plan["resume_sha256"]: raise RuntimeError("Resume checkpoint changed after campaign creation") import torch saved = torch.load(pilot,map_location='cpu',weights_only=False) config = saved['config'] config['selection_metric'] = plan.get('selection_metric',config.get('selection_metric','raw_nll')) del saved train_output = campaign / "training" code = child_stage(campaign,state,"training",trainer_command( config,train_output,pilot,training_deadline),training_deadline + 60) state["training_output"] = str(train_output) state["resume_checkpoint"] = str(train_output / "checkpoint.pt") if STOP: state.update(status="interrupted",stage="stopped") write_json(campaign / "state.json",state) return 0 if code != 0: raise RuntimeError(f"Training exited {code}; durable checkpoint retained") if not (train_output / "summary.json").exists(): raise RuntimeError("Training summary missing") summary = json.loads((train_output / "summary.json").read_text()) state["training_summary"] = {k:summary[k] for k in ('completed_steps','best_validation_macro_nll','status','peak_cuda_allocated_bytes')} state["training_summary"].update({k:summary[k] for k in ('selection_metric','best_validation_selection_score') if k in summary}) if plan.get('train_only', False): state.update(status='training_complete', stage='training_complete', train_only=True, best_checkpoint=str(train_output / 'best.pt'), best_sha256=checksum(train_output / 'best.pt'), finished_utc=datetime.now(timezone.utc).isoformat(), selection='Await cross-candidate validation selection before calibration/test/holdout') predictions = train_output / 'best_validation_predictions.json' if predictions.exists(): state['best_validation_predictions'] = str(predictions) state['best_validation_predictions_sha256'] = checksum(predictions) write_json(campaign / 'state.json',state) return 0 eval_output = campaign / "evaluation" eval_command = [sys.executable,str(ROOT / 'experiment.py'),'finalize','--checkpoint', str(train_output / 'best.pt'),'--dataset',config['dataset'],'--output',str(eval_output)] code = child_stage(campaign,state,'evaluation',eval_command,final_deadline - 120) if STOP: state.update(status='interrupted',stage='stopped') write_json(campaign / 'state.json',state) return 0 if code != 0: raise RuntimeError(f"Evaluation exited {code}; test completion not established") result = json.loads((eval_output / 'metrics.json').read_text()) if result['status'] != 'complete': raise RuntimeError('Evaluation metrics are incomplete') deployed = Path.home() / 'ai/opensysone/deploy' deployed.mkdir(parents=True,exist_ok=True) model = eval_output / 'model.pt' pointer = {'model':str(model),'sha256':checksum(model),'campaign':str(campaign), 'evaluation':str(eval_output / 'metrics.json'), 'inference_max_tokens':plan.get('inference_max_tokens'), 'model_id':'opensysone-'+config['model'].split('/')[-1].rsplit('-',1)[0].lower(), 'finalized_utc':datetime.now(timezone.utc).isoformat()} write_json(deployed / 'current.json',pointer) # Source-controlled result copy contains small evidence, never checkpoints. small = ROOT / 'results' / campaign.name small.mkdir(exist_ok=False) import shutil for directory,names in ((train_output,['manifest.json','summary.json','correctness_initial.json','correctness_final.json']), (eval_output,['manifest.json','metrics.json','correctness.json','data_filter.json'])): for name in names: if (directory / name).exists(): shutil.copy2(directory / name,small / (directory.name + '_' + name)) write_json(small / 'deployment.json',pointer) # Validate the actual inference payload before starting a long-lived service. request_command = [sys.executable,str(ROOT / 'jev_harness.py'),'--backend','local', '--checkpoint',str(model),'--request',str(ROOT / 'examples/jev_request.json')] if plan.get('inference_max_tokens'): request_command.extend(['--max-tokens',str(plan['inference_max_tokens'])]) if child_stage(campaign,state,'harness_check',request_command,final_deadline - 60) != 0: raise RuntimeError('Trained-model harness check failed') if STOP: state.update(status='interrupted',stage='stopped') write_json(campaign / 'state.json',state) return 0 with (campaign / 'api.log').open('a') as api_log: command = [sys.executable,str(ROOT / 'jev_harness.py'),'--backend','serve', '--checkpoint',str(model),'--port',str(plan['port'])] if plan.get('inference_max_tokens'): command.extend(['--max-tokens',str(plan['inference_max_tokens'])]) api = subprocess.Popen(command,cwd=ROOT,stdin=subprocess.DEVNULL, stdout=api_log,stderr=subprocess.STDOUT,start_new_session=True) state.update(api_pid=api.pid,api_command=command,model=str(model), status='complete',stage='deploying',url=f"http://127.0.0.1:{plan['port']}") write_json(campaign / 'state.json',state) import urllib.request ready = False for _ in range(12): if STOP: if api.poll() is None: api.terminate() try: api.wait(timeout=5) except subprocess.TimeoutExpired: api.kill() state.update(status='interrupted',stage='stopped',api_ready=False) write_json(campaign/'state.json',state) return 0 if api.poll() is not None: break try: headers = {} if os.environ.get('OPENSYSONE_API_KEY'): headers['Authorization'] = 'Bearer ' + os.environ['OPENSYSONE_API_KEY'] with urllib.request.urlopen(urllib.request.Request(state['url']+'/health',headers=headers),timeout=5) as response: health = json.load(response) ready = health['status']=='ready' and health.get('checkpoint')==str(model) if ready: break except Exception: time.sleep(5) state['stage'] = 'serving' if ready else 'api_start_failed' state['api_ready'] = ready write_json(campaign / 'state.json',state) if not ready: if api.poll() is None: api.terminate() raise RuntimeError('API did not become ready; finalized model is retained') return 0 except Exception as error: state.update(status='failed',error=str(error)) write_json(campaign / 'state.json',state) print(json.dumps(state),flush=True) return 1 def resume_correctness(parent): summary_path = parent / 'summary.json' initial_path = parent / 'correctness_initial.json' checks = None if summary_path.exists(): summary = json.loads(summary_path.read_text()) checks = summary.get('correctness') if (checks is None and summary.get('final_correctness_status') == 'skipped_on_stop' and not (parent / 'correctness_final.json').exists() and initial_path.exists()): checks = json.loads(initial_path.read_text()) elif initial_path.exists(): checks = json.loads(initial_path.read_text()) expected = {'branch_chunks_1_probability_max_abs', 'branch_chunks_2_probability_max_abs', 'branch_chunks_4_probability_max_abs', 'question_isolation_probability_max_abs', 'candidate_permutation_probability_max_abs', 'repeat_probability_max_abs'} if (not isinstance(checks,dict) or not expected.issubset(checks) or any(not isinstance(v,(int,float)) or not math.isfinite(v) or not 0 <= v <= 1e-4 for k,v in checks.items() if k.endswith('max_abs'))): raise ValueError('Pilot must pass all FP32 correctness gates') return checks def detach(args): parent = Path(args.pilot).resolve() checkpoint = parent if parent.suffix == '.pt' else parent / 'checkpoint.pt' if not checkpoint.exists() or not (checkpoint.parent / 'best.pt').exists(): raise ValueError('Pilot/resume run must have a checkpoint and best artifact') resume_correctness(checkpoint.parent) deadline = datetime.fromisoformat(args.deadline.replace('Z','+00:00')) if args.deadline else datetime.now(timezone.utc)+timedelta(hours=24) if deadline.timestamp() - time.time() > 24 * 3600 + 10: raise ValueError('Campaign deadline cannot exceed 24 hours from launch') name = datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%SZ') + '-24h' campaign = RUN_ROOT / name reserve, estimate = evaluation_budget(checkpoint.parent) requested = (datetime.fromisoformat(args.training_deadline.replace('Z','+00:00')).timestamp() if args.training_deadline else None) training_deadline = training_cutoff(deadline.timestamp(),reserve,requested) campaign.mkdir(parents=True,exist_ok=False) plan = {'resume_checkpoint':str(checkpoint),'resume_sha256':checksum(checkpoint), 'final_deadline':deadline.isoformat(),'eval_reserve_seconds':reserve,'evaluation_estimate':estimate,'port':args.port, 'training_deadline':datetime.fromtimestamp(training_deadline,timezone.utc).isoformat(), 'train_only':args.train_only, 'selection_metric':getattr(args,'selection_metric',SELECTION_METRIC), 'inference_max_tokens':args.inference_max_tokens, 'created_utc':datetime.now(timezone.utc).isoformat(), 'source_commit':subprocess.check_output(['git','rev-parse','HEAD'],cwd=ROOT,text=True).strip()} write_json(campaign / 'plan.json',plan) environment = os.environ.copy() environment.update(PYTHONUNBUFFERED='1',TOKENIZERS_PARALLELISM='false',HF_HUB_DISABLE_XET='1', PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True') with (campaign / 'supervisor.log').open('a') as log: process = subprocess.Popen([sys.executable,str(Path(__file__).resolve()),'--campaign',str(campaign)], cwd=ROOT,stdin=subprocess.DEVNULL,stdout=log,stderr=subprocess.STDOUT, env=environment,start_new_session=True) (RUN_ROOT / 'LAST_CAMPAIGN').write_text(str(campaign)+'\n') print(json.dumps({'campaign':str(campaign),'supervisor_pid':process.pid,'final_deadline':plan['final_deadline']})) def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--pilot') parser.add_argument('--deadline') parser.add_argument('--training-deadline',help='Optional earlier training cutoff; must preserve the measured evaluation reserve') parser.add_argument('--port',type=int,default=18081) parser.add_argument('--inference-max-tokens',type=int,default=1024) parser.add_argument('--selection-metric',choices=('raw_nll',SELECTION_METRIC),default=SELECTION_METRIC, help='Checkpoint selection on validation only; deployment calibration remains separate') parser.add_argument('--train-only',action='store_true', help='Stop with validation-selected artifacts before calibration/test/holdout or API deployment') parser.add_argument('--campaign',help=argparse.SUPPRESS) args = parser.parse_args() if args.campaign: code = run_campaign(args.campaign) (Path(args.campaign) / 'exit_code').write_text(str(code)+'\n') sys.exit(code) if not args.pilot: parser.error('--pilot is required') detach(args) if __name__ == '__main__': main()