Blame for src/nida/jobqueue.py

"""Sequential job queue for one Colab runtime (one Drive mount, many commands).

The queue is a text file (usually on Drive, so it can be edited while the queue runs):

* one ``nida.cli`` command per line, written exactly like the runner's ``COMMAND_ARGS``
  (without ``--runs-root/--run-id/--git-commit``; the queue adds them);
* ``@checkout <sha>``: fetch ``<bundle_dir>/nida-<sha>.bundle`` and check that commit
  out in the repository (the package is installed editable, so later jobs run it);
* ``@stop``: end the queue here;
* blank lines and lines starting with ``#`` are ignored.

The file is re-read before every job, so lines can be appended at any time. **Every line
must end with a newline**: the last line of a file without a trailing newline may be
half written (Drive sync, an editor mid-save), so it is ignored until the newline appears.
A read that fails (``OSError``, undecodable bytes, a line that does not parse) is retried
``read_retries`` times, ``poll_seconds`` apart, before the error is raised.

Lines must be unique after canonicalisation (``" ".join(shlex.split(line))``): a duplicate
is a ``QueueError``, so that deleting an earlier line can never shift a pending job onto
the key of a finished one. Each line is done at most once: its key is ``<canonical>#0``
and finished keys are stored next to the queue file in ``<queue>.state.json``, so a queue
restarted on a new runtime skips them. A restarted queue first re-establishes the code
state: the last successful ``@checkout`` in the state file is checked out again when HEAD
is something else (the notebook checked out its own ``GIT_REF``).

A job cut off by a dying runtime (exit code in ``INTERRUPTED_EXIT_CODES``: SIGTERM,
SIGINT, ``nida train``'s interrupted code) is not recorded and the queue stops with
``QueueSummary.interrupted``: it runs again on the next queue start, so give resumable
commands, e.g. ``train ... --resume``. Any other failed job is recorded and the queue
moves on; it is not retried. When nothing is pending the queue polls the file until
``wait_minutes`` have passed since the last job finished, then returns.
"""

from __future__ import annotations

import json
import logging
import re
import shlex
import shutil
import sys
import time
from collections.abc import Callable, Mapping, Sequence
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Literal

from nida.cli import EXIT_TRAIN_INTERRUPTED
from nida.runs import new_run_id
from nida.utils.io import atomic_write_json

__all__ = [
    "INTERRUPTED_EXIT_CODES",
    "Job",
    "JobOutcome",
    "QueueError",
    "QueueSummary",
    "child_command",
    "parse_queue",
    "run_queue",
    "state_path_for",
]

log = logging.getLogger(__name__)

STATE_SUFFIX = ".state.json"
_SHA_RE = re.compile(r"[0-9a-f]{40}")
_RESERVED_FLAGS = frozenset({"--runs-root", "--run-id", "--git-commit"})
_REQUIRED_STATE_KEYS = ("kind", "returncode", "git_commit")

#: Exit codes of a child that was stopped from outside (the runtime is going away): the
#: job is not recorded as done. ``EXIT_TRAIN_INTERRUPTED`` is ``nida train`` after a
#: SIGTERM + checkpoint; 130 / 143 are the shell codes of SIGINT / SIGTERM and -2 / -15
#: the same signals as reported by ``subprocess``.
INTERRUPTED_EXIT_CODES = frozenset({EXIT_TRAIN_INTERRUPTED, 130, -2, -15, 143})

#: ``runner(cmd, log_path) -> exit code``: runs ``cmd``, streaming its output to the log.
Runner = Callable[[Sequence[str], Path], int]
#: ``git(args, cwd) -> stdout``: runs git, raises on a non-zero exit.
Git = Callable[[Sequence[str], Path], str]


class QueueError(ValueError):
    """A malformed queue line (reported with its line number)."""


@dataclass(frozen=True)
class Job:
    """One queue line. ``key`` identifies it across re-reads of the file."""

    line_no: int
    text: str
    kind: Literal["run", "checkout", "stop"]
    args: tuple[str, ...]
    key: str


@dataclass(frozen=True)
class JobOutcome:
    key: str
    kind: str
    run_id: str | None
    returncode: int
    git_commit: str
    started_at: float
    finished_at: float


@dataclass
class QueueSummary:
    outcomes: list[JobOutcome] = field(default_factory=list)
    stopped: bool = False
    #: A child exited with a code in ``INTERRUPTED_EXIT_CODES``; the queue stopped and
    #: that job was not recorded (it runs again on the next start).
    interrupted: bool = False
    interrupted_key: str | None = None

    @property
    def failed(self) -> list[JobOutcome]:
        return [o for o in self.outcomes if o.returncode != 0]


