feat(gdpval): allow partial_completion on the final stage

benchmarks/gdpval/README.md

@@ -188,11 +188,24 @@ balanced before dispatch. When the policy is added while resuming an existing
 strict run, the recorded assignment remains authoritative so the successful
 rows are reused exactly as collected.
 
-Partial completion is deliberately unavailable on the final stage. Its accepted
-row set and omitted keys are frozen in the stage journal, so resume cannot absorb
-late rows and silently change downstream reference selection. Enabling this
-policy while resuming an incomplete stage reuses its successful rollout evidence
-instead of invalidating the rollout cache.
+A stage's accepted row set and omitted keys are frozen in the stage journal, so
+resume cannot absorb late rows and silently change downstream reference
+selection. Enabling this policy while resuming an incomplete stage reuses its
+successful rollout evidence instead of invalidating the rollout cache.
+
+### Accepting a partial final stage
+
+The same `partial_completion` block may be set on the final stage. There it sets
+the coverage floor for the headline: a final stage the policy accepts reports
+`comparison/final_stage_partial_accepted: 1` and `comparison/final_stage_degraded: 0`,
+and emits the flat `comparison/eval_elo` / `comparison/normalized_elo`, even when
+`comparison/observed_final_stage_row_count` is below the expected count.
+`comparison/final_stage_complete` stays exact. A final stage below the floor is
+still scored, but reports `final_stage_degraded: 1` and no flat headline.
+
+```bash
+++multistage.stages='[{num_tasks: 45, partial_completion: {...}}, {num_tasks: 220, num_models: 4, partial_completion: {min_success_fraction: 0.95, min_per_reference_success_fraction: 0.5, min_successful_rows_per_reference: 1}}]'
+```
 
 ### Fresh vs. cached deliverables
 
@@ -236,7 +249,7 @@ full run:
 
 | Key | Default | Meaning |
 |-----|---------|---------|
