trainmeter 0.0.2__py3-none-any.whl

trainmeter/supervisor/launcher.py

@@ -0,0 +1,90 @@
+"""Run a command as a child process and stay out of its way.
+
+Contract: the child inherits stdio and is not touched; its exit status is returned as a shell
+would report it (128+N when killed by signal N); SIGTERM, SIGHUP and SIGQUIT sent to `tm` alone
+are forwarded to the child; Ctrl-C reaches both processes through the terminal, so SIGINT is not
+forwarded and `tm` never exits before the child.
+"""
+
+from __future__ import annotations
+
+import signal
+import subprocess
+import time
+import warnings
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from typing import Any
+
+FORWARDED = ("SIGTERM", "SIGHUP", "SIGQUIT")
+
+
+@dataclass(frozen=True)
+class Outcome:
+    status: int
+    signal: int | None
+    elapsed_s: float
+    pid: int | None = None
+    error: str | None = None
+
+
+def exit_status(returncode: int) -> tuple[int, int | None]:
+    """(shell-style status, killing signal or None) for a `Popen.returncode`."""
+    if returncode < 0:
+        return 128 - returncode, -returncode
+    return returncode, None
+
+
+def launch(
+    argv: list[str],
+    *,
+    env: Mapping[str, str] | None = None,
+    on_spawn: Callable[[int], None] | None = None,
+    clock: Callable[[], float] = time.monotonic,
+) -> Outcome:
+    state: dict[str, Any] = {"proc": None, "pending": []}
+
+    def forward(signum: int, _frame: Any) -> None:
+        proc = state["proc"]
+        if proc is None:
+            state["pending"].append(signum)
+            return
+        try:
+            proc.send_signal(signum)
+        except OSError:
+            pass
+
+    def ignore(_signum: int, _frame: Any) -> None:
+        # A Python handler, not SIG_IGN: an ignored signal would be inherited by the child.
+        return None
+
+    previous: dict[int, Any] = {}
+    handlers = {getattr(signal, n): forward for n in FORWARDED if hasattr(signal, n)}
+    handlers[signal.SIGINT] = ignore
+    try:
+        for signum, handler in handlers.items():
+            previous[signum] = signal.signal(signum, handler)
+    except ValueError:  # not the main thread: no handlers, plain child
+        previous.clear()
+
+    start = clock()
+    try:
+        try:
+            proc = subprocess.Popen(argv, env=None if env is None else dict(env))
+        except FileNotFoundError as exc:
+            return Outcome(127, None, clock() - start, error=str(exc))
+        except OSError as exc:
+            return Outcome(126, None, clock() - start, error=str(exc))
+        state["proc"] = proc
+        for signum in state["pending"]:
+            forward(signum, None)
+        if on_spawn is not None:
+            try:
+                on_spawn(proc.pid)
+            except Exception as exc:  # noqa: BLE001 - monitoring must not touch the child
+                warnings.warn(f"trainmeter monitoring failed to start: {exc!r}", stacklevel=2)
+        status, killed_by = exit_status(proc.wait())
+        return Outcome(status, killed_by, clock() - start, pid=proc.pid)
+    finally:
+        for signum, handler in previous.items():
+            signal.signal(signum, handler)

trainmeter/supervisor/live.py

