source/scripts/launch_24h.py

"""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()