rpx_benchmark.cell_log

  1"""Per-(scene, phase) metric log — the canonical experiment artifact.
  2
  3The runner produces a flat list of per-sample metric rows. This module
  4groups them into per-(scene, phase) cells and writes them to a typed
  5log file. The log is the *single source of truth* for downstream
  6analysis: Φ, JEDI, ESD-stratified summaries, bootstrap CIs, sensitivity
  7sweeps, and the paper's per-task results tables all read from it.
  8
  9Why a separate log?
 10-------------------
 11
 12- **JEDI bounds must be frozen across the full model zoo** before
 13  individual model JEDI values are comparable. The cell log lets us
 14  collect every model first, then freeze bounds once, then recompute
 15  JEDI without re-running models.
 16- **Sensitivity to scoring choices** (α, ε, Wilks vs Pillai, ESD
 17  weighting tier, etc.) reduces to a re-read of the same cells.
 18- **Bootstrap CIs on Φ** resample scenes from the cell log; no need to
 19  re-run any model.
 20- **Bug fixes** in metric calculators can be partially recovered by
 21  re-running on cached predictions; the cell log records what the
 22  numbers *were* at run time so the diff is interpretable.
 23
 24Storage format
 25--------------
 26
 27Parquet by default (typed columns, fast pandas/DuckDB reads). When
 28``pyarrow`` is unavailable, falls back to JSON Lines — same schema,
 29slightly larger files.
 30
 31Schema
 32------
 33
 34One row per (model, task, scene, phase, frame_budget) cell:
 35
 36  model_name        str       e.g. "DA3 Metric-L"
 37  task              str       TaskType.value, e.g. "monocular_depth"
 38  scene_id          str       canonical scene identifier
 39  phase             str       "clu" | "int" | "cln"
 40  difficulty        str|null  "easy" | "medium" | "hard" (when known)
 41  n_samples         int       number of frames averaged into this cell
 42  frame_budget      int       0 = full clip / per-frame task; >0 = the
 43                              subsampled clip size for the Video Depth temporal-
 44                              resolution ablation (typical values 25,
 45                              75, 150). Defaulting to 0 keeps the column
 46                              meaningful for Image Depth (always 0) while
 47                              giving Video Depth a single column to slice on.
 48  latency_ms_mean   float|n   mean per-sample latency
 49  toolkit_version   str       rpx_benchmark.__version__
 50  git_sha           str|null  caller-supplied (e.g. adapter version)
 51  timestamp_utc     str       ISO-8601, when the row was written
 52  metric:<name>     float|n   one column per metric key (prefixed
 53                              ``metric:`` so the metric columns are
 54                              namespaced and never collide with the
 55                              fixed-schema columns above)
 56"""
 57
 58from __future__ import annotations
 59
 60import json
 61import os
 62from dataclasses import dataclass
 63from datetime import datetime, timezone
 64from pathlib import Path
 65from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence
 66
 67import numpy as np
 68
 69from .exceptions import ConfigTypeError, DatasetError
 70from .logging_utils import get_logger
 71
 72log = get_logger(__name__)
 73
 74
 75# Columns that live in the fixed schema. Metric columns are prefixed
 76# ``metric:`` so they're always namespaced apart.
 77FIXED_COLUMNS: tuple[str, ...] = (
 78    "model_name",
 79    "task",
 80    "scene_id",
 81    "phase",
 82    "difficulty",
 83    "n_samples",
 84    "frame_budget",
 85    "latency_ms_mean",
 86    # Hardware context — required so cross-host aggregation can tell
 87    # apart cells produced on different GPUs / precisions. The runner
 88    # auto-detects via SystemCard.auto_detect() and threads the dict
 89    # into cells_from_per_sample as ``system_card``. None for any field
 90    # the host couldn't detect (e.g. running on CPU has gpu_name="").
 91    "gpu_name",
 92    "gpu_memory_gb",
 93    "precision",
 94    "batch_size",
 95    "toolkit_version",
 96    "git_sha",
 97    "timestamp_utc",
 98)
 99