@@ -0,0 +1,133 @@
+"""Live state behind the dashboard: the engine and the timeline, fed under one lock.
+
+Records come from the ingest thread and the samplers; the web threads only read. A ticker takes
+one snapshot per second, samples the derived metrics into the timeline and keeps the snapshot
+for readers, so a request never waits on the engine.
+"""
+
+from __future__ import annotations
+
+import threading
+import time
+import warnings
+from collections.abc import Callable
+from typing import Any
+
+from ..engine import Engine
+from ..timeline import Timeline
+
+
+class LiveState:
+    def __init__(
+        self,
+        engine: Engine,
+        header: dict[str, Any] | None = None,
+        mono: Callable[[], float] = time.monotonic,
+    ) -> None:
+        self.engine = engine
+        self.timeline = Timeline()
+        self.header = header or {}
+        self._mono = mono
+        self.t0 = mono()
+        self.lock = threading.Lock()
+        self.running = True
+        self.exit: dict[str, Any] | None = None
+        self.replay: str | None = None  # the run directory name when this is a replay
+        self._latest: dict[str, Any] = self._envelope({})
+        self._stop = threading.Event()
+        self._thread = threading.Thread(target=self._run, name="trainmeter-live", daemon=True)
+
+    @classmethod
+    def from_replay(cls, replayed: Any) -> LiveState:
+        """A finished run, served as it ended: no ticker, no clock, the log's own numbers."""
+        live = cls(replayed.engine, replayed.header, mono=lambda: replayed.t_end)
+        live.t0 = replayed.t_start
+        live.timeline = replayed.timeline
+        live.replay = replayed.path.name
+        live.running = False
+        live.exit = replayed.exit
+        live._latest = live._envelope(replayed.snapshot)
+        return live
+
+    def feed(self, record: dict[str, Any]) -> None:
+        with self.lock:
+            self.engine.feed(record)
+            self.timeline.feed(record, record["t"] - self.t0, self.engine.position())
+
+    def _envelope(self, snap: dict[str, Any]) -> dict[str, Any]:
+        elapsed = self._mono() - self.t0
+        out: dict[str, Any] = {
+            "snapshot": snap,
+            "elapsed_s": elapsed,
+            "running": self.running,
+            "exit": self.exit,
+            "replay": self.replay,
+        }
+        n_gpus = (snap.get("n_gpus") or {}).get("value")
+        if n_gpus:
+            out["gpu_hours"] = elapsed * n_gpus / 3600  # view-time; the price stays in the browser
+        return out
+
+    def tick(self) -> None:
+        """One snapshot: kept for readers and sampled into the timeline. Never raises."""
+        try:
+            with self.lock:
+                snap = self.engine.snapshot(self._mono())
+                t = self._mono() - self.t0
+                self.timeline.add_snapshot(snap, t, self.engine.position())
+                self._latest = self._envelope(snap)
+        except Exception as exc:  # noqa: BLE001 - the dashboard must not cost a run
+            warnings.warn(f"trainmeter: snapshot failed: {exc!r}", stacklevel=2)
+
+    def latest(self) -> dict[str, Any]:
+        with self.lock:
+            return self._latest
+
+    def facts(self) -> dict[str, Any]:
+        with self.lock:
+            return {
+                k: {"value": f.value, "source": f.source, "note": f.note}
+                for k, f in self.engine.facts.items().items()
+            }
+
+    def passport(self) -> dict[str, Any]:
+        """The run passport of the latest snapshot, plus its Markdown."""
+        from .. import passport
+
+        with self.lock:
+            snap = self._latest["snapshot"]
+            data = passport.build(
+                snap, self.header, self.facts_unlocked(), self.engine.hello, self.exit
+            )
+        return {"passport": data, "markdown": passport.to_markdown(data)}
+
+    def facts_unlocked(self) -> dict[str, Any]:
+        return {
+            k: {"value": f.value, "source": f.source, "note": f.note}
+            for k, f in self.engine.facts.items().items()
+        }
+
+    def series(self, names: list[str], x: str, since: float | None) -> dict[str, Any]:
+        with self.lock:
+            return self.timeline.query(names, x, since)
+
+    def series_names(self) -> list[str]:
+        with self.lock:
+            return self.timeline.names()
+
+    def finish(self, exit_info: dict[str, Any]) -> None:
+        self.running = False
+        self.exit = exit_info
+        self.tick()
+
+    def start(self) -> None:
+        self._thread.start()
+
+    def stop(self) -> None:
+        self._stop.set()
+        if self._thread.is_alive():
+            self._thread.join(timeout=2.0)
+
+    def _run(self) -> None:
+        while not self._stop.wait(1.0):
+            self.tick()

trainmeter/supervisor/store.py

@@ -0,0 +1,118 @@
+"""The run directory and its single writer.
+
+`events.jsonl` is appended by the supervisor only. Every failure here is reported once and
+then swallowed: a full disk must never cost a training run.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import threading
+import time
+import warnings
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any, TextIO
+
+from ..records import encode, make_record
+
+RUNS_DIR = Path(".trainmeter") / "runs"
+
+
+def _safe_name(name: str) -> str:
+    return re.sub(r"[^A-Za-z0-9._-]+", "-", name).strip("-.") or "run"
+
+
+class RunStore:
+    def __init__(
+        self,
+        path: Path,
+        mono: Callable[[], float] = time.monotonic,
+        wall: Callable[[], float] = time.time,
+    ) -> None:
+        self.path = path
+        self.run_id = path.name
+        self._mono = mono
+        self.mono0 = mono()
+        self.wall0 = wall()
+        self._lock = threading.Lock()
+        self._warned = False
+        self._file: TextIO | None = open(path / "events.jsonl", "a", encoding="utf-8")  # noqa: SIM115
+
+    @classmethod
+    def create(
+        cls,
+        base: Path,
+        name: str,
+        mono: Callable[[], float] = time.monotonic,
+        wall: Callable[[], float] = time.time,
+    ) -> RunStore:
+        """Create `<base>/.trainmeter/runs/<UTC timestamp>-<name>/`, adding a suffix on a clash."""
+        stamp = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime(wall()))
+        runs = Path(base) / RUNS_DIR
+        runs.mkdir(parents=True, exist_ok=True)
+        stem = f"{stamp}-{_safe_name(name)}"
+        for n in range(1, 1000):
+            path = runs / (stem if n == 1 else f"{stem}-{n}")
+            try:
+                path.mkdir()
+            except FileExistsError:
+                continue
+            return cls(path, mono, wall)
+        raise FileExistsError(f"no free run directory for {stem}")
+
+    def _warn_once(self, message: str) -> None:
+        if not self._warned:
+            self._warned = True
+            warnings.warn(f"trainmeter run log: {message}", stacklevel=3)
+
+    def emit(
+        self, kind: str, src: str, d: dict[str, Any] | None = None, *, pid: int | None = None
+    ) -> None:
+        """Append one record. Never raises."""
+        try:
+            line = encode(make_record(kind, src, self._mono(), d, pid=pid))
+        except (TypeError, ValueError) as exc:
+            self._warn_once(f"dropped a record that is not valid JSON: {exc}")
+            return
+        self._write_line(line)
+
+    def _write_line(self, line: str) -> None:
+        with self._lock:
+            if self._file is None:
+                return
+            try:
+                self._file.write(line + "\n")
+                self._file.flush()
+            except OSError as exc:
+                self._warn_once(f"stopped writing: {exc}")
+                self._file = None
+
+    def append_record(self, record: dict[str, Any]) -> None:
+        """Append a record that already has its envelope, such as one received from an agent."""
+        try:
+            line = encode(record)
+        except (TypeError, ValueError) as exc:
+            self._warn_once(f"dropped a record that is not valid JSON: {exc}")
+            return
+        self._write_line(line)
+
+    def write_json(self, name: str, obj: dict[str, Any]) -> None:
+        """Write `name` atomically. Never raises."""
+        try:
+            tmp = self.path / f".{name}.tmp"
+            tmp.write_text(json.dumps(obj, indent=2) + "\n", encoding="utf-8")
+            os.replace(tmp, self.path / name)
+        except (OSError, TypeError, ValueError) as exc:
+            self._warn_once(f"could not write {name}: {exc}")
+
+    def close(self) -> None:
+        with self._lock:
+            if self._file is not None:
+                try:
+                    self._file.close()
+                except OSError:
+                    pass
+                self._file = None

trainmeter/timeline.py

@@ -0,0 +1,102 @@
+"""Chart history in memory: named series of points that can be read against any x axis.
+
+Pure. It is fed from the same records as the engine, so a replay of a run log rebuilds it.
+The x axes are `step`, `tokens`, `flops` (modeled FLOPs spent) and `time` (seconds since the
+run started). A point whose x is unknown, such as `flops` before a FLOPs per token exists, is
+left out of that axis instead of being placed at zero.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from . import catalog
+
+AXES = ("step", "tokens", "flops", "time")
+_LADDER = ("gpu_util", "sm_active", "tensor_active")
+
+
+class Timeline:
+    def __init__(self, max_points: int = 4000) -> None:
+        self._max = max_points
+        # name -> list of (time, step, tokens, flops, value)
+        self._series: dict[str, list[tuple[float, Any, Any, Any, float]]] = {}
+        self._explicit_train = False
+
+    def add(self, name: str, value: Any, t: float, pos: dict[str, Any]) -> None:
+        if isinstance(value, bool) or not isinstance(value, (int, float)) or value != value:
+            return
+        if value in (float("inf"), float("-inf")):
+            return
+        points = self._series.setdefault(name, [])
+        points.append((t, pos.get("step"), pos.get("tokens"), pos.get("flops"), float(value)))
+        if len(points) > self._max:
+            del points[1:-1:2]  # halve the resolution, keep the first and the last point
+
+    # ---- feeding -----------------------------------------------------------------------
+
+    def feed(self, record: dict[str, Any], t: float, pos: dict[str, Any]) -> None:
+        """`t` is seconds since the run started, `pos` the engine's position after the record."""
+        d = record.get("d") or {}
+        kind = record.get("kind")
+        if kind == "metric" and isinstance(d.get("key"), str):
+            here = dict(pos)
+            if isinstance(d.get("step"), int):
+                here["step"] = d["step"]
+            key = d["key"]
+            if catalog.canonical(key) == "train_loss":
+                self._explicit_train = True
+            self.add(key, d.get("value"), t, here)
+        elif kind == "sample.gpu":
+            for row in d.get("gpus") or []:
+                for field, *_ in catalog.GPU_FIELDS:
+                    if field in row:
+                        self.add(f"gpu.{row.get('index')}.{field}", row[field], t, pos)
+
+    def add_snapshot(self, snap: dict[str, Any], t: float, pos: dict[str, Any]) -> None:
+        """The derived metrics, sampled once per tick."""
+        for name in ("tokens_per_s", "step_time_s", "progress", "flops_spent"):
+            self.add(name, snap.get(name), t, pos)
+        for name in ("mfu", "ofu", "mfu_ofu_gap", "arithmetic_intensity", "flops_observed"):
+            self.add(name, (snap.get(name) or {}).get("value"), t, pos)
+        self.add("goodput", (snap.get("time") or {}).get("goodput"), t, pos)
+        host = snap.get("host") or {}
+        self.add("host_cpu_cores", host.get("cpu_cores"), t, pos)
+        self.add("host_rss_b", host.get("rss_b"), t, pos)
+        mean = (snap.get("ladder") or {}).get("mean") or {}
+        for name in _LADDER:
+            self.add(name, mean.get(name), t, pos)
+
+    # ---- reading -----------------------------------------------------------------------
+
+    def _keys_for(self, name: str) -> str | None:
+        """The stored key behind a canonical name: the name itself, or the first alias seen."""
+        if name in self._series:
+            return name
+        for key in self._series:
+            if catalog.canonical(key) == name:
+                return key
+        if name == "train_loss" and not self._explicit_train and catalog.BARE_LOSS in self._series:
+            return catalog.BARE_LOSS
+        return None
+
+    def names(self) -> list[str]:
+        return sorted(self._series)
+
+    def query(
+        self, names: list[str], x: str = "step", since: float | None = None
+    ) -> dict[str, list[list[float]]]:
+        if x not in AXES:
+            raise ValueError(f"x must be one of {', '.join(AXES)}")
+        column = {"time": 0, "step": 1, "tokens": 2, "flops": 3}[x]
+        out: dict[str, list[list[float]]] = {}
+        for name in names:
+            key = self._keys_for(name)
+            points = self._series.get(key, []) if key else []
+            rows = [
+                [p[column], p[4]]
+                for p in points
+                if p[column] is not None and (since is None or p[column] > since)
+            ]
+            out[name] = rows
+        return out