diff --git a/oddish/alembic/versions/tvm_metrics_001_add_task_version_model_metrics.py b/oddish/alembic/versions/tvm_metrics_001_add_task_version_model_metrics.py new file mode 100644 index 000000000..c034e5ad4 --- /dev/null +++ b/oddish/alembic/versions/tvm_metrics_001_add_task_version_model_metrics.py @@ -0,0 +1,126 @@ +"""Add per-(task version, agent, model) trial metrics. + +Sibling of ``task_version_browse_summaries`` one grain finer, so pass rate and +trajectory length can be selected per model instead of only filtered on. +Backfill is deliberately NOT done here -- the table is inert until +``backfill_task_version_model_metrics`` populates it, which keeps this +migration fast on a hot ``trials`` table. + +Revision ID: tvm_metrics_001 +Revises: agentcap01 +""" + +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + +revision: str = "tvm_metrics_001" +down_revision: Union[str, Sequence[str], None] = "agentcap01" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + +_TABLE = "task_version_model_metrics" + + +def _table_exists(name: str) -> bool: + return name in sa.inspect(op.get_bind()).get_table_names() + + +def _index_exists(table: str, name: str) -> bool: + return name in { + item["name"] for item in sa.inspect(op.get_bind()).get_indexes(table) + } + + +def _counter(name: str) -> sa.Column: + return sa.Column(name, sa.Integer(), nullable=False, server_default="0") + + +def upgrade() -> None: + if not _table_exists(_TABLE): + op.create_table( + _TABLE, + sa.Column("task_version_id", sa.String(160), nullable=False), + sa.Column("agent", sa.String(128), nullable=False), + # NULL model is a real state on older trials; the empty string keeps + # it inside the primary key without inventing a fake model name. + sa.Column("model", sa.String(256), nullable=False, server_default=""), + sa.Column("task_id", sa.String(128), nullable=False), + # scored population + _counter("n_pass"), + _counter("n_partial"), + _counter("n_fail"), + # terminal but unscored, attributed by harbor_stage + _counter("n_unscored_agent"), + _counter("n_unscored_env"), + _counter("n_unscored_verify"), + _counter("n_cancelled_user"), + _counter("n_cancelled_reaped"), + _counter("n_cancelled_other"), + _counter("n_skipped"), + _counter("n_scoreless"), + _counter("n_unscored_unknown"), + # never in any denominator + _counter("n_inflight"), + # metric sums + sa.Column( + "sum_reward", sa.Float(), nullable=False, server_default="0" + ), + _counter("n_reward_present"), + sa.Column( + "sum_runtime", sa.Float(), nullable=False, server_default="0" + ), + _counter("n_runtime_present"), + _counter("n_with_trajectory"), + # step distribution -- nullable: NULL means "never measured", which + # is distinct from a measured zero. See spec 4.6. + _counter("n_steps_present"), + sa.Column("sum_steps", sa.BigInteger(), nullable=False, server_default="0"), + sa.Column("steps_all_min", sa.Integer()), + sa.Column("steps_all_p50", sa.Integer()), + sa.Column("steps_all_max", sa.Integer()), + _counter("steps_pass_n"), + sa.Column("steps_pass_min", sa.Integer()), + sa.Column("steps_pass_p50", sa.Integer()), + sa.Column("steps_pass_max", sa.Integer()), + _counter("steps_fail_n"), + sa.Column("steps_fail_min", sa.Integer()), + sa.Column("steps_fail_p50", sa.Integer()), + sa.Column("steps_fail_max", sa.Integer()), + _counter("steps_partial_n"), + sa.Column( + "updated_at", + sa.DateTime(timezone=True), + nullable=False, + server_default=sa.text("NOW()"), + ), + sa.PrimaryKeyConstraint("task_version_id", "agent", "model"), + sa.ForeignKeyConstraint( + ["task_version_id"], ["task_versions.id"], ondelete="CASCADE" + ), + sa.ForeignKeyConstraint(["task_id"], ["tasks.id"], ondelete="CASCADE"), + ) + + # Ratios are computed at read time rather than stored: Postgres generated + # columns cannot reference other generated columns, and every denominator + # here is a sum of counters that a partial index can cover just as well. + if not _index_exists(_TABLE, "idx_tvm_metrics_model"): + op.create_index("idx_tvm_metrics_model", _TABLE, ["model", "agent"]) + if not _index_exists(_TABLE, "idx_tvm_metrics_task_id"): + op.create_index("idx_tvm_metrics_task_id", _TABLE, ["task_id"]) + # Anomaly triage (spec 4.7) seeks on a low passing-step minimum. + if not _index_exists(_TABLE, "idx_tvm_metrics_steps_pass_min"): + op.create_index( + "idx_tvm_metrics_steps_pass_min", + _TABLE, + ["steps_pass_min"], + postgresql_where=sa.text("steps_pass_n > 0"), + ) + + +def downgrade() -> None: + op.drop_index("idx_tvm_metrics_steps_pass_min", table_name=_TABLE) + op.drop_index("idx_tvm_metrics_task_id", table_name=_TABLE) + op.drop_index("idx_tvm_metrics_model", table_name=_TABLE) + op.drop_table(_TABLE) diff --git a/oddish/src/oddish/core/backfill_task_version_model_metrics.py b/oddish/src/oddish/core/backfill_task_version_model_metrics.py new file mode 100644 index 000000000..fd4b8de54 --- /dev/null +++ b/oddish/src/oddish/core/backfill_task_version_model_metrics.py @@ -0,0 +1,103 @@ +"""Populate task_version_model_metrics for task versions that have no row yet. + +Batched and resumable: each batch commits on its own and the next run picks up +where this one stopped, because progress is the presence of rows rather than a +stored cursor. Safe to re-run -- a recomputed group produces an identical row. + + uv run python -m oddish.core.backfill_task_version_model_metrics --limit 500 +""" + +from __future__ import annotations + +import argparse +import asyncio +import logging + +from sqlalchemy import func, select + +from oddish.core.task_version_model_metrics import ( + metrics_query, + refresh_task_version_model_metrics, +) +from oddish.db import TaskVersionModel, TaskVersionModelMetricsModel, get_session + +logger = logging.getLogger(__name__) + +DEFAULT_BATCH = 200 + + +async def _next_version_ids(session, after_id: str | None, batch: int) -> list[str]: + """Next page of task versions in id order. + + Keyset, not "versions missing a row": a version whose trials are all out of + scope legitimately produces no rows, so a missing-row cursor would hand back + the same page forever. + """ + query = select(TaskVersionModel.id).order_by(TaskVersionModel.id.asc()) + if after_id is not None: + query = query.where(TaskVersionModel.id > after_id) + rows = await session.execute(query.limit(batch)) + return [str(value) for value in rows.scalars().all()] + + +async def backfill( + batch: int = DEFAULT_BATCH, + limit: int | None = None, + after_id: str | None = None, +) -> int: + """Recompute metrics for every task version. Returns the number processed.""" + processed = 0 + cursor = after_id + while limit is None or processed < limit: + size = batch if limit is None else min(batch, limit - processed) + async with get_session() as session: + version_ids = await _next_version_ids(session, cursor, size) + if not version_ids: + break + await refresh_task_version_model_metrics(session, version_ids) + written = ( + await session.execute( + select(func.count()) + .select_from(TaskVersionModelMetricsModel) + .where( + TaskVersionModelMetricsModel.task_version_id.in_(version_ids) + ) + ) + ).scalar_one() + await session.commit() + cursor = version_ids[-1] + processed += len(version_ids) + logger.info( + "%d task versions (%d metric rows); %d total; resume after %s", + len(version_ids), + written, + processed, + cursor, + ) + if len(version_ids) < size: + break + return processed + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--batch", type=int, default=DEFAULT_BATCH) + parser.add_argument( + "--limit", type=int, default=None, help="stop after N task versions" + ) + parser.add_argument( + "--after-id", default=None, help="resume after this task version id" + ) + args = parser.parse_args() + logging.basicConfig(level=logging.INFO, format="%(message)s") + total = asyncio.run( + backfill(batch=args.batch, limit=args.limit, after_id=args.after_id) + ) + logger.info("done: %d task versions", total) + + +if __name__ == "__main__": + main() + + +__all__ = ["backfill", "metrics_query"] diff --git a/oddish/src/oddish/core/task_browse_summary.py b/oddish/src/oddish/core/task_browse_summary.py index 7a3dc51c2..2ff283dc4 100644 --- a/oddish/src/oddish/core/task_browse_summary.py +++ b/oddish/src/oddish/core/task_browse_summary.py @@ -11,6 +11,7 @@ browse_trial_scope, trial_bucket_label, ) +from oddish.core.task_version_model_metrics import refresh_task_version_model_metrics from oddish.db import ( TaskBrowseSummaryModel, TaskVersionModel, @@ -174,3 +175,8 @@ async def refresh_task_browse_summaries( | {"updated_at": func.now()}, ) ) + + # The per-model rollup is invalidated by exactly the same events, so it + # rides this refresh rather than duplicating the hook at all 18 call sites. + # Runs under the advisory locks taken above. + await refresh_task_version_model_metrics(session, list(summaries)) diff --git a/oddish/src/oddish/core/task_version_model_metrics.py b/oddish/src/oddish/core/task_version_model_metrics.py new file mode 100644 index 000000000..637a6cd4a --- /dev/null +++ b/oddish/src/oddish/core/task_version_model_metrics.py @@ -0,0 +1,284 @@ +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +from sqlalchemy import Integer, case, delete, func, select, text +from sqlalchemy.dialects.postgresql import insert as pg_insert +from sqlalchemy.ext.asyncio import AsyncSession + +from oddish.core.task_browse_metrics import browse_trial_scope, trial_bucket_label +from oddish.db import ( + TaskVersionModel, + TaskVersionModelMetricsModel, + TrialModel, + TrialStatus, +) + +_INFLIGHT = ( + TrialStatus.PENDING, + TrialStatus.QUEUED, + TrialStatus.RUNNING, + TrialStatus.RETRYING, +) + +_ENV_STAGES = ( + "starting", + "trial_started", + "environment_setup", + "image_build_failed", +) + +_CANCELLED_USER = "Cancelled by user" +_CANCELLED_REAPED = "Worker heartbeat stalled%" + +# Written by the recompute; everything else in the row is part of the key. +_VALUE_FIELDS = ( + "task_id", + "n_pass", + "n_partial", + "n_fail", + "n_unscored_agent", + "n_unscored_env", + "n_unscored_verify", + "n_cancelled_user", + "n_cancelled_reaped", + "n_cancelled_other", + "n_skipped", + "n_scoreless", + "n_unscored_unknown", + "n_inflight", + "sum_reward", + "n_reward_present", + "sum_runtime", + "n_runtime_present", + "n_with_trajectory", + "n_steps_present", + "sum_steps", + "steps_all_min", + "steps_all_p50", + "steps_all_max", + "steps_pass_n", + "steps_pass_min", + "steps_pass_p50", + "steps_pass_max", + "steps_fail_n", + "steps_fail_min", + "steps_fail_p50", + "steps_fail_max", + "steps_partial_n", +) + + +def _count_where(condition: Any) -> Any: + return func.count(case((condition, 1))) + + +def _median(condition: Any = None) -> Any: + """Observed median step count. + + ``percentile_disc`` rather than ``_cont``: most groups hold four trials or + fewer, where interpolation returns a step count no trial ever had. + """ + expr = func.percentile_disc(0.5).within_group(TrialModel.total_steps.asc()) + return expr.filter(condition) if condition is not None else expr + + +def _step_block(prefix: str, condition: Any) -> list[Any]: + steps = TrialModel.total_steps + return [ + # count() of a column skips NULLs, which is exactly the population the + # min/median/max below are computed over. + func.count(steps).filter(condition).label(f"{prefix}_n"), + func.min(steps).filter(condition).label(f"{prefix}_min"), + func.cast(_median(condition), Integer).label(f"{prefix}_p50"), + func.max(steps).filter(condition).label(f"{prefix}_max"), + ] + + +def metric_columns() -> list[Any]: + """Aggregate expressions producing one metrics row per group.""" + bucket = trial_bucket_label() + stage = TrialModel.harbor_stage + error = TrialModel.error_message + reward = TrialModel.reward + steps = TrialModel.total_steps + + # Trials carry no duration column; wall-clock comes from the timestamps. + runtime = func.extract( + "epoch", TrialModel.finished_at - TrialModel.started_at + ) + + cancelled = stage == "cancelled" + scored = reward.isnot(None) + is_pass = bucket == "pass" + is_fail = bucket == "fail" + + terminal_unscored = case( + (stage == "agent_running", "agent"), + (stage.in_(_ENV_STAGES), "env"), + (stage == "verification", "verify"), + (cancelled, "cancelled"), + else_="unknown", + ) + + return [ + func.count(case((is_pass, 1))).label("n_pass"), + func.count(case((bucket == "partial", 1))).label("n_partial"), + func.count(case((is_fail, 1))).label("n_fail"), + # Unscored attribution only applies to trials that never got a reward + # and are not in flight; the bucket taxonomy owns everything else. + _count_where( + (~scored) + & (terminal_unscored == "agent") + & (~TrialModel.status.in_(_INFLIGHT)) + & (bucket != "skipped") + ).label("n_unscored_agent"), + _count_where( + (~scored) + & (terminal_unscored == "env") + & (~TrialModel.status.in_(_INFLIGHT)) + & (bucket != "skipped") + ).label("n_unscored_env"), + _count_where( + (~scored) + & (terminal_unscored == "verify") + & (~TrialModel.status.in_(_INFLIGHT)) + & (bucket != "skipped") + ).label("n_unscored_verify"), + _count_where(cancelled & (error == _CANCELLED_USER)).label("n_cancelled_user"), + _count_where(cancelled & error.like(_CANCELLED_REAPED)).label( + "n_cancelled_reaped" + ), + _count_where( + cancelled + & ~func.coalesce(error, "").like(_CANCELLED_REAPED) + & (func.coalesce(error, "") != _CANCELLED_USER) + & (bucket != "skipped") + ).label("n_cancelled_other"), + _count_where(bucket == "skipped").label("n_skipped"), + _count_where(bucket == "scoreless").label("n_scoreless"), + _count_where( + (~scored) + & (terminal_unscored == "unknown") + & (~TrialModel.status.in_(_INFLIGHT)) + & (bucket.notin_(("skipped", "scoreless"))) + ).label("n_unscored_unknown"), + _count_where(TrialModel.status.in_(_INFLIGHT)).label("n_inflight"), + func.coalesce(func.sum(reward), 0.0).label("sum_reward"), + func.count(reward).label("n_reward_present"), + # Cancelled and reaped trials died early by definition; folding their + # runtimes in drags the average toward zero. + func.coalesce(func.sum(runtime).filter(~cancelled), 0.0).label("sum_runtime"), + func.count(runtime).filter(~cancelled).label("n_runtime_present"), + _count_where(TrialModel.has_trajectory.is_(True)).label("n_with_trajectory"), + func.count(steps).label("n_steps_present"), + func.coalesce(func.sum(steps), 0).label("sum_steps"), + func.min(steps).label("steps_all_min"), + func.cast(_median(), Integer).label("steps_all_p50"), + func.max(steps).label("steps_all_max"), + *_step_block("steps_pass", is_pass), + *_step_block("steps_fail", is_fail), + func.count(steps).filter(bucket == "partial").label("steps_partial_n"), + ] + + +def metrics_query(version_ids: list[str]) -> Any: + """Grouped metrics over the browse trial population for these versions.""" + return ( + select( + TrialModel.task_version_id.label("task_version_id"), + TrialModel.agent.label("agent"), + func.coalesce(TrialModel.model, "").label("model"), + func.min(TrialModel.task_id).label("task_id"), + *metric_columns(), + ) + .where(TrialModel.task_version_id.in_(version_ids), *browse_trial_scope()) + # Group on the raw column, not the coalesce: the SELECT and GROUP BY + # coalesces compile to separate bind parameters, which Postgres will not + # recognise as the same expression. Grouping on `model` is equivalent -- + # SQL groups NULLs together -- and the coalesce above only renames that + # group to "" so it fits a non-nullable primary key. + .group_by(TrialModel.task_version_id, TrialModel.agent, TrialModel.model) + ) + + +async def refresh_task_version_model_metrics( + session: AsyncSession, task_version_ids: Iterable[str | None] +) -> None: + """Rebuild affected (version, agent, model) rows inside the caller transaction. + + Recompute, never increment: a retried trial is reset to RUNNING with its + reward and ``total_steps`` nulled (``trial_handler.py``), so trials move + backwards out of terminal buckets and delta arithmetic would drift. + """ + version_ids = sorted({str(value) for value in task_version_ids if value}) + if not version_ids: + return + await session.flush() + + # Lock here rather than relying on the caller. refresh_task_browse_summaries + # already holds these, but the backfill calls this function directly, and an + # unlocked backfill batch racing a live refresh can overwrite a fresh row + # with the snapshot it aggregated moments earlier -- silent staleness, not a + # crash. pg_advisory_xact_lock is re-entrant within a transaction, so taking + # an already-held lock is free; the same sorted order keeps both paths in a + # single global ordering and cannot deadlock against each other. + for version_id in version_ids: + await session.execute( + text( + "SELECT pg_advisory_xact_lock(" + "hashtextextended(CAST(:version_id AS text), 0))" + ), + {"version_id": version_id}, + ) + + known = ( + await session.execute( + select(TaskVersionModel.id) + .where(TaskVersionModel.id.in_(version_ids)) + .order_by(TaskVersionModel.id) + ) + ).scalars().all() + if not known: + return + + rows = (await session.execute(metrics_query([str(v) for v in known]))).mappings() + values = [dict(row) for row in rows] + + # A version whose last trial was deleted or superseded keeps no row, and a + # group that lost its final trial must not survive as a stale one. + present = {(v["task_version_id"], v["agent"], v["model"]) for v in values} + existing = ( + await session.execute( + select( + TaskVersionModelMetricsModel.task_version_id, + TaskVersionModelMetricsModel.agent, + TaskVersionModelMetricsModel.model, + ).where(TaskVersionModelMetricsModel.task_version_id.in_(known)) + ) + ).all() + stale = [key for key in existing if tuple(key) not in present] + for version_id, agent, model in stale: + await session.execute( + delete(TaskVersionModelMetricsModel).where( + TaskVersionModelMetricsModel.task_version_id == version_id, + TaskVersionModelMetricsModel.agent == agent, + TaskVersionModelMetricsModel.model == model, + ) + ) + + if not values: + return + insert = pg_insert(TaskVersionModelMetricsModel).values(values) + await session.execute( + insert.on_conflict_do_update( + index_elements=[ + TaskVersionModelMetricsModel.task_version_id, + TaskVersionModelMetricsModel.agent, + TaskVersionModelMetricsModel.model, + ], + set_={field: getattr(insert.excluded, field) for field in _VALUE_FIELDS} + | {"updated_at": func.now()}, + ) + ) diff --git a/oddish/src/oddish/db/__init__.py b/oddish/src/oddish/db/__init__.py index e0ddaa29d..d5637f39d 100644 --- a/oddish/src/oddish/db/__init__.py +++ b/oddish/src/oddish/db/__init__.py @@ -46,6 +46,7 @@ TaskModel, TaskVersionModel, TaskBrowseSummaryModel, + TaskVersionModelMetricsModel, TrialEventModel, TrialFacetModel, TrialModel, @@ -137,6 +138,7 @@ "TaskModel", "TaskVersionModel", "TaskBrowseSummaryModel", + "TaskVersionModelMetricsModel", "TrialEventModel", "TrialFacetModel", "TrialModel", diff --git a/oddish/src/oddish/db/models.py b/oddish/src/oddish/db/models.py index 7bcf5ab26..193cdf65e 100644 --- a/oddish/src/oddish/db/models.py +++ b/oddish/src/oddish/db/models.py @@ -991,6 +991,87 @@ class TaskBrowseSummaryModel(Base): ) +def _counter() -> Mapped[int]: + return mapped_column(Integer, nullable=False, default=0, server_default="0") + + +class TaskVersionModelMetricsModel(Base): + """Trial metrics for one task version under one agent and model. + + A finer grain than ``TaskBrowseSummaryModel``: that table answers "how did + this task version do", this one answers "how did this model do on it", which + cannot be recovered by splitting the coarser row. + """ + + __tablename__ = "task_version_model_metrics" + + task_version_id: Mapped[str] = mapped_column( + String(160), + ForeignKey("task_versions.id", ondelete="CASCADE"), + primary_key=True, + ) + agent: Mapped[str] = mapped_column(String(128), primary_key=True) + # Older trials carry no model; "" keeps them addressable in the primary key + # rather than dropping them or inventing a name. + model: Mapped[str] = mapped_column( + String(256), primary_key=True, server_default="" + ) + task_id: Mapped[str] = mapped_column( + String(128), ForeignKey("tasks.id", ondelete="CASCADE"), nullable=False + ) + + n_pass: Mapped[int] = _counter() + n_partial: Mapped[int] = _counter() + n_fail: Mapped[int] = _counter() + + n_unscored_agent: Mapped[int] = _counter() + n_unscored_env: Mapped[int] = _counter() + n_unscored_verify: Mapped[int] = _counter() + n_cancelled_user: Mapped[int] = _counter() + n_cancelled_reaped: Mapped[int] = _counter() + n_cancelled_other: Mapped[int] = _counter() + n_skipped: Mapped[int] = _counter() + n_scoreless: Mapped[int] = _counter() + n_unscored_unknown: Mapped[int] = _counter() + n_inflight: Mapped[int] = _counter() + + sum_reward: Mapped[float] = mapped_column( + Float, nullable=False, default=0.0, server_default="0" + ) + n_reward_present: Mapped[int] = _counter() + sum_runtime: Mapped[float] = mapped_column( + Float, nullable=False, default=0.0, server_default="0" + ) + n_runtime_present: Mapped[int] = _counter() + n_with_trajectory: Mapped[int] = _counter() + + # NULL on every distribution column means "never measured", which is not the + # same as a measured zero -- total_steps is absent on most pre-July trials. + n_steps_present: Mapped[int] = _counter() + sum_steps: Mapped[int] = mapped_column( + BigInteger, nullable=False, default=0, server_default="0" + ) + steps_all_min: Mapped[int | None] = mapped_column(Integer, nullable=True) + steps_all_p50: Mapped[int | None] = mapped_column(Integer, nullable=True) + steps_all_max: Mapped[int | None] = mapped_column(Integer, nullable=True) + steps_pass_n: Mapped[int] = _counter() + steps_pass_min: Mapped[int | None] = mapped_column(Integer, nullable=True) + steps_pass_p50: Mapped[int | None] = mapped_column(Integer, nullable=True) + steps_pass_max: Mapped[int | None] = mapped_column(Integer, nullable=True) + steps_fail_n: Mapped[int] = _counter() + steps_fail_min: Mapped[int | None] = mapped_column(Integer, nullable=True) + steps_fail_p50: Mapped[int | None] = mapped_column(Integer, nullable=True) + steps_fail_max: Mapped[int | None] = mapped_column(Integer, nullable=True) + steps_partial_n: Mapped[int] = _counter() + + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + nullable=False, + default=utcnow, + server_default=text("NOW()"), + ) + + class TrialModel(TimestampedMixin, Base): """Trial database model.""" diff --git a/oddish/tests/test_task_version_model_metrics.py b/oddish/tests/test_task_version_model_metrics.py new file mode 100644 index 000000000..fb6bf56b1 --- /dev/null +++ b/oddish/tests/test_task_version_model_metrics.py @@ -0,0 +1,415 @@ +"""Per-(task version, agent, model) metric rollups. + +Exercises the recompute against a real database rather than a compiled query, +because the properties that matter -- percentile_disc's observed median, NULL +step handling, and the pass/fail split -- are Postgres semantics, not ORM ones. +""" + +from __future__ import annotations + +import uuid + +import pytest +from sqlalchemy import select + +from oddish.core.task_version_model_metrics import ( + refresh_task_version_model_metrics, +) +from oddish.db import ( + ExperimentModel, + TaskModel, + TaskStatus, + TaskVersionModel, + TaskVersionModelMetricsModel, + TrialModel, + TrialStatus, +) + + +async def _seed(session, trials: list[dict]) -> tuple[str, str]: + """Create a task + version and the given trials. Returns (task, version).""" + suffix = uuid.uuid4().hex[:8] + task_id = f"task-{suffix}" + version_id = f"tv-{suffix}" + experiment_id = f"exp-{suffix}" + + session.add( + ExperimentModel(id=experiment_id, name=experiment_id, org_id="org-1") + ) + session.add( + TaskModel( + id=task_id, + name=task_id, + org_id="org-1", + user="tester", + task_path="s3://test-bucket/tvm", + status=TaskStatus.COMPLETED, + ) + ) + # No relationship() links these, so SQLAlchemy will not order the inserts + # by the foreign key on its own. + await session.flush() + session.add( + TaskVersionModel( + id=version_id, + task_id=task_id, + version=1, + task_path="s3://test-bucket/tvm", + ) + ) + await session.flush() + + for index, spec in enumerate(trials): + trial_id = f"{task_id}-{index}" + session.add( + TrialModel( + id=trial_id, + name=trial_id, + task_id=task_id, + task_version_id=version_id, + experiment_id=experiment_id, + org_id="org-1", + agent=spec.pop("agent", "claude-code"), + model=spec.pop("model", "claude-opus-4-8"), + provider=spec.pop("provider", "anthropic"), + queue_key="anthropic/claude-opus-4-8", + is_probe=False, + status=spec.pop("status", TrialStatus.SUCCESS), + **spec, + ) + ) + await session.flush() + return task_id, version_id + + +async def _row(session, version_id: str, agent="claude-code", model="claude-opus-4-8"): + result = await session.execute( + select(TaskVersionModelMetricsModel).where( + TaskVersionModelMetricsModel.task_version_id == version_id, + TaskVersionModelMetricsModel.agent == agent, + TaskVersionModelMetricsModel.model == model, + ) + ) + return result.scalar_one_or_none() + + +@pytest.mark.asyncio +async def test_buckets_and_step_distribution(session): + """Pass/fail split, and min/median/max over each outcome separately.""" + _, version_id = await _seed( + session, + [ + {"reward": 1.0, "total_steps": 10}, + {"reward": 1.0, "total_steps": 30}, + {"reward": 1.0, "total_steps": 20}, + {"reward": 0.0, "total_steps": 200}, + {"reward": 0.0, "total_steps": 300}, + {"reward": 0.5, "total_steps": 77}, + ], + ) + await refresh_task_version_model_metrics(session, [version_id]) + row = await _row(session, version_id) + + assert (row.n_pass, row.n_fail, row.n_partial) == (3, 2, 1) + + # percentile_disc returns an observed value: 10/20/30 -> 20. + assert (row.steps_pass_n, row.steps_pass_min, row.steps_pass_max) == (3, 10, 30) + assert row.steps_pass_p50 == 20 + + # Even-sized group: _disc takes the lower of 200/300, never 250. + assert (row.steps_fail_n, row.steps_fail_min, row.steps_fail_max) == (2, 200, 300) + assert row.steps_fail_p50 == 200 + + # Partial carries a count only -- no distribution columns (decision 11). + assert row.steps_partial_n == 1 + + # The `all` block spans every outcome including the partial. + assert (row.steps_all_min, row.steps_all_max) == (10, 300) + assert row.n_steps_present == 6 + + +@pytest.mark.asyncio +async def test_null_steps_are_excluded_not_zeroed(session): + """A missing total_steps must not read as a zero-step trial.""" + _, version_id = await _seed( + session, + [ + {"reward": 1.0, "total_steps": 40}, + {"reward": 1.0, "total_steps": None}, + {"reward": 1.0, "total_steps": None}, + ], + ) + await refresh_task_version_model_metrics(session, [version_id]) + row = await _row(session, version_id) + + assert row.n_pass == 3 + # Three passing trials, one measured: stats describe the one. + assert row.steps_pass_n == 1 + assert row.steps_pass_min == row.steps_pass_max == row.steps_pass_p50 == 40 + + +@pytest.mark.asyncio +async def test_group_with_no_steps_yields_null_not_zero(session): + """Never-measured must be NULL, so a filter excludes it rather than ranking it.""" + _, version_id = await _seed( + session, [{"reward": 1.0, "total_steps": None}, {"reward": 0.0}] + ) + await refresh_task_version_model_metrics(session, [version_id]) + row = await _row(session, version_id) + + assert row.n_steps_present == 0 + assert row.steps_all_min is None + assert row.steps_all_p50 is None + assert row.steps_pass_min is None + assert row.steps_pass_n == 0 + + +@pytest.mark.asyncio +async def test_inflight_trials_stay_out_of_scored_counts(session): + """A running trial must not move any pass/fail number.""" + _, version_id = await _seed( + session, + [ + {"reward": 1.0, "total_steps": 5}, + {"status": TrialStatus.RUNNING, "harbor_stage": "agent_running"}, + {"status": TrialStatus.QUEUED, "harbor_stage": "starting"}, + ], + ) + await refresh_task_version_model_metrics(session, [version_id]) + row = await _row(session, version_id) + + assert row.n_pass == 1 + assert row.n_fail == 0 + assert row.n_inflight == 2 + assert row.n_unscored_agent == 0 + assert row.n_unscored_env == 0 + + +@pytest.mark.asyncio +async def test_gate_skipped_trial_is_not_also_cancelled_other(session): + """Gate-skipped trials own the skipped bucket despite their cancelled stage.""" + _, version_id = await _seed( + session, + [ + { + "status": TrialStatus.SKIPPED, + "harbor_stage": "cancelled", + "error_message": "Skipped by gate", + } + ], + ) + await refresh_task_version_model_metrics(session, [version_id]) + row = await _row(session, version_id) + + assert row.n_skipped == 1 + assert row.n_cancelled_other == 0 + + +@pytest.mark.asyncio +async def test_grain_splits_by_agent_and_model(session): + """Same task version, two models -> two rows, not one blended one.""" + _, version_id = await _seed( + session, + [ + {"model": "claude-opus-4-8", "reward": 1.0, "total_steps": 10}, + {"model": "claude-opus-4-8", "reward": 1.0, "total_steps": 12}, + {"model": "gpt-5.5", "reward": 0.0, "total_steps": 400}, + ], + ) + await refresh_task_version_model_metrics(session, [version_id]) + + opus = await _row(session, version_id, model="claude-opus-4-8") + gpt = await _row(session, version_id, model="gpt-5.5") + assert (opus.n_pass, opus.n_fail) == (2, 0) + assert (gpt.n_pass, gpt.n_fail) == (0, 1) + assert opus.steps_pass_max == 12 + assert gpt.steps_fail_min == 400 + + +@pytest.mark.asyncio +async def test_recompute_is_idempotent(session): + """Running twice yields identical rows -- the property deltas would break.""" + _, version_id = await _seed( + session, [{"reward": 1.0, "total_steps": 8}, {"reward": 0.0, "total_steps": 9}] + ) + await refresh_task_version_model_metrics(session, [version_id]) + first = await _row(session, version_id) + snapshot = (first.n_pass, first.n_fail, first.steps_pass_p50, first.sum_steps) + + await refresh_task_version_model_metrics(session, [version_id]) + second = await _row(session, version_id) + assert ( + second.n_pass, + second.n_fail, + second.steps_pass_p50, + second.sum_steps, + ) == snapshot + + +@pytest.mark.asyncio +async def test_trial_moving_backwards_leaves_its_bucket(session): + """A retried trial is reset to RUNNING with reward and steps nulled.""" + task_id, version_id = await _seed( + session, [{"reward": 1.0, "total_steps": 15}, {"reward": 1.0, "total_steps": 25}] + ) + await refresh_task_version_model_metrics(session, [version_id]) + assert (await _row(session, version_id)).n_pass == 2 + + trial = await session.get(TrialModel, f"{task_id}-0") + trial.status = TrialStatus.RUNNING + trial.reward = None + trial.total_steps = None + await session.flush() + + await refresh_task_version_model_metrics(session, [version_id]) + row = await _row(session, version_id) + assert row.n_pass == 1 + assert row.n_inflight == 1 + # The retired trial's steps must leave the distribution too. + assert row.steps_pass_n == 1 + assert row.steps_pass_min == 25 + + +@pytest.mark.asyncio +async def test_soft_deleted_and_superseded_trials_drop_out(session): + """Scope predicate is applied by the recompute's own SQL, not an ORM listener.""" + task_id, version_id = await _seed( + session, + [ + {"reward": 1.0, "total_steps": 3}, + {"reward": 1.0, "total_steps": 4}, + {"reward": 1.0, "total_steps": 5}, + ], + ) + from datetime import datetime, timezone + + deleted = await session.get(TrialModel, f"{task_id}-1") + deleted.deleted_at = datetime.now(timezone.utc) + superseded = await session.get(TrialModel, f"{task_id}-2") + superseded.superseded_by_trial_id = f"{task_id}-0" + await session.flush() + + await refresh_task_version_model_metrics(session, [version_id]) + row = await _row(session, version_id) + assert row.n_pass == 1 + assert row.steps_pass_n == 1 + assert row.steps_pass_min == 3 + + +@pytest.mark.asyncio +async def test_group_losing_its_last_trial_is_deleted(session): + """A stale row must not survive its population going empty.""" + task_id, version_id = await _seed(session, [{"reward": 1.0, "total_steps": 6}]) + await refresh_task_version_model_metrics(session, [version_id]) + assert await _row(session, version_id) is not None + + from datetime import datetime, timezone + + trial = await session.get(TrialModel, f"{task_id}-0") + trial.deleted_at = datetime.now(timezone.utc) + await session.flush() + + await refresh_task_version_model_metrics(session, [version_id]) + assert await _row(session, version_id) is None + + +# --- backfill --------------------------------------------------------------- +# +# The backfill drives its own sessions and commits, so these use committed rows +# rather than the rolled-back `session` fixture and clean up after themselves. + + +@pytest.mark.asyncio +async def test_backfill_covers_versions_and_skips_empty_ones(session): + """A version with no in-scope trials yields no row and must not stall the loop. + + A "select versions missing a row" cursor would hand back the trial-less + version forever; this is the guard on the keyset pagination that replaced it. + """ + from oddish.core.backfill_task_version_model_metrics import backfill + + _, with_trials = await _seed(session, [{"reward": 1.0, "total_steps": 11}]) + _, without_trials = await _seed(session, []) + await session.commit() + + try: + processed = await backfill(batch=5) + assert processed > 0 + + assert await _row(session, with_trials) is not None + assert await _row(session, without_trials) is None + finally: + await session.rollback() + + +@pytest.mark.asyncio +async def test_backfill_is_idempotent(session): + """Re-running must not duplicate or change rows.""" + from sqlalchemy import func as sa_func + + from oddish.core.backfill_task_version_model_metrics import backfill + + await _seed(session, [{"reward": 1.0, "total_steps": 21}, {"reward": 0.0}]) + await session.commit() + + async def _count() -> int: + return ( + await session.execute( + select(sa_func.count()).select_from(TaskVersionModelMetricsModel) + ) + ).scalar_one() + + try: + await backfill(batch=5) + first = await _count() + assert first > 0 + + await backfill(batch=5) + session.expire_all() + assert await _count() == first + finally: + await session.rollback() + + +@pytest.mark.asyncio +async def test_recompute_takes_the_version_advisory_lock(session): + """The backfill calls this directly, so it cannot rely on a caller's lock. + + Without a lock, a backfill batch racing a live refresh overwrites a fresh + row with the snapshot it aggregated moments earlier. Proven by holding the + lock from a second connection and asserting the recompute blocks on it. + """ + import asyncio + + import oddish.db.connection as conn + from sqlalchemy import text as sa_text + + _, version_id = await _seed(session, [{"reward": 1.0, "total_steps": 12}]) + await session.commit() + + holder = conn.async_session_maker() + try: + # Same hash expression as refresh_task_browse_summaries, so the two + # paths contend on one lock rather than two. + await holder.execute( + sa_text( + "SELECT pg_advisory_xact_lock(" + "hashtextextended(CAST(:v AS text), 0))" + ), + {"v": version_id}, + ) + + blocked = conn.async_session_maker() + try: + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for( + refresh_task_version_model_metrics(blocked, [version_id]), + timeout=2.0, + ) + finally: + await blocked.rollback() + await blocked.close() + finally: + await holder.rollback() + await holder.close() + await session.rollback()