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