training_deadline.py

1"""Persistent total-job deadline and Linux owned-process-group watchdog.
2
3The deadline starts before Popen (including imports, data setup and checkpointing).
4Resume reuses the original file; wall-clock rollback cannot extend this invocation
5because its watchdog also uses a monotonic deadline. No automatic restart.
6"""
7import argparse
8import hashlib
9import json
10import math
11import os
12from pathlib import Path
13import signal
14import subprocess
15import sys
16import time
17import uuid
18
19
20def read_deadline(path):
21    value=json.loads(Path(path).read_text())
22    if value.get('format')!='chess-total-deadline-v1':raise ValueError('deadline format')
23    for key in ('started_unix','deadline_unix','budget_seconds','checkpoint_margin_seconds'):
24        if isinstance(value.get(key),bool) or not isinstance(value.get(key),(int,float)) or not math.isfinite(value[key]):
25            raise ValueError('invalid deadline number')
26    if not 0<value['budget_seconds']<=14400:raise ValueError('budget exceeds four hours')
27    if not 0<value['checkpoint_margin_seconds']<value['budget_seconds']:raise ValueError('invalid checkpoint margin')
28    if abs(value['deadline_unix']-value['started_unix']-value['budget_seconds'])>1e-5:raise ValueError('inconsistent deadline')
29    return value
30
31
32def create_or_resume(path,identity,budget=14400,margin=120):
33    if isinstance(budget,bool) or not isinstance(budget,(int,float)) or not math.isfinite(budget) or not 0<budget<=14400:raise ValueError('invalid four-hour budget')
34    if isinstance(margin,bool) or not isinstance(margin,(int,float)) or not math.isfinite(margin) or not 0<margin<budget:raise ValueError('invalid checkpoint margin')
35    path=Path(path)
36    if path.exists():
37        value=read_deadline(path)
38        if value['identity']!=identity or value['budget_seconds']!=budget or value['checkpoint_margin_seconds']!=margin:
39            raise ValueError('deadline identity/config mismatch; never extend an existing job')
40    else:
41        now=time.time()
42        value={'format':'chess-total-deadline-v1','started_unix':now,'deadline_unix':now+budget,
43               'budget_seconds':budget,'checkpoint_margin_seconds':margin,'identity':identity}
44        # Exclusive creation prevents concurrent supervisors resetting the budget.
45        with path.open('x') as f:json.dump(value,f,sort_keys=True);f.flush();os.fsync(f.fileno())
46        value=read_deadline(path)
47    if time.time()>=value['deadline_unix']:raise TimeoutError('original job deadline expired')
48    return value
49
50
51class Deadline:
52    def __init__(self,value):
53        self.value=value
54        self.monotonic_end=time.monotonic()+max(0,value['deadline_unix']-time.time())
55    def remaining(self):
56        return min(self.value['deadline_unix']-time.time(),self.monotonic_end-time.monotonic())
57    def soft_stop(self):return self.remaining()<=self.value['checkpoint_margin_seconds']
58
59
60def _guard_job(command,value,env=None,owner_pid=None):
61    """Linux only. TERM at 5s remaining, KILL at deadline; clean own group on failure.
62
63    OS scheduling can delay signal delivery; this is not a real-time guarantee.
64    Group containment covers torchrun workers; children must not detach sessions.
65    """
66    if os.name!='posix':raise RuntimeError('watchdog requires Linux process groups')
67    clock=Deadline(value)
68    if clock.soft_stop():raise TimeoutError('insufficient original budget to start/resume')
69    process=None;term=False;expired=False
70    try:
71        if owner_pid is not None and os.getppid()!=owner_pid:raise RuntimeError('supervisor died before launch')
72        process=subprocess.Popen(command,start_new_session=True,env=env,stdout=sys.stderr)
73        while True:
74            if owner_pid is not None and os.getppid()!=owner_pid:
75                return {'returncode':125,'deadline_expired':False,'supervisor_lost':True}
76            remaining=clock.remaining()
77            if remaining<=0:
78                expired=True;break
79            result=process.poll()
80            if result is not None:return {'returncode':result,'deadline_expired':False}
81            if remaining<=5 and not term:
82                os.killpg(process.pid,signal.SIGTERM);term=True
83            time.sleep(min(.1,max(.001,remaining)))
84    finally:
85        if process is not None:
86            # Also remove orphaned workers after leader failure/success.
87            try:os.killpg(process.pid,signal.SIGKILL)
88            except ProcessLookupError:pass
89            process.wait()
90    return {'returncode':process.returncode,'deadline_expired':expired}
91
92
93def supervise(command,value,env=None):
94    """A separate guardian owns/reaps workers and survives supervisor SIGKILL.
95
96    It also notices loss of its supervisor and immediately kills its owned group.
97    The guardian has a separate session; no global job/process enumeration occurs.
98    """
99    if os.name!='posix':raise RuntimeError('watchdog requires Linux process groups')
100    if Deadline(value).soft_stop():raise TimeoutError('insufficient original budget to start/resume')
101    lock_fd=int(env['CHESS_JOB_LOCK_FD']) if env and 'CHESS_JOB_LOCK_FD' in env else None
102    guardian=subprocess.Popen([sys.executable,str(Path(__file__).resolve()),'--guardian'],
103        stdin=subprocess.PIPE,stdout=subprocess.PIPE,text=True,start_new_session=True,
104        pass_fds=() if lock_fd is None else (lock_fd,))
105    try:
106        output,_=guardian.communicate(json.dumps({'command':command,'value':value,'env':env,'owner_pid':os.getpid()}))
107        if guardian.returncode:raise RuntimeError('deadline guardian failed: '+output)
108        result=json.loads(output)
109        if result.get('error_type')=='FileNotFoundError':raise FileNotFoundError(result['error'])
110        if 'error' in result:raise RuntimeError(result['error'])
111        return result
112    finally:
113        if guardian.poll() is None:
114            guardian.terminate()
115            guardian.wait(timeout=5)
116
117
118def guardian_main():
119    # Linux subreaper: adopt and reap orphaned torchrun workers after killing the
120    # owned group. PR_SET_CHILD_SUBREAPER needs no elevated privilege.
121    import ctypes
122    if ctypes.CDLL(None,use_errno=True).prctl(36,1,0,0,0)!=0:raise OSError(ctypes.get_errno(),'subreaper unavailable')
123    def terminate(signum,frame):raise SystemExit(128+signum)
124    signal.signal(signal.SIGTERM,terminate);signal.signal(signal.SIGINT,terminate)
125    request=json.load(sys.stdin)
126    try:
127        try:result=_guard_job(**request)
128        except Exception as error:result={'error':repr(error),'error_type':type(error).__name__}
129    finally:
130        # All owned processes have already been killed; only reap dead children.
131        until=time.monotonic()+2
132        while time.monotonic()<until:
133            try:pid,_=os.waitpid(-1,os.WNOHANG)
134            except ChildProcessError:break
135            if pid==0:time.sleep(.01)
136    if os.getppid()==request['owner_pid']:print(json.dumps(result),flush=True)
137
138
139def main():
140    p=argparse.ArgumentParser();p.add_argument('--deadline',required=True);p.add_argument('--approval',required=True)
141    p.add_argument('--resume-checkpoint')
142    p.add_argument('--dataset-manifest',required=True);p.add_argument('command',nargs=argparse.REMAINDER)
143    args=p.parse_args()
144    root=Path(__file__).parent
145    identity={'dataset_manifest_sha256':hashlib.sha256(Path(args.dataset_manifest).read_bytes()).hexdigest(),
146              'code_sha256':{name:hashlib.sha256((root/name).read_bytes()).hexdigest() for name in
147                             ('training_deadline.py','training_core.py','train_chess.py','verified_training_data.py')}}
148    if not args.command:raise ValueError('missing command')
149    approval=json.loads(Path(args.approval).read_text())
150    run_id=approval.get('run_id')
151    try:uuid.UUID(run_id)
152    except (ValueError,TypeError,AttributeError):raise ValueError('approval must name immutable run UUID')
153    deadline_path=str(Path(args.deadline).resolve())
154    resume_path=approval.get('resume_checkpoint_path')
155    if resume_path is not None and (not isinstance(resume_path,str) or str(Path(resume_path).resolve())!=resume_path):
156        raise ValueError('approved resume checkpoint must be canonical absolute path')
157    if approval!={'approved':True,'identity':identity,'budget_seconds':14400,'command':args.command,
158                  'run_id':run_id,'deadline_path':deadline_path,'resume_checkpoint_path':resume_path}:raise ValueError('missing/mismatched review release or alternate deadline path')
159    command=args.command
160    if args.resume_checkpoint:
161        if resume_path is None or str(Path(args.resume_checkpoint).resolve())!=resume_path:raise ValueError('unapproved resume checkpoint')
162        if not Path(resume_path).is_file():raise ValueError('missing resume checkpoint')
163        command=[*command,'--resume',resume_path]
164    identity={**identity,'run_id':run_id,'deadline_path':deadline_path,'command':args.command,'resume_checkpoint_path':resume_path}
165    # Lock held for entire supervisor lifetime, including deadline-file creation.
166    import fcntl
167    with open(deadline_path+'.lock','a') as lock:
168        fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB)
169        if args.resume_checkpoint and not Path(deadline_path).is_file():
170            raise ValueError('explicit resume requires the original deadline file; no fresh budget')
171        value=create_or_resume(deadline_path,identity)
172        env=dict(os.environ,CHESS_JOB_DEADLINE=deadline_path,CHESS_JOB_LOCK_FD=str(lock.fileno()))
173        result=supervise(command,value,env)
174        Path(deadline_path+'.result.json').write_text(json.dumps(result,indent=2))
175        raise SystemExit(result['returncode'] or (124 if result['deadline_expired'] else 0))
176
177
178if __name__=='__main__':
179    if sys.argv[1:]==['--guardian']:guardian_main()
180    else:main()
181