Fix sharded-eval poll loop: stop killing healthy QUEUED shards, reconcile before failing

The 45-min poll deadline cancelled shards that were merely QUEUED waiting
for a10g capacity — throwing away their GPU queue position — then, after
the retry rounds, marked the row failed while leaving jobs running, so
their outputs completed late and stranded in staging. This failed
otherwise-successful submissions under GPU load.

- Poll loop never cancels healthy QUEUED/RUNNING shards; only ERROR shards
re-dispatch. Each job's EVAL_JOB_TIMEOUT is the real RUNNING bound.
- Replace the per-round reaper with one generous env-tunable backstop
(default 6h) that only stops an unbounded loop.
- Reconcile-before-fail: always attempt the merge from whatever landed;
a shard flagged ERROR/timeout may still have written all its fixtures
(outputs are the source of truth). Only fail on genuinely-missing ones.
- Cancel orphan jobs on the true-failure path so nothing keeps billing.
- Boot sweep now reconciles stranded pending rows (finalize if artifacts
are complete) instead of blindly failing them — self-heals across Space
restarts, which is what stranded the recovered submission.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

submit.py

@@ -234,19 +234,20 @@ SHARD_BUCKET_PREFIX = os.getenv(
 # overwrites its upload prefix), so one cheap retry absorbs a transient
 # job/runtime blip without re-running the shards that already passed.
 SHARD_MAX_RETRIES = 1
-# Whole-fan-out poll deadline. Each shard job carries its own
-# ``EVAL_JOB_TIMEOUT``; this guards the Space-side poll loop. Generous
-# vs. the per-shard ceiling because queued shards (past the ~8
-# concurrent slots) wait their turn before their own timeout starts.
-SHARD_POLL_DEADLINE_SECONDS = 45 * 60
-# When the poll window elapses with shards still non-terminal — typically the
-# tail shard stuck QUEUED waiting for GPU (a10g-large) capacity rather than a
-# compute failure — re-dispatch those stragglers and reset the window, up to
-# this many rounds, before giving up. A fresh dispatch after the window can land
-# a freed slot; shard uploads are idempotent (each rewrites its own staging
-# prefix), so a re-dispatch is safe. Worst-case total wait is roughly
-# SHARD_POLL_DEADLINE_SECONDS * (1 + SHARD_DEADLINE_RETRY_ROUNDS).
-SHARD_DEADLINE_RETRY_ROUNDS = 2
+# Whole-fan-out poll backstop. This is NOT a per-round reaper: each shard
+# job carries its own ``EVAL_JOB_TIMEOUT`` (which bounds RUNNING work and
+# turns a genuinely-hung job into ERROR), and a shard sitting QUEUED is
+# healthy — it is just waiting for a10g-large capacity past the account's
+# ~8 concurrent slots. Cancelling a QUEUED shard only throws away its place
+# in the GPU queue and makes a capacity crunch worse, so the poll loop never
+# does that. This constant is a single generous absolute ceiling (hours),
+# purely so the poll loop can't spin forever if a shard never reaches a
+# terminal stage; on hitting it the caller still reconciles against durable
+# staging artifacts before declaring failure. Env-tunable for very large
+# fan-outs behind a deep GPU queue.
+SHARD_POLL_DEADLINE_SECONDS = int(
+    os.getenv("CADGENBENCH_SHARD_POLL_DEADLINE_SECONDS", str(6 * 60 * 60))
+)
 
 # One HfApi client per process. HF_TOKEN is picked up from the env at
 # construction time and reused for every call.
@@ -1164,7 +1165,42 @@ def _run_worker_sharded(
     failures = _poll_shards_until_done(
         submission_id, submission_blob_url, shards,
     )
+
+    # Reconcile-before-fail: outputs are the source of truth. Even when a
+    # shard was flagged ERROR or caught by the poll backstop, its fixtures
+    # may all be in staging (it completed late, or wrote results before
+    # erroring). Always attempt the merge from whatever landed; only fail
+    # if the merged set genuinely can't cover every expected fixture.
     if failures:
+        logger.warning(
+            "Sharded eval for %s: %d shard(s) did not COMPLETE (%s); "
+            "reconciling from staging before deciding.",
+            submission_id, len(failures), "; ".join(failures),
+        )
+        progress.publish(
+            submission_id,
+            progress.RUNNING,
+            "Some chunks reported issues — checking for late results…",
+        )
+    else:
+        progress.publish(
+            submission_id,
+            progress.RUNNING,
+            "All chunks evaluated — merging results…",
+        )
+
+    try:
+        summary = _merge_shards_and_publish(
+            submission_id, list(shards.keys()), fixture_names,
+        )
+    except Exception as merge_err:  # noqa: BLE001 - map to a failed row
+        if not failures:
+            # Every shard COMPLETED but the merge itself broke: a real
+            # error, surfaced to _run_worker's handler.
+            raise
+        # Genuinely incomplete. Tear down any shard job still running so
+        # it stops billing now that the submission is failing, then fail.
+        _cancel_orphan_shard_jobs(shards)
         reason = ("sharded eval failed: " + "; ".join(failures))[
             :FAILURE_REASON_MAX_CHARS
         ]
@@ -1172,17 +1208,18 @@ def _run_worker_sharded(
         progress.publish(
             submission_id, progress.FAILED, _failed_progress_message(reason),
         )
-        logger.warning("Sharded eval for %s failed: %s", submission_id, reason)
+        logger.warning(
+            "Sharded eval for %s failed; reconcile could not finalize "
+            "(%s): %s", submission_id, merge_err, reason,
+        )
         return
 
-    progress.publish(
-        submission_id,
-        progress.RUNNING,
-        "All chunks evaluated — merging results…",
-    )
-    summary = _merge_shards_and_publish(
-        submission_id, list(shards.keys()), fixture_names,
-    )
+    if failures:
+        logger.info(
+            "Sharded eval for %s reconciled: all fixtures present in "
+            "staging despite %d non-completed shard(s).",
+            submission_id, len(failures),
+        )
     _flip_row_to_completed(submission_id, summary)
     progress.publish(
         submission_id, progress.COMPLETED, _completed_progress_message(summary),
@@ -1291,13 +1328,13 @@ def _dispatch_shard(
 
 
 def _cancel_shard_job(state: dict[str, Any]) -> None:
-    """Best-effort cancel of a shard's in-flight job before re-dispatch.
+    """Best-effort cancel of a shard's in-flight job.
 
-    Used on the poll-deadline retry path so a straggler that is still
-    QUEUED/RUNNING releases its slot and does not keep writing its staging
-    prefix once a replacement is dispatched. Best-effort: a failure is
-    logged and ignored, since shard uploads are idempotent (a stale job
-    only ever overwrites its own prefix with an equivalent result).
+    Used on the final failure path (via :func:`_cancel_orphan_shard_jobs`)
+    so a shard still QUEUED/RUNNING when the submission is being failed is
+    torn down instead of left to run and keep billing. Best-effort: a
+    failure is logged and ignored, since a lingering job only ever
+    overwrites its own idempotent staging prefix.
     """
     job_id = state.get("job_id")
     if not job_id:
@@ -1310,11 +1347,24 @@ def _cancel_shard_job(state: dict[str, Any]) -> None:
             namespace=EVAL_JOB_NAMESPACE,
             token=_jobs_token(),
         )
-        logger.info("Cancelled straggler shard job %s before retry", job_id)
+        logger.info("Cancelled orphan shard job %s", job_id)
     except Exception as e:  # noqa: BLE001 - cancel is best-effort
         logger.warning("Could not cancel shard job %s: %s", job_id, e)
 
 
+def _cancel_orphan_shard_jobs(shards: dict[str, dict[str, Any]]) -> None:
+    """Cancel every shard job that never reached COMPLETED.
+
+    Called once, on the final failure path, after reconcile has confirmed
+    the submission cannot be finalized. A shard still QUEUED/RUNNING (or
+    one flagged ERROR/FAILED whose HF job lingers) is cancelled so it stops
+    consuming GPU/billing after the submission has already been failed.
+    """
+    for st in shards.values():
+        if st.get("stage") != "COMPLETED" and st.get("job_id"):
+            _cancel_shard_job(st)
+
+
 def _poll_shards_until_done(
     submission_id: str,
     submission_blob_url: str,
@@ -1322,24 +1372,29 @@ def _poll_shards_until_done(
 ) -> list[str]:
     """Poll every shard to terminal, retrying only ERROR shards.
 
-    Mirrors the orchestrator's eval poll loop: a single thread sweeps
-    all running shards each tick (``inspect_job`` calls are cheap), an
-    ERROR shard re-dispatches up to :data:`SHARD_MAX_RETRIES` times,
-    and a non-terminal stage just waits. Returns a list of
-    ``"<shard_id>: <reason>"`` strings for shards that stayed ERROR
-    after their retries (empty list means every shard COMPLETED).
+    A single thread sweeps all shards each tick (``inspect_job`` calls
+    are cheap). A shard that is QUEUED (waiting for a10g-large capacity
+    past the account's concurrent-job cap) or RUNNING is *healthy* and
+    is simply waited on — its own job ``EVAL_JOB_TIMEOUT`` bounds RUNNING
+    work, and HF turns a genuinely-hung job into ERROR. Only an ERROR
+    shard is re-dispatched, up to :data:`SHARD_MAX_RETRIES`. The loop
+    never cancels a healthy QUEUED/RUNNING shard: doing so throws away
+    its place in the GPU queue and makes a capacity crunch strictly
+    worse (the bug that stranded otherwise-successful runs).
+
+    Returns a list of ``"<shard_id>: <reason>"`` for shards that did not
+    COMPLETE (empty means every shard COMPLETED). This is a *signal, not
+    a verdict*: the caller reconciles against durable staging artifacts
+    before declaring failure, because a shard flagged ERROR or caught by
+    the backstop may still have written all of its per-fixture outputs.
     Transient ``inspect_job`` failures retry up to
     :data:`JOB_POLL_MAX_CONSECUTIVE_ERRORS` before raising.
 
-    If the :data:`SHARD_POLL_DEADLINE_SECONDS` window elapses with shards
-    still non-terminal (the GPU-capacity-starvation case, where a tail
-    shard sits QUEUED), those stragglers are cancelled and re-dispatched
-    and the window resets, up to :data:`SHARD_DEADLINE_RETRY_ROUNDS`
-    rounds, before the submission is finally failed. All-or-nothing is
-    preserved: the list is non-empty unless every shard COMPLETED.
+    :data:`SHARD_POLL_DEADLINE_SECONDS` is a single generous absolute
+    backstop (hours), not a per-round reaper: it only stops the loop
+    spinning forever if a shard never reaches a terminal stage.
     """
     deadline = time.monotonic() + SHARD_POLL_DEADLINE_SECONDS
-    deadline_rounds_left = SHARD_DEADLINE_RETRY_ROUNDS
     consecutive_errors = 0
     last_done = -1
     total = len(shards)
@@ -1402,46 +1457,22 @@ def _poll_shards_until_done(
                         "Shard %s FAILED after %d attempt(s): %s",
                         shard_id, st["attempts"], st["message"],
                     )
+            # QUEUED / RUNNING (or any other non-terminal stage): healthy,
+            # just not done. Wait — the job's own EVAL_JOB_TIMEOUT bounds
+            # it. No Space-side cancellation or re-dispatch.
 
         if time.monotonic() >= deadline:
-            stragglers = [
-                sid for sid, st in shards.items()
-                if st["stage"] not in ("COMPLETED", "FAILED")
-            ]
-            if stragglers and deadline_rounds_left > 0:
-                deadline_rounds_left -= 1
-                logger.warning(
-                    "Poll deadline (%ds) hit for %s with %d straggler shard(s) "
-                    "%s; re-dispatching (%d round(s) left).",
-                    SHARD_POLL_DEADLINE_SECONDS, submission_id,
-                    len(stragglers), stragglers, deadline_rounds_left,
-                )
-                for sid in stragglers:
-                    st = shards[sid]
-                    _cancel_shard_job(st)
-                    # Give the replacement a fresh ERROR-retry budget too.
-                    st["attempts"] = 0
-                    _dispatch_shard(
-                        submission_id, submission_blob_url, sid, st,
-                    )
-                progress.publish(
-                    submission_id,
-                    progress.RUNNING,
-                    f"GPU capacity was tight — retrying {len(stragglers)} "
-                    f"straggler chunk(s) (round "
-                    f"{SHARD_DEADLINE_RETRY_ROUNDS - deadline_rounds_left} of "
-                    f"{SHARD_DEADLINE_RETRY_ROUNDS})…",
-                )
-                deadline = time.monotonic() + SHARD_POLL_DEADLINE_SECONDS
-                last_done = -1  # force a progress republish on the next sweep
-                continue
+            # Absolute backstop reached. Do NOT cancel here — a shard may
+            # be finishing right now, and its outputs are the source of
+            # truth. Mark the non-terminal ones so the caller reconciles
+            # against staging; genuinely-orphaned jobs are torn down on
+            # the failure path once reconcile confirms they're incomplete.
             for shard_id, st in shards.items():
                 if st["stage"] not in ("COMPLETED", "FAILED"):
                     st["stage"] = "FAILED"
                     st["message"] = (
-                        f"Space-side poll deadline exceeded "
-                        f"({SHARD_POLL_DEADLINE_SECONDS}s) after "
-                        f"{SHARD_DEADLINE_RETRY_ROUNDS} retry round(s)"
+                        f"shard did not reach a terminal stage within the "
+                        f"{SHARD_POLL_DEADLINE_SECONDS}s poll backstop"
                     )
             break
         time.sleep(JOB_POLL_INTERVAL_SECONDS)
@@ -1806,21 +1837,94 @@ def _flip_row_to_failed(submission_id: str, reason: str) -> None:
     )
 
 
+# ---------------------------------------------------------------------------
+# Reconcile: finalize a submission from durable artifacts alone
+# ---------------------------------------------------------------------------
+
+
+def _expected_fixtures(submission_id: str) -> list[str]:
+    """The fixture (sample) names a submission was expected to cover.
+
+    Read from the submitted ``submissions/<id>.zip`` top-level sample
+    dirs — the same set the eval fans out over. Used to verify staging
+    shards are complete before finalizing a reconstructed run, without
+    threading the in-memory fixture list through a dead worker.
+    """
+    zip_path = hf_hub_download(
+        repo_id=HF_SUBMISSIONS_REPO,
+        repo_type="dataset",
+        filename=f"{SUBMISSIONS_DIR}/{submission_id}.zip",
+        force_download=True,
+    )
+    with zipfile.ZipFile(zip_path) as zf:
+        tops = {n.split("/")[0] for n in zf.namelist() if "/" in n}
+    return sorted(
+        t for t in tops
+        if t and t != "meta.json" and not t.startswith("__")
+    )
+
+
+def _reconcile_submission(submission_id: str) -> tuple[str, str]:
+    """Try to finalize a stranded submission from artifacts alone.
+
+    Outputs are the source of truth: a submission whose eval actually
+    produced results — a published report bundle, or staging shards that
+    cover every expected fixture — is completed even though its worker
+    died and left the row ``pending``. Only a submission with no usable
+    artifacts is failed. Idempotent and safe to re-run.
+
+    Returns ``(outcome, detail)`` with ``outcome`` in
+    ``{"completed", "failed"}``; the row flip is performed here.
+    """
+    # 1. Single-job (or already-merged) path: a complete report bundle
+    #    means the eval finished and only the row flip was lost.
+    try:
+        summary = _fetch_run_summary_from_report(submission_id)
+    except Exception:  # noqa: BLE001 - no usable report; try shard merge
+        summary = None
+    if summary is not None:
+        _flip_row_to_completed(submission_id, summary)
+        return "completed", "recovered from an existing report bundle"
+
+    # 2. Sharded path: merge the staging shards iff they cover every
+    #    expected fixture. _merge_shards_and_publish raises otherwise.
+    try:
+        expected = _expected_fixtures(submission_id)
+    except Exception as e:  # noqa: BLE001 - can't even read the submission
+        reason = f"reconcile could not read submission fixtures: {e}"
+        _flip_row_to_failed(submission_id, reason[:FAILURE_REASON_MAX_CHARS])
+        return "failed", reason
+    try:
+        summary = _merge_shards_and_publish(submission_id, [], expected)
+    except Exception as e:  # noqa: BLE001 - no complete artifacts
+        reason = f"{STUCK_PENDING_REASON}; no complete artifacts to reconcile"
+        _flip_row_to_failed(submission_id, reason[:FAILURE_REASON_MAX_CHARS])
+        return "failed", f"{reason} ({e})"
+
+    _flip_row_to_completed(submission_id, summary)
+    _cleanup_shard_artifacts(submission_id)
+    return "completed", f"merged {len(expected)} fixtures from staging"
+
+
 # ---------------------------------------------------------------------------
 # Boot-time stuck-pending sweep
 # ---------------------------------------------------------------------------
 
 
 def _sweep_stuck_pending() -> None:
-    """Flip pending rows older than the threshold to failed.
-
-    A ``pending`` row whose worker died (Space restart, OOM, crash)
-    has no one to flip it; without this sweep it stays pending in
-    the leaderboard forever. The check is "submitted_at older than
-    30 min" - well above the real eval ceiling (~5 min on
-    cpu-upgrade), so any genuinely-still-running submission is safe.
-    Runs once per process at module-import time inside a daemon
-    thread so app boot doesn't block on the Hub read.
+    """Reconcile pending rows whose worker died, then decide their fate.
+
+    A ``pending`` row whose worker died (Space restart, OOM, crash) has
+    no one to finalize it. At boot no worker thread from the previous
+    process survives, so any such stale row is orphaned. Rather than
+    blindly failing it — which is exactly how a fully-successful eval got
+    marked failed with its outputs stranded — each stale row is run
+    through :func:`_reconcile_submission`: if its report bundle or
+    staging shards are complete it is finalized to ``completed``, else it
+    is failed. The "submitted_at older than 30 min" gate avoids racing a
+    submission whose worker is legitimately starting in *this* process.
+    Runs once per process at module-import time inside a daemon thread so
+    app boot doesn't block on the Hub read.
     """
     try:
         body = _download_results_jsonl()
@@ -1862,19 +1966,22 @@ def _sweep_stuck_pending() -> None:
             stuck_ids.append(sid)
 
     if not stuck_ids:
-        logger.info("Stuck-pending sweep: nothing stale")
+        logger.info("Boot reconcile sweep: nothing stale")
         return
 
     logger.warning(
-        "Stuck-pending sweep: flipping %d row(s) to failed: %s",
+        "Boot reconcile sweep: reconciling %d stale pending row(s): %s",
         len(stuck_ids), stuck_ids,
     )
     for sid in stuck_ids:
         try:
-            _flip_row_to_failed(sid, STUCK_PENDING_REASON)
+            outcome, detail = _reconcile_submission(sid)
+            logger.warning(
+                "Boot reconcile: %s -> %s (%s)", sid, outcome, detail,
+            )
         except Exception as e:  # noqa: BLE001 - log + carry on per-row
             logger.exception(
-                "Stuck-pending flip failed for %s (%s: %s)",
+                "Boot reconcile failed for %s (%s: %s)",
                 sid, type(e).__name__, e,
             )