def parse_queue(text: str) -> list[Job]:
    """Parse a queue file.

    Raises:
        QueueError: an unknown ``@`` directive, a bad ``@checkout`` SHA, a line that sets
            a flag the queue owns, a line that is not valid shell quoting, or a line that
            duplicates an earlier one (after canonicalisation).
    """
    jobs: list[Job] = []
    seen: dict[str, int] = {}  # canonical line -> line number of its first occurrence
    for line_no, raw in enumerate(text.splitlines(), start=1):
        line = raw.strip()
        if not line or line.startswith("#"):
            continue
        try:
            args = tuple(shlex.split(line))
        except ValueError as exc:
            raise QueueError(f"line {line_no}: {exc}") from exc
        kind: Literal["run", "checkout", "stop"]
        if args[0] == "@checkout":
            if len(args) != 2 or not _SHA_RE.fullmatch(args[1]):
                raise QueueError(f"line {line_no}: '@checkout' needs one 40-char lowercase SHA")
            kind = "checkout"
        elif args[0] == "@stop":
            if len(args) != 1:
                raise QueueError(f"line {line_no}: '@stop' takes no arguments")
            kind = "stop"
        elif args[0].startswith("@"):
            raise QueueError(f"line {line_no}: unknown directive {args[0]!r}")
        else:
            owned = sorted(
                a for a in args if a in _RESERVED_FLAGS or a.split("=", 1)[0] in _RESERVED_FLAGS
            )
            if owned:
                raise QueueError(f"line {line_no}: {owned} are set by the queue")
            kind = "run"
        canonical = " ".join(args)
        if canonical in seen:
            raise QueueError(
                f"line {line_no}: duplicate of line {seen[canonical]} ({canonical!r}); "
                "every queue line must be unique"
            )
        seen[canonical] = line_no
        # '#0' is kept from the days duplicates were allowed: existing state files use it.
        jobs.append(Job(line_no, canonical, kind, args, f"{canonical}#0"))
    return jobs


def child_command(
    args: Sequence[str], *, runs_root: Path, run_id: str, git_commit: str
) -> list[str]:
    """The ``python -m nida.cli`` command line for one queue job."""
    return [
        sys.executable,
        "-m",
        "nida.cli",
        *args,
        "--runs-root",
        str(runs_root),
        "--run-id",
        run_id,
        "--git-commit",
        git_commit,
    ]


def _load_state(path: Path) -> dict[str, dict[str, object]]:
    """Read the finished-job records.

    Raises:
        ValueError: the file is not a JSON object of job records (named with ``path``).
    """
    if not path.exists():
        return {}
    try:
        loaded = json.loads(path.read_text(encoding="utf-8"))
    except ValueError as exc:  # JSONDecodeError and UnicodeDecodeError
        raise ValueError(f"{path}: queue state is not valid JSON: {exc}") from exc
    if not isinstance(loaded, dict):
        raise ValueError(f"{path}: queue state must be a JSON object")
    state: dict[str, dict[str, object]] = {}
    for key, record in loaded.items():
        if not isinstance(record, dict):
            raise ValueError(f"{path}: record {key!r} must be a JSON object")
        missing = [k for k in _REQUIRED_STATE_KEYS if k not in record]
        if missing:
            raise ValueError(f"{path}: record {key!r} lacks {missing}")
        returncode = record["returncode"]
        if not isinstance(returncode, int) or isinstance(returncode, bool):
            raise ValueError(f"{path}: record {key!r}: returncode must be an integer")
        if not isinstance(record["kind"], str) or not isinstance(record["git_commit"], str):
            raise ValueError(f"{path}: record {key!r}: kind and git_commit must be strings")
        state[str(key)] = dict(record)
    return state


def _last_checkout_sha(state: Mapping[str, Mapping[str, object]]) -> str | None:
    """The commit of the most recently finished successful ``@checkout`` in ``state``.

    Ordered by ``finished_at``, not by key order: the state file is written with sorted
    keys, so its order is alphabetical (a restart once re-checked out an older SHA that
    sorted after the newer one).
    """
    sha: str | None = None
    latest = float("-inf")
    for done in state.values():
        if done["kind"] != "checkout" or done["returncode"] != 0:
            continue
        finished = done.get("finished_at")
        if not isinstance(finished, (int, float)) or isinstance(finished, bool):
            raise ValueError(f"checkout record for {done['git_commit']} has no finished_at")
        if finished >= latest:
            latest = float(finished)
            sha = str(done["git_commit"])
    return sha


