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']<=28800:raise ValueError('budget exceeds eight 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=28800,margin=120):
33    if isinstance(budget,bool) or not isinstance(budget,(int,float)) or not math.isfinite(budget) or not 0<budget<=28800:raise ValueError('invalid eight-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
139
140if __name__=='__main__':
141    if sys.argv[1:]==['--guardian']:guardian_main()
142    else:raise SystemExit('Use reviewed curriculum owner; direct admission disabled')
143