| """Progress lines for every long Phase 11 job, so a slow job is never mistaken for a stuck one. |
| |
| Every line is one JSON object on stderr, which the launcher writes to the run's `log.txt` |
| and `make status` shows. Three shapes: `stage_start`/`stage_end` around a named piece of |
| work, `progress` through a count of units with elapsed time and an estimate of what is left, |
| and `heartbeat` from a background thread every minute, with GPU memory and the job's own |
| status, for the stretches where no unit completes. |
| """ |
| from __future__ import annotations |
| import json |
| import sys |
| import threading |
| import time |
| from collections.abc import Callable, Iterator |
| from contextlib import contextmanager |
| from types import TracebackType |
| from typing import Final, Self |
| PROGRESS_INTERVAL_S: Final = 30.0 |
| HEARTBEAT_S: Final = 60.0 |
|
|
| def progress(event: str, **values: object) -> None: |
| """One JSON line on stderr.""" |
| line = json.dumps({'event': event, 'time': round(time.time(), 1), **values}, default=str) |
| sys.stderr.write(line + '\n') |
| sys.stderr.flush() |
|
|
| @contextmanager |
| def stage(name: str, **values: object) -> Iterator[None]: |
| """`stage_start` and `stage_end` with the seconds between; `stage_failed` if it raises.""" |
| started = time.perf_counter() |
| progress('stage_start', stage=name, **values) |
| try: |
| yield |
| except BaseException as error: |
| elapsed = round(time.perf_counter() - started, 1) |
| progress('stage_failed', stage=name, seconds=elapsed, error=repr(error)[:300]) |
| raise |
| progress('stage_end', stage=name, seconds=round(time.perf_counter() - started, 1)) |
|
|
| class Tracker: |
| """Progress through `total` units, or an open count when `total` is None: a line on the |
| first unit, at most every `interval_s` seconds after it, and on the last.""" |
|
|
| def __init__(self, task: str, total: int | None, interval_s: float=PROGRESS_INTERVAL_S) -> None: |
| self.task = task |
| self.total = total |
| self.interval_s = interval_s |
| self.done = 0 |
| self.started = time.perf_counter() |
| self.last: float | None = None |
|
|
| def advance(self, units: int=1, **values: object) -> None: |
| self.done += units |
| now = time.perf_counter() |
| finished = self.total is not None and self.done >= self.total |
| if self.last is not None and now - self.last < self.interval_s and (not finished): |
| return |
| self.last = now |
| elapsed = now - self.started |
| remaining = elapsed / self.done * (self.total - self.done) if self.total is not None and self.done else None |
| progress('progress', task=self.task, done=self.done, total=self.total, elapsed_s=round(elapsed, 1), eta_s=round(remaining, 1) if remaining is not None else None, **values) |
|
|
| class Heartbeat: |
| """A line every `interval_s` seconds from a background thread while a block runs.""" |
|
|
| def __init__(self, label: str, status: Callable[[], dict[str, object]] | None=None, interval_s: float=HEARTBEAT_S) -> None: |
| self.label = label |
| self.status = status |
| self.interval_s = interval_s |
| self.stopped = threading.Event() |
| self.thread = threading.Thread(target=self.beat, daemon=True) |
| self.started = time.perf_counter() |
|
|
| def beat(self) -> None: |
| while not self.stopped.wait(self.interval_s): |
| extra = self.status() if self.status is not None else {} |
| progress('heartbeat', label=self.label, elapsed_s=round(time.perf_counter() - self.started, 1), gpu_allocated_gib=gpu_allocated_gib(), **extra) |
|
|
| def __enter__(self) -> Self: |
| self.thread.start() |
| return self |
|
|
| def __exit__(self, kind: type[BaseException] | None, error: BaseException | None, trace: TracebackType | None) -> None: |
| self.stopped.set() |
| self.thread.join() |
|
|
| def gpu_allocated_gib() -> float | None: |
| """GPU memory this process holds, or None where torch or CUDA is absent.""" |
| if 'torch' not in sys.modules: |
| return None |
| import torch |
| if not torch.cuda.is_available(): |
| return None |
| return round(torch.cuda.memory_allocated() / 2 ** 30, 2) |
|
|
