Defer baseline gates while retry jobs are active (#1051)

Co-authored-by: Charles Huang <77707222+GeeseGoo@users.noreply.github.com>
Co-authored-by: kyle-compute <kylehvdbnl@gmail.com>

Author: charlesyhuang

oddish/src/oddish/queue.py

@@ -77,6 +77,12 @@ class TrialSupersedeConflict(RuntimeError):
     TrialStatus.RUNNING,
     TrialStatus.RETRYING,
 )
+ACTIVE_WORKER_JOB_STATUSES = (
+    WorkerJobStatus.QUEUED,
+    WorkerJobStatus.RUNNING,
+    WorkerJobStatus.RETRYING,
+    WorkerJobStatus.BLOCKED,
+)
 ACTIVE_PIPELINE_STATUSES = (
     AnalysisStatus.PENDING,
     AnalysisStatus.QUEUED,
@@ -1487,13 +1493,14 @@ async def maybe_gate_llm_trials(session: AsyncSession, trial_id: str) -> bool:
 
     Fires only when *trial_id* is a nop/oracle baseline. The decision is scoped
     to the baseline's **(task version, experiment)**: when every baseline trial
-    for that task version in that experiment is terminal, evaluates them and —
-    if they validate the task (oracle passes, nop fails) — releases that scope's
-    BLOCKED LLM trials to QUEUED; otherwise cancels them and mirrors them to
-    FAILED so the task can advance. Scoping by experiment keeps concurrent sweeps
-    in different experiments from sharing each other's gate timing or verdict;
-    scoping by task version keeps an older version's baselines from validating a
-    newer version's (different code) LLM trials.
+    and its authoritative worker job for that task version in that experiment
+    are terminal, evaluates them and — if they validate the task (oracle passes,
+    nop fails) — releases that scope's BLOCKED LLM trials to QUEUED; otherwise
+    cancels them and mirrors them to FAILED so the task can advance. Scoping by
+    experiment keeps concurrent sweeps in different experiments from sharing
+    each other's gate timing or verdict; scoping by task version keeps an older
+    version's baselines from validating a newer version's (different code) LLM
+    trials.
 
     A no-op when there are no BLOCKED LLM trials in this scope (the gate was
     never armed) or other baselines are still running. Uses SELECT FOR UPDATE on
@@ -1595,6 +1602,23 @@ async def _resolve_baseline_gate_for_scope(
     if not blocked_trial_ids:
         return None
 
+    # worker_jobs is authoritative for scheduling state. During failure
+    # settlement the trial mirror can briefly read FAILED before the runner
+    # records the still-live job as RETRYING. Treat either active mirror as
+    # pending so that window cannot turn a retryable baseline into a permanent
+    # faulty gate verdict.
+    active_baseline_job = (
+        select(WorkerJobModel.id)
+        .where(
+            and_(
+                WorkerJobModel.kind == WorkerJobKind.TRIAL,
+                WorkerJobModel.subject_table == "trials",
+                WorkerJobModel.subject_id == TrialModel.id,
+                WorkerJobModel.status.in_(ACTIVE_WORKER_JOB_STATUSES),
+            )
+        )
+        .exists()
+    )
     pending_baselines = await session.scalar(
         select(func.count(TrialModel.id)).where(
             and_(
@@ -1603,7 +1627,10 @@ async def _resolve_baseline_gate_for_scope(
                 TrialModel.task_version_id == task_version_id,
                 TrialModel.queue_key == NOP_ORACLE_QUEUE_KEY,
                 TrialModel.superseded_by_trial_id.is_(None),
-                TrialModel.status.in_(ACTIVE_TRIAL_STATUSES),
+                or_(
+                    TrialModel.status.in_(ACTIVE_TRIAL_STATUSES),
+                    active_baseline_job,
+                ),
             )
         )
     )

oddish/src/oddish/worker/local_runner.py

@@ -43,6 +43,7 @@
     get_session,
 )
 from oddish.core.harbor_artifacts import cache_write_tokens_from_trajectory
+from oddish.core.cost_basis import CANCELLED_HARBOR_STAGE
 from oddish.core.llm_key_fingerprint import platform_key_hash_for_provider
 from oddish.db.models import WorkerJobKind, WorkerJobModel, WorkerJobStatus
 from oddish.db.storage import resolve_task_directory
@@ -314,7 +315,8 @@ async def run_trial_locally(trial_id: str, *, dry_run: bool = False) -> None:
     """Execute a probe trial in-process and mirror status to the DB.
 
     Status transitions: ``QUEUED`` -> ``RUNNING`` -> ``SUCCESS``
-    (or ``FAILED`` on exception, with ``error_message`` populated).
+    (or ``FAILED`` on exception, with ``error_message`` populated) on both the
+    trial and its worker job.
 
     When ``dry_run`` is True, skips the actual Harbor call. Used in
     tests to exercise the status-transition path without spinning up
@@ -326,8 +328,9 @@ async def run_trial_locally(trial_id: str, *, dry_run: bool = False) -> None:
     # and a gated (BLOCKED) LLM trial is skipped until the baseline gate
     # releases it. ``run_trial_locally`` is the only dispatch entrypoint, so
     # this claim is the single choke point that prevents double-dispatch.
+    claimed_at = datetime.now(timezone.utc)
     async with get_session() as session:
-        claimed = (
+        claimed_trial_id = (
             await session.execute(
                 update(TrialModel)
                 .where(
@@ -344,19 +347,47 @@ async def run_trial_locally(trial_id: str, *, dry_run: bool = False) -> None:
                 )
                 .values(
                     status=TrialStatus.RUNNING,
-                    started_at=datetime.now(timezone.utc),
+                    started_at=claimed_at,
                 )
-                .returning(TrialModel.org_id, TrialModel.billed_user_id)
+                .returning(TrialModel.id)
             )
-        ).one_or_none()
-    if claimed is None:
+        ).scalar_one_or_none()
+        if claimed_trial_id is not None:
+            org_id, billed_user_id = (
+                await session.execute(
+                    select(TrialModel.org_id, TrialModel.billed_user_id).where(
+                        TrialModel.id == trial_id
+                    )
+                )
+            ).one()
+            # Local mode bypasses the unified dispatcher, so it must mirror the
+            # scheduling row itself. The baseline gate treats worker_jobs as
+            # authoritative and must never mistake a completed local baseline's
+            # original QUEUED row for work that can still retry.
+            await session.execute(
+                update(WorkerJobModel)
+                .where(
+                    WorkerJobModel.kind == WorkerJobKind.TRIAL,
+                    WorkerJobModel.subject_table == "trials",
+                    WorkerJobModel.subject_id == trial_id,
+                    WorkerJobModel.status.in_(
+                        (WorkerJobStatus.QUEUED, WorkerJobStatus.RETRYING)
+                    ),
+                )
+                .values(
+                    status=WorkerJobStatus.RUNNING,
+                    claimed_at=claimed_at,
+                    started_at=claimed_at,
+                    heartbeat_at=claimed_at,
+                )
+            )
+    if claimed_trial_id is None:
         logger.info(
             "local_runner: trial %s not claimable (already dispatched, gated, "
             "or gone), skipping",
             trial_id,
         )
         return
-    org_id, billed_user_id = claimed
     logger.info("local_runner: trial %s -> RUNNING", trial_id)
 
     failure: Exception | None = None
@@ -368,6 +399,7 @@ async def run_trial_locally(trial_id: str, *, dry_run: bool = False) -> None:
         failure = exc
 
     completed = False
+    finished_at = datetime.now(timezone.utc)
     async with get_session() as session:
         trial = await session.get(TrialModel, trial_id, with_for_update=True)
         if trial is None:
@@ -382,7 +414,7 @@ async def run_trial_locally(trial_id: str, *, dry_run: bool = False) -> None:
             else:
                 trial.status = TrialStatus.SUCCESS
                 logger.info("local_runner: trial %s -> SUCCESS", trial_id)
-            trial.finished_at = datetime.now(timezone.utc)
+            trial.finished_at = finished_at
             completed = True
         else:
             logger.info(
@@ -391,6 +423,47 @@ async def run_trial_locally(trial_id: str, *, dry_run: bool = False) -> None:
                 trial.status.value,
             )
 
+        # A cancellation or another terminal writer can win while Harbor is
+        # still exiting. Preserve that trial outcome, but do not leave local
+        # mode's scheduling row RUNNING forever: the baseline gate treats an
+        # active worker job as authoritative retry evidence.
+        if trial is not None and trial.status in (
+            TrialStatus.SUCCESS,
+            TrialStatus.FAILED,
+            TrialStatus.SKIPPED,
+        ):
+            cancelled = (
+                trial.status == TrialStatus.SKIPPED
+                or trial.harbor_stage == CANCELLED_HARBOR_STAGE
+            )
+            if cancelled:
+                worker_job_status = WorkerJobStatus.CANCELLED
+            elif trial.status == TrialStatus.SUCCESS:
+                worker_job_status = WorkerJobStatus.SUCCESS
+            else:
+                worker_job_status = WorkerJobStatus.FAILED
+            settled_at = trial.finished_at or finished_at
+            await session.execute(
+                update(WorkerJobModel)
+                .where(
+                    WorkerJobModel.kind == WorkerJobKind.TRIAL,
+                    WorkerJobModel.subject_table == "trials",
+                    WorkerJobModel.subject_id == trial_id,
+                    WorkerJobModel.status == WorkerJobStatus.RUNNING,
+                )
+                .values(
+                    status=worker_job_status,
+                    finished_at=settled_at,
+                    heartbeat_at=settled_at,
+                    next_retry_at=None,
+                    error_message=(
+                        None
+                        if worker_job_status == WorkerJobStatus.SUCCESS
+                        else trial.error_message
+                    ),
+                )
+            )
+
     from oddish.core.quota_enforcement import enforce_trial_quotas_until_checked
 
     # Enforce settled spend, but run hooks only for the winning completion.

oddish/src/oddish/workers/queue/cleanup.py

@@ -883,15 +883,14 @@ async def _advance_running_tasks_to_analysis(
 
     # -----------------------------------------------------------------
     # 2b. Baseline gate backstop: (task_version, experiment) groups whose
-    #     nop/oracle baselines are all terminal but whose LLM trials are
-    #     still BLOCKED. Normally the last baseline's handler resolves the
-    #     gate; this re-drives it if that handler was killed first. The gate
-    #     is (task version, experiment)-scoped, so group + match BLOCKED LLM
-    #     trials by (task_id, task_version_id, experiment_id) and hand it one
-    #     representative baseline trial id per group. ``IS NOT DISTINCT
-    #     FROM`` so a NULL version/experiment still matches itself (plain
-    #     ``=`` would drop those scopes, unlike the ORM push path). Skipped
-    #     entirely when the gate is off so it never touches the hot path.
+    #     nop/oracle trial mirrors and worker jobs are all terminal but whose
+    #     LLM trials are still BLOCKED. Normally the last baseline's handler
+    #     resolves the gate; this re-drives it if that handler was killed first.
+    #     The gate is (task version, experiment)-scoped, so group + match BLOCKED
+    #     LLM trials by (task_id, task_version_id, experiment_id) and hand it one
+    #     representative baseline trial id per group. ``IS NOT DISTINCT FROM``
+    #     lets a NULL version/experiment match itself (plain ``=`` would drop
+    #     those scopes, unlike the ORM push path).
     # -----------------------------------------------------------------
     # Only run the heavy grouped scan when something is actually BLOCKED.
     # Runs regardless of the feature flag so a flag rollback can't strand
@@ -929,8 +928,19 @@ async def _advance_running_tasks_to_analysis(
                     GROUP BY base.task_id, base.task_version_id,
                              base.experiment_id
                     HAVING COUNT(*) FILTER (
-                        WHERE base.status
-                            IN ('PENDING', 'QUEUED', 'RUNNING', 'RETRYING')
+                        WHERE base.status IN (
+                            'PENDING', 'QUEUED', 'RUNNING', 'RETRYING'
+                        )
+                        OR EXISTS (
+                            SELECT 1
+                            FROM worker_jobs baseline_job
+                            WHERE baseline_job.subject_table = 'trials'
+                              AND baseline_job.kind::text = 'TRIAL'
+                              AND baseline_job.subject_id = base.id
+                              AND baseline_job.status::text IN (
+                                  'QUEUED', 'RUNNING', 'RETRYING', 'BLOCKED'
+                              )
+                        )
                     ) = 0
                     """
                 ),