src/fpgm/utils/timing.py

"""Per-step wall-clock timing, structured for ``meta.json`` and logged as it runs.

Every stage in this repo already records *what* it produced and *why* a value is
missing (``StageCache``'s fingerprint/limitations, ``DepthMethod``'s reason
codes). What none of them record is *how long each step took* -- which is the
first question asked the moment a batch is slower than expected, and the one
question that cannot be answered after the fact from the artifacts on disk.

The physics stages this module was written for make that acute: a single
posterior run is a loop of build-spec / simulate-N-particles / score, repeated
per object per episode, with the simulate step running in a *different conda
env via subprocess*. Without per-step numbers, "the run took 40 minutes" is
indistinguishable between "MuJoCo is slow" (buy fewer particles), "the
subprocess handshake is slow" (batch harder), and "reading poses.npz is slow"
(cache it) -- three different fixes.

Two properties worth stating because they shaped the design:

**A repeated label accumulates rather than overwrites.** ``simulate`` entered
once per episode across 15 episodes is one row with ``count=15`` and a total,
not 15 rows or (worse) only the last one. That is what makes the summary
readable on a batch run instead of only on a single-episode run.

**Steps nest, and a nested step's time is *also* counted in its parent.**
Labels are joined with ``/`` (``posterior/simulate``), so a summary can be read
top-down. Fractions are therefore reported against the timer's own total, not
against a sum of rows -- summing the rows would double-count nesting, and a
"fractions add up to 240%" table is worse than no table.

Usage::

    timer = StepTimer("s9.physics")
    with timer.step("load_poses", uuid=uuid):
        track = load_track(...)
    with timer.step("simulate", n=len(particles)):
        rollouts = simulate(spec, particles)
    logger.info("%s", timer.report())
    cache.write_meta(..., payload={..., "timings": timer.summary()}, ...)

``timer.summary()`` is plain JSON-able types, so it drops straight into a
``StageCache`` payload with no conversion.
"""

from __future__ import annotations

import time
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Any

from fpgm.utils.logging import get_logger

__all__ = ["StepRecord", "StepTimer"]

logger = get_logger(__name__)


@dataclass
class StepRecord:
    """Accumulated wall-clock for one label.

    Attributes:
        label: Full ``parent/child`` path of the step.
        seconds: Total wall-clock across every entry of this label.
        count: How many times the label was entered.
        units: Sum of the ``n=`` keyword across entries, when given. This is
            what makes ``per_unit_ms`` meaningful for a step whose cost scales
            with a batch size (particles simulated, frames scored) rather than
            with the number of calls.
        fields: The last entry's extra keyword fields, kept for context in the
            summary (e.g. which uuid was slowest). Deliberately last-wins rather
            than a list: this is a timing record, not an event log.
    """

    label: str
    seconds: float = 0.0
    count: int = 0
    units: int = 0
    fields: dict[str, Any] = field(default_factory=dict)

    def as_dict(self, total_seconds: float) -> dict[str, Any]:
        out: dict[str, Any] = {
            "label": self.label,
            "seconds": round(self.seconds, 4),
            "count": self.count,
            "fraction": round(self.seconds / total_seconds, 4) if total_seconds > 0 else 0.0,
        }
        if self.count > 1:
            out["mean_seconds"] = round(self.seconds / self.count, 4)
        if self.units:
            out["units"] = self.units
            out["per_unit_ms"] = round(1000.0 * self.seconds / self.units, 4)
        if self.fields:
            out["last_fields"] = self.fields
        return out


