File size: 3,325 Bytes
0662d8e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | """Run upstream `read` to a CHARACTER target, resuming after crashes.
python harness/supervise.py --target-chars N | --match runs/other/weights
-- <run_upstream.py args...>
Polls the manifest (written at every evaluation). At the target it sends
SIGTERM, which upstream handles by finishing the chunk and checkpointing. If
the process dies before the target, it is restarted and resumes from the last
checkpoint - upstream's normal resume path. With --match, the target is the
other run's read_chars once that run has stopped (until then, keep reading).
Every session start/stop is appended to <weights>/../supervise.log.
"""
import argparse
import json
import os
import signal
import subprocess
import sys
import time
def read_chars(wdir):
try:
with open(os.path.join(wdir, "manifest.json")) as f:
return int(json.load(f).get("read_chars") or 0)
except (OSError, ValueError):
return 0
def running(pattern):
r = subprocess.run(["pgrep", "-f", pattern], capture_output=True, text=True)
return bool(r.stdout.strip())
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--target-chars", type=int, default=0)
ap.add_argument("--match", default=None, help="weights dir of the run to match")
ap.add_argument("--match-pattern", default=None, help="pgrep pattern: other run still going")
ap.add_argument("--max-restarts", type=int, default=5)
ap.add_argument("rest", nargs=argparse.REMAINDER)
a = ap.parse_args()
cmd = a.rest[1:] if a.rest and a.rest[0] == "--" else a.rest
wdir = cmd[cmd.index("--weights-dir") + 1]
log = open(os.path.join(os.path.dirname(wdir.rstrip("/")), "supervise.log"), "a")
out = open(os.path.join(os.path.dirname(wdir.rstrip("/")), "train.log"), "a")
def note(s):
log.write(f"{time.strftime('%F %T')} {s}\n")
log.flush()
def target():
if a.target_chars:
return a.target_chars
if a.match_pattern and running(a.match_pattern):
return None
return read_chars(a.match)
restarts = 0
while True:
t = target()
if t is not None and read_chars(wdir) >= t:
note(f"target {t:,} reached ({read_chars(wdir):,})")
return 0
note(f"session start at {read_chars(wdir):,} chars")
p = subprocess.Popen([sys.executable, "harness/run_upstream.py"] + cmd,
stdout=out, stderr=subprocess.STDOUT)
stopped = False
while p.poll() is None:
time.sleep(10)
t = target()
if not stopped and t is not None and read_chars(wdir) >= t:
note(f"target {t:,} reached - SIGTERM")
p.send_signal(signal.SIGTERM)
stopped = True
note(f"session exit {p.returncode} at {read_chars(wdir):,} chars")
if stopped:
return 0
if p.returncode == 0 and t is None:
continue # minutes ran out before the target was known
if p.returncode != 0:
restarts += 1
if restarts > a.max_restarts:
note("too many restarts")
return 1
elif t is not None and read_chars(wdir) >= t:
return 0
if __name__ == "__main__":
sys.exit(main())
|