100METRIC_PREFIX = "metric:"
101
102
103def metric_col(name: str) -> str:
104    """Namespaced column name for a metric: ``metric:absrel``, etc."""
105    return f"{METRIC_PREFIX}{name}"
106
107
108def metric_keys(row: Mapping[str, Any]) -> List[str]:
109    """Return the metric keys present in a cell row (without prefix)."""
110    return [k[len(METRIC_PREFIX) :] for k in row if k.startswith(METRIC_PREFIX)]
111
112
113# --------------------------------------------------------------------------- #
114# Grouping
115# --------------------------------------------------------------------------- #
116
117
118def _toolkit_version() -> str:
119    try:
120        from . import __version__
121
122        return str(__version__)
123    except Exception:  # noqa: BLE001 — best-effort, never fail the log write
124        return "unknown"
125
126
127def cells_from_per_sample(
128    per_sample_metrics: Sequence[Mapping[str, Any]],
129    *,
130    model_name: str,
131    task: str,
132    git_sha: Optional[str] = None,
133    metric_keys: Optional[Sequence[str]] = None,
134    system_card: Optional[Any] = None,
135) -> List[Dict[str, Any]]:
136    """Group per-sample metric rows into per-(scene, phase) cells.
137
138    Each output cell carries the *mean* of every metric across the frames
139    that fell in that (scene, phase) bucket. Non-numeric fields (id,
140    difficulty) are taken from the first row in the bucket; difficulty
141    is assumed stable per (scene, phase).
142
143    Parameters
144    ----------
145    per_sample_metrics
146        Same list the runner builds: each dict has metric keys + ``id``
147        / ``phase`` / ``scene`` / ``difficulty`` / ``latency_ms``.
148    model_name, task
149        Stamped onto every emitted row.
150    git_sha
151        Optional adapter version stamp.
152    metric_keys
153        If given, only these metric columns are emitted. Otherwise all
154        numeric keys that aren't fixed-schema metadata are kept.
155    system_card
156        Optional :class:`~rpx_benchmark.profiler.SystemCard` dict (as
157        produced by ``card.to_dict()``). When provided, ``gpu_name``,
158        ``gpu_memory_gb``, ``precision``, and ``batch_size`` from the
159        card are stamped onto every emitted cell so cross-host /
160        cross-precision aggregation can disambiguate them. None means
161        those columns will be empty / None in the output — fine for
162        synthetic tests but not for production sweeps.
163
164    Returns
165    -------
166    list of dict
167        One row per cell, ready to feed :func:`write_cells`.
168    """
169    metadata_keys = {"id", "phase", "difficulty", "scene", "latency_ms"}
170
171    # Auto-detect metric keys from the first row if not given.
172    if metric_keys is None and per_sample_metrics:
173        first = per_sample_metrics[0]
174        metric_keys = [
175            k
176            for k, v in first.items()
177            if k not in metadata_keys and isinstance(v, (int, float)) and not isinstance(v, bool)
178        ]
179    metric_keys = list(metric_keys or [])
180
181    # Bucket by (scene, phase).
182    buckets: Dict[tuple[Any, Any], List[Mapping[str, Any]]] = {}
183    for row in per_sample_metrics:
184        scene = row.get("scene")
185        phase = row.get("phase")
186        if scene is None or phase is None:
187            continue
188        buckets.setdefault((scene, phase), []).append(row)
189
190    timestamp = datetime.now(timezone.utc).isoformat()
191    version = _toolkit_version()
192
193    # Accept either a SystemCard dataclass or its .to_dict() form for
194    # ergonomics — third-party users on the public API typically have
195    # the dataclass, internal pipelines pass dicts.
196    if system_card is None:
197        sc: Mapping[str, Any] = {}
198    elif hasattr(system_card, "to_dict"):
199        sc = system_card.to_dict()
200    elif isinstance(system_card, Mapping):
201        sc = system_card
202    else:
203        raise ConfigTypeError(
204            f"system_card must be a SystemCard, a dict, or None; "
205            f"got {type(system_card).__name__}",
206            hint="pass SystemCard.auto_detect(), its to_dict() value, or None",
207        )
208    sc_gpu_name = sc.get("gpu_name") or None
209    sc_gpu_mem = sc.get("gpu_memory_gb")
210    sc_precision = sc.get("precision") or None
211    sc_batch_size = sc.get("batch_size")
212
213    cells: List[Dict[str, Any]] = []
214    for (scene, phase), rows in buckets.items():
215        # frame_budget is the first-seen value for the cell (per-sample
216        # rows for one (scene, phase) should always carry the same
217        # budget — they come from one VideoDepthDataset run). Defaults to 0
218        # for per-frame tasks that don't subsample.
219        budget = next(
220            (int(r["frame_budget"]) for r in rows if isinstance(r.get("frame_budget"), int)),
221            0,
222        )
223        cell: Dict[str, Any] = {
224            "model_name": model_name,
225            "task": task,
226            "scene_id": str(scene),
227            "phase": str(phase.value if hasattr(phase, "value") else phase),
228            "difficulty": _first_str(rows, "difficulty"),
229            "n_samples": len(rows),
230            "frame_budget": budget,
231            "latency_ms_mean": _mean_or_none(rows, "latency_ms"),
232            # Hardware context from the runner's SystemCard. Columns
233            # always present in the schema; values are None when the
234            # caller didn't pass a system_card.
235            "gpu_name": sc_gpu_name,
236            "gpu_memory_gb": sc_gpu_mem,
237            "precision": sc_precision,
238            "batch_size": sc_batch_size,
239            "toolkit_version": version,
240            "git_sha": git_sha,
241            "timestamp_utc": timestamp,
242        }
243        for k in metric_keys:
244            vals = [r[k] for r in rows if isinstance(r.get(k), (int, float))]
245            cell[metric_col(k)] = float(np.mean(vals)) if vals else None
246        cells.append(cell)
247
248    return cells
249
250
251def _first_str(rows: Sequence[Mapping[str, Any]], key: str) -> Optional[str]:
252    """First non-null string-like value for ``key`` across rows."""
253    for r in rows:
254        v = r.get(key)
255        if v is None:
256            continue
257        return str(v.value if hasattr(v, "value") else v)
258    return None
259
260
261def _mean_or_none(rows: Sequence[Mapping[str, Any]], key: str) -> Optional[float]:
262    vals = [r[key] for r in rows if isinstance(r.get(key), (int, float))]
263    return float(np.mean(vals)) if vals else None
264
265
266def cell_from_metrics(
267    metrics: Mapping[str, Any],
268    *,
269    model_name: str,
270    task: str,
271    scene_id: str,
272    phase: Any,
273    n_samples: int,
274    difficulty: Any = None,
275    latency_ms_mean: Optional[float] = None,
276    git_sha: Optional[str] = None,
277    metric_keys: Optional[Sequence[str]] = None,
278    timestamp_utc: Optional[str] = None,
279    frame_budget: int = 0,
280    system_card: Optional[Any] = None,
281) -> Dict[str, Any]:
282    """Build one canonical (scene, phase) cell from *pre-aggregated* metrics.
283
284    Companion to :func:`cells_from_per_sample`.  Use this for
285    clip/sequence-level tasks (Video Depth, tracking) where one model forward
286    already yields a single cell's worth of numbers — there are no
287    per-sample rows to group.  The emitted row has the **identical
288    schema** (:data:`FIXED_COLUMNS` + ``metric:`` columns), so per-frame
289    and per-clip cells merge into one ``cells.parquet`` and feed the same
290    cross-phase / phase×difficulty Φ analysis.
291
292    Parameters
293    ----------
294    metrics
295        Already-aggregated metric name → scalar (e.g. the output of
296        ``evaluate_video_depth_clip``). ``NaN`` values are preserved (the metric
297        was attempted but undefined — e.g. OPW with no flow backend);
298        the downstream MANOVA drops them.
299    model_name, task, scene_id, phase, n_samples
300        Fixed-schema identity of the cell. ``phase`` / ``difficulty`` may
301        be enums (``.value`` is taken) or plain strings.
302    difficulty
303        The scene's ESD tier (easy/medium/hard) — required for
304        cross-split Φ; ``None`` when unknown.
305    """
306    if metric_keys is None:
307        metric_keys = [
308            k for k, v in metrics.items() if isinstance(v, (int, float)) and not isinstance(v, bool)
309        ]
310
311    def _label(v: Any) -> Optional[str]:
312        return None if v is None else str(v.value if hasattr(v, "value") else v)
313
314    if system_card is None:
315        sc = {}
316    elif hasattr(system_card, "to_dict"):
317        sc = system_card.to_dict()
318    elif isinstance(system_card, Mapping):
319        sc = system_card
320    else:
321        raise ConfigTypeError("system_card must be a SystemCard, a dict, or None")
322
323    cell: Dict[str, Any] = {
324        "gpu_name": sc.get("gpu_name") or None,
325        "gpu_memory_gb": sc.get("gpu_memory_gb"),
326        "precision": sc.get("precision") or None,
327        "batch_size": sc.get("batch_size"),
328        "model_name": model_name,
329        "task": task,
330        "scene_id": str(scene_id),
331        "phase": _label(phase),
332        "difficulty": _label(difficulty),
333        "n_samples": int(n_samples),
334        "frame_budget": int(frame_budget),
335        "latency_ms_mean": latency_ms_mean,
336        "toolkit_version": _toolkit_version(),
337        "git_sha": git_sha,
338        "timestamp_utc": timestamp_utc or datetime.now(timezone.utc).isoformat(),
339    }
340    for k in metric_keys:
341        v = metrics.get(k)
342        cell[metric_col(k)] = (
343            float(v) if isinstance(v, (int, float)) and not isinstance(v, bool) else None
344        )
345    return cell
346
347
348# --------------------------------------------------------------------------- #
349# Writers
350# --------------------------------------------------------------------------- #
351
352
353@dataclass
354class WriteResult:
355    """Outcome of a :func:`write_cells` call."""
356
357    path: Path
358    format: str  # "parquet" | "jsonl"
359    n_rows: int
360    metric_columns: List[str]
361
362    def to_dict(self) -> Dict[str, Any]:
363        return {
364            "path": str(self.path),
365            "format": self.format,
366            "n_rows": self.n_rows,
367            "metric_columns": list(self.metric_columns),
368        }
369
370
371def _try_parquet() -> bool:
372    try:
373        import pyarrow  # noqa: F401
374
375        return True
376    except ImportError:
377        return False
378
379
380def write_cells(
381    cells: Sequence[Mapping[str, Any]],
382    path: os.PathLike,
383    *,
384    format: str = "auto",
385) -> WriteResult:
386    """Write cells to disk.
387
388    Parameters
389    ----------
390    cells
391        Output of :func:`cells_from_per_sample`.
392    path
393        File path. If ``format='auto'`` and the extension is ``.parquet``,
394        Parquet is preferred; otherwise JSON Lines. Missing extension
395        defaults to Parquet when ``pyarrow`` is available, JSONL otherwise.
396    format
397        Override: ``"parquet"`` | ``"jsonl"`` | ``"auto"``.
398
399    Returns
400    -------
401    WriteResult
402        Path, chosen format, and the metric-column names that were written.
403    """
404    out = Path(path)
405    out.parent.mkdir(parents=True, exist_ok=True)
406
407    if format == "auto":
408        if out.suffix.lower() == ".jsonl":
409            format = "jsonl"
410        elif out.suffix.lower() in (".parquet", ".pq"):
411            format = "parquet"
412        elif _try_parquet():
413            format = "parquet"
414            out = out.with_suffix(".parquet")
415        else:
416            format = "jsonl"
417            out = out.with_suffix(".jsonl")
418
419    if format == "parquet":
420        if not _try_parquet():
421            raise DatasetError(
422                "Parquet write requested but pyarrow is not installed.",
423                hint="pip install 'rpx-benchmark[hub]' (or pyarrow directly).",
424            )
425        _write_parquet(cells, out)
426    elif format == "jsonl":
427        _write_jsonl(cells, out)
428    else:
429        raise DatasetError(
430            f"unknown cell-log format {format!r}",
431            hint="Use 'parquet', 'jsonl', or 'auto'.",
432        )
433
434    metric_columns = sorted({k for row in cells for k in row if k.startswith(METRIC_PREFIX)})
435    return WriteResult(
436        path=out,
437        format=format,
438        n_rows=len(cells),
439        metric_columns=metric_columns,
440    )
441
442
443def _write_parquet(cells: Sequence[Mapping[str, Any]], path: Path) -> None:
444    import pyarrow as pa
445    import pyarrow.parquet as pq
446
447    if not cells:
448        # Write an empty file with just the fixed schema so readers don't
449        # crash on a missing column. Metric columns are unknown until we
450        # see data.
451        schema = pa.schema(
452            [
453                pa.field("model_name", pa.string()),
454                pa.field("task", pa.string()),
455                pa.field("scene_id", pa.string()),
456                pa.field("phase", pa.string()),
457                pa.field("difficulty", pa.string()),
458                pa.field("n_samples", pa.int64()),
459                pa.field("frame_budget", pa.int64()),
460                pa.field("latency_ms_mean", pa.float64()),
461                pa.field("gpu_name", pa.string()),
462                pa.field("gpu_memory_gb", pa.float64()),
463                pa.field("precision", pa.string()),
464                pa.field("batch_size", pa.int64()),
465                pa.field("toolkit_version", pa.string()),
466                pa.field("git_sha", pa.string()),
467                pa.field("timestamp_utc", pa.string()),
468            ]
469        )
470        pq.write_table(pa.Table.from_arrays([[] for _ in schema], schema=schema), path)
471        return
472
473    # Union of every column key across rows so missing values become null.
474    all_keys: List[str] = list(FIXED_COLUMNS)
475    seen = set(all_keys)
476    for row in cells:
477        for k in row:
478            if k not in seen and k.startswith(METRIC_PREFIX):
479                all_keys.append(k)
480                seen.add(k)
481
482    columns: Dict[str, List[Any]] = {k: [] for k in all_keys}
483    for row in cells:
484        for k in all_keys:
485            columns[k].append(row.get(k))
486
487    table = pa.table(columns)
488    pq.write_table(table, path)
489
490
491def _write_jsonl(cells: Sequence[Mapping[str, Any]], path: Path) -> None:
492    with path.open("w") as f:
493        for row in cells:
494            f.write(json.dumps(row) + "\n")
495
496
497# --------------------------------------------------------------------------- #
498# Readers
499# --------------------------------------------------------------------------- #
500
501
502def read_cells(path: os.PathLike) -> List[Dict[str, Any]]:
503    """Read cells back into a list of dicts.
504
505    Format detected from extension: ``.parquet`` / ``.pq`` → Parquet,
506    anything else → JSON Lines.
507    """
508    p = Path(path)
509    if not p.is_file():
510        raise DatasetError(f"cell log not found: {p}")
511
512    if p.suffix.lower() in (".parquet", ".pq"):
513        try:
514            import pyarrow.parquet as pq
515        except ImportError as e:
516            raise DatasetError(
517                f"can't read {p} — pyarrow not installed.",
518                hint="pip install 'rpx-benchmark[hub]'",
519            ) from e
520        table = pq.read_table(p)
521        return table.to_pylist()
522
523    rows: List[Dict[str, Any]] = []
524    with p.open() as f:
525        for line in f:
526            line = line.strip()
527            if line:
528                rows.append(json.loads(line))
529    return rows
530
531
532def merge_cells(paths: Iterable[os.PathLike]) -> List[Dict[str, Any]]:
533    """Concatenate multiple cell logs into one list.
534
535    Per-log column sets may differ (different tasks have different
536    metric columns); the union schema is preserved on output.
537    """
538    merged: List[Dict[str, Any]] = []
539    for p in paths:
540        merged.extend(read_cells(p))
541    return merged
542
543
544__all__ = [
545    "FIXED_COLUMNS",
546    "METRIC_PREFIX",
547    "WriteResult",
548    "cells_from_per_sample",
549    "merge_cells",
550    "metric_col",
551    "metric_keys",
552    "read_cells",
553    "write_cells",
554]