rpx_benchmark.tasks._pipeline

  1"""Shared pipeline: download -> load -> run -> report.
  2
  3Every task runner calls :func:`run_pipeline`. The task module only
  4owns the task-specific bits: a config subclass, a ``TASK_SPEC``, and
  5a one-line ``run_<task>`` wrapper around :func:`run_pipeline`.
  6
  7The toolkit does not ship models. Users supply a
  8:class:`BenchmarkableModel` (usually built via
  9``rpx.make_numpy_<task>_model(fn)`` or a custom ``InputAdapter`` /
 10``OutputAdapter`` pair) and set it on the config's ``model`` field.
 11"""
 12
 13from __future__ import annotations
 14
 15from dataclasses import dataclass
 16from pathlib import Path
 17from typing import Dict, Optional, Tuple
 18
 19from ..adapters import BenchmarkableModel
 20from ..api import Difficulty, TaskType
 21from ..cell_log import cells_from_per_sample, write_cells
 22from ..deployment import DeploymentReadinessReport
 23from ..exceptions import ConfigError
 24from ..hub import DEFAULT_REPO_ID, DEFAULT_REVISION, download_split
 25from ..loader import RPXDataset
 26from ..logging_utils import get_logger
 27from ..metrics.registry import BenchmarkResult, MetricSuite
 28from ..profiler import EfficiencyMetadata, SystemCard, count_parameters
 29from ..reports import format_markdown_summary, write_json
 30from ..runner import BenchmarkRunner, ProgressCallback
 31
 32log = get_logger(__name__)
 33
 34PipelineResult = Tuple[BenchmarkResult, DeploymentReadinessReport, Dict[str, Path]]
 35
 36
 37@dataclass
 38class TaskRunConfig:
 39    """Base config for every task runner.
 40
 41    Subclass with ``@dataclass`` when a task wants extra fields. The
 42    user supplies ``model`` — the toolkit doesn't ship any.
 43    """
 44
 45    model: Optional[BenchmarkableModel] = None
 46
 47    split: Difficulty | str = Difficulty.HARD
 48    repo_id: Optional[str] = None
 49    cache_dir: Optional[str] = None
 50    revision: Optional[str] = DEFAULT_REVISION
 51    batch_size: int = 1
 52    #: Refuse CPU fallback when a CUDA run was requested. Production smoke
 53    #: and benchmark launchers set this True; lightweight library users keep
 54    #: the historical best-effort fallback by default.
 55    require_cuda: bool = False
 56    #: Skip the first-sample torch FLOP counter. Useful for smoke gates and
 57    #: memory-constrained GPUs where instrumentation can OOM independently of
 58    #: the real forward pass.
 59    skip_flops: bool = False
 60    #: Keep only the first N manifest entries. This is a sample count for
 61    #: frame tasks and a clip count for Video Depth.
 62    max_samples: Optional[int] = None
 63
 64    #: When set, the pipeline skips :func:`download_split` and loads
 65    #: the manifest from disk. Use for runs against a locally-staged
 66    #: lossless dataset (e.g. ``dataset_hub.cli manifest`` output)
 67    #: before the v2-webp tree has been uploaded to HF. Path must point
 68    #: at the JSON file (typically ``<staging>/manifests/<task>/<split>.json``);
 69    #: the ``root`` inside that JSON resolves modality file paths.
 70    manifest_path: Optional[str] = None
 71
 72    device: str = "cuda"
 73    output_dir: Optional[str] = None
 74    progress: Optional[ProgressCallback] = None
 75    #: Optional protocol-specific metric suite. The default remains the
 76    #: registered public suite for the task.
 77    metric_suite: Optional[MetricSuite] = None
 78
 79    #: When True, mirror the per-run output directory to UTD Box at
 80    #: ``<box_folder_id>/<task>/<model>/<split>/`` after the run finishes.
 81    #: Requires ``BOX_DEVELOPER_TOKEN`` in the environment. Idempotent
 82    #: re-uploads (size-matched skip).
 83    upload_to_box: bool = False
 84    #: Box folder id to root the upload under. Defaults to the team's
 85    #: RPX-Outputs folder. Unused unless ``upload_to_box=True``.
 86    box_folder_id: Optional[str] = None
 87
 88    def __post_init__(self) -> None:
 89        if self.model is None:
 90            raise ConfigError(
 91                f"{type(self).__name__}.model is required. Wrap your "
 92                "model via rpx.make_numpy_<task>_model(fn) or build a "
 93                "BenchmarkableModel with your own input/output adapters.",
 94            )
 95        if isinstance(self.split, str):
 96            try:
 97                self.split = Difficulty(self.split)
 98            except ValueError as e:
 99                raise ConfigError(
100                    f"Unknown split {self.split!r}. Use one of: {[d.value for d in Difficulty]}",
101                ) from e
102        if self.batch_size < 1:
103            raise ConfigError(f"batch_size must be >= 1, got {self.batch_size}")
104        if self.max_samples is not None and self.max_samples < 1:
105            raise ConfigError(f"max_samples must be >= 1, got {self.max_samples}")
106
107
108def resolve_device(requested: str, *, require_cuda: bool = False) -> str:
109    """Fall back to CPU when CUDA was requested but isn't available."""
110    if require_cuda and not requested.startswith("cuda"):
111        raise ConfigError(
112            f"GPU-only run requires device='cuda', got {requested!r}",
113            hint="Use --device cuda, or explicitly opt into CPU diagnostics.",
114        )
115    if not requested.startswith("cuda"):
116        return requested
117    try:
118        import torch
119    except ImportError as exc:
120        if require_cuda:
121            raise ConfigError(
122                "GPU-only run requires a CUDA-enabled PyTorch installation.",
123                hint="Install the PyTorch wheel matching this host, then rerun.",
124            ) from exc
125        return "cpu"
126    if torch.cuda.is_available():
127        return requested
128    if require_cuda:
129        raise ConfigError(
130            "CUDA was required but torch.cuda.is_available() is false.",
131            hint="Check the NVIDIA driver, CUDA-capable PyTorch wheel, and GPU visibility.",
132        )
133    log.warning("CUDA requested but unavailable; falling back to CPU.")
134    return "cpu"
135
136
137def _count_params(model: BenchmarkableModel) -> EfficiencyMetadata:
138    raw = getattr(model, "model", None)
139    if raw is None or not hasattr(raw, "parameters"):
140        return EfficiencyMetadata(model_type="local", notes="non-torch callable")
141    return EfficiencyMetadata(params_m=count_parameters(raw), model_type="local")
142
143
144def run_pipeline(
145    *,
146    task: TaskType,
147    primary_metric: str,
148    cfg: TaskRunConfig,
149    compute_ts: bool = True,
150    compute_sgc: bool = False,
151) -> PipelineResult:
152    """Run a task end-to-end: download, load, predict, score, report.
153
154    Parameters
155    ----------
156    task : TaskType
157    primary_metric : str
158        Metric key that drives ESD-weighted phase scoring.
159    cfg : TaskRunConfig
160        Must already have ``model`` set.
161    compute_ts : bool
162        Compute Temporal Stability. Disable for tasks where adjacent
163        samples are unrelated (e.g. keypoint matching).
164    compute_sgc : bool
165        Compute Stack Geometric Coherence. Only meaningful for
166        segmentation with depth available.
167    """
168    cfg.device = resolve_device(cfg.device, require_cuda=cfg.require_cuda)
169    split_name = cfg.split.value if isinstance(cfg.split, Difficulty) else str(cfg.split)
170    repo_id = cfg.repo_id or DEFAULT_REPO_ID
171
172    log.info("pipeline: task=%s split=%s device=%s", task.value, split_name, cfg.device)
173
174    if cfg.manifest_path:
175        manifest_path = Path(cfg.manifest_path)
176        if not manifest_path.exists():
177            raise ConfigError(
178                f"manifest_path does not exist: {manifest_path}",
179                hint="Generate per-task manifests with "
180                "`python -m rpx_benchmark.dataset_hub.cli manifest "
181                "--src <raw> --staging <out> --splits "
182                "benchmark/data/splits/scene_splits.json` and pass the "
183                "resulting <out>/manifests/<task>/<split>.json path.",
184            )
185        log.info("using local manifest: %s", manifest_path)
186    else:
187        manifest_path = download_split(
188            task=task,
189            split=cfg.split,
190            repo_id=repo_id,
191            cache_dir=cfg.cache_dir,
192            revision=cfg.revision,
193            max_samples=cfg.max_samples,
194        )
195    dataset = RPXDataset.from_manifest(
196        manifest_path,
197        batch_size=cfg.batch_size,
198        max_samples=cfg.max_samples,
199    )
200    log.info("loaded %d samples from %s", len(dataset), manifest_path)
201
202    model = cfg.model
203    name = getattr(model, "name", "model")
204
205    model.setup()
206    efficiency = _count_params(model)
207    if efficiency.params_m is not None:
208        log.info("model %s: %.2f M params", name, efficiency.params_m)
209    # SystemCard auto-detects GPU / precision / batch_size and is
210    # stamped onto every cell row so cross-host aggregation can
211    # disambiguate cells produced on different hardware.
212    efficiency.system_card = SystemCard.auto_detect(
213        precision=getattr(model, "native_precision", "fp32"),
214        batch_size=cfg.batch_size,
215    )
216
217    runner = BenchmarkRunner(
218        model=model,
219        dataset=dataset,
220        metric_suite=cfg.metric_suite or MetricSuite.for_task(task),
221        call_setup=False,
222    )
223    result, dr_report = runner.run_with_report(
224        primary_metric=primary_metric,
225        model_name=name,
226        efficiency=efficiency,
227        compute_ts=compute_ts,
228        compute_sgc_flag=compute_sgc,
229        skip_flops=cfg.skip_flops,
230        progress=cfg.progress,
231    )
232
233    safe_name = name.replace("/", "__")
234    out_dir = Path(cfg.output_dir or f"./rpx_results/{safe_name}/{split_name}")
235    out_dir.mkdir(parents=True, exist_ok=True)
236    json_path = out_dir / "result.json"
237    md_path = out_dir / "summary.md"
238    cells_path = out_dir / "cells.parquet"
239    per_sample_path = out_dir / "per_sample_metrics.parquet"
240
241    try:
242        import pyarrow as pa
243        import pyarrow.parquet as pq
244
245        pq.write_table(pa.Table.from_pylist(result.per_sample), per_sample_path)
246    except ImportError as exc:
247        raise ConfigError(
248            "Writing canonical per-sample metrics requires pyarrow.",
249            hint="Install 'rpx-benchmark[hub]' or use the RPX Docker image.",
250        ) from exc
251
252    # cells.parquet is the canonical artefact downstream Φ / 𝒥 / paper-
253    # table fills read from. Bucket per_sample (one row per frame) into
254    # per-(scene, phase) cells with the mean of every numeric metric.
255    # Mirrors the per-clip pipeline (_video_pipeline.py) so both tasks
256    # feed the same downstream path.
257    cells = cells_from_per_sample(
258        result.per_sample,
259        model_name=name,
260        task=task.value,
261        system_card=efficiency.system_card,
262    )
263    if cells:
264        write_cells(cells, cells_path)
265        log.info("wrote %d cell rows to %s", len(cells), cells_path)
266    else:
267        log.warning(
268            "no cells emitted for %s/%s — per_sample lacks scene/phase fields. "
269            "This usually means the dataset manifest's samples don't carry "
270            "scene_id + phase metadata, or the runner's metric_suite stripped "
271            "them. Downstream Φ/J aggregation will be unable to bucket.",
272            name, split_name,
273        )
274
275    write_json(
276        json_path,
277        task=task.value,
278        model_name=name,
279        split=split_name,
280        repo_id=repo_id,
281        result=result,
282        dr_report=dr_report,
283    )
284    md_path.write_text(
285        format_markdown_summary(
286            task=task.value,
287            model_name=name,
288            split=split_name,
289            repo_id=repo_id,
290            result=result,
291            dr_report=dr_report,
292        ),
293        encoding="utf-8",
294    )
295    log.info("wrote %s and %s", json_path, md_path)
296
297    paths: Dict[str, Path] = {
298        "json": json_path,
299        "markdown": md_path,
300        "out_dir": out_dir,
301        "per_sample_metrics": per_sample_path,
302    }
303    if cells:
304        paths["cells"] = cells_path
305
306    if cfg.upload_to_box:
307        # Lazy import keeps Box deps (requests) out of the import path
308        # for users who never upload.
309        from ..box_upload import DEFAULT_BOX_FOLDER_ID, upload_run_dir
310
311        try:
312            summary = upload_run_dir(
313                out_dir,
314                task=task.value,
315                model_name=name,
316                split=split_name,
317                root_folder_id=cfg.box_folder_id or DEFAULT_BOX_FOLDER_ID,
318            )
319            log.info(
320                "box: %d uploaded, %d skipped → box:%s",
321                summary["uploaded"],
322                summary["skipped"],
323                summary["remote_path"],
324            )
325            paths["box_remote"] = summary["remote_path"]
326        except Exception as e:  # noqa: BLE001
327            # Don't lose the result.json on the floor just because the
328            # token expired — log loudly and continue.
329            log.error(
330                "Box upload failed (%s: %s) — local artefacts at %s are intact; "
331                "use `scripts/sync_results_to_box.py` to recover with a fresh token.",
332                type(e).__name__,
333                e,
334                out_dir,
335            )
336
337    return result, dr_report, paths
338
339
340__all__ = ["TaskRunConfig", "PipelineResult", "resolve_device", "run_pipeline"]