def _read_queue(
    queue_path: Path,
    *,
    read_retries: int,
    poll_seconds: float,
    sleep: Callable[[float], None],
) -> list[Job]:
    """Read and parse the queue file, retrying transient failures.

    A final line without a newline is ignored: it may be half written.

    Raises:
        OSError / UnicodeDecodeError / QueueError: the last of ``read_retries`` attempts
            failed.
    """
    for attempt in range(1, read_retries + 1):
        try:
            text = queue_path.read_text(encoding="utf-8")
            return parse_queue(text[: text.rfind("\n") + 1])
        except (OSError, UnicodeDecodeError, QueueError) as exc:
            if attempt == read_retries:
                raise
            log.warning(
                "queue: reading %s failed (attempt %d/%d): %s",
                queue_path,
                attempt,
                read_retries,
                exc,
            )
            sleep(poll_seconds)
    raise AssertionError("unreachable: run_queue validates read_retries >= 1")


def _checkout(
    sha: str,
    repo_dir: Path,
    bundle_dir: Path,
    git: Git,
    *,
    wait_minutes: float,
    poll_seconds: float,
    clock: Callable[[], float],
    sleep: Callable[[float], None],
) -> None:
    """Check out ``sha`` from its bundle (idempotent), waiting for the bundle to sync."""
    bundle = bundle_dir / f"nida-{sha}.bundle"
    deadline = clock() + wait_minutes * 60
    while not bundle.is_file():
        if clock() >= deadline:
            raise FileNotFoundError(f"bundle not found after {wait_minutes:g} min: {bundle}")
        log.warning("queue: bundle %s not there yet; waiting", bundle)
        sleep(poll_seconds)
    git(["bundle", "verify", "--quiet", str(bundle)], repo_dir)
    git(["fetch", "--quiet", str(bundle), "HEAD"], repo_dir)
    git(["checkout", "--quiet", "--force", "--detach", sha], repo_dir)
    head = git(["rev-parse", "HEAD"], repo_dir).strip()
    if head != sha:
        raise RuntimeError(f"checkout mismatch: {head} != {sha}")


def state_path_for(queue_path: Path) -> Path:
    """Where the finished-job keys of ``queue_path`` are stored."""
    return queue_path.with_name(queue_path.name + STATE_SUFFIX)