class StepTimer:
    """Wall-clock per named step, nestable, accumulating, JSON-summarisable.

    Not thread-safe and deliberately so: the nesting stack is a plain list, and
    a timer shared across threads would interleave labels into nonsense. Give
    each worker its own timer and merge the summaries if that is ever needed.

    Args:
        name: Shown in log lines and carried into the summary; use the stage
            name (``"s9.physics"``) so a grep of a batch log separates stages.
        log_each: Emit a DEBUG line as each step closes. Left on by default --
            a long-running stage that prints nothing until it finishes is
            indistinguishable from a hung one, which is the exact situation
            per-step timing is supposed to make diagnosable.
        clock: Injectable time source, for tests. Must be monotonic;
            ``time.perf_counter`` is used rather than ``time.time`` so an NTP
            step mid-run cannot produce a negative duration.
    """

    def __init__(
        self,
        name: str,
        *,
        log_each: bool = True,
        clock: Any = time.perf_counter,
    ) -> None:
        self.name = name
        self._log_each = log_each
        self._clock = clock
        self._records: dict[str, StepRecord] = {}
        self._stack: list[str] = []
        self._t0 = clock()
        self._total: float | None = None

    # -- measuring ---------------------------------------------------------- #

    @contextmanager
    def step(self, label: str, *, n: int | None = None, **fields: Any) -> Iterator[StepRecord]:
        """Time the body, filed under ``label`` (nested under any enclosing step).

        Args:
            label: Short name for this step. Must not contain ``/`` -- nesting
                builds the path, so a label with a slash in it would produce a
                row indistinguishable from a real parent/child pair.
            n: Batch size this step processed, if it has one. Enables
                ``per_unit_ms`` in the summary.
            **fields: Extra context recorded with the step (uuid, label,
                particle count...). Kept last-wins; see :class:`StepRecord`.

        Yields:
            The :class:`StepRecord` being accumulated, so a body that only
            learns its own ``n`` partway through can set ``rec.units += k``.

        The step is recorded even if the body raises: a crash in the slow step
        is precisely when the timing is most wanted, and an exception path that
        silently dropped the measurement would hide it.
        """
        if "/" in label:
            raise ValueError(f"step label must not contain '/': {label!r}")
        self._stack.append(label)
        full = "/".join(self._stack)
        rec = self._records.setdefault(full, StepRecord(label=full))
        start = self._clock()
        try:
            yield rec
        finally:
            elapsed = self._clock() - start
            rec.seconds += elapsed
            rec.count += 1
            if n is not None:
                rec.units += int(n)
            if fields:
                rec.fields = {str(k): v for k, v in fields.items()}
            self._stack.pop()
            if self._log_each:
                suffix = ""
                if n:
                    suffix = f" ({n} units, {1000.0 * elapsed / max(n, 1):.2f} ms/unit)"
                logger.debug("%s: %s took %.3fs%s", self.name, full, elapsed, suffix)

    def mark(self, label: str, seconds: float, *, n: int | None = None, **fields: Any) -> None:
        """Record a duration measured elsewhere (e.g. reported by a subprocess).

        The MuJoCo worker runs in another conda env and times its own stepping
        loop; that number is more useful than the parent's view of it (which
        includes process spawn and JSON I/O), and both are worth keeping. This
        is how the child's self-report gets into the same table.
        """
        full = "/".join([*self._stack, label])
        rec = self._records.setdefault(full, StepRecord(label=full))
        rec.seconds += float(seconds)
        rec.count += 1
        if n is not None:
            rec.units += int(n)
        if fields:
            rec.fields = {str(k): v for k, v in fields.items()}

    def stop(self) -> float:
        """Freeze the total. Idempotent; :meth:`summary` calls it if you did not."""
        if self._total is None:
            self._total = self._clock() - self._t0
        return self._total

    # -- reporting ---------------------------------------------------------- #

    @property
    def total_seconds(self) -> float:
        """Elapsed since construction, or since :meth:`stop` if it was called."""
        return self._total if self._total is not None else self._clock() - self._t0

    def summary(self) -> dict[str, Any]:
        """JSON-able timing summary, ready for a ``StageCache`` payload.

        Rows are ordered slowest-first: the reason anyone opens this table is to
        find the bottleneck, and sorting by name would bury it.
        """
        total = self.stop()
        rows = sorted(self._records.values(), key=lambda r: -r.seconds)
        return {
            "name": self.name,
            "total_seconds": round(total, 4),
            "steps": [r.as_dict(total) for r in rows],
        }

    def report(self) -> str:
        """Human-readable one-block summary for a log line."""
        s = self.summary()
        lines = [f"{self.name}: total {s['total_seconds']:.2f}s"]
        for row in s["steps"]:
            bits = f"  {row['label']:<38s} {row['seconds']:8.2f}s  {row['fraction']*100:5.1f}%"
            if row["count"] > 1:
                bits += f"  x{row['count']}"
            if "per_unit_ms" in row:
                bits += f"  {row['per_unit_ms']:.2f} ms/unit"
            lines.append(bits)
        return "\n".join(lines)