Download src/nida/jobqueue.py from Bilal7Dev/nida-tts: direct link, hf CLI and curl.
- Browser
- Download file 18.3 kB
-
https://huggingface.co/Bilal7Dev/nida-tts/resolve/main/src/nida/jobqueue.py
- Command line
-
hf download hf://Bilal7Dev/nida-tts/src/nida/jobqueue.py
-
curl -L -o jobqueue.py https://huggingface.co/Bilal7Dev/nida-tts/resolve/main/src/nida/jobqueue.py
18.3 kB
| """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).""" | |
| 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 | |
| class JobOutcome: | |
| key: str | |
| kind: str | |
| run_id: str | None | |
| returncode: int | |
| git_commit: str | |
| started_at: float | |
| finished_at: float | |
| 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 | |
| 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" | |