trainmeter 0.0.2__py3-none-any.whl

trainmeter/emit.py

@@ -0,0 +1,64 @@
+"""`emit(...)`: the one optional line a training loop can add.
+
+    from trainmeter import emit
+    emit(train_loss=loss.item())
+
+It sends one datagram and returns. It is a no-op outside `tm`, returns None, never raises, and
+never converts a tensor that lives on an accelerator: numbers and 0-dimensional CPU tensors only.
+"""
+
+from __future__ import annotations
+
+import numbers
+import os
+import sys
+from typing import Any
+
+from .agent.sender import get_sender
+
+# Keys that are facts about the run, not series.
+FACT_KEYS = frozenset({"tokens_per_step"})
+
+_warned: set[str] = set()
+
+
+def as_number(value: Any) -> float | None:
+    """A float for Python and NumPy numbers and 0-dim CPU tensors, else None. Never syncs."""
+    if isinstance(value, bool):
+        return None
+    if isinstance(value, numbers.Real):
+        return float(value)
+    try:
+        device = getattr(value, "device", None)
+        if getattr(value, "ndim", None) == 0 and getattr(device, "type", None) == "cpu":
+            detach = getattr(value, "detach", None)  # a loss with requires_grad warns otherwise
+            return float(detach() if detach is not None else value)
+    except Exception:  # noqa: BLE001
+        return None
+    return None
+
+
+def emit(**values: Any) -> None:
+    try:
+        sender = get_sender()
+        if sender is None:
+            return
+        rank = int(os.environ["RANK"]) if "RANK" in os.environ else 0
+        for key, value in values.items():
+            number = as_number(value)
+            if number is None:
+                if key not in _warned:
+                    _warned.add(key)
+                    print(
+                        f"trainmeter: emit({key}=...) skipped, not a number on the host",
+                        file=sys.stderr,
+                    )
+                continue
+            if key in FACT_KEYS:
+                sender.send(
+                    "fact", "emit", {"key": key, "value": number, "source": "emitted"}, rank
+                )
+            else:
+                sender.send("metric", "emit", {"key": key, "value": number, "via": "emit"}, rank)
+    except Exception:  # noqa: BLE001 - invariant 1: never raise into the training loop
+        return

trainmeter/engine.py