def run_queue(
    queue_path: Path,
    *,
    runs_root: Path,
    git_commit: str,
    repo_dir: Path,
    bundle_dir: Path,
    log_dir: Path,
    runner: Runner,
    git: Git,
    wait_minutes: float = 0.0,
    poll_seconds: float = 60.0,
    read_retries: int = 10,
    checkout_wait_minutes: float = 10.0,
    clock: Callable[[], float] = time.time,
    sleep: Callable[[float], None] = time.sleep,
    make_run_id: Callable[[str], str] = new_run_id,
) -> QueueSummary:
    """Run the queue until it is empty (after waiting), ``@stop`` or an interruption.

    Args:
        git_commit: the commit the repository is at when the queue starts (the notebook's
            ``GIT_REF``); used when the state file has no ``@checkout``.
        read_retries: attempts to read and parse the queue file (>= 1).
        checkout_wait_minutes: how long a missing ``@checkout`` bundle is waited for.

    Returns:
        What ran this time (jobs finished by an earlier queue run are skipped silently,
        they are in the state file). ``interrupted`` is set when a child was stopped from
        outside; that job is not recorded.

    Raises:
        ValueError: ``read_retries < 1`` or a malformed state file.
        QueueError: the file has a malformed or duplicate line (checked before every job,
            after ``read_retries`` attempts).
        FileNotFoundError / RuntimeError: a ``@checkout`` failed, or the repository is not
            at ``git_commit`` although the state has no checkout; later jobs would run the
            wrong code, so the queue stops.
    """
    if read_retries < 1:
        raise ValueError(f"read_retries must be >= 1, got {read_retries}")
    state_path = state_path_for(queue_path)
    state = _load_state(state_path)
    summary = QueueSummary()

    def checkout(sha: str) -> None:
        _checkout(
            sha,
            repo_dir,
            bundle_dir,
            git,
            wait_minutes=checkout_wait_minutes,
            poll_seconds=poll_seconds,
            clock=clock,
            sleep=sleep,
        )

    commit = git_commit
    head = git(["rev-parse", "HEAD"], repo_dir).strip()
    last_sha = _last_checkout_sha(state)
    if last_sha is None:
        if head != git_commit:
            raise RuntimeError(
                f"repository HEAD {head} is not git_commit {git_commit} and the queue state "
                "has no @checkout to restore: the caller passed the wrong commit"
            )
    else:
        commit = last_sha  # a restarted queue continues on the last checkout
        if head != last_sha:
            log.info("queue: restart: HEAD %s, re-checking out %s", head, last_sha)
            checkout(last_sha)
    last_finished = clock()
    while True:
        jobs = _read_queue(
            queue_path, read_retries=read_retries, poll_seconds=poll_seconds, sleep=sleep
        )
        pending = [j for j in jobs if j.key not in state]
        if not pending:
            if clock() - last_finished >= wait_minutes * 60:
                return summary
            sleep(poll_seconds)
            continue
        job = pending[0]
        started = clock()
        run_id: str | None = None
        if job.kind == "stop":
            code = 0
        elif job.kind == "checkout":
            log.info("queue: checkout %s", job.args[1])
            checkout(job.args[1])
            commit = job.args[1]
            code = 0
        else:
            run_id = make_run_id(job.args[0])
            cmd = child_command(job.args, runs_root=runs_root, run_id=run_id, git_commit=commit)
            log_path = log_dir / f"{run_id}.log"
            log_dir.mkdir(parents=True, exist_ok=True)
            log.info("queue: line %d -> %s", job.line_no, run_id)
            code = runner(cmd, log_path)
            if (runs_root / run_id).is_dir() and log_path.is_file():
                try:
                    shutil.copyfile(log_path, runs_root / run_id / "log.txt")
                except OSError as exc:
                    log.warning("queue: could not copy %s into the run dir: %s", log_path, exc)
            log.info("queue: %s exited with %d", run_id, code)
            if code in INTERRUPTED_EXIT_CODES:
                log.warning(
                    "queue: %s was interrupted (exit %d); not recording it, stopping the queue",
                    run_id,
                    code,
                )
                summary.interrupted = True
                summary.interrupted_key = job.key
                return summary
        last_finished = clock()
        outcome = JobOutcome(job.key, job.kind, run_id, code, commit, started, last_finished)
        state[job.key] = asdict(outcome)
        atomic_write_json(state_path, state)
        summary.outcomes.append(outcome)
        if job.kind == "stop":
            summary.stopped = True
            return summary


def stream_to_log(cmd: Sequence[str], log_path: Path) -> int:
    """Default runner: run ``cmd``, echo its output and write it to ``log_path``."""
    import subprocess

    with (
        log_path.open("w", encoding="utf-8") as out,
        subprocess.Popen(
            list(cmd), stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1
        ) as proc,
    ):
        assert proc.stdout is not None  # stdout=PIPE guarantees a stream
        for line in proc.stdout:
            print(line, end="", flush=True)
            out.write(line)
            out.flush()
        return proc.wait()


def run_git(args: Sequence[str], cwd: Path) -> str:
    """Default git: raise with the output on a non-zero exit."""
    import subprocess

    proc = subprocess.run(["git", *args], cwd=cwd, capture_output=True, text=True, check=False)
    if proc.returncode != 0:
        raise RuntimeError(f"git {' '.join(args[:2])} failed: {proc.stdout}{proc.stderr}")
    return proc.stdout


def summary_markdown(summary: QueueSummary, state: Mapping[str, Mapping[str, object]]) -> str:
    """``report.md`` body: every job ever finished by this queue run dir."""
    lines = ["# queue", "", "| job | run_id | exit | commit |", "|---|---|---|---|"]
    for key, done in state.items():
        lines.append(
            f"| `{key}` | {done.get('run_id') or '-'} | {done.get('returncode')} | "
            f"{str(done.get('git_commit'))[:12]} |"
        )
    lines += ["", f"this session: {len(summary.outcomes)} jobs, {len(summary.failed)} failed"]
    if summary.stopped:
        lines.append("stopped by @stop")
    if summary.interrupted:
        lines.append(
            f"INTERRUPTED: `{summary.interrupted_key}` was stopped from outside and is not "
            "recorded; it runs again when the queue is restarted"
        )
    return "\n".join(lines) + "\n"