-| `stages` | *(required)* | List of `{num_tasks?, num_models?, seed?, partial_completion?}` (or `"[num_tasks]:[num_models]:seed"` strings). `num_tasks` omitted ⇒ full task set; `num_models` omitted ⇒ all references. `partial_completion` is an opt-in non-final calibration policy (see [above](#accepting-partial-calibration-after-timeouts)) with keys `min_success_fraction` (default `1.0`), `min_per_reference_success_fraction` (`1.0`), `min_successful_rows_per_reference` (`1`), `waivable_failure_classes` (`[timeout_exceeded]`) and `tolerate_unresolved` (`false`). |
+| `stages` | *(required)* | List of `{num_tasks?, num_models?, seed?, partial_completion?}` (or `"[num_tasks]:[num_models]:seed"` strings). `num_tasks` omitted ⇒ full task set; `num_models` omitted ⇒ all references. `partial_completion` is an opt-in policy with overall/per-reference success floors (see [above](#accepting-partial-calibration-after-timeouts)); on the final stage it sets the floor below which the headline is reported as degraded. Keys: `min_success_fraction` (default `1.0`), `min_per_reference_success_fraction` (`1.0`), `min_successful_rows_per_reference` (`1`), `waivable_failure_classes` (`[timeout_exceeded]`) and `tolerate_unresolved` (`false`). |
 | `column` | `[occupation]` | Dataset column(s) the task sample is drawn proportionally over. |
 | `distribution_path` | *(auto)* | Reuse/write the task-distribution JSON here; built from the dataset when absent. |
 | `dataset_path` | *(prepared dataset)* | Dataset the distribution is built from. |

resources_servers/gdpval/app.py

@@ -1460,6 +1460,7 @@ def _fit_mle(ref_totals: Dict[str, List[Any]]) -> tuple[Optional[float], Optiona
         # appears once per stage — distinguished only by ``stage_index``.
         staged: Dict[int, List[Dict[str, Any]]] = {}
         expected_stage_row_counts: Dict[int, Set[int]] = {}
+        accepted_stage_row_counts: Dict[int, Set[Optional[int]]] = {}
         expected_final_stage_values: Set[int] = set()
         expected_final_stage_rows = 0
         for vr in body.verify_responses:
@@ -1472,6 +1473,10 @@ def _fit_mle(ref_totals: Dict[str, List[Any]]) -> tuple[Optional[float], Optiona
                     expected_stage_row_counts.setdefault(normalized_stage_index, set()).add(
                         int(expected_stage_row_count)
                     )
+                accepted_stage_row_count = vr.get("accepted_stage_row_count")
+                accepted_stage_row_counts.setdefault(normalized_stage_index, set()).add(
+                    int(accepted_stage_row_count) if accepted_stage_row_count is not None else None
+                )
             expected_final_stage_index = vr.get("expected_final_stage_index")
             if expected_final_stage_index is not None:
                 expected_final_stage_values.add(int(expected_final_stage_index))
@@ -1594,15 +1599,27 @@ def _fit_mle(ref_totals: Dict[str, List[Any]]) -> tuple[Optional[float], Optiona
                     and expected_final_stage_count is not None
                     and observed_final_stage_count == expected_final_stage_count
                 )
+                # A final stage the orchestrator accepted under its
+                # partial-completion policy stamps the accepted row count on
+                # every row it kept; the count must match what is observed.
+                accepted_count_values = accepted_stage_row_counts.get(expected_final_stage_index, set())
+                final_stage_partial_accepted = (
+                    final_stage_present
+                    and not final_stage_complete
+                    and len(accepted_count_values) == 1
+                    and next(iter(accepted_count_values)) == observed_final_stage_count
+                )
+                final_stage_accepted = final_stage_complete or final_stage_partial_accepted
                 candidate = stage_fits.get(expected_final_stage_index)
                 final_stage_fit = candidate is not None and candidate[0] is not None
-                if final_stage_complete and final_stage_fit:
+                if final_stage_accepted and final_stage_fit:
                     headline_stage_index = expected_final_stage_index
                     headline = candidate
                 extra["comparison/final_stage_present"] = int(final_stage_present)
                 extra["comparison/final_stage_complete"] = int(final_stage_complete)
+                extra["comparison/final_stage_partial_accepted"] = int(final_stage_partial_accepted)
                 extra["comparison/final_stage_fit"] = int(final_stage_fit)
-                extra["comparison/final_stage_degraded"] = int(not (final_stage_complete and final_stage_fit))
+                extra["comparison/final_stage_degraded"] = int(not (final_stage_accepted and final_stage_fit))
                 extra["comparison/observed_final_stage_row_count"] = observed_final_stage_count
                 extra["comparison/expected_final_stage_row_count_consistent"] = int(expected_count_consistent)
                 if expected_final_stage_count is not None:

resources_servers/gdpval/multistage_elo.py

@@ -61,7 +61,7 @@
 
 @dataclass
 class PartialStagePolicy:
-    """Opt-in policy for accepting incomplete non-final calibration stages.
+    """Opt-in policy for accepting incomplete stages.
 
     ``min_success_fraction`` gates successful rollout coverage across the whole
     stage. ``min_per_reference_success_fraction`` and
@@ -109,9 +109,10 @@ class StageSpec:
     available references" (used for the first, broad stage). Each task is judged
     against **one** reference sampled uniformly from the included set. ``seed``
     makes this stage's task sampling and per-task reference assignment
-    reproducible. ``partial_completion`` is an explicit, non-final-stage-only
-    escape hatch for accepting a calibrated ELO with bounded missing evidence;
-    it is disabled by default.
+    reproducible. ``partial_completion`` is an explicit escape hatch for
+    accepting a stage ELO with bounded missing evidence; it is disabled by
+    default. On the final stage it also sets the coverage floor below which
+    aggregation reports the headline as degraded.
     """
 
     num_tasks: Optional[int] = None

resources_servers/gdpval/multistage_orchestrator.py

@@ -216,6 +216,10 @@ def restart_from(self, index: int) -> None:
 # rather than reproducing the same outcome at full rollout cost.
 _IN_PROCESS_RETRYABLE_CLASSES = frozenset({"timeout_exceeded", "transient"})
 
+# Stamped on every row of a stage accepted by its partial-completion policy, so
+# aggregation can tell an accepted partial final stage from a degraded one.
+ACCEPTED_STAGE_ROW_COUNT_KEY = "accepted_stage_row_count"
+
 
 def _is_in_process_retryable(row: Mapping[str, Any]) -> bool:
     return (
@@ -773,8 +777,6 @@ def parse_multistage_config(raw: Mapping[str, Any]) -> MultiStageRunConfig:
             "multistage.enabled=true but no stages were configured. Set "
             "multistage.stages, e.g. ++multistage.stages='[{num_tasks: 110, num_models: 12}, {num_models: 4}]'."
         )
-    if stages[-1].partial_completion is not None:
-        raise ValueError("partial_completion is allowed only on non-final calibration stages")
 
     column = raw.get("column") or raw.get("columns") or ["occupation"]
     if isinstance(column, str):
@@ -978,8 +980,6 @@ async def run_multistage_stages(
     for stage in multistage_config.stages:
         if stage.partial_completion is not None:
             _validate_partial_stage_policy(stage.partial_completion)
-    if total_stages and multistage_config.stages[-1].partial_completion is not None:
-        raise ValueError("partial_completion is allowed only on non-final calibration stages")
 
     def _emit(name: str, **data: object) -> None:
         if on_event is not None:
@@ -1151,7 +1151,7 @@ def _emit(name: str, **data: object) -> None:
         # persisted timeouts.  Evaluate that frozen evidence before dispatch so
         # the already-timed-out rows are not forced through another long attempt.
         pre_dispatch_partial_outcome: Optional[Dict[str, Any]] = None
-        if resume is not None and index < total_stages - 1 and stage.partial_completion is not None and pending_rows:
+        if resume is not None and stage.partial_completion is not None and pending_rows:
             pending_keys_for_policy = {(row[TASK_INDEX_KEY_NAME], row[ROLLOUT_INDEX_KEY_NAME]) for row in pending_rows}
             latest_failures = resume.latest_failures_by_stage.get(index, {})
             latest_dispositions = resume.latest_attempt_dispositions_by_stage.get(index, {})
@@ -1321,7 +1321,7 @@ def _emit(name: str, **data: object) -> None:
 
         # Outcomes contain no authoritative ELO; it is re-fit from rows on
         # resume. A retryable failure or drained row leaves the stage open unless
-        # an explicit non-final partial-completion policy accepts its evidence.
+        # an explicit partial-completion policy accepts its evidence.
         returned_keys: set[Tuple[Any, Any]] = set()
         prior_attempts = resume.attempts_by_stage.get(index, {}) if resume is not None else {}
         for result in new_tagged:
@@ -1344,43 +1344,46 @@ def _emit(name: str, **data: object) -> None:
         successful_keys = {_stage_key(row) for row in tagged}
         planned_keys = {_stage_key(row) for row in stage_rows}
         missing_success_keys = planned_keys - successful_keys
-        if index < total_stages - 1:
-            if stage.partial_completion is not None:
-                coverage_outcome = partial_outcome or _partial_stage_outcome(
-                    stage.partial_completion,
-                    stage_rows,
-                    tagged,
-                    new_tagged,
-                    unresolved_keys,
-                    reference_ids,
-                    stage_elo,
-                    num_references,
-                )
-                if missing_success_keys:
-                    partial_outcome = coverage_outcome
-                    if partial_outcome is not None:
-                        if "stage_index" not in partial_outcome:
-                            partial_outcome = {"stage_index": index, **partial_outcome}
-                        stage_complete = True
-                    else:
-                        stage_complete = False
-                        coverage_rejected = not unresolved_keys
-                elif coverage_outcome is None:
-                    # Even a persisted "success" must contain usable battle
-                    # evidence for every configured coverage gate.
-                    stage_complete = False
-                    coverage_rejected = True
-            else:
-                # Terminal/max-attempt means "do not retry", not "safe adaptive
-                # calibration".  Without an explicit partial policy every
-                # planned non-final row must contribute usable battle evidence.
-                if (
-                    stage_elo is None
-                    or not math.isfinite(stage_elo)
-                    or _fit_eligible_stage_keys(stage_rows, tagged) != planned_keys
-                ):
+        is_final_stage = index == total_stages - 1
+        if stage.partial_completion is not None:
+            coverage_outcome = partial_outcome or _partial_stage_outcome(
+                stage.partial_completion,
+                stage_rows,
+                tagged,
+                new_tagged,
+                unresolved_keys,
+                reference_ids,
+                stage_elo,
+                num_references,
+            )
+            if missing_success_keys:
+                partial_outcome = coverage_outcome
+                if partial_outcome is not None:
+                    if "stage_index" not in partial_outcome:
+                        partial_outcome = {"stage_index": index, **partial_outcome}
+                    stage_complete = True
+                elif not is_final_stage:
                     stage_complete = False
                     coverage_rejected = not unresolved_keys
+            elif coverage_outcome is None and not is_final_stage:
+                # Even a persisted "success" must contain usable battle
+                # evidence for every configured coverage gate.
+                stage_complete = False
+                coverage_rejected = True
+        elif not is_final_stage:
+            # Terminal/max-attempt means "do not retry", not "safe adaptive
+            # calibration".  Without an explicit partial policy every
+            # planned non-final row must contribute usable battle evidence.
+            if (
+                stage_elo is None
+                or not math.isfinite(stage_elo)
+                or _fit_eligible_stage_keys(stage_rows, tagged) != planned_keys
+            ):
+                stage_complete = False
+                coverage_rejected = not unresolved_keys
+        if partial_outcome is not None:
+            for row in tagged:
+                row[ACCEPTED_STAGE_ROW_COUNT_KEY] = len(partial_outcome["included_keys"])
         if resume is not None and stage_complete:
             resume.on_outcome(
                 index,
@@ -1593,6 +1596,9 @@ def _resume_complete_stage(
         for r in resume.rows_by_stage.get(index, [])
         if included_keys is None or _stage_key(r) in included_keys
     ]
+    if included_keys is not None:
+        for row in cached_rows:
+            row[ACCEPTED_STAGE_ROW_COUNT_KEY] = len(included_keys)
     plan = resume.plans.get(index, {})
     reference_ids = list(plan.get("reference_ids", []))
     task_ids = list(plan.get("task_ids", []))

resources_servers/gdpval/tests/test_app.py

@@ -1879,6 +1879,54 @@ def test_partial_expected_final_stage_does_not_emit_headline(self) -> None:
         assert metrics["comparison/observed_final_stage_row_count"] == 1
         assert metrics["comparison/expected_final_stage_row_count"] == 2
 
+    @pytest.mark.parametrize(
+        ("accepted_count", "accepted"),
+        [
+            (1, True),
+            (None, False),
+            (2, False),
+        ],
+    )
+    def test_policy_accepted_partial_final_stage_emits_headline(
+        self, accepted_count: int | None, accepted: bool
+    ) -> None:
+        """A final stage accepted by its partial-completion policy is not degraded."""
+        import asyncio as _asyncio
+
+        from nemo_gym.config_types import AggregateMetricsRequest
+
+        server = _server(
+            reward_mode="comparison",
+            reference_models={"low": {"deliverables_dir": "/tmp/low", "elo": 1000.0}},
+        )
+        final_row: dict = {
+            "_ng_task_index": 1,
+            "_ng_rollout_index": 0,
+            "stage_index": 1,
+            "expected_final_stage_index": 1,
+            "expected_stage_row_count": 2,
+            "task_id": "t1",
+            "reward": 1.0,
+            "total_wins": 8,
+            "total_losses": 2,
+            "total_ties": 0,
+            "per_reference": {
+                "low": {"wins": 8, "losses": 2, "ties": 0, "reference_elo": 1000.0},
+            },
+            "response": {},
+        }
+        if accepted_count is not None:
+            final_row["accepted_stage_row_count"] = accepted_count
+
+        metrics = _asyncio.run(
+            server.aggregate_metrics(AggregateMetricsRequest(verify_responses=[final_row]))
+        ).agent_metrics
+
+        assert metrics["comparison/final_stage_complete"] == 0
+        assert metrics["comparison/final_stage_partial_accepted"] == int(accepted)
+        assert metrics["comparison/final_stage_degraded"] == int(not accepted)
+        assert ("comparison/eval_elo" in metrics) is accepted
+
     def test_aggregate_metrics_handles_repeated_task_across_stages(self) -> None:
         """The same ``(task_index, rollout_index)`` may recur across stages (one
         rollout judged against a different reference subset per stage). The