From 4f05a4265e9d3c1f62f27c107bdc06e57ebf4ade Mon Sep 17 00:00:00 2001 From: Alex Gronskiy Date: Tue, 29 Sep 2026 17:47:21 +0200 Subject: [PATCH] feat(gdpval): allow partial_completion on the final stage The final stage had no coverage floor: a run either reported exactly N/N rows or was flagged final_stage_degraded, with nothing in between. Three 0.16.0 runs landed 215-217/220 and were blocked from certification. partial_completion may now be set on the final stage. When the policy accepts it, the orchestrator stamps accepted_stage_row_count on the kept rows, and aggregation reports final_stage_partial_accepted=1, final_stage_degraded=0, and emits the flat headline. Below the floor the stage is still scored but degraded, as before. final_stage_complete stays exact. Refs #3341 Co-Authored-By: Claude Opus 5.5 Signed-off-by: Alex Gronskiy --- benchmarks/gdpval/README.md | 25 +++- resources_servers/gdpval/app.py | 21 +++- resources_servers/gdpval/multistage_elo.py | 9 +- .../gdpval/multistage_orchestrator.py | 88 +++++++------- resources_servers/gdpval/tests/test_app.py | 48 ++++++++ .../tests/test_multistage_orchestrator.py | 112 ++++++++++++++---- 6 files changed, 227 insertions(+), 76 deletions(-) diff --git a/benchmarks/gdpval/README.md b/benchmarks/gdpval/README.md index 307c52b9c16..3bc92710043 100644 --- a/benchmarks/gdpval/README.md +++ b/benchmarks/gdpval/README.md @@ -160,11 +160,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 @@ -208,7 +221,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 with overall/per-reference success floors. | +| `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; on the final stage it sets the floor below which the headline is reported as degraded. | | `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. | diff --git a/resources_servers/gdpval/app.py b/resources_servers/gdpval/app.py index 707f5950d4a..d4f78e4cc45 100644 --- a/resources_servers/gdpval/app.py +++ b/resources_servers/gdpval/app.py @@ -1429,6 +1429,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: @@ -1441,6 +1442,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)) @@ -1563,15 +1568,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: diff --git a/resources_servers/gdpval/multistage_elo.py b/resources_servers/gdpval/multistage_elo.py index bc0c58e8d5c..bc7def9d305 100644 --- a/resources_servers/gdpval/multistage_elo.py +++ b/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 diff --git a/resources_servers/gdpval/multistage_orchestrator.py b/resources_servers/gdpval/multistage_orchestrator.py index 1f8494cd7a1..a850d71d8fe 100644 --- a/resources_servers/gdpval/multistage_orchestrator.py +++ b/resources_servers/gdpval/multistage_orchestrator.py @@ -178,6 +178,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 ( @@ -729,8 +733,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): @@ -934,8 +936,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: @@ -1107,7 +1107,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, {}) @@ -1277,7 +1277,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: @@ -1300,43 +1300,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, @@ -1534,6 +1537,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", [])) diff --git a/resources_servers/gdpval/tests/test_app.py b/resources_servers/gdpval/tests/test_app.py index cf918ee0f37..c54ac05cc92 100644 --- a/resources_servers/gdpval/tests/test_app.py +++ b/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 diff --git a/resources_servers/gdpval/tests/test_multistage_orchestrator.py b/resources_servers/gdpval/tests/test_multistage_orchestrator.py index 01e1acda7b5..389bf7e86d2 100644 --- a/resources_servers/gdpval/tests/test_multistage_orchestrator.py +++ b/resources_servers/gdpval/tests/test_multistage_orchestrator.py @@ -39,6 +39,7 @@ ) from resources_servers.gdpval.multistage_elo import PartialStagePolicy from resources_servers.gdpval.multistage_orchestrator import ( + ACCEPTED_STAGE_ROW_COUNT_KEY, MultiStageRunConfig, StageResume, _cached_partial_snapshot_is_valid, @@ -194,17 +195,18 @@ def test_rejects_unsafe_partial_completion_policy(self, partial_completion: Dict } ) - def test_rejects_partial_completion_on_final_stage(self) -> None: - with pytest.raises(ValueError, match="non-final calibration stages"): - parse_multistage_config( - { - "enabled": True, - "stages": [ - {"num_tasks": 5}, - {"num_tasks": 10, "partial_completion": {"min_success_fraction": 0.9}}, - ], - } - ) + def test_accepts_partial_completion_on_final_stage(self) -> None: + cfg = parse_multistage_config( + { + "enabled": True, + "stages": [ + {"num_tasks": 5}, + {"num_tasks": 10, "partial_completion": {"min_success_fraction": 0.9}}, + ], + } + ) + assert cfg.stages[0].partial_completion is None + assert cfg.stages[-1].partial_completion == PartialStagePolicy(min_success_fraction=0.9) class TestFindReferenceElos: @@ -2112,19 +2114,83 @@ async def omit_one_result(rows_in: List[Dict[str, Any]]): assert len(summaries) == 1 assert resume.completed == [] - async def test_direct_final_stage_policy_is_rejected(self) -> None: - task_ids = ["t0"] - cfg = self._config(enabled=True) - cfg.stages[-1].partial_completion = cfg.stages[0].partial_completion + @staticmethod + def _final_stage_only_config(min_success_fraction: float) -> MultiStageRunConfig: + return MultiStageRunConfig( + enabled=True, + stages=parse_multistage_config( + { + "enabled": True, + "stages": [ + { + "num_tasks": 10, + "partial_completion": { + "min_success_fraction": min_success_fraction, + "min_per_reference_success_fraction": 0.5, + }, + } + ], + } + ).stages, + seed=0, + ) - with pytest.raises(ValueError, match="non-final calibration stages"): - await run_multistage_stages( - cfg, - {"a": 1000.0}, - _distribution(task_ids), - _materialized_rows(task_ids), - _fake_run_rollouts_factory(), - ) + async def test_final_stage_policy_stamps_accepted_row_count(self) -> None: + task_ids = [f"t{i}" for i in range(10)] + resume = self._balanced_resume(task_ids) + + results, summaries = await run_multistage_stages( + self._final_stage_only_config(0.9), + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + self._runner({0}, failure_class="skipped", terminal=True), + resume=resume, + ) + + assert len(results) == 9 + assert {row[ACCEPTED_STAGE_ROW_COUNT_KEY] for row in results} == {9} + assert summaries[0]["partial"] is True + [(index, outcome)] = resume.completed + assert index == 0 + assert outcome["status"] == "partial_complete" + assert len(outcome["omitted_keys"]) == 1 + + cached = RecordingResume( + plans=resume.plans, + outcomes={0: outcome}, + rows_by_stage={0: [row for row in resume.appended[0] if NG_FAILURE_CLASS_KEY not in row]}, + gated_keys={0: {(i, 0) for i in range(10)}}, + ) + dispatched: List[int] = [] + resumed, _ = await run_multistage_stages( + self._final_stage_only_config(0.9), + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + self._runner(set(), dispatched_stages=dispatched), + resume=cached, + ) + assert dispatched == [] + assert {row[ACCEPTED_STAGE_ROW_COUNT_KEY] for row in resumed} == {9} + + async def test_final_stage_below_policy_floor_is_not_stamped(self) -> None: + task_ids = [f"t{i}" for i in range(10)] + resume = self._balanced_resume(task_ids) + + results, summaries = await run_multistage_stages( + self._final_stage_only_config(0.9), + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + self._runner({0, 1}, failure_class="skipped", terminal=True), + resume=resume, + ) + + assert len(results) == 8 + assert all(ACCEPTED_STAGE_ROW_COUNT_KEY not in row for row in results) + assert "incomplete" not in summaries[0] + assert resume.completed == [(0, {"stage_index": 0, "status": "complete"})] async def test_invalid_directly_constructed_policy_is_rejected(self) -> None: task_ids = ["t0"]