@@ -0,0 +1,621 @@
+"""The engine: records in, snapshots out. A pure function of the records it is fed.
+
+It holds no clock and does no I/O, so a live run and a replay of `events.jsonl` give the same
+numbers. Every number that cannot be computed is absent and named in `absent`, never a zero.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from . import catalog, metrics
+from .facts import FactStore, rank
+from .flops import CONVENTIONS, ModelShape, all_conventions, flops_per_token
+from .peaks import lookup
+
+SHAPE_KEYS = ("n_params", "n_embedding_params", "n_layers", "d_model", "seq_len")
+
+
+class Engine:
+    def __init__(self) -> None:
+        self.facts = FactStore()
+        self._devices: dict[tuple[Any, Any], dict[str, Any]] = {}
+        self._world_size = 1
+        self._prev_t: float | None = None
+        self._ema = metrics.Ema()
+        self._steps = 0
+        self._tokens_seen = 0
+        self._detected_tokens_per_step: float | None = None
+        self._series: dict[str, dict[str, Any]] = {}
+        self._seq = 0
+        # time split
+        self._t_start: float | None = None
+        self._t_end: float | None = None
+        self._t_last = 0.0
+        self._first_step_t: float | None = None
+        self._last_step_t: float | None = None
+        self._phase_s: dict[str, float] = {}
+        self._data_before_first = 0.0
+        self._probes: set[str] = set()
+        self.hello: dict[str, Any] = {}
+        self._host: dict[str, Any] = {}
+        self._host_prev: tuple[float, float] | None = None  # (t, cpu_s) of the last sample
+        self._node_uptime: float | None = None
+        self._rank_prev: dict[int, float] = {}
+        self._rank_dt: dict[int, metrics.Ema] = {}
+        # GPU sampler state
+        self._gpu_devices: dict[Any, dict[str, Any]] = {}
+        self._gpu_selected_by: str | None = None
+        self._gpu_versions: dict[str, str] = {}
+        self._gpu_unavailable: str | None = None
+        self._gpm_unavailable: str | None = None
+        self._ladder: dict[int, dict[str, float]] = {}
+        self._ofu_ema = metrics.Ema(alpha=0.2, warmup=0)
+        self._dram_ema = metrics.Ema(alpha=0.2, warmup=0)
+        self._observed_flops = 0.0
+        self._observed_unknown: str | None = None
+
+    # ---- input -------------------------------------------------------------------------
+
+    def feed(self, record: dict[str, Any]) -> None:
+        kind, d = record.get("kind"), record.get("d") or {}
+        if isinstance(record.get("t"), (int, float)):
+            self._t_last = max(self._t_last, record["t"])
+        if kind == "run.start":
+            self._t_start = record["t"]
+            self._node_uptime = d.get("node_uptime_s")
+        elif kind == "run.end":
+            self._t_end = record["t"]
+        elif kind == "phase":
+            self._phase(record, d)
+        elif kind == "fact":
+            self._fact(record, d)
+        elif kind == "hello":
+            self._world_size = max(self._world_size, int(d.get("world_size") or 1))
+            if (record.get("rank") or 0) == 0:
+                self._probes |= set(d.get("probes") or [])
+                self.hello = {**d, **self.hello}
+        elif kind == "step":
+            self._step(record, d)
+        elif kind == "metric":
+            self._metric(record, d)
+        elif kind == "sample.gpu":
+            self._sample_gpu(d)
+        elif kind == "sample.host":
+            self._sample_host(record, d)
+
+    def _fact(self, record: dict[str, Any], d: dict[str, Any]) -> None:
+        key, source = d.get("key"), d.get("source", "detected")
+        if key == "device":
+            dev = d.get("value") or {}
+            ident = (record.get("rank"), dev.get("uuid") or dev.get("index"))
+            self._devices[ident] = dev
+        elif isinstance(key, str):
+            self.facts.offer(key, d.get("value"), source, d.get("note"))
+
+    def _tokens_in(self, d: dict[str, Any], k: int) -> float | None:
+        """Tokens in a step record. Detection always runs, so a declared value can be checked."""
+        per_rank = d.get("tokens")
+        if per_rank is not None:
+            dp = self.facts.get("dp")
+            dp_value = int(dp.value) if dp is not None else self._world_size
+            self._detected_tokens_per_step = float(per_rank) * dp_value / k
+        declared = self.facts.get("tokens_per_step")
+        if declared is not None and declared.source in ("declared", "emitted"):
+            return k * float(declared.value)
+        return None if per_rank is None else k * self._detected_tokens_per_step  # type: ignore[operator]
+
+    def _rank_step_time(self, rank: int, t: float, k: int) -> None:
+        """Per-rank seconds per step, for straggler detection. Rank 0 also gives the step time."""
+        prev = self._rank_prev.get(rank)
+        self._rank_prev[rank] = t
+        if prev is not None and t > prev:
+            self._rank_dt.setdefault(rank, metrics.Ema(alpha=0.2, warmup=1)).update((t - prev) / k)
+
+    def _step(self, record: dict[str, Any], d: dict[str, Any]) -> None:
+        rank, k, t = record.get("rank") or 0, int(d.get("k") or 1), record["t"]
+        self._rank_step_time(rank, t, k)
+        if rank != 0:
+            return  # rank 0 is the step clock
+        if self._first_step_t is None:
+            self._first_step_t = t
+            self._data_before_first = self._phase_s.get("data", 0.0)
+        self._last_step_t = t
+        self._steps += k
+        tokens = self._tokens_in(d, k)
+        if tokens is not None:
+            self._tokens_seen += tokens
+        if self._prev_t is not None and t > self._prev_t and tokens is not None:
+            self._ema.update(tokens / (t - self._prev_t))
+        self._prev_t = t
+
+    def _sample_host(self, record: dict[str, Any], d: dict[str, Any]) -> None:
+        """Process-tree use from /proc. CPU is cores busy between two samples."""
+        host: dict[str, Any] = {}
+        if isinstance(d.get("procs"), int):
+            host["procs"] = d["procs"]
+        if isinstance(d.get("rss_kb"), (int, float)):
+            host["rss_b"] = d["rss_kb"] * 1024
+        cpu, t = d.get("cpu_s"), record.get("t")
+        if isinstance(cpu, (int, float)) and isinstance(t, (int, float)):
+            prev = self._host_prev
+            self._host_prev = (t, cpu)
+            # CPU time of processes that already exited drops out, so a fall is not a rate
+            if prev is not None and t > prev[0] and cpu >= prev[1]:
+                host["cpu_cores"] = (cpu - prev[1]) / (t - prev[0])
+            elif "cpu_cores" in self._host:
+                host["cpu_cores"] = self._host["cpu_cores"]
+        self._host = host
+
+    def _sample_gpu(self, d: dict[str, Any]) -> None:
+        if "unavailable" in d:
+            self._gpu_unavailable = str(d["unavailable"])
+            return
+        self._gpu_selected_by = d.get("selected_by", self._gpu_selected_by)
+        if isinstance(d.get("versions"), dict):
+            self._gpu_versions = d["versions"]
+        self._gpm_unavailable = d.get("gpm_unavailable")
+        ofus: list[float] = []
+        drams: list[float] = []
+        for row in d.get("gpus") or []:
+            index = row.get("index")
+            ident = row.get("uuid") or index
+            self._gpu_devices[ident] = {"name": row.get("name"), "index": index}
+            ladder = {k: row[k] for k in ("gpu_util", "sm_active", "tensor_active") if k in row}
+            ofu = self._row_ofu(row)
+            if ofu is not None:
+                ladder["ofu"] = ofu
+                ofus.append(ofu)
+                self._observe(row, ofu)
+            if "dram_active" in row:
+                drams.append(row["dram_active"])
+            self._ladder[index] = ladder
+        if ofus:
+            self._ofu_ema.update(sum(ofus) / len(ofus))
+        if drams:
+            self._dram_ema.update(sum(drams) / len(drams))
+
+    @staticmethod
+    def _row_ofu(row: dict[str, Any]) -> float | None:
+        need = ("tensor_active", "sm_clock_mhz", "max_sm_clock_mhz")
+        if (
+            not all(isinstance(row.get(k), (int, float)) for k in need)
+            or not row["max_sm_clock_mhz"]
+        ):
+            return None
+        return metrics.ofu(row["tensor_active"], row["sm_clock_mhz"], row["max_sm_clock_mhz"])
+
+    def _observe(self, row: dict[str, Any], ofu: float) -> None:
+        """Observed FLOPs: this GPU's OFU x its own dense peak x the window it was measured over."""
+        window = row.get("window_s")
+        peak = lookup(row["name"]) if row.get("name") else None
+        if peak is None:
+            self._observed_unknown = f"{row.get('name') or 'a GPU'} has no dense peak in the table"
+        elif window and window > 0:
+            self._observed_flops += ofu * peak.bf16_dense_flops * window
+
+    def _phase(self, record: dict[str, Any], d: dict[str, Any]) -> None:
+        """Seconds spent in a named phase. Rank 0 is the clock, like for steps."""
+        name, seconds = d.get("name"), d.get("s")
+        if (record.get("rank") or 0) == 0 and isinstance(name, str):
+            if isinstance(seconds, (int, float)) and seconds >= 0:
+                self._phase_s[name] = self._phase_s.get(name, 0.0) + float(seconds)
+
+    def _metric(self, record: dict[str, Any], d: dict[str, Any]) -> None:
+        key, value = d.get("key"), d.get("value")
+        if isinstance(key, str) and isinstance(value, (int, float)):
+            via = d.get("via")
+            emitted = via == "emit" if via else record.get("src") == "emit"
+            self._seq += 1
+            count = self._series[key]["count"] if key in self._series else 0
+            self._series[key] = {
+                "value": float(value),
+                "count": count + 1,
+                "source": "emitted" if emitted else "tapped",
+                "via": via or ("emit" if emitted else None),
+                "seq": self._seq,
+            }
+
+    # ---- output ------------------------------------------------------------------------
+
+    def position(self) -> dict[str, Any]:
+        """Where the run is now, the x values of every chart. Cheap: no snapshot is built."""
+        fpt = self._flops_per_token()
+        value = fpt["value"] if fpt else None
+        return {
+            "step": self._steps,
+            "tokens": self._tokens_seen,
+            "flops": None if value is None else value * self._tokens_seen,
+        }
+
+    def _visible(self) -> dict[Any, dict[str, Any]]:
+        """The devices the training process reported, else the ones the sampler was watching."""
+        return self._devices or self._gpu_devices
+
+    def _devices_source(self) -> str:
+        if self._devices:
+            return "detected"
+        if self._gpu_selected_by == "process":
+            return "detected"  # the GPUs NVML lists for the watched process
+        return "declared" if self._gpu_selected_by == "flag" else "assumed"
+
+    def _n_gpus(self) -> int | None:
+        return len(self._visible()) or None
+
+    def _peak(self):
+        names = {dev.get("name") for dev in self._visible().values() if dev.get("name")}
+        return lookup(next(iter(names))) if len(names) == 1 else None
+
+    def _flops_per_token(self) -> dict[str, Any] | None:
+        declared = self.facts.get("flops_per_token")
+        conv = self.facts.get("convention")
+        if declared is not None:
+            label = conv.value if conv is not None else "declared value"
+            return {"value": float(declared.value), "convention": label, "source": declared.source}
+        parts = [self.facts.get(k) for k in SHAPE_KEYS]
+        if any(p is None for p in parts):
+            return None
+        shape = ModelShape(*(int(p.value) for p in parts))  # type: ignore[union-attr]
+        source = max((p.source for p in parts), key=rank)  # type: ignore[union-attr]
+        if conv is not None and conv.value in CONVENTIONS:
+            return {
+                "value": flops_per_token(shape, conv.value),
+                "convention": conv.value,
+                "source": source,
+            }
+        values = all_conventions(shape)
+        return {
+            "value": None,
+            "range": (min(values.values()), max(values.values())),
+            "convention": "not chosen",
+            "source": source,
+        }
+
+    def _series_snapshot(self) -> dict[str, Any]:
+        """One entry per canonical name; for two keys, emitted beats tapped, then the latest."""
+        out: dict[str, dict[str, Any]] = {}
+        explicit_train = any(catalog.canonical(k) == "train_loss" for k in self._series)
+        for key, entry in self._series.items():
+            name = catalog.canonical(key)
+            if name is None and key == catalog.BARE_LOSS and not explicit_train:
+                name = "train_loss"
+            shown = {
+                "value": entry["value"],
+                "key": key,
+                "count": entry["count"],
+                "mapped_by_name": name is not None and name != key,
+                "source": entry["source"],
+                "via": entry["via"],
+            }
+            better = (rank(entry["source"]), -entry["seq"])
+            old = out.get(name or key)
+            if old is None or better < old["_order"]:
+                out[name or key] = {**shown, "_order": better}
+        return {k: {f: v for f, v in e.items() if f != "_order"} for k, e in out.items()}
+
+    def snapshot(self, now: float | None = None) -> dict[str, Any]:
+        absent: dict[str, str] = {}
+        warnings: list[str] = []
+        snap: dict[str, Any] = {"steps": self._steps, "tokens_seen": self._tokens_seen}
+        tps = self._ema.value
+        if tps:
+            snap["tokens_per_s"] = tps
+        step_time = self._rank_dt.get(0)
+        if step_time is not None and step_time.value:
+            snap["step_time_s"] = step_time.value
+        self._ranks(snap)
+        fpt = self._flops_per_token()
+        peak, n_gpus = self._peak(), self._n_gpus()
+        if fpt is not None:
+            snap["flops_per_token"] = fpt
+            if fpt["value"] is not None:
+                snap["flops_spent"] = fpt["value"] * self._tokens_seen
+        else:
+            absent["flops_per_token"] = (
+                "declare --params, --layers, --d-model and --seq-len, or --flops-per-token"
+            )
+        self._budget(snap, fpt, tps, absent)
+        self._mfu(snap, fpt, tps, peak, n_gpus, absent)
+        self._counters(snap, peak, absent)
+        if peak is not None:
+            snap["peak"] = {
+                "name": peak.name,
+                "bf16_dense_flops": peak.bf16_dense_flops,
+                "hbm_bytes_per_s": peak.hbm_bytes_per_s,
+                "ridge": peak.ridge,
+                "source": peak.source,
+            }
+        if self._visible():
+            snap["devices"] = [
+                {"name": dev.get("name"), "index": dev.get("index")}
+                for dev in self._visible().values()
+            ]
+        if self._gpu_versions:
+            snap["gpu_versions"] = self._gpu_versions
+        shape = self._shape()
+        if shape is not None:
+            snap["conventions"] = all_conventions(shape)
+        if n_gpus is not None:
+            source = self._devices_source()
+            snap["n_gpus"] = {"value": n_gpus, "source": source}
+            if source == "assumed":
+                warnings.append(
+                    "the training process did not report its GPUs, so the GPU count and model "
+                    "are the visible devices"
+                )
+        if self._host:
+            snap["host"] = dict(self._host)
+        snap["series"] = self._series_snapshot()
+        self._time(snap, now, absent)
+        self._counted(snap, peak, n_gpus, tps, absent)
+        self._plan(snap, peak, n_gpus, fpt, absent)
+        if self.facts.get("moe_experts") is not None:
+            warnings.append(
+                f"the model config lists {self.facts.value('moe_experts')} experts: FLOPs per "
+                "token from total parameters overstate a mixture-of-experts model, declare the "
+                "active parameters with --params"
+            )
+        declared = self.facts.get("tokens_per_step")
+        det = self._detected_tokens_per_step
+        if declared is not None and det and abs(det - float(declared.value)) > 0.01 * det:
+            warnings.append(
+                f"declared tokens per step {float(declared.value):g} disagrees with "
+                f"detected {det:g}; the declared value is used"
+            )
+        if self.facts.get("observe_only") is not None:
+            warnings.append(
+                "attached from outside: there is no agent in the program, so steps, tokens, loss "
+                "and MFU are absent; only GPU and host counters are shown"
+            )
+        snap["warnings"] = warnings
+        snap["absent"] = absent
+        return snap
+
+    def _counters(self, snap, peak, absent) -> None:
+        """OFU, observed FLOPs, arithmetic intensity, the gap and the ladder: all from counters."""
+        reason = self._gpu_unavailable or self._gpm_unavailable
+        if self._gpu_unavailable is None and self._ladder:
+            snap["ladder"] = {
+                "gpus": {str(i): dict(v) for i, v in sorted(self._ladder.items())},
+                "mean": self._ladder_mean(),
+                "mfu": snap.get("mfu", {}).get("value"),
+            }
+        ofu = self._ofu_ema.value
+        if ofu is None:
+            absent["ofu"] = reason or "no GPM sample yet"
+            absent["observed_flops"] = absent["ofu"]
+        else:
+            snap["ofu"] = {"value": ofu, "convention": "counters", "source": "measured"}
+            if self._observed_unknown is not None:
+                absent["observed_flops"] = self._observed_unknown
+            else:
+                snap["flops_observed"] = {
+                    "value": self._observed_flops,
+                    "convention": "counters",
+                    "note": "approximate; includes recompute and padding",
+                }
+            mfu = snap.get("mfu", {})
+            if "value" in mfu:
+                snap["mfu_ofu_gap"] = {
+                    "value": ofu - mfu["value"],
+                    "ofu": "counters",
+                    "mfu": mfu["convention"],
+                }
+        self._intensity(snap, peak, absent)
+
+    def _ladder_mean(self) -> dict[str, float]:
+        out: dict[str, float] = {}
+        for rung in ("gpu_util", "sm_active", "tensor_active", "ofu"):
+            values = [v[rung] for v in self._ladder.values() if rung in v]
+            if values:
+                out[rung] = sum(values) / len(values)
+        return out
+
+    def _intensity(self, snap, peak, absent) -> None:
+        mfu, dram = snap.get("mfu", {}), self._dram_ema.value
+        if "value" not in mfu:
+            absent["arithmetic_intensity"] = "needs MFU under a chosen convention"
+        elif peak is None or dram is None:
+            absent["arithmetic_intensity"] = (
+                self._gpu_unavailable or self._gpm_unavailable or "needs DRAM activity and a peak"
+            )
+        elif dram <= 0:
+            absent["arithmetic_intensity"] = "no DRAM traffic in the window"
+        else:
+            snap["arithmetic_intensity"] = {
+                "value": metrics.step_intensity(mfu["value"], dram, peak.ridge),
+                "unit": "FLOP/byte",
+                "note": "aggregate over the window, not per kernel",
+                "convention": mfu["convention"],
+            }
+
+    # (phase, probe that measures it, label)
+    _MEASURED = (
+        ("data", "dataloader", "Data wait"),
+        ("checkpoint", "checkpoint", "Checkpoints"),
+        ("eval", "eval", "Eval"),
+    )
+
+    def _time(self, snap, now, absent) -> None:
+        """Where the wall time went. Measured parts are named; the rest is the remainder.
+
+        elapsed = startup + data + checkpoint + eval + compute + after, where startup runs to the
+        first step, and compute is what is left of the training window after the measured parts.
+        A part whose probe never attached is absent, and its time stays in compute.
+        """
+        if self._t_start is None:
+            absent["time"] = "no run start record"
+            return
+        end = self._t_end if self._t_end is not None else max(now or 0.0, self._t_last)
+        elapsed = max(end - self._t_start, 0.0)
+        time_split: dict[str, Any] = {"elapsed_s": elapsed, "finished": self._t_end is not None}
+        if self._node_uptime is not None:
+            time_split["node_uptime_s"] = self._node_uptime
+        first, last = self._first_step_t, self._last_step_t
+        if first is None:
+            absent["time_split"] = "needs the agent's steps: no optimizer step has been seen"
+            snap["time"] = time_split
+            return
+        window_end = last if self._t_end is not None else end
+        window = max(window_end - first, 0.0)
+        measured = {
+            name: self._phase_s.get(name, 0.0)
+            for name, probe, _ in self._MEASURED
+            if probe in self._probes
+        }
+        in_window = {k: v for k, v in measured.items() if k != "data"}
+        data_rest = max(measured.get("data", 0.0) - self._data_before_first, 0.0)
+        startup = max(
+            first - self._t_start - min(self._data_before_first, first - self._t_start), 0.0
+        )
+        compute = max(window - data_rest - sum(in_window.values()), 0.0)
+        parts = [
+            {"name": "startup", "label": "Startup and compile", "s": startup, "measured": True}
+        ]
+        for name, _probe, label in self._MEASURED:
+            parts.append(
+                {
+                    "name": name,
+                    "label": label,
+                    "s": measured.get(name),
+                    "measured": name in measured,
+                }
+            )
+        unmeasured = [p["label"] for p in parts if not p["measured"]]
+        parts.append(
+            {
+                "name": "compute",
+                "label": "Compute and the rest" if unmeasured else "Compute",
+                "s": compute,
+                "measured": False,
+            }
+        )
+        if self._t_end is not None:
+            parts.append(
+                {
+                    "name": "after",
+                    "label": "After the last step",
+                    "s": max(end - window_end, 0.0),
+                    "measured": True,
+                }
+            )
+        time_split["parts"] = parts
+        time_split["remainder_includes"] = unmeasured
+        if elapsed > 0:
+            time_split["goodput"] = (compute + data_rest) / elapsed
+        snap["time"] = time_split
+
+    def _counted(self, snap, peak, n_gpus, tps, absent) -> None:
+        """The `counted` convention: executed FLOPs, shown beside MFU and never in its place."""
+        fact = self.facts.get("counted_flops_per_token")
+        if fact is None:
+            absent["counted"] = (
+                "no sampled step was counted (agent flops probe, or FlopCounterMode saw nothing)"
+            )
+            return
+        per_token = float(fact.value)
+        entry: dict[str, Any] = {
+            "per_token": per_token,
+            "convention": "counted",
+            "source": fact.source,
+            "note": "executed FLOPs from operator shapes: includes recompute, not model FLOPs",
+            "flops_spent": per_token * self._tokens_seen,
+        }
+        if peak is not None and n_gpus and tps:
+            entry["utilization"] = metrics.mfu(per_token, tps, n_gpus, peak.bf16_dense_flops)
+        shape = self._shape()
+        if shape is not None:
+            entry["against_conventions"] = {
+                name: per_token / value for name, value in all_conventions(shape).items()
+            }
+        snap["counted"] = entry
+
+    def _shape(self) -> ModelShape | None:
+        parts = [self.facts.get(k) for k in SHAPE_KEYS]
+        if any(p is None for p in parts):
+            return None
+        return ModelShape(*(int(p.value) for p in parts))  # type: ignore[union-attr]
+
+    def _plan(self, snap, peak, n_gpus, fpt, absent) -> None:
+        """What was declared next to what was measured."""
+        plan_mfu, plan_hours = self.facts.value("plan_mfu"), self.facts.value("plan_hours")
+        if plan_mfu is None and plan_hours is None:
+            absent["plan"] = "declare --plan-mfu and --plan-hours to compare the plan with the run"
+            return
+        plan: dict[str, Any] = {}
+        mfu = snap.get("mfu") or {}
+        if plan_mfu is not None:
+            plan["mfu"] = {
+                "planned": plan_mfu,
+                "measured": mfu.get("value"),
+                "convention": mfu.get("convention"),
+            }
+            if peak is not None and n_gpus and fpt and fpt["value"]:
+                plan["tokens_per_s"] = {
+                    "planned": plan_mfu * n_gpus * peak.bf16_dense_flops / fpt["value"],
+                    "measured": snap.get("tokens_per_s"),
+                }
+        if plan_hours is not None:
+            elapsed = (snap.get("time") or {}).get("elapsed_s")
+            eta = snap.get("eta_s")
+            plan["hours"] = {
+                "planned": plan_hours,
+                "projected": (elapsed + eta) / 3600
+                if elapsed is not None and eta is not None
+                else None,
+            }
+        if "budget_tokens" in snap:
+            plan["budget_tokens"] = snap["budget_tokens"]
+        snap["plan"] = plan
+
+    def _ranks(self, snap) -> None:
+        times = {r: e.value for r, e in self._rank_dt.items() if e.value}
+        if len(times) > 1:
+            slow, fast = max(times.values()), min(times.values())
+            snap["ranks"] = {
+                "step_time_s": {str(r): v for r, v in sorted(times.items())},
+                "slowest": str(max(times, key=times.get)),
+                "spread": slow / fast - 1,
+            }
+
+    def _budget(self, snap, fpt, tps, absent) -> None:
+        tokens_budget = self.facts.value("budget_tokens")
+        if tokens_budget is None and self.facts.get("budget_steps") is not None:
+            per_step = self.facts.value("tokens_per_step") or self._detected_tokens_per_step
+            if per_step:
+                tokens_budget = self.facts.value("budget_steps") * per_step
+        flops_budget = self.facts.value("budget_flops")
+        value = fpt["value"] if fpt else None
+        if tokens_budget is None and flops_budget is not None and value:
+            tokens_budget = flops_budget / value
+        if tokens_budget is None:
+            absent["progress"] = "declare --budget-tokens, --budget-steps or --budget-flops"
+            return
+        snap["budget_tokens"] = tokens_budget
+        snap["progress"] = self._tokens_seen / tokens_budget
+        if tps:
+            snap["eta_s"] = metrics.eta_seconds(tokens_budget - self._tokens_seen, tps)
+        if value:
+            snap["flops_required"] = value * tokens_budget
+            snap["flops_remaining"] = max(value * (tokens_budget - self._tokens_seen), 0.0)
+
+    def _mfu(self, snap, fpt, tps, peak, n_gpus, absent) -> None:
+        if fpt is None or not tps:
+            return
+        if peak is None or n_gpus is None:
+            absent["mfu"] = "device is unknown, so there is no dense peak to divide by"
+            return
+        entry: dict[str, Any] = {
+            "convention": fpt["convention"],
+            "peak": peak.name,
+            "peak_source": peak.source,
+        }
+        if fpt["value"] is not None:
+            entry["value"] = metrics.mfu(fpt["value"], tps, n_gpus, peak.bf16_dense_flops)
+        else:
+            lo, hi = fpt["range"]
+            entry["range"] = (
+                metrics.mfu(lo, tps, n_gpus, peak.bf16_dense_flops),
+                metrics.mfu(hi, tps, n_gpus, peak.bf16_dense_flops),
+            )
+        snap["mfu"] = entry