code/harness/supervise.py

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