rpx_benchmark.tasks.registry

  1"""Task plugin registry.
  2
  3Each task module registers a :class:`TaskSpec` describing the primary
  4metric, required modalities, and two optional deployment-readiness
  5hooks. The runner looks these up and dispatches through them — it
  6never branches on task identity.
  7"""
  8
  9from __future__ import annotations
 10
 11from dataclasses import dataclass
 12from typing import Any, Callable, Dict, List, Optional, Tuple
 13
 14from ..api import TaskType
 15from ..deployment import DeploymentReadinessReport
 16from ..exceptions import ConfigError
 17from ..logging_utils import get_logger
 18from ..metrics.registry import BenchmarkResult
 19
 20log = get_logger(__name__)
 21
 22
 23TaskRunResult = Tuple[BenchmarkResult, Optional[DeploymentReadinessReport], Dict[str, Any]]
 24
 25# Hook signatures. Kept as ``Callable[..., Any]`` so importing this
 26# module does not drag the ``deployment`` result types through — the
 27# return types are documented on the attribute docstrings.
 28TemporalStabilityFn = Callable[[List[Any], List[Any], List[Any]], Any]
 29"""Signature: ``(predictions, samples, camera_poses) -> TemporalStabilityResult | None``."""
 30GeometricCoherenceFn = Callable[[List[Any], List[Any]], Any]
 31"""Signature: ``(predictions, samples) -> StackGeometricCoherenceResult | None``."""
 32
 33
 34@dataclass
 35class TaskSpec:
 36    """Per-task plugin record.
 37
 38    Attributes
 39    ----------
 40    task, display_name, description
 41        Identification.
 42    primary_metric
 43        Metric key driving the ESD-weighted phase score.
 44    required_modalities
 45        Hint for hub download scope (``rgb``, ``depth``, ``mask`` …).
 46    higher_is_better
 47        Direction of the primary metric.
 48    run
 49        End-to-end callable ``(cfg) -> (BenchmarkResult, report, paths)``.
 50    temporal_stability_fn, geometric_coherence_fn
 51        Optional deployment-readiness hooks. The runner calls them if
 52        set; tasks without a hook skip that metric silently.
 53    """
 54
 55    task: TaskType
 56    display_name: str
 57    primary_metric: str
 58    required_modalities: List[str]
 59    run: Callable[[Any], TaskRunResult]
 60    higher_is_better: bool = False
 61    description: str = ""
 62    temporal_stability_fn: Optional[TemporalStabilityFn] = None
 63    geometric_coherence_fn: Optional[GeometricCoherenceFn] = None
 64
 65
 66_TASK_REGISTRY: Dict[TaskType, TaskSpec] = {}
 67
 68
 69def register_task(spec: TaskSpec) -> TaskSpec:
 70    if spec.task in _TASK_REGISTRY:
 71        raise ConfigError(
 72            f"Task {spec.task.value!r} is already registered.",
 73            hint="Call unregister_task() first to replace it.",
 74        )
 75    _TASK_REGISTRY[spec.task] = spec
 76    log.debug("registered task %s (primary=%s)", spec.task.value, spec.primary_metric)
 77    return spec
 78
 79
 80def unregister_task(task: TaskType) -> bool:
 81    return _TASK_REGISTRY.pop(task, None) is not None
 82
 83
 84def get_task_spec(task: TaskType) -> TaskSpec:
 85    if task not in _TASK_REGISTRY:
 86        raise ConfigError(
 87            f"No runner registered for task {task.value!r}.",
 88            hint="Registered: " + ", ".join(sorted(t.value for t in _TASK_REGISTRY)),
 89        )
 90    return _TASK_REGISTRY[task]
 91
 92
 93def available_tasks() -> List[TaskType]:
 94    return sorted(_TASK_REGISTRY.keys(), key=lambda t: t.value)
 95
 96
 97def iter_task_specs() -> List[TaskSpec]:
 98    return [_TASK_REGISTRY[t] for t in available_tasks()]
 99
100
101__all__ = [
102    "TaskSpec",
103    "TaskRunResult",
104    "register_task",
105    "unregister_task",
106    "get_task_spec",
107    "available_tasks",
108    "iter_task_specs",
109]