diff --git a/resources_servers/gdpval/configs/gdpval_minimax_multistage_judge.yaml b/resources_servers/gdpval/configs/gdpval_minimax_multistage_judge.yaml new file mode 100644 index 00000000000..a26cde5520a --- /dev/null +++ b/resources_servers/gdpval/configs/gdpval_minimax_multistage_judge.yaml @@ -0,0 +1,193 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# GDPVal overlay: MULTI-STAGE ELO comparison mode judged by a LOCAL, SELF-HOSTED +# MiniMax-M3. You serve MiniMax-M3 yourself (vendor `vllm/vllm-openai:minimax-m3` +# container); gym does NOT spawn it — gym talks to your OpenAI-compatible /v1 +# endpoint. This is the self-hosted judge overlay (see +# gdpval_minimax_selfhosted_judge.yaml) PLUS comparison mode with the reference +# models + ELO anchors the multi-stage driver fits against. +# +# ── STEP 1: serve MiniMax-M3 yourself (see "Local multimodal judge +# (MiniMax-M3)" in benchmarks/gdpval/README.md), then: +# export MINIMAX_BASE_URL=http://$HEAD_IP:8000/v1 # must include /v1 +# export MINIMAX_MODEL=minimax-m3 # match --served-model-name +# curl -s "$MINIMAX_BASE_URL/models" # confirm it's up +# +# ── STEP 2: run multi-stage ELO through the benchmark's driver. `--benchmark +# gdpval` loads the multistage orchestrator + policy model + dataset; +# this overlay swaps in comparison mode + the MiniMax judge; the +# ++multistage.* flags opt into staged ELO: +# +# export PERSIST_DELIVERABLES_DIR=/abs/path/to/eval_deliverables_cache +# gym eval run \ +# --model-type vllm_model --benchmark gdpval --split benchmark \ +# --config resources_servers/gdpval/configs/gdpval_minimax_multistage_judge.yaml \ +# --output results/gdpval_minimax_multistage.jsonl \ +# ++multistage.enabled=true \ +# ++multistage.stages='[{num_tasks: 5}, {num_tasks: 15, num_models: 4}]' \ +# ++num_repeats=2 +# +# This example is a small, smoke-sized shape: stage 1 samples 5 tasks, each judged +# against one reference drawn from ALL the references below, for a coarse ELO; +# stage 2 judges 15 tasks against the 4 anchors nearest the stage-1 estimate for +# the final fit. Grow num_tasks for a reportable run (see +# benchmarks/gdpval/README.md). ++num_repeats=2 gives 2 independent rollouts per +# task in EVERY stage, pooled into each stage's BT fit (there is NO per-stage +# repeats key — one inside a stage dict is silently ignored). The headline +# comparison/eval_elo is the last stage's fit (per-stage values under +# comparison/stage_/*). Add ++multistage.seed= for reproducible task +# sampling / reference selection. +# +# Stage 1 here is strict: every planned row must be judged before stage 2 is +# planned. A timeout or other retryable failure stops the run after stage 1 +# (exit 0, no comparison/eval_elo); a row that can no longer be retried fails +# the run. To accept stage 1 with a bounded number of missing rows, give it a +# partial_completion policy (see benchmarks/gdpval/README.md). +# +# To resume, rerun the same command with --resume, adding RERUN_INCOMPLETE=true +# to reuse finished deliverables and cached judgements. The resume journal is +# discarded (renamed *.stale.) when the dataset, stage shape, seed, +# num_repeats, references or any model/agent/judge setting other than +# endpoints, API keys, headers, RERUN_INCOMPLETE and agent concurrency changes. +# +# NOTE ON MEDIA: MiniMax-M3 reads images + video but has NO audio tower, so audio +# deliverables are dropped with a warning (everything else is still graded). + +# --- Repoint the judge proxy at your self-hosted MiniMax-M3 endpoint --------- +gdpval_judge_model: + responses_api_models: + openai_model: + entrypoint: app.py + # Your MiniMax-M3 server's OpenAI-compatible endpoint. Must include /v1. + openai_base_url: ${oc.env:MINIMAX_BASE_URL,http://127.0.0.1:8000/v1} + openai_api_key: ${oc.env:MINIMAX_API_KEY,sk-local} + # Must match the server's --served-model-name (defaults to the HF id). + openai_model: ${oc.env:MINIMAX_MODEL,MiniMaxAI/MiniMax-M3} + # Ride out overload 502/503s instead of failing after 3 quick retries. + max_http_attempts: 1800 + openai_default_headers: {} + +# --- Comparison mode + references + single MiniMax-M3 judge ------------------- +gdpval_resources_server: + resources_servers: + gdpval: + verified: false + reward_mode: comparison + # Reference models the anchored Bradley-Terry MLE ELO is fit against. The + # multi-stage driver picks the `num_models` nearest anchors per stage. + # + # ANCHOR VINTAGE: Artificial Analysis GDPval-AA v2 live board, read + # 2026-09-09 (nine anchors), 2026-09-11 (deepseek_v4_flash_0731, + # inkling_small, qwen35_122b) and 2026-09-17 (the six under the 2026-09-17 + # note below). Anchors are held FIXED in the fit, so a stale anchor shifts + # every candidate score by roughly the anchor error; re-read them from the + # current board before a reportable run. + # Each anchor is one exact AA row: deepseek_v4_pro is the plain "DeepSeek + # V4 Pro (Reasoning, Max Effort)" row (1223), and deepseek_v4_pro_0813_max + # is the "DeepSeek V4 Pro 0813" row (1493). Use deliverables from the exact + # row variant (release and effort) whose elo is listed, and check the + # variant, not just the family name, when refreshing. + # Leave out references whose deliverables are incomplete: missing + # deliverables inflate the fit. Keep the anchors spread across the board's + # range: a candidate rated above the top anchor is fitted by extrapolation + # past the ceiling. + # The reference deliverables are NOT shipped with Gym: you must provide + # each reference model's per-task `task_/` deliverables yourself + # (the layout the Stirrup agent persists), produced by the exact model + # variant whose `elo` is listed. /gdpval/refs is only a placeholder. + # Each deliverables_dir resolves to $GDPVAL_REFS_ROOT/, defaulting to + # /gdpval/refs/. Three ways to point at your copies: + # - refs share a parent: export GDPVAL_REFS_ROOT=/abs/path/to/refs + # - refs are scattered (e.g. per-run output dirs, which share no parent): + # edit these paths directly in your local checkout, or leave the file + # as-is and override them from `env.yaml` in the repo root — it is + # gitignored and merges over this overlay, so nothing site-specific + # gets committed and no extra flags are needed per run: + # + # gdpval_resources_server: + # resources_servers: + # gdpval: + # reference_models: + # kimi_k26: {deliverables_dir: /abs/path/to/deliverables_cache} + # + # - one-off: ++...gdpval.reference_models..deliverables_dir=/abs/path + reference_models: + deepseek_v4_flash_0731: + deliverables_dir: ${oc.env:GDPVAL_REFS_ROOT,/gdpval/refs}/deepseek_v4_flash_0731 + elo: 1442 + deepseek_v4_pro: + deliverables_dir: ${oc.env:GDPVAL_REFS_ROOT,/gdpval/refs}/deepseek_v4_pro + elo: 1223 + inkling_small: + deliverables_dir: ${oc.env:GDPVAL_REFS_ROOT,/gdpval/refs}/inkling_small + elo: 1191 + glm51_fp8: + deliverables_dir: ${oc.env:GDPVAL_REFS_ROOT,/gdpval/refs}/glm51_fp8 + elo: 1181 + kimi_k26: + deliverables_dir: ${oc.env:GDPVAL_REFS_ROOT,/gdpval/refs}/kimi_k26 + elo: 1096 + nemotron3_ultra: + deliverables_dir: ${oc.env:GDPVAL_REFS_ROOT,/gdpval/refs}/nemotron3_ultra + elo: 1091 + qwen36_35b: + deliverables_dir: ${oc.env:GDPVAL_REFS_ROOT,/gdpval/refs}/qwen36_35b + elo: 992 + qwen35_122b: + deliverables_dir: ${oc.env:GDPVAL_REFS_ROOT,/gdpval/refs}/qwen35_122b + elo: 925 + qwen35_397b: + deliverables_dir: ${oc.env:GDPVAL_REFS_ROOT,/gdpval/refs}/qwen35_397b + elo: 905 + gptoss_120b: + deliverables_dir: ${oc.env:GDPVAL_REFS_ROOT,/gdpval/refs}/gptoss_120b + elo: 745 + gemma4_26b: + deliverables_dir: ${oc.env:GDPVAL_REFS_ROOT,/gdpval/refs}/gemma4_26b + elo: 693 + qwen3_30b_thinking: + deliverables_dir: ${oc.env:GDPVAL_REFS_ROOT,/gdpval/refs}/qwen3_30b_thinking + elo: 274 + # 2026-09-17 additions. Effort is part of the AA row, so the key + # carries it where AA scores efforts separately. + kimi_k3_max: + deliverables_dir: ${oc.env:GDPVAL_REFS_ROOT,/gdpval/refs}/kimi_k3_max + elo: 1548 + qwen38_27b_medium: + deliverables_dir: ${oc.env:GDPVAL_REFS_ROOT,/gdpval/refs}/qwen38_27b_medium + elo: 1430 + gpt55_high: + deliverables_dir: ${oc.env:GDPVAL_REFS_ROOT,/gdpval/refs}/gpt55_high + elo: 1372 + claude_sonnet5_high: + deliverables_dir: ${oc.env:GDPVAL_REFS_ROOT,/gdpval/refs}/claude_sonnet5_high + elo: 1308 + minimax_m3: + deliverables_dir: ${oc.env:GDPVAL_REFS_ROOT,/gdpval/refs}/minimax_m3 + elo: 1304 + # 0813 release at max effort, matching AA's only 0813 row (1493). Not a + # duplicate of deepseek_v4_pro (1223), which is the earlier release. + deepseek_v4_pro_0813_max: + deliverables_dir: ${oc.env:GDPVAL_REFS_ROOT,/gdpval/refs}/deepseek_v4_pro_0813_max + elo: 1493 + # Pairwise trials per (task x reference x repeat). 4 debiases position with + # swap/no-swap; lower to 2 for a faster smoke test. + num_comparison_trials: 4 + # Single self-hosted MiniMax-M3 judge (no frontier panel). + judge_model_server: + type: responses_api_models + name: gdpval_judge_model + judge_responses_create_params_overrides: + model: ${oc.env:MINIMAX_MODEL,MiniMaxAI/MiniMax-M3} + judge_panel: null + # A VLM can't decode an application/pdf data URL, so rasterize PDFs/Office + # docs to page images + extracted text. + judge_media_mode: images_and_text + judge_pdf_render_dpi: 144 + judge_pdf_max_pages: 30 + judge_pdf_include_text: true + # MiniMax-M3 reads video (native video_url blocks) but has no audio tower. + judge_handles_video: true + judge_handles_audio: false + on_missing_av_judge: warn diff --git a/resources_servers/gdpval/multistage_elo.py b/resources_servers/gdpval/multistage_elo.py index 6a8feb18a1c..33cd665ece1 100644 --- a/resources_servers/gdpval/multistage_elo.py +++ b/resources_servers/gdpval/multistage_elo.py @@ -20,8 +20,9 @@ ``responses_api_agents.stirrup_agent.task_distribution``); ``T`` is configurable per stage and **defaults to the full task set**, 2. Includes a set of ``M`` reference models and assigns **each task a single - reference** drawn uniformly (equal weight) from that set — so a task is - judged against one reference, not all ``M``, + reference** from that set — an independent uniform draw by default, with an + opt-in seeded balanced assignment — so a task is judged against one reference, + not all ``M``, 3. Fits an anchored Bradley-Terry MLE ELO from that stage's win/loss/tie battles pooled per reference (reusing ``comparison.calculate_mle_elo``), and 4. Uses that estimate to choose the ``M`` references for the next stage. @@ -47,7 +48,7 @@ import random from dataclasses import dataclass from pathlib import Path -from typing import Any, Dict, List, Mapping, Optional, Sequence +from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple from resources_servers.gdpval.comparison import calculate_mle_elo @@ -58,6 +59,45 @@ PerReferenceTotals = Dict[str, Dict[str, float]] +@dataclass +class PartialStagePolicy: + """Opt-in policy for accepting incomplete non-final calibration stages. + + ``min_success_fraction`` gates successful rollout coverage across the whole + stage. ``min_per_reference_success_fraction`` and + ``min_successful_rows_per_reference`` apply independently to every selected + reference model, so an anchor cannot disappear entirely from a partial fit. + ``waivable_failure_classes`` lists the persisted failure classes that may be + newly waived; it defaults to task timeouts only. Rows already resolved by the standard terminal/max-attempt rules + remain eligible, but every omission still counts against the configured + coverage floors. Drained/no-persist and missing rows keep the stage open. + + ``tolerate_unresolved`` relaxes *only* the failure-class test: an unresolved + row that came back with any persisted failure class is accepted, subject to + the coverage floors. It never waives a row that never returned, that was + drained/no-persist, or that carries no failure class at all -- their absence + is the resume signal that lets the next allocation re-dispatch them intact, + so those keep the stage open as they do without the flag. Set it on a + calibration stage whose omissions are structural -- a missing reference + deliverable does not appear on retry unless the reference tree is repaired, + so holding the stage open for it only costs further launches until the row + reaches ``NEMO_GYM_MAX_ROLLOUT_ATTEMPTS``. + + Enabling it accepts more than structural gaps: *every* persisted class + becomes waivable, including the infrastructure-wide ones + (``agent_request_failed``, ``agent_run_error``, ...). A whole agent-server + outage mid-stage is therefore indistinguishable from scattered structural + omissions, and the coverage floors above are the only remaining guard. Set + them deliberately tight before turning this on. + """ + + min_success_fraction: float = 1.0 + min_per_reference_success_fraction: float = 1.0 + min_successful_rows_per_reference: int = 1 + waivable_failure_classes: Tuple[str, ...] = ("timeout_exceeded",) + tolerate_unresolved: bool = False + + @dataclass class StageSpec: """Configuration for a single stage. @@ -69,12 +109,15 @@ 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. + 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. """ num_tasks: Optional[int] = None num_models: Optional[int] = None seed: Optional[int] = None + partial_completion: Optional[PartialStagePolicy] = None # --------------------------------------------------------------------------- @@ -146,13 +189,16 @@ def assign_task_references( reference_ids: Sequence[str], *, rng: random.Random, + balanced: bool = False, ) -> Dict[str, str]: - """Assign each task a single reference model, sampled with equal probability. + """Assign each task a single reference model. Rather than judging every task against every included reference, each task is - compared against **one** reference drawn uniformly (each included reference - weighted equally) from ``reference_ids``. Returns a ``{task_id: ref_id}`` map; - empty when no references are included. + compared against **one** reference from ``reference_ids``. By default each + assignment is an independent uniform draw. With ``balanced=True``, assignment + slots are shuffled while ensuring per-reference task counts differ by at most + one. Returns a ``{task_id: ref_id}`` map; empty when no references are + included. Deterministic given *rng*, so a stage's assignment replays identically on resume (it is also recorded in the stage journal). @@ -160,7 +206,16 @@ def assign_task_references( refs = list(reference_ids) if not refs: return {} - return {str(tid): rng.choice(refs) for tid in task_ids} + if not balanced: + return {str(tid): rng.choice(refs) for tid in task_ids} + + tasks = [str(tid) for tid in task_ids] + full_cycles, remainder = divmod(len(tasks), len(refs)) + slots = refs * full_cycles + if remainder: + slots.extend(rng.sample(refs, remainder)) + rng.shuffle(slots) + return dict(zip(tasks, slots)) # --------------------------------------------------------------------------- diff --git a/resources_servers/gdpval/multistage_orchestrator.py b/resources_servers/gdpval/multistage_orchestrator.py index 0726833f1eb..b53c7eb7b69 100644 --- a/resources_servers/gdpval/multistage_orchestrator.py +++ b/resources_servers/gdpval/multistage_orchestrator.py @@ -26,10 +26,11 @@ * Each stage is one pass of the standard rollout collection over the stage's sampled ``T`` tasks (``T`` is configurable per stage and defaults to the full task distribution). The stage includes a set of reference models (chosen - adaptively — see below) and assigns **each task a single reference** sampled - uniformly (equal weight) from that set; the task's row is tagged with that one - ``reference_ids=[ref]`` (honored by the GDPVal verifier's per-request reference - filter) and a ``stage_index``. + adaptively — see below) and assigns **each task a single reference** from that + set; the default is an independent uniform sample, while partial-completion + stages use a seeded balanced assignment. The task's row is tagged with that + one ``reference_ids=[ref]`` (honored by the GDPVal verifier's per-request + reference filter) and a ``stage_index``. * Between stages we fit the stage's anchored Bradley-Terry MLE ELO (the same math the server's ``aggregate_metrics`` uses) — pooling each reference's win/loss/tie counts over the tasks assigned to it — to pick the next stage's @@ -54,7 +55,13 @@ from __future__ import annotations import hashlib +import math +import os import random +import secrets +import tempfile +import time +from collections import Counter from copy import deepcopy from dataclasses import dataclass, field from pathlib import Path @@ -62,15 +69,24 @@ import orjson -from nemo_gym.global_config import AGENT_REF_KEY_NAME, ROLLOUT_INDEX_KEY_NAME, TASK_INDEX_KEY_NAME +from nemo_gym.global_config import ( + AGENT_REF_KEY_NAME, + ATTEMPT_INDEX_KEY_NAME, + ROLLOUT_INDEX_KEY_NAME, + TASK_INDEX_KEY_NAME, +) from nemo_gym.path_utils import failures_path_for from nemo_gym.rollout_collection import ( NG_FAILURE_CLASS_KEY, NG_NO_PERSIST_KEY, - NG_TERMINAL_KEY, - _get_max_rollout_attempts, + DispatchLatencyTracker, + get_max_rollout_attempts, + is_terminal_failure, + migrate_invalid_judge_main_rows, + observed_elapsed, ) from resources_servers.gdpval.multistage_elo import ( + PartialStagePolicy, PerReferenceTotals, StageSpec, assign_task_references, @@ -83,10 +99,52 @@ ) +# GDPVal opts into both parts of the shared failure contract, which rollout +# collection keeps off by default: legacy terminal timeouts are retried, and +# rows carrying an invalid judge response are retried instead of scored. The +# single-pass delegate applies each key the config leaves unset; the multistage +# bookkeeping below applies the same contract directly and so refuses a config +# that explicitly turns either key off. +GDPVAL_FAILURE_CONTRACT = {"retry_terminal_timeouts": True, "retry_invalid_judge_responses": True} + +# Top-level server-config fields the resume fingerprint ignores (see +# ``compute_fingerprint``). Connection fields say how a server is reached; +# operational fields are knobs a resume is expected to change without +# affecting any result: ``rerun_incomplete`` (``RERUN_INCOMPLETE``) is how a +# resume reuses cached work, and ``concurrency`` only bounds parallelism. +_CONNECTION_SERVER_FIELDS = frozenset({"host", "port", "base_url", "api_key", "default_headers"}) +_OPERATIONAL_SERVER_FIELDS = frozenset({"rerun_incomplete", "concurrency"}) + + +def _failure_contract_updates(rollout_collection_config: Any) -> Dict[str, bool]: + """The failure-contract keys to apply: every key the config leaves unset.""" + explicitly_set = getattr(rollout_collection_config, "model_fields_set", None) or set() + return {name: value for name, value in GDPVAL_FAILURE_CONTRACT.items() if name not in explicitly_set} + + +def _reject_failure_contract_opt_out(rollout_collection_config: Any) -> None: + """Refuse an explicit ``false`` the multistage bookkeeping cannot honour.""" + explicitly_set = getattr(rollout_collection_config, "model_fields_set", None) or set() + opted_out = sorted( + name + for name in GDPVAL_FAILURE_CONTRACT + if name in explicitly_set and getattr(rollout_collection_config, name, None) is False + ) + if opted_out: + raise ValueError( + f"multistage GDPVal runs always apply the failure contract, but {opted_out} " + "is explicitly set to false; remove the override or run without multistage.enabled" + ) + + # A rollout runner: given a list of fully-formed rollout rows, run them and # return ``(row, result)`` pairs (result == the agent's /run response, i.e. the # GDPVal verify response). Injected so tests can avoid real servers. RolloutRunner = Callable[[List[Dict[str, Any]]], Awaitable[List[Tuple[Dict[str, Any], Dict[str, Any]]]]] +AssignmentRepair = Callable[ + [int, Sequence[str], Mapping[str, str]], + Tuple[Dict[str, str], Dict[str, Any]], +] @dataclass @@ -99,6 +157,8 @@ class StageResume: ``(task_index, rollout_index)`` not to re-dispatch (success, terminal, or max-attempt). ``on_plan``/``on_outcome``/``on_rows`` persist newly produced state (``on_rows`` routes success/failure/kill_shaped like ``run_from_config``). + ``restart_from(i)`` invokes ``on_restart`` and preserves stage ``i``'s plan + and evidence while invalidating its outcome and every dependent later stage. """ plans: Mapping[int, dict] @@ -108,24 +168,400 @@ class StageResume: on_plan: Callable[[int, dict], None] on_outcome: Callable[[int, dict], None] on_rows: Callable[[int, List[Dict[str, Any]]], None] + on_restart: Callable[[int], None] = field(default=lambda _index: None, repr=False) + # Longest observed failed-attempt duration, keyed by stage then + # (task_index, rollout_index). This is the stage-aware equivalent of the + # single-pass dispatcher's resume timing map. + elapsed_by_stage: Mapping[int, Mapping[Tuple[Any, Any], float]] = field(default_factory=dict) + # Failed judging attempts whose persisted policy artifact should be reused + # on retry, keyed by stage then (task_index, rollout_index). + reuse_cached_keys: Mapping[int, AbstractSet[Tuple[Any, Any]]] = field(default_factory=dict) + # Number of sidecar attempts already consumed, keyed stage/task/rollout. + attempts_by_stage: Mapping[int, Mapping[Tuple[Any, Any], int]] = field(default_factory=dict) + # Latest persisted failure per stage/task/rollout. An explicit partial-stage + # policy may accept an already-recorded timeout without dispatching it again. + latest_failures_by_stage: Mapping[int, Mapping[Tuple[Any, Any], Dict[str, Any]]] = field(default_factory=dict) + # Latest attempt disposition written by this orchestrator, including + # drained/no-persist attempts that deliberately never enter the sidecar. + latest_attempt_dispositions_by_stage: Mapping[int, Mapping[Tuple[Any, Any], Dict[str, Any]]] = field( + default_factory=dict + ) + + def restart_from(self, index: int) -> None: + """Durably invalidate a stale outcome and every dependent later stage.""" + self.on_restart(index) + self.plans = {stage: value for stage, value in self.plans.items() if stage <= index} + self.outcomes = {stage: value for stage, value in self.outcomes.items() if stage < index} + for name in ( + "rows_by_stage", + "gated_keys", + "elapsed_by_stage", + "reuse_cached_keys", + "attempts_by_stage", + "latest_failures_by_stage", + "latest_attempt_dispositions_by_stage", + ): + values = getattr(self, name) + setattr(self, name, {stage: value for stage, value in values.items() if stage <= index}) + + +# Classes worth re-dispatching inside the same process. Deliberately narrower +# than "not terminal": is_terminal_failure also spares reference_missing / +# eval_missing / transport_ineligible on the grounds that the deliverable tree +# may be repaired *between* runs. Nothing repairs it mid-process, so retrying +# those here only burns dispatch budget. +# `incomplete`, `judge_invalid` and `legitimate` are retried across resumes up to +# NEMO_GYM_MAX_ROLLOUT_ATTEMPTS but are excluded here too: they are a recorded +# decision, pending evidence that an *immediate* re-dispatch recovers them +# rather than reproducing the same outcome at full rollout cost. +_IN_PROCESS_RETRYABLE_CLASSES = frozenset({"timeout_exceeded", "transient"}) + + +def _is_in_process_retryable(row: Mapping[str, Any]) -> bool: + return ( + not _is_success_row(row) + and not row.get(NG_NO_PERSIST_KEY) + and not is_terminal_failure(row, retry_terminal_timeouts=True) + and row.get(NG_FAILURE_CLASS_KEY) in _IN_PROCESS_RETRYABLE_CLASSES + ) def _is_success_row(row: Mapping[str, Any]) -> bool: """A row is a success iff it carries neither a failure class nor no-persist.""" - return row.get(NG_FAILURE_CLASS_KEY) is None and not row.get(NG_NO_PERSIST_KEY) + return ( + row.get(NG_FAILURE_CLASS_KEY) is None + and not row.get(NG_NO_PERSIST_KEY) + and not row.get("invalid_judge_response") + ) + + +def _stage_key(row: Mapping[str, Any]) -> Tuple[Any, Any]: + return (row.get(TASK_INDEX_KEY_NAME), row.get(ROLLOUT_INDEX_KEY_NAME)) + + +def _fit_eligible_stage_keys( + stage_rows: Sequence[Mapping[str, Any]], successful_rows: Sequence[Mapping[str, Any]] +) -> set[Tuple[Any, Any]]: + """Keys whose persisted success contains a usable assigned-reference battle.""" + planned_by_key = {_stage_key(row): row for row in stage_rows} + successful_by_key = {_stage_key(row): row for row in successful_rows if _stage_key(row) in planned_by_key} + fit_eligible: set[Tuple[Any, Any]] = set() + for key, success in successful_by_key.items(): + refs = list(planned_by_key[key].get("reference_ids") or []) + if not refs: + continue + counts = (success.get("per_reference") or {}).get(str(refs[0])) or {} + if sum(float(counts.get(name, 0) or 0) for name in ("wins", "losses", "ties")) > 0: + fit_eligible.add(key) + return fit_eligible + + +def _recorded_key_set(record: Mapping[str, Any], field_name: str) -> set[Tuple[Any, Any]]: + keys: set[Tuple[Any, Any]] = set() + for value in record.get(field_name, []) or []: + if isinstance(value, (list, tuple)) and len(value) == 2: + keys.add((value[0], value[1])) + return keys + + +def _partial_policy_record(policy: PartialStagePolicy) -> Dict[str, Any]: + return { + "min_success_fraction": policy.min_success_fraction, + "min_per_reference_success_fraction": policy.min_per_reference_success_fraction, + "min_successful_rows_per_reference": policy.min_successful_rows_per_reference, + "newly_waivable_failure_classes": sorted(policy.waivable_failure_classes), + "tolerate_unresolved": policy.tolerate_unresolved, + } + + +def _elo_evidence_sha256(rows: Sequence[Mapping[str, Any]]) -> str: + """Hash exactly the pooled battle evidence that determines a stage ELO.""" + pooled = pool_per_reference(rows) + normalized = { + str(reference_id): { + "wins": int(counts.get("wins", 0) or 0), + "losses": int(counts.get("losses", 0) or 0), + "ties": int(counts.get("ties", 0) or 0), + "reference_elo": counts.get("reference_elo"), + } + for reference_id, counts in pooled.items() + } + return hashlib.sha256(orjson.dumps(normalized, option=orjson.OPT_SORT_KEYS)).hexdigest() + + +def _partial_stage_outcome( + policy: PartialStagePolicy, + stage_rows: Sequence[Mapping[str, Any]], + successful_rows: Sequence[Mapping[str, Any]], + new_results: Sequence[Mapping[str, Any]], + unresolved_keys: AbstractSet[Tuple[Any, Any]], + reference_ids: Sequence[str], + stage_elo: Optional[float], + num_references: int, +) -> Optional[Dict[str, Any]]: + """Return a durable partial outcome when every configured gate passes.""" + if ( + stage_elo is None + or not math.isfinite(stage_elo) + or not stage_rows + or not reference_ids + or num_references != len(reference_ids) + ): + return None + + planned_by_key = {_stage_key(row): row for row in stage_rows} + if len(planned_by_key) != len(stage_rows): + return None + successful_by_key = {_stage_key(row): row for row in successful_rows if _stage_key(row) in planned_by_key} + + planned_per_reference: Counter[str] = Counter() + successful_per_reference: Counter[str] = Counter() + judged_per_reference: Counter[str] = Counter() + fit_eligible_keys: set[Tuple[Any, Any]] = set() + for key, row in planned_by_key.items(): + refs = list(row.get("reference_ids") or []) + if refs: + planned_per_reference[str(refs[0])] += 1 + success = successful_by_key.get(key) + if success is None or not refs: + continue + reference_id = str(refs[0]) + successful_per_reference[reference_id] += 1 + counts = (success.get("per_reference") or {}).get(reference_id) or {} + if sum(float(counts.get(name, 0) or 0) for name in ("wins", "losses", "ties")) > 0: + judged_per_reference[reference_id] += 1 + fit_eligible_keys.add(key) + + # A main-file row without a usable battle is not calibration evidence. Do + # not let it masquerade as a success or become part of the frozen snapshot. + if set(successful_by_key) - fit_eligible_keys: + return None + + success_fraction = len(fit_eligible_keys) / len(planned_by_key) + if success_fraction < policy.min_success_fraction: + return None + + per_reference_success_fractions: Dict[str, float] = {} + for reference_id in reference_ids: + planned = planned_per_reference[reference_id] + if planned <= 0: + # Balanced assignment gives a reference zero tasks whenever + # num_tasks < num_models. A reference with nothing planned has no + # coverage floor to violate; it contributes no evidence. + per_reference_success_fractions[reference_id] = 1.0 + continue + success_fraction_for_reference = judged_per_reference[reference_id] / planned + per_reference_success_fractions[reference_id] = success_fraction_for_reference + if ( + judged_per_reference[reference_id] < policy.min_successful_rows_per_reference + or success_fraction_for_reference < policy.min_per_reference_success_fraction + ): + return None + + planned_keys = list(planned_by_key) + included_keys = [key for key in planned_keys if key in fit_eligible_keys] + omitted_keys = [key for key in planned_keys if key not in fit_eligible_keys] + included_rows = [successful_by_key[key] for key in included_keys] + already_resolved_omitted_keys = set(omitted_keys) - set(unresolved_keys) + result_by_key = {_stage_key(result): result for result in new_results} + for key in unresolved_keys: + result = result_by_key.get(key) + # A row that never returned, or that was drained before its POST, is the + # resume signal for the next allocation: its absence is what makes the + # sample re-dispatchable intact. `tolerate_unresolved` never waives those. + if result is None or result.get(NG_NO_PERSIST_KEY): + return None + # A persisted failure class is required either way: `tolerate_unresolved` + # widens *which* classes are waivable, it never waives a row carrying no + # class at all. + failure_class = result.get(NG_FAILURE_CLASS_KEY) + if failure_class is None: + return None + if not policy.tolerate_unresolved and failure_class not in policy.waivable_failure_classes: + return None + + per_reference = { + reference_id: { + "planned": planned_per_reference[reference_id], + "successful": successful_per_reference[reference_id], + "judged": judged_per_reference[reference_id], + "success_fraction": per_reference_success_fractions[reference_id], + } + for reference_id in reference_ids + } + return { + "status": "partial_complete", + "included_keys": [list(key) for key in included_keys], + "omitted_keys": [list(key) for key in omitted_keys], + "accepted_unresolved_keys": [list(key) for key in planned_keys if key in unresolved_keys], + "already_resolved_omitted_keys": [list(key) for key in planned_keys if key in already_resolved_omitted_keys], + "evidence_sha256": _elo_evidence_sha256(included_rows), + "success_fraction": success_fraction, + "persisted_success_fraction": len(successful_by_key) / len(planned_by_key), + "per_reference": per_reference, + "policy": _partial_policy_record(policy), + } + + +def _cached_partial_snapshot_is_valid( + outcome: Mapping[str, Any], + expected_policy: Optional[PartialStagePolicy], + stage_rows: Sequence[Mapping[str, Any]], + persisted_rows: Sequence[Mapping[str, Any]], + reference_ids: Sequence[str], + reference_elos: Mapping[str, float], +) -> bool: + """Revalidate a frozen partial outcome from its included evidence.""" + omitted_keys = _recorded_key_set(outcome, "omitted_keys") + accepted_unresolved_keys = _recorded_key_set(outcome, "accepted_unresolved_keys") + already_resolved_omitted_keys = _recorded_key_set(outcome, "already_resolved_omitted_keys") + if ( + accepted_unresolved_keys & already_resolved_omitted_keys + or accepted_unresolved_keys | already_resolved_omitted_keys != omitted_keys + ): + return False + + policy_record = outcome.get("policy") + required_policy_fields = { + "min_success_fraction", + "min_per_reference_success_fraction", + "min_successful_rows_per_reference", + "newly_waivable_failure_classes", + } + if expected_policy is None or not isinstance(policy_record, Mapping): + return False + if not required_policy_fields <= set(policy_record): + return False + # `tolerate_unresolved` is optional in the record and reads as the off value + # when absent. + recorded_tolerate_unresolved = policy_record.get("tolerate_unresolved", False) + # It is also non-material to validity and so excluded from the comparison: + # the snapshot is revalidated by re-deriving the outcome from its recorded + # evidence, and this flag does not change that evidence -- only whether a + # stage would be accepted. Flipping it mid-run (its intended use, to break a + # hang) must therefore not invalidate an already-frozen snapshot. + expected_policy_record = _partial_policy_record(expected_policy) + recorded_policy_fields = {field: policy_record[field] for field in required_policy_fields} + if recorded_policy_fields != {field: expected_policy_record[field] for field in required_policy_fields}: + return False + # The durable record spells the waiver set `newly_waivable_failure_classes`; + # the config field is `waivable_failure_classes`. Map it back before parsing. + replayed_policy_fields = { + field: policy_record[field] for field in required_policy_fields if field != "newly_waivable_failure_classes" + } + replayed_policy_fields["waivable_failure_classes"] = policy_record["newly_waivable_failure_classes"] + replayed_policy_fields["tolerate_unresolved"] = recorded_tolerate_unresolved + try: + policy = _parse_partial_stage_policy(replayed_policy_fields) + except (TypeError, ValueError): + return False + if policy is None: + return False + + included_keys = _recorded_key_set(outcome, "included_keys") + frozen_rows = [row for row in persisted_rows if _stage_key(row) in included_keys and _is_success_row(row)] + stage_elo, _, num_references = fit_stage_elo(pool_per_reference(frozen_rows), reference_elos) + candidate = _partial_stage_outcome( + policy, + stage_rows, + frozen_rows, + [], + set(), + reference_ids, + stage_elo, + num_references, + ) + return bool( + candidate is not None + and _recorded_key_set(candidate, "included_keys") == included_keys + and _recorded_key_set(candidate, "omitted_keys") == omitted_keys + and candidate.get("evidence_sha256") == outcome.get("evidence_sha256") + ) def compute_fingerprint( multistage_config: MultiStageRunConfig, reference_elos: Mapping[str, float], distribution: Mapping[str, Mapping[str, object]], + *, + materialized_rows: Optional[Sequence[Mapping[str, Any]]] = None, + rollout_collection_config: Optional[Any] = None, + resolved_global_config: Optional[Mapping[str, Any]] = None, ) -> str: - """Stable hash of everything that affects stage planning. + """Stable hash of the configured inputs that affect stage planning or judged results. A mismatch between the current run's fingerprint and a journal's marks the journal stale: the plans/outcomes it records were produced under a different configuration or task distribution and cannot be safely replayed. + + Not hashed: server endpoints, credentials and headers, the operational + agent knobs in ``_OPERATIONAL_SERVER_FIELDS``, ``partial_completion``, + ``retry_inprocess``, ``NEMO_GYM_MAX_ROLLOUT_ATTEMPTS``, and the contents + of the reference deliverable directories. """ + + def jsonable(value: Any) -> Any: + if hasattr(value, "model_dump"): + return value.model_dump(mode="json") + if isinstance(value, Mapping): + return {str(key): jsonable(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [jsonable(item) for item in value] + if isinstance(value, Path): + return str(value) + if value is None or isinstance(value, (bool, int, float, str)): + return value + return repr(value) + + component_types = {"responses_api_agents", "responses_api_models", "resources_servers"} + + # Connection endpoints, credentials and headers change across allocations + # (head-node IPs, rotated keys) without affecting any judgment, and secrets + # should not be hashed into persisted state at all. Operational knobs that a + # resume is expected to change are dropped for the same reason. Only fields + # at the top level of a server config are stripped; nested config remains + # fingerprint-relevant. + def is_unhashed_field(name: str) -> bool: + # ``endpoint_*`` covers the vLLM proxy's endpoint-file liveness knobs + # (stale grace, connection retries, poll interval): how a replica is + # reached, never what it generates. + return ( + name in _CONNECTION_SERVER_FIELDS + or name in _OPERATIONAL_SERVER_FIELDS + or name.endswith("_api_key") + or name.endswith("_base_url") + or name.endswith("_default_headers") + or name.startswith("endpoint_") + ) + + def jsonable_runtime_block(block: Mapping[str, Any]) -> Dict[str, Any]: + """Serialize a component block without its unhashed top-level server fields.""" + result: Dict[str, Any] = {} + for key, value in block.items(): + key = str(key) + if key not in component_types: + result[key] = jsonable(value) + continue + + servers = jsonable(value) + if not isinstance(servers, Mapping): + result[key] = servers + continue + + result[key] = { + str(server_name): ( + { + str(field): field_value + for field, field_value in server_config.items() + if not is_unhashed_field(str(field)) + } + if isinstance(server_config, Mapping) + else server_config + ) + for server_name, server_config in servers.items() + } + return result + payload = { "stages": [(s.num_tasks, s.num_models, s.seed) for s in multistage_config.stages], "seed": multistage_config.seed, @@ -141,6 +577,52 @@ def compute_fingerprint( for grp in sorted(distribution) }, } + # Keep partial completion out of the rollout-cache fingerprint so an + # interrupted strict run can opt in without discarding successful rows. Its + # recorded plan remains authoritative; only a fresh policy-enabled plan uses + # balanced assignment. A persisted partial outcome records the exact policy + # and is revalidated on every resume. + if materialized_rows is not None: + materialized_hasher = hashlib.sha256() + for row in materialized_rows: + materialized_hasher.update(orjson.dumps(row, option=orjson.OPT_SORT_KEYS)) + materialized_hasher.update(b"\n") + payload["materialized_rows"] = { + "count": len(materialized_rows), + "sha256": materialized_hasher.hexdigest(), + } + if rollout_collection_config is not None: + # ``input_jsonl_fpath`` is deliberately absent: the materialized rows + # above hash the input by content, and the CLI's generated preprocessed + # path moves between releases without a byte of input changing. + result_affecting_fields = ( + "agent_name", + "limit", + "num_repeats", + "num_repeats_add_seed", + "responses_create_params", + "prompt_config", + "skills", + ) + + def config_value(name: str) -> Any: + if isinstance(rollout_collection_config, Mapping): + return rollout_collection_config.get(name) + return getattr(rollout_collection_config, name, None) + + payload["rollout_collection_config"] = {name: jsonable(config_value(name)) for name in result_affecting_fields} + if resolved_global_config is not None: + payload["runtime_components"] = { + str(name): jsonable_runtime_block(block) + for name, block in resolved_global_config.items() + if isinstance(block, Mapping) and component_types.intersection(block) + } + # Transport-assignment repair rewrites which reference each task is + # judged against, so its settings must invalidate journals like any + # other planner input. An absent block adds nothing to the payload. + transport_repair = (resolved_global_config.get("multistage") or {}).get("transport_assignment_repair") + if transport_repair is not None: + payload["transport_assignment_repair"] = jsonable(transport_repair) encoded = orjson.dumps(payload, option=orjson.OPT_SORT_KEYS) return hashlib.sha256(encoded).hexdigest() @@ -169,12 +651,96 @@ class MultiStageRunConfig: # Judge a task's cached deliverable in later stages instead of re-running the # policy. Falls back to a fresh rollout when the deliverable is missing. reuse_cached_deliverables: bool = True + # Re-dispatch rows that failed with `timeout_exceeded` or `transient` inside + # the same process, up to NEMO_GYM_MAX_ROLLOUT_ATTEMPTS. Other retryable + # classes wait for a resume. Off by default: without it such a row keeps its + # stage open, and the run stops after that stage until it is resumed. + retry_inprocess: bool = False + + +# Failure classes a partial-calibration policy may newly waive. A waived row is +# omitted from the fit and still counts against every coverage floor, so this is +# a bound on which *causes* are acceptable, not on how much may be missing. +# `skipped` is excluded: an unusable sample is not evidence of anything. +_WAIVABLE_FAILURE_CLASSES = frozenset({"timeout_exceeded", "transient"}) + + +def _validate_partial_stage_policy(policy: PartialStagePolicy) -> None: + """Validate parsed and directly constructed partial policies.""" + for name in ("min_success_fraction", "min_per_reference_success_fraction"): + value = getattr(policy, name) + if isinstance(value, bool) or not isinstance(value, (int, float)) or not 0 < value <= 1: + raise ValueError(f"partial_completion.{name} must be in (0, 1]") + minimum_rows = policy.min_successful_rows_per_reference + if isinstance(minimum_rows, bool) or not isinstance(minimum_rows, int) or minimum_rows <= 0: + raise ValueError("partial_completion.min_successful_rows_per_reference must be a positive integer") + if not isinstance(policy.tolerate_unresolved, bool): + raise ValueError("partial_completion.tolerate_unresolved must be a boolean") + waivable = policy.waivable_failure_classes + if isinstance(waivable, (str, bytes)) or not isinstance(waivable, Sequence) or not waivable: + raise ValueError("partial_completion.waivable_failure_classes must be a non-empty sequence of strings") + unknown_classes = sorted(set(map(str, waivable)) - _WAIVABLE_FAILURE_CLASSES) + if unknown_classes: + raise ValueError( + "partial_completion.waivable_failure_classes may only contain " + f"{sorted(_WAIVABLE_FAILURE_CLASSES)}; got {unknown_classes}" + ) + + +def _parse_partial_stage_policy(raw: Any) -> Optional[PartialStagePolicy]: + """Parse and validate an opt-in partial-completion mapping.""" + if raw is None: + return None + if not isinstance(raw, Mapping): + raise ValueError("stage partial_completion must be a mapping") + + allowed_fields = { + "min_success_fraction", + "min_per_reference_success_fraction", + "min_successful_rows_per_reference", + "waivable_failure_classes", + "tolerate_unresolved", + } + unknown_fields = sorted(set(raw) - allowed_fields) + if unknown_fields: + raise ValueError(f"unknown partial_completion field(s): {', '.join(map(str, unknown_fields))}") + + numeric_fields = allowed_fields - {"waivable_failure_classes", "tolerate_unresolved"} + if any(isinstance(raw.get(field), bool) for field in numeric_fields if field in raw): + raise ValueError("partial_completion thresholds must be numeric, not boolean") + + raw_waivable = raw.get("waivable_failure_classes", ("timeout_exceeded",)) + if isinstance(raw_waivable, (str, bytes)) or not isinstance(raw_waivable, Sequence): + raise ValueError("partial_completion.waivable_failure_classes must be a sequence of strings") + waivable = tuple(str(value) for value in raw_waivable) + + raw_tolerate = raw.get("tolerate_unresolved", False) + if not isinstance(raw_tolerate, bool): + raise ValueError("partial_completion.tolerate_unresolved must be a boolean") + + try: + raw_minimum_rows = raw.get("min_successful_rows_per_reference", 1) + minimum_rows = int(raw_minimum_rows) + policy = PartialStagePolicy( + min_success_fraction=float(raw.get("min_success_fraction", 1.0)), + min_per_reference_success_fraction=float(raw.get("min_per_reference_success_fraction", 1.0)), + min_successful_rows_per_reference=minimum_rows, + waivable_failure_classes=waivable, + tolerate_unresolved=raw_tolerate, + ) + except (TypeError, ValueError) as exc: + raise ValueError("partial_completion thresholds must be numeric") from exc + if minimum_rows != raw_minimum_rows: + raise ValueError("partial_completion.min_successful_rows_per_reference must be a positive integer") + _validate_partial_stage_policy(policy) + return policy def parse_multistage_config(raw: Mapping[str, Any]) -> MultiStageRunConfig: """Build a :class:`MultiStageRunConfig` from a raw config mapping. - Accepts stages as a list of mappings (``{num_tasks?, num_models?, seed?}``) + Accepts stages as a list of mappings (``{num_tasks?, num_models?, seed?, + partial_completion?}``) or as a list of ``"[num_tasks][:num_models[:seed]]"`` strings (handy for CLI overrides). ``num_tasks`` is the per-stage task count and defaults to the full task set when omitted (empty leading field in the string form). Raises @@ -192,6 +758,7 @@ def parse_multistage_config(raw: Mapping[str, Any]) -> MultiStageRunConfig: num_models=int(num_models) if num_models is not None else None, seed=int(seed) if seed is not None else None, num_tasks=int(num_tasks) if num_tasks is not None else None, + partial_completion=_parse_partial_stage_policy(entry.get("partial_completion")), ) ) else: @@ -206,6 +773,8 @@ 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): @@ -220,6 +789,7 @@ def parse_multistage_config(raw: Mapping[str, Any]) -> MultiStageRunConfig: nested_tasks=bool(raw.get("nested_tasks", False)), seed=raw.get("seed"), reuse_cached_deliverables=bool(raw.get("reuse_cached_deliverables", True)), + retry_inprocess=bool(raw.get("retry_inprocess", False)), ) @@ -311,6 +881,8 @@ def build_stage_rows( def tag_results( pairs: Sequence[Tuple[Mapping[str, Any], Mapping[str, Any]]], stage_index: int, + expected_final_stage_index: Optional[int] = None, + expected_stage_row_count: Optional[int] = None, ) -> List[Dict[str, Any]]: """Attach rollout identity + ``stage_index`` to each stage result row. @@ -318,6 +890,9 @@ def tag_results( result (task/rollout indices, agent ref) so the merged rollouts file and the standard ``_call_aggregate_metrics`` see well-formed rows, and stamps ``stage_index``/``task_id`` so the stage-aware aggregation can group by stage. + Multi-stage callers also provide ``expected_final_stage_index`` and + ``expected_stage_row_count`` so aggregation can distinguish a complete run + from a missing or partially drained final stage. """ tagged: List[Dict[str, Any]] = [] for row, result in pairs: @@ -325,7 +900,15 @@ def tag_results( out[TASK_INDEX_KEY_NAME] = row[TASK_INDEX_KEY_NAME] out[ROLLOUT_INDEX_KEY_NAME] = row[ROLLOUT_INDEX_KEY_NAME] out[AGENT_REF_KEY_NAME] = row[AGENT_REF_KEY_NAME] + if ATTEMPT_INDEX_KEY_NAME in row: + out[ATTEMPT_INDEX_KEY_NAME] = row[ATTEMPT_INDEX_KEY_NAME] + if "verify_cache_namespace" in row: + out["verify_cache_namespace"] = row["verify_cache_namespace"] out["stage_index"] = stage_index + if expected_final_stage_index is not None: + out["expected_final_stage_index"] = expected_final_stage_index + if expected_stage_row_count is not None: + out["expected_stage_row_count"] = expected_stage_row_count if out.get("task_id") is None: tid = row_task_id(row) if tid is not None: @@ -349,17 +932,21 @@ async def run_multistage_stages( rng: Optional[random.Random] = None, on_event: Optional[Callable[[str, dict], None]] = None, resume: Optional[StageResume] = None, + dispatch_longest_first: bool = False, + assignment_repair: Optional[AssignmentRepair] = None, ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: """Run every stage and return ``(all_result_rows, stage_summaries)``. For each stage: sample the stage's ``T`` tasks (``T`` defaults to the full task set), select the included references (closest known ELO to the running - estimate), assign each sampled task one of them uniformly at random, build the - stage's rollout rows, execute them via ``run_rollouts``, tag the results, pool - the per-reference votes, and fit the stage ELO (threaded into the next stage's - selection). ``all_result_rows`` is the concatenation of every stage's tagged - results (ready to write as the standard rollouts file); ``stage_summaries`` is - one dict per stage for logging. + estimate), assign each sampled task one reference, build the stage's rollout + rows, execute them via ``run_rollouts``, tag the results, pool the + per-reference votes, and fit the stage ELO (threaded into the next stage's + selection). Fresh partial-completion stages balance assignments across the + selected references; other fresh stages use independent uniform draws. + ``all_result_rows`` is the concatenation of every stage's tagged results + (ready to write as the standard rollouts file); ``stage_summaries`` is one + dict per stage for logging. ``rng`` seeds task sampling (defaults to ``multistage_config.seed``); per-task reference assignment is seeded independently per stage via @@ -370,8 +957,8 @@ async def run_multistage_stages( interrupted stages re-dispatch only the ``(task, rollout)`` rows without a persisted success, and recorded plans (including each stage's per-task reference assignment) are replayed so selection is identical even when - ``multistage.seed`` is ``None``. ``resume=None`` is byte-for-byte the - pre-resume behavior. + ``multistage.seed`` is ``None``. ``resume=None`` preserves the same execution + semantics without file-backed persistence. """ base_rng = rng or ( random.Random(multistage_config.seed) if multistage_config.seed is not None else random.Random() @@ -388,6 +975,11 @@ async def run_multistage_stages( nested=multistage_config.nested_tasks, ) total_stages = len(multistage_config.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: @@ -401,23 +993,116 @@ def _emit(name: str, **data: object) -> None: # (task_id, rollout_index) deliverables already produced by earlier stages. # Later stages reuse these instead of re-running the policy. produced: set[Tuple[str, int]] = set() + sidecar_produced_by_stage: Dict[int, set[Tuple[str, int]]] = {} + if resume is not None and resume.reuse_cached_keys: + task_id_by_key: Dict[Tuple[Any, Any], str] = {} + for row in materialized_rows: + task_id = row_task_id(row) + if task_id is not None: + task_id_by_key[(row.get(TASK_INDEX_KEY_NAME), row.get(ROLLOUT_INDEX_KEY_NAME))] = task_id + for stage_index, keys in resume.reuse_cached_keys.items(): + for key in keys: + task_id = task_id_by_key.get(key) + if task_id is not None: + sidecar_produced_by_stage.setdefault(stage_index, set()).add((task_id, int(key[1] or 0))) + + max_attempts = get_max_rollout_attempts() + retry_inprocess = bool(getattr(multistage_config, "retry_inprocess", False)) for index, stage in enumerate(multistage_config.stages): + # A sidecar reuse flag is emitted only after the policy artifact exists. + # Treat it as produced from this stage onward even when the judging row + # is terminal/max-attempt gated and therefore never enters the main file. + produced.update(sidecar_produced_by_stage.get(index, set())) if resume is not None and index in resume.outcomes: - eval_elo = _resume_complete_stage( - index, - total_stages, - resume, - reference_elos, - produced, - all_results, - stage_summaries, - _emit, - rows_by_task=rows_by_task, + outcome = resume.outcomes[index] + plan = resume.plans.get(index, {}) + expected_rows = [row for task_id in plan.get("task_ids", []) for row in rows_by_task.get(str(task_id), [])] + expected_keys = {(row.get(TASK_INDEX_KEY_NAME), row.get(ROLLOUT_INDEX_KEY_NAME)) for row in expected_rows} + resolved_keys = set(resume.gated_keys.get(index, set())) + completion_is_covered = len(expected_keys) == len(expected_rows) and expected_keys <= resolved_keys + task_reference_ids = { + str(key): str(value) for key, value in (plan.get("task_reference_ids") or {}).items() + } + if not task_reference_ids and plan.get("reference_ids"): + assignment_rng = stage_assignment_rng(multistage_config.seed, plan.get("seed"), index) + task_reference_ids = assign_task_references( + plan.get("task_ids", []), + plan.get("reference_ids", []), + rng=assignment_rng, + ) + planned_stage_rows = build_stage_rows(rows_by_task, task_reference_ids, index) + if outcome.get("status") == "partial_complete": + included_keys = _recorded_key_set(outcome, "included_keys") + omitted_keys = _recorded_key_set(outcome, "omitted_keys") + persisted_success_keys = {_stage_key(row) for row in resume.rows_by_stage.get(index, [])} + completion_is_covered = ( + completion_is_covered + and not (included_keys & omitted_keys) + and included_keys | omitted_keys == expected_keys + and included_keys <= persisted_success_keys + and _cached_partial_snapshot_is_valid( + outcome, + stage.partial_completion, + planned_stage_rows, + resume.rows_by_stage.get(index, []), + plan.get("reference_ids", []), + reference_elos, + ) + ) + if not completion_is_covered: + raise RuntimeError( + f"cached partial stage {index} snapshot is invalid; refusing to recompute " + "adaptive downstream stages from changed evidence" + ) + elif index < total_stages - 1: + # Older runs used "complete" to mean terminal/max-attempt + # resolved. Such an outcome is not safe to reuse for adaptive + # reference selection unless every planned row has usable + # persisted battle evidence. + persisted_rows = resume.rows_by_stage.get(index, []) + completion_is_covered = completion_is_covered and ( + _fit_eligible_stage_keys(planned_stage_rows, persisted_rows) == expected_keys + and len(planned_stage_rows) == len(expected_rows) + ) + if completion_is_covered: + resumed_elo = _resume_complete_stage( + index, + total_stages, + len(expected_rows), + resume, + reference_elos, + produced, + all_results, + stage_summaries, + _emit, + ) + # Match uninterrupted execution: a completed stage with no + # usable battles does not erase the last fitted ELO. Later + # adaptive stages continue from the most recent valid fit. + if resumed_elo is not None: + eval_elo = resumed_elo + continue + _emit( + "stage_completion_stale", + index=index, + total_stages=total_stages, + expected_rows=len(expected_rows), + resolved_rows=len(expected_keys & resolved_keys), ) - continue + resume.restart_from(index) + sidecar_produced_by_stage = { + stage_index: keys for stage_index, keys in sidecar_produced_by_stage.items() if stage_index <= index + } reference_ids, task_ids, task_reference_ids, replayed = _plan_stage( - index, stage, reference_elos, eval_elo, stage_task_sets, multistage_config, resume + index, + stage, + reference_elos, + eval_elo, + stage_task_sets, + multistage_config, + resume, + assignment_repair, ) stage_rows = build_stage_rows( @@ -432,6 +1117,70 @@ def _emit(name: str, **data: object) -> None: # failures from the sidecar (mirrors ``_load_from_cache``, stage-keyed). gated_keys = set(resume.gated_keys.get(index, set())) if resume is not None else set() pending_rows = [r for r in stage_rows if (r[TASK_INDEX_KEY_NAME], r[ROLLOUT_INDEX_KEY_NAME]) not in gated_keys] + if resume is not None: + reuse_cached_keys = resume.reuse_cached_keys.get(index, set()) + attempts_by_key = resume.attempts_by_stage.get(index, {}) + for row in pending_rows: + key = (row[TASK_INDEX_KEY_NAME], row[ROLLOUT_INDEX_KEY_NAME]) + if key in reuse_cached_keys: + row["reuse_cached_deliverable"] = True + prior_attempts = attempts_by_key.get(key, 0) + if prior_attempts > 0: + row[ATTEMPT_INDEX_KEY_NAME] = prior_attempts + + if dispatch_longest_first and resume is not None: + elapsed_by_key = resume.elapsed_by_stage.get(index, {}) + if elapsed_by_key: + known = [ + row + for row in pending_rows + if (row[TASK_INDEX_KEY_NAME], row[ROLLOUT_INDEX_KEY_NAME]) in elapsed_by_key + ] + unknown = [ + row + for row in pending_rows + if (row[TASK_INDEX_KEY_NAME], row[ROLLOUT_INDEX_KEY_NAME]) not in elapsed_by_key + ] + known.sort( + key=lambda row: elapsed_by_key[(row[TASK_INDEX_KEY_NAME], row[ROLLOUT_INDEX_KEY_NAME])], + reverse=True, + ) + pending_rows = known + unknown + + # A run may opt into partial calibration after an allocation ended with + # 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: + 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, {}) + # Persisted failures are fsynced before their journal disposition, + # while no-persist tombstones are journaled first. A tombstone must + # therefore override an older sidecar; otherwise the sidecar is at + # least as new as the journal. If a later persisted attempt crashed + # before journaling, preferring its sidecar fails closed correctly. + latest_attempts = { + key: latest_dispositions[key] + if latest_dispositions.get(key, {}).get(NG_NO_PERSIST_KEY) + else latest_failures[key] + for key in pending_keys_for_policy + if latest_dispositions.get(key, {}).get(NG_NO_PERSIST_KEY) or key in latest_failures + } + persisted_failures = [latest_attempts[key] for key in pending_keys_for_policy if key in latest_attempts] + cached_elo, _, cached_num_references = fit_stage_elo(pool_per_reference(cached_rows), reference_elos) + candidate = _partial_stage_outcome( + stage.partial_completion, + stage_rows, + cached_rows, + persisted_failures, + pending_keys_for_policy, + reference_ids, + cached_elo, + cached_num_references, + ) + if candidate is not None: + pre_dispatch_partial_outcome = {"stage_index": index, **candidate} num_reused = sum(1 for r in stage_rows if r.get("reuse_cached_deliverable")) _emit( @@ -447,17 +1196,120 @@ def _emit(name: str, **data: object) -> None: replayed=replayed, ) - pairs = await run_rollouts(pending_rows) - new_tagged = tag_results(pairs, index) + pairs = [] if pre_dispatch_partial_outcome is not None else await run_rollouts(pending_rows) + expected_stage_row_count = len(stage_rows) + new_tagged = tag_results( + pairs, + index, + expected_final_stage_index=total_stages - 1, + expected_stage_row_count=expected_stage_row_count, + ) if resume is not None: resume.on_rows(index, new_tagged) + + # In-process retry: re-dispatch rows that came back retryable, applying the + # same gate a resume would (success / terminal / max-attempt), until the cap + # is reached or the dispatch budget runs out. Without this, max_attempts only + # ever advances across resumes, so a timed-out row is lost whenever a run + # exits cleanly -- the whole point of the cap never fires inside a leg. + # ``resume.attempts_by_stage`` is a startup snapshot that never mutates in + # process, so the cap has to be tracked locally. The total consumed for a + # key is: attempts recorded at startup + this leg's initial dispatch + the + # in-process retries made below. + inprocess_attempts: Dict[Tuple[Any, Any], int] = {} + if retry_inprocess and pairs and resume is not None: + startup_attempts = resume.attempts_by_stage.get(index, {}) + # Sticky for the whole stage: once a key is known to have a reusable + # deliverable on disk it stays reusable, even if a later attempt comes + # back stamped on the rollout path without the flag. + reuse_cached_keys_for_stage = set(resume.reuse_cached_keys.get(index, frozenset())) + pending_by_key = {(r[TASK_INDEX_KEY_NAME], r[ROLLOUT_INDEX_KEY_NAME]): r for r in pending_rows} + while True: + retry_rows = [] + for row in new_tagged: + if not _is_in_process_retryable(row): + continue + key = (row[TASK_INDEX_KEY_NAME], row[ROLLOUT_INDEX_KEY_NAME]) + attempts_used = startup_attempts.get(key, 0) + 1 + inprocess_attempts.get(key, 0) + if attempts_used >= max_attempts: + continue + source = pending_by_key.get(key) + if source is None: + continue + retry_row = dict(source) + # 0-based index of the attempt about to be made, so each retry + # gets its own rollout capture key. + retry_row[ATTEMPT_INDEX_KEY_NAME] = attempts_used + # A failed judging attempt keeps its policy artifact: rejudge it + # rather than rerunning the agent, exactly as resume does. The + # flag may have been set on this very row moments ago, so honour + # both it and the startup snapshot. + if row.get("reuse_cached_deliverable"): + reuse_cached_keys_for_stage.add(key) + if key in reuse_cached_keys_for_stage: + retry_row["reuse_cached_deliverable"] = True + retry_rows.append(retry_row) + inprocess_attempts[key] = inprocess_attempts.get(key, 0) + 1 + if not retry_rows: + break + retry_pairs = await run_rollouts(retry_rows) + retried = tag_results( + retry_pairs, + index, + expected_final_stage_index=total_stages - 1, + expected_stage_row_count=expected_stage_row_count, + ) + # Once the dispatch budget is gone every row comes back drained + # without a POST. A drained row means "we never got to try", not + # "the earlier persisted attempt is void": it must neither become + # the row of record nor journal a no-persist disposition over a + # key whose sidecar already holds a real failure -- that tombstone + # is durable and would permanently disable the pre-dispatch + # partial-outcome shortcut for that key. + dispatched = [r for r in retried if not r.get(NG_NO_PERSIST_KEY)] + num_drained = len(retried) - len(dispatched) + if dispatched: + resume.on_rows(index, dispatched) + _emit( + "stage_retry", + index=index, + total_stages=total_stages, + num_selected=len(retry_rows), + num_dispatched=len(dispatched), + num_drained=num_drained, + num_recovered=sum(1 for r in dispatched if _is_success_row(r)), + ) + # Later *dispatched* attempt wins for a key; earlier attempts stay + # in the sidecar. Drained rows are discarded, leaving the persisted + # attempt as the row of record; the attempt counter still advanced + # above, so the loop cannot spin on them. + superseded = {(r[TASK_INDEX_KEY_NAME], r[ROLLOUT_INDEX_KEY_NAME]) for r in dispatched} + new_tagged = [ + r for r in new_tagged if (r[TASK_INDEX_KEY_NAME], r[ROLLOUT_INDEX_KEY_NAME]) not in superseded + ] + dispatched + # A verify-side failure can still carry a complete, reusable policy + # artifact. Make it available immediately to later stages in this same + # process; startup-only sidecar loading covers only resumed runs. + for row in new_tagged: + if not row.get("reuse_cached_deliverable"): + continue + tid = row_task_id(row) + if tid is not None: + produced.add((tid, int(row.get(ROLLOUT_INDEX_KEY_NAME, 0) or 0))) # Only successful rows feed pooling / aggregate. new_successes = [r for r in new_tagged if _is_success_row(r)] - tagged = list(cached_rows) + new_successes + tagged = [ + dict( + r, + expected_final_stage_index=total_stages - 1, + expected_stage_row_count=expected_stage_row_count, + ) + for r in cached_rows + ] + new_successes all_results.extend(tagged) # Record this stage's deliverables so later stages can reuse them. - for row in stage_rows: + for row in tagged: tid = row_task_id(row) if tid is not None: produced.add((tid, int(row.get(ROLLOUT_INDEX_KEY_NAME, 0) or 0))) @@ -467,9 +1319,73 @@ def _emit(name: str, **data: object) -> None: if stage_elo is not None: eval_elo = stage_elo - # Completion marker only; eval_elo is re-fit from rows on resume. - if resume is not None: - resume.on_outcome(index, {"stage_index": index, "status": "complete"}) + # 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. + 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: + key = (result[TASK_INDEX_KEY_NAME], result[ROLLOUT_INDEX_KEY_NAME]) + resolved = _is_success_row(result) or is_terminal_failure(result, retry_terminal_timeouts=True) + if ( + not resolved + and result.get(NG_FAILURE_CLASS_KEY) is not None + and not result.get(NG_NO_PERSIST_KEY) + and prior_attempts.get(key, 0) + 1 + inprocess_attempts.get(key, 0) >= max_attempts + ): + resolved = True + if resolved: + returned_keys.add(key) + pending_keys = {(r[TASK_INDEX_KEY_NAME], r[ROLLOUT_INDEX_KEY_NAME]) for r in pending_rows} + unresolved_keys = pending_keys - returned_keys + stage_complete = not unresolved_keys + partial_outcome = pre_dispatch_partial_outcome + coverage_rejected = False + 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 + ): + stage_complete = False + coverage_rejected = not unresolved_keys + if resume is not None and stage_complete: + resume.on_outcome( + index, + partial_outcome or {"stage_index": index, "status": "complete"}, + ) _emit( "stage_end", @@ -479,22 +1395,103 @@ def _emit(name: str, **data: object) -> None: normalized_elo=normalized, num_references=num_references, ) - stage_summaries.append( - { - "stage_index": index, - "num_tasks": len(task_ids), - "num_rollouts": len(stage_rows), - "num_reused": num_reused, - "reference_ids": list(reference_ids), - "eval_elo": stage_elo, - "normalized_elo": normalized, - "num_references": num_references, - } - ) + summary = { + "stage_index": index, + "num_tasks": len(task_ids), + "num_rollouts": len(stage_rows), + "num_reused": num_reused, + "reference_ids": list(reference_ids), + "eval_elo": stage_elo, + "normalized_elo": normalized, + "num_references": num_references, + } + if partial_outcome is not None: + summary.update( + partial=True, + success_fraction=partial_outcome["success_fraction"], + num_successful=len(partial_outcome["included_keys"]), + num_omitted=len(partial_outcome["omitted_keys"]), + ) + _emit( + "stage_partial_complete", + index=index, + total_stages=total_stages, + success_fraction=partial_outcome["success_fraction"], + num_omitted=len(partial_outcome["omitted_keys"]), + ) + stage_summaries.append(summary) + if not stage_complete: + # Recorded on the summary so the caller can tell a dead end from a + # stop a resume would clear (see ``_require_final_stage_rows``). + summary["incomplete"] = True + summary["coverage_blocked"] = coverage_rejected + _emit( + "stage_incomplete", + index=index, + total_stages=total_stages, + num_pending=len(unresolved_keys), + num_omitted=len(missing_success_keys), + coverage_blocked=coverage_rejected, + ) + # Later reference selection depends on an accepted ELO. Stop before + # creating downstream plans; retryable rows can resume, while a + # coverage-only rejection requires an explicit policy/data change. + break return all_results, stage_summaries +def _require_final_stage_rows( + all_results: Sequence[Mapping[str, Any]], + total_stages: int, + stage_summaries: Sequence[Mapping[str, Any]], +) -> None: + """Fail the run when a coverage rejection left the final stage empty. + + Every row is stamped ``expected_final_stage_index = total_stages - 1``, so a + run that stops early still declares a final stage it never reached. + Aggregation records that as ``comparison/final_stage_degraded`` and returns + normally, which makes an abandoned run indistinguishable from a finished one + for anything reading the exit status. + + Only a coverage rejection is terminal here. When the stage loop stopped with + rows still retryable, a resume picks them up and the run is expected to + continue in a later process; that path only logs how to continue, and the + run exits normally without a ``comparison/eval_elo``. Likewise a final stage + with *some* rows is not fatal: aggregation reports its per-stage fit and + marks it degraded, with no headline ``comparison/eval_elo`` unless every + expected final-stage row is present. + """ + if not total_stages: + return + final_stage_index = total_stages - 1 + if any(int(row.get("stage_index", -1) or 0) == final_stage_index for row in all_results): + return + blocked = [s for s in stage_summaries if s.get("coverage_blocked")] + if not blocked: + stopped = [int(s.get("stage_index", -1) or 0) for s in stage_summaries if s.get("incomplete")] + if stopped: + import sys + + print( + f"[multistage-elo] stopped after stage {stopped[-1] + 1}/{total_stages} with retryable or " + "undispatched rows; no final-stage result and no comparison/eval_elo will be reported. " + "Rerun the same command with --resume to continue (multistage.retry_inprocess=true retries " + "timeouts and transient failures in the same process).", + file=sys.stderr, + flush=True, + ) + return + reached = sorted({int(s.get("stage_index", -1) or 0) for s in stage_summaries}) + raise RuntimeError( + f"multi-stage run declared stage {final_stage_index} as its final stage but produced no rows for it " + f"(stages reached: {reached}, {len(all_results)} row(s) total). " + f"Stage {int(blocked[-1].get('stage_index', -1) or 0)} was rejected on coverage with nothing left to " + "retry, so no downstream stage was planned and there is no final-stage result to report. " + "This needs an explicit policy or data change, not a resume." + ) + + def _plan_stage( index: int, stage: StageSpec, @@ -503,6 +1500,7 @@ def _plan_stage( stage_task_sets: Sequence[Sequence[str]], multistage_config: MultiStageRunConfig, resume: Optional[StageResume], + assignment_repair: Optional[AssignmentRepair] = None, ) -> Tuple[List[str], List[str], Dict[str, str], bool]: """Return ``(reference_ids, task_ids, task_reference_ids, replayed)`` for a stage. @@ -510,7 +1508,8 @@ def _plan_stage( the ``num_models`` closest to the running ELO estimate). ``task_ids`` is the stage's sampled ``T`` tasks (the full task set when ``num_tasks`` is unset). ``task_reference_ids`` maps each task to the single reference it is judged - against, sampled uniformly (equal weight) from ``reference_ids``. + against. Fresh partial-completion stages use a seeded balanced assignment; + other fresh stages use independent uniform draws. ``replayed`` is True when the recorded plan was returned from ``resume.plans[index]`` (deterministic replay, even with ``seed=None``), False @@ -532,34 +1531,49 @@ def _plan_stage( reference_ids = select_references(reference_elos, eval_elo, stage.num_models) task_ids = list(stage_task_sets[index]) rng = stage_assignment_rng(multistage_config.seed, stage.seed, index) - task_reference_ids = assign_task_references(task_ids, reference_ids, rng=rng) - if resume is not None: - resume.on_plan( + task_reference_ids = assign_task_references( + task_ids, + reference_ids, + rng=rng, + balanced=stage.partial_completion is not None, + ) + repair_receipt: Optional[Dict[str, Any]] = None + if assignment_repair is not None: + task_reference_ids, repair_receipt = assignment_repair( index, - { - "stage_index": index, - "status": "planned", - "reference_ids": list(reference_ids), - "task_ids": list(task_ids), - "task_reference_ids": task_reference_ids, - "seed": stage.seed, - "prior_eval_elo": eval_elo, - }, + reference_ids, + task_reference_ids, ) + if set(task_reference_ids) != set(task_ids): + raise ValueError("transport assignment repair changed the stage task set") + if any(reference_id not in reference_ids for reference_id in task_reference_ids.values()): + raise ValueError("transport assignment repair selected a reference outside the stage") + if resume is not None: + plan = { + "stage_index": index, + "status": "planned", + "reference_ids": list(reference_ids), + "task_ids": list(task_ids), + "task_reference_ids": task_reference_ids, + "seed": stage.seed, + "prior_eval_elo": eval_elo, + } + if repair_receipt is not None: + plan["transport_assignment_repair"] = repair_receipt + resume.on_plan(index, plan) return list(reference_ids), task_ids, task_reference_ids, False def _resume_complete_stage( index: int, total_stages: int, + expected_stage_row_count: int, resume: StageResume, reference_elos: Mapping[str, float], produced: set[Tuple[str, int]], all_results: List[Dict[str, Any]], stage_summaries: List[Dict[str, Any]], emit: Callable[..., None], - *, - rows_by_task: Mapping[str, List[Dict[str, Any]]], ) -> Optional[float]: """Reuse a completed stage's cached rows without dispatch; return threaded ELO. @@ -567,10 +1581,21 @@ def _resume_complete_stage( truth) rather than trusting the recorded ``eval_elo`` field, so the value threaded to later stages is always consistent with the persisted rows. """ - cached_rows = list(resume.rows_by_stage.get(index, [])) - plan = resume.plans[index] + outcome = resume.outcomes.get(index, {}) + partial = outcome.get("status") == "partial_complete" + included_keys = _recorded_key_set(outcome, "included_keys") if partial else None + cached_rows = [ + dict( + r, + expected_final_stage_index=total_stages - 1, + expected_stage_row_count=expected_stage_row_count, + ) + for r in resume.rows_by_stage.get(index, []) + if included_keys is None or _stage_key(r) in included_keys + ] + plan = resume.plans.get(index, {}) reference_ids = list(plan.get("reference_ids", [])) - task_ids = list(plan["task_ids"]) + task_ids = list(plan.get("task_ids", [])) for row in cached_rows: tid = row_task_id(row) @@ -589,21 +1614,29 @@ def _resume_complete_stage( normalized_elo=normalized, num_references=num_references, num_rollouts=len(cached_rows), + partial=partial, ) - stage_summaries.append( - { - "stage_index": index, - "num_tasks": len(task_ids), - # Cached results can be incomplete even when the stage was marked complete. - "num_rollouts": sum(len(rows_by_task[task_id]) for task_id in task_ids), - "num_reused": 0, - "reference_ids": reference_ids, - "eval_elo": stage_elo, - "normalized_elo": normalized, - "num_references": num_references, - "cached": True, - } - ) + summary = { + "stage_index": index, + "num_tasks": len(task_ids), + # Cached results can be incomplete even when the stage was marked complete. + "num_rollouts": expected_stage_row_count, + "num_reused": 0, + "reference_ids": reference_ids, + "eval_elo": stage_elo, + "normalized_elo": normalized, + "num_references": num_references, + "cached": True, + } + if partial: + summary.update( + partial=True, + num_rollouts=expected_stage_row_count, + num_successful=len(cached_rows), + success_fraction=outcome.get("success_fraction"), + num_omitted=len(_recorded_key_set(outcome, "omitted_keys")), + ) + stage_summaries.append(summary) return stage_elo @@ -625,9 +1658,35 @@ def write_rollouts(all_results: Sequence[Mapping[str, Any]], output_fpath: str | deduped.values(), key=lambda r: (r.get("stage_index", 0), r.get(TASK_INDEX_KEY_NAME, 0), r.get(ROLLOUT_INDEX_KEY_NAME, 0)), ) - with output_fpath.open("wb") as handle: - for row in ordered: - handle.write(orjson.dumps(row) + b"\n") + prior_mode = output_fpath.stat().st_mode & 0o7777 if output_fpath.exists() else None + temp_path: Optional[Path] = None + try: + # NamedTemporaryFile forces 0600, which would silently make a fresh + # shared rollout file owner-only. Create the staging file with the same + # 0666+umask semantics as a normal ``open(..., 'wb')`` instead. + for _ in range(100): + candidate = output_fpath.parent / f".{output_fpath.name}.merge-{secrets.token_hex(8)}" + try: + fd = os.open(candidate, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o666) + except FileExistsError: + continue + temp_path = candidate + break + else: # pragma: no cover - cryptographically-random collisions + raise FileExistsError(f"could not allocate staging file beside {output_fpath}") + + with os.fdopen(fd, "wb") as handle: + for row in ordered: + handle.write(orjson.dumps(row) + b"\n") + handle.flush() + os.fsync(handle.fileno()) + if prior_mode is not None: + os.chmod(temp_path, prior_mode) + os.replace(temp_path, output_fpath) + temp_path = None + finally: + if temp_path is not None: + temp_path.unlink(missing_ok=True) return output_fpath @@ -642,6 +1701,12 @@ def journal_path_for(output_fpath: str | Path) -> Path: return output_fpath.with_name(f"{output_fpath.stem}_multistage_state.jsonl") +def aggregate_metrics_path_for(output_fpath: str | Path) -> Path: + """Standard aggregate-metrics path associated with a rollout JSONL.""" + output_fpath = Path(output_fpath) + return output_fpath.with_stem(output_fpath.stem + "_aggregate_metrics").with_suffix(".json") + + def read_journal(journal_fpath: str | Path) -> Tuple[Dict[int, dict], Dict[int, dict], Optional[str]]: """Read the append-only journal; latest record per ``stage_index`` wins. @@ -664,13 +1729,130 @@ def read_journal(journal_fpath: str | Path) -> Tuple[Dict[int, dict], Dict[int, index = record.get("stage_index") if index is None: continue - if record.get("status") == "planned": + status = record.get("status") + if status == "restart_from_stage": + restart_stage = int(index) + # The marker is durable before downstream files are pruned. It + # therefore invalidates old journal state immediately even if + # the process dies during cleanup; newer records later in the + # journal can repopulate the recomputed stages. + plans = {i: plan for i, plan in plans.items() if i <= restart_stage} + outcomes = {i: outcome for i, outcome in outcomes.items() if i < restart_stage} + elif status == "planned": plans[int(index)] = record - elif record.get("status") == "complete": + elif status in {"complete", "partial_complete"}: outcomes[int(index)] = record return plans, outcomes, fingerprint +def load_latest_attempt_dispositions( + journal_fpath: str | Path, +) -> Dict[int, Dict[Tuple[Any, Any], Dict[str, Any]]]: + """Load latest failure/no-persist attempt state recorded in the journal.""" + latest: Dict[int, Dict[Tuple[Any, Any], Dict[str, Any]]] = {} + path = Path(journal_fpath) + if not path.exists(): + return latest + with path.open("rb") as handle: + for line in handle: + line = line.strip() + if not line: + continue + record = orjson.loads(line) + status = record.get("status") + if status == "restart_from_stage": + restart_stage = int(record["stage_index"]) + latest = {index: rows for index, rows in latest.items() if index <= restart_stage} + continue + if status != "attempt_dispositions": + continue + stage_index = int(record["stage_index"]) + stage_latest = latest.setdefault(stage_index, {}) + for disposition in record.get("attempts", []) or []: + if TASK_INDEX_KEY_NAME not in disposition or ROLLOUT_INDEX_KEY_NAME not in disposition: + continue + key = (disposition[TASK_INDEX_KEY_NAME], disposition[ROLLOUT_INDEX_KEY_NAME]) + stage_latest[key] = disposition + return latest + + +def _pending_restart_stage(journal_fpath: str | Path) -> Optional[int]: + """Return a restart marker whose downstream file cleanup was not acknowledged.""" + pending: Optional[int] = None + path = Path(journal_fpath) + if not path.exists(): + return None + with path.open("rb") as handle: + for line in handle: + line = line.strip() + if not line: + continue + record = orjson.loads(line) + status = record.get("status") + if status == "restart_from_stage": + pending = int(record["stage_index"]) + elif status == "restart_cleanup_complete" and pending == int(record["stage_index"]): + pending = None + return pending + + +def _atomic_filter_jsonl(path: Path, keep: Callable[[Mapping[str, Any]], bool]) -> None: + """Atomically rewrite ``path`` with only records accepted by ``keep``.""" + if not path.exists(): + return + mode = path.stat().st_mode & 0o7777 + temp_path: Optional[Path] = None + try: + with tempfile.NamedTemporaryFile( + mode="wb", dir=path.parent, prefix=f".{path.name}.prune-", delete=False + ) as out: + temp_path = Path(out.name) + with path.open("rb") as source: + for line in source: + stripped = line.strip() + if stripped and keep(orjson.loads(stripped)): + out.write(stripped + b"\n") + out.flush() + os.fsync(out.fileno()) + os.chmod(temp_path, mode) + os.replace(temp_path, path) + temp_path = None + finally: + if temp_path is not None: + temp_path.unlink(missing_ok=True) + + +def _prune_downstream_files(output_fpath: str | Path, restart_stage: int) -> None: + """Move persisted main/sidecar rows after ``restart_stage`` to a quarantine file. + + Pruned rows are completed provider evidence; a restart decision must never + destroy them. Each prune appends the dropped rows to a timestamped + ``.pruned.`` sibling before atomically rewriting the source. + """ + output_fpath = Path(output_fpath) + + def keep(record: Mapping[str, Any]) -> bool: + return int(record.get("stage_index", 0) or 0) <= restart_stage + + stale_suffix = f".pruned.{time.time_ns()}" + for path in (output_fpath, failures_path_for(output_fpath)): + if not path.exists(): + continue + dropped: list[bytes] = [] + with path.open("rb") as handle: + for line in handle: + stripped = line.strip() + if stripped and not keep(orjson.loads(stripped)): + dropped.append(stripped) + if dropped: + quarantine = path.with_name(path.name + stale_suffix) + with quarantine.open("wb") as out: + out.write(b"\n".join(dropped) + b"\n") + out.flush() + os.fsync(out.fileno()) + _atomic_filter_jsonl(path, keep) + + def load_persisted_rows(output_fpath: str | Path) -> Dict[int, List[Dict[str, Any]]]: """Group the main-jsonl success rows by ``stage_index``. @@ -702,7 +1884,7 @@ def load_gated_keys( Mirrors ``_load_from_cache`` with the stage dimension added: a stage-row is gated if it is a success (main jsonl), a terminal sidecar failure - (``_ng_failure_terminal``), or has hit ``_get_max_rollout_attempts`` attempts + (``_ng_failure_terminal``), or has hit ``get_max_rollout_attempts`` attempts in the sidecar. Everything else is re-dispatched. """ gated: Dict[int, set[Tuple[Any, Any]]] = { @@ -714,7 +1896,7 @@ def load_gated_keys( if not failures_fpath.exists(): return gated - max_attempts = _get_max_rollout_attempts() + max_attempts = get_max_rollout_attempts() attempts: Dict[Tuple[int, Any, Any], int] = {} terminal: set[Tuple[int, Any, Any]] = set() with failures_fpath.open("rb") as handle: @@ -727,7 +1909,7 @@ def load_gated_keys( continue key = (int(fr.get("stage_index", 0) or 0), fr[TASK_INDEX_KEY_NAME], fr[ROLLOUT_INDEX_KEY_NAME]) attempts[key] = attempts.get(key, 0) + 1 - if fr.get(NG_TERMINAL_KEY): + if is_terminal_failure(fr, retry_terminal_timeouts=True): terminal.add(key) for key in attempts: @@ -737,6 +1919,130 @@ def load_gated_keys( return gated +def load_failure_timings(output_fpath: str | Path) -> Dict[int, Dict[Tuple[Any, Any], float]]: + """Load longest failed-attempt duration per stage/task/rollout key.""" + elapsed_by_stage: Dict[int, Dict[Tuple[Any, Any], float]] = {} + failures_fpath = failures_path_for(Path(output_fpath)) + if not failures_fpath.exists(): + return elapsed_by_stage + + with failures_fpath.open("rb") as handle: + for line in handle: + line = line.strip() + if not line: + continue + row = orjson.loads(line) + if TASK_INDEX_KEY_NAME not in row or ROLLOUT_INDEX_KEY_NAME not in row: + continue + elapsed = observed_elapsed(row) + if elapsed is None: + continue + stage_index = int(row.get("stage_index", 0) or 0) + key = (row[TASK_INDEX_KEY_NAME], row[ROLLOUT_INDEX_KEY_NAME]) + stage_timings = elapsed_by_stage.setdefault(stage_index, {}) + stage_timings[key] = max(stage_timings.get(key, 0.0), elapsed) + return elapsed_by_stage + + +def load_failure_attempts(output_fpath: str | Path) -> Dict[int, Dict[Tuple[Any, Any], int]]: + """Load prior sidecar attempt counts per stage/task/rollout key.""" + attempts_by_stage: Dict[int, Dict[Tuple[Any, Any], int]] = {} + failures_fpath = failures_path_for(Path(output_fpath)) + if not failures_fpath.exists(): + return attempts_by_stage + + with failures_fpath.open("rb") as handle: + for line in handle: + line = line.strip() + if not line: + continue + row = orjson.loads(line) + if TASK_INDEX_KEY_NAME not in row or ROLLOUT_INDEX_KEY_NAME not in row: + continue + stage_index = int(row.get("stage_index", 0) or 0) + key = (row[TASK_INDEX_KEY_NAME], row[ROLLOUT_INDEX_KEY_NAME]) + stage_attempts = attempts_by_stage.setdefault(stage_index, {}) + stage_attempts[key] = stage_attempts.get(key, 0) + 1 + return attempts_by_stage + + +def load_latest_failures(output_fpath: str | Path) -> Dict[int, Dict[Tuple[Any, Any], Dict[str, Any]]]: + """Load the newest persisted sidecar failure for each stage-row key.""" + latest_by_stage: Dict[int, Dict[Tuple[Any, Any], Dict[str, Any]]] = {} + failures_fpath = failures_path_for(Path(output_fpath)) + if not failures_fpath.exists(): + return latest_by_stage + + with failures_fpath.open("rb") as handle: + for line in handle: + line = line.strip() + if not line: + continue + row = orjson.loads(line) + if TASK_INDEX_KEY_NAME not in row or ROLLOUT_INDEX_KEY_NAME not in row: + continue + stage_index = int(row.get("stage_index", 0) or 0) + key = (row[TASK_INDEX_KEY_NAME], row[ROLLOUT_INDEX_KEY_NAME]) + latest_by_stage.setdefault(stage_index, {})[key] = row + return latest_by_stage + + +def load_retryable_failure_stages( + output_fpath: str | Path, + gated_keys: Mapping[int, AbstractSet[Tuple[Any, Any]]], + rows_by_stage: Mapping[int, Sequence[Mapping[str, Any]]], +) -> set[int]: + """Stages with a persisted failure that still needs another attempt.""" + retryable: set[int] = set() + failures_fpath = failures_path_for(Path(output_fpath)) + if not failures_fpath.exists(): + return retryable + + completed_keys = { + stage_index: {(row.get(TASK_INDEX_KEY_NAME), row.get(ROLLOUT_INDEX_KEY_NAME)) for row in rows} + for stage_index, rows in rows_by_stage.items() + } + + with failures_fpath.open("rb") as handle: + for line in handle: + line = line.strip() + if not line: + continue + row = orjson.loads(line) + if TASK_INDEX_KEY_NAME not in row or ROLLOUT_INDEX_KEY_NAME not in row: + continue + stage_index = int(row.get("stage_index", 0) or 0) + key = (row[TASK_INDEX_KEY_NAME], row[ROLLOUT_INDEX_KEY_NAME]) + # Historical failures remain in the append-only sidecar after a + # later attempt succeeds. Such a key is resolved, not retryable. + if key not in completed_keys.get(stage_index, set()) and key not in gated_keys.get(stage_index, set()): + retryable.add(stage_index) + return retryable + + +def load_reuse_cached_keys(output_fpath: str | Path) -> Dict[int, set[Tuple[Any, Any]]]: + """Load stage-aware retry keys that should reuse a persisted deliverable.""" + reuse_cached_keys: Dict[int, set[Tuple[Any, Any]]] = {} + failures_fpath = failures_path_for(Path(output_fpath)) + if not failures_fpath.exists(): + return reuse_cached_keys + + with failures_fpath.open("rb") as handle: + for line in handle: + line = line.strip() + if not line: + continue + row = orjson.loads(line) + if not row.get("reuse_cached_deliverable"): + continue + if TASK_INDEX_KEY_NAME not in row or ROLLOUT_INDEX_KEY_NAME not in row: + continue + stage_index = int(row.get("stage_index", 0) or 0) + key = (row[TASK_INDEX_KEY_NAME], row[ROLLOUT_INDEX_KEY_NAME]) + reuse_cached_keys.setdefault(stage_index, set()).add(key) + return reuse_cached_keys + + def append_journal_record(journal_fpath: str | Path, record: Mapping[str, Any], fingerprint: str) -> None: """Append a single journal record stamped with the run fingerprint.""" path = Path(journal_fpath) @@ -745,6 +2051,8 @@ def append_journal_record(journal_fpath: str | Path, record: Mapping[str, Any], out["fingerprint"] = fingerprint with path.open("ab") as handle: handle.write(orjson.dumps(out) + b"\n") + handle.flush() + os.fsync(handle.fileno()) def route_stage_rows(output_fpath: str | Path, rows: Sequence[Mapping[str, Any]]) -> None: @@ -767,13 +2075,78 @@ def route_stage_rows(output_fpath: str | Path, rows: Sequence[Mapping[str, Any]] fail_handle.write(orjson.dumps(row) + b"\n") else: main_handle.write(orjson.dumps(row) + b"\n") + main_handle.flush() + fail_handle.flush() + os.fsync(main_handle.fileno()) + os.fsync(fail_handle.fileno()) + + +def _gate_partial_outcome_omissions( + outcomes: Mapping[int, Mapping[str, Any]], + gated_keys: Dict[int, set[Tuple[Any, Any]]], +) -> None: + """Keep journaled partial omissions closed on deterministic resume.""" + for stage_index, outcome in outcomes.items(): + if outcome.get("status") == "partial_complete": + gated_keys.setdefault(stage_index, set()).update(_recorded_key_set(outcome, "omitted_keys")) def build_file_resume(output_fpath: str | Path, journal_fpath: str | Path, fingerprint: str) -> StageResume: """Build a file-backed :class:`StageResume` from the journal + rollout files.""" + output_fpath = Path(output_fpath) + journal_fpath = Path(journal_fpath) + migrate_invalid_judge_main_rows(output_fpath) + + # Finish an interrupted cleanup before reading any reusable stage state. + recovered_restart = _pending_restart_stage(journal_fpath) + if recovered_restart is not None: + aggregate_metrics_path_for(output_fpath).unlink(missing_ok=True) + _prune_downstream_files(output_fpath, recovered_restart) + append_journal_record( + journal_fpath, + {"stage_index": recovered_restart, "status": "restart_cleanup_complete"}, + fingerprint, + ) + plans, outcomes, _ = read_journal(journal_fpath) rows_by_stage = load_persisted_rows(output_fpath) gated_keys = load_gated_keys(output_fpath, rows_by_stage) + _gate_partial_outcome_omissions(outcomes, gated_keys) + elapsed_by_stage = load_failure_timings(output_fpath) + reuse_cached_keys = load_reuse_cached_keys(output_fpath) + attempts_by_stage = load_failure_attempts(output_fpath) + latest_failures_by_stage = load_latest_failures(output_fpath) + latest_attempt_dispositions_by_stage = load_latest_attempt_dispositions(journal_fpath) + # Older multistage runs wrote a completion marker even when a retryable + # sidecar failure remained. The sidecar is authoritative for retry state; + # discard only those stale outcomes so the recorded plan is replayed. + retryable_failure_stages = load_retryable_failure_stages(output_fpath, gated_keys, rows_by_stage) + if retryable_failure_stages and recovered_restart is None: + restart_stage = min(retryable_failure_stages) + # Persist the logical invalidation first. If cleanup is interrupted, the + # next process sees this marker, refuses stale downstream journal state, + # and idempotently finishes pruning the data files. + append_journal_record( + journal_fpath, + {"stage_index": restart_stage, "status": "restart_from_stage"}, + fingerprint, + ) + aggregate_metrics_path_for(output_fpath).unlink(missing_ok=True) + _prune_downstream_files(output_fpath, restart_stage) + append_journal_record( + journal_fpath, + {"stage_index": restart_stage, "status": "restart_cleanup_complete"}, + fingerprint, + ) + plans, outcomes, _ = read_journal(journal_fpath) + rows_by_stage = load_persisted_rows(output_fpath) + gated_keys = load_gated_keys(output_fpath, rows_by_stage) + _gate_partial_outcome_omissions(outcomes, gated_keys) + elapsed_by_stage = load_failure_timings(output_fpath) + reuse_cached_keys = load_reuse_cached_keys(output_fpath) + attempts_by_stage = load_failure_attempts(output_fpath) + latest_failures_by_stage = load_latest_failures(output_fpath) + latest_attempt_dispositions_by_stage = load_latest_attempt_dispositions(journal_fpath) def on_plan(index: int, plan: dict) -> None: append_journal_record(journal_fpath, plan, fingerprint) @@ -782,7 +2155,56 @@ def on_outcome(index: int, outcome: dict) -> None: append_journal_record(journal_fpath, outcome, fingerprint) def on_rows(index: int, rows: List[Dict[str, Any]]) -> None: + dispositions = [ + { + TASK_INDEX_KEY_NAME: row[TASK_INDEX_KEY_NAME], + ROLLOUT_INDEX_KEY_NAME: row[ROLLOUT_INDEX_KEY_NAME], + NG_FAILURE_CLASS_KEY: row.get(NG_FAILURE_CLASS_KEY), + NG_NO_PERSIST_KEY: bool(row.get(NG_NO_PERSIST_KEY)), + } + for row in rows + if not _is_success_row(row) and TASK_INDEX_KEY_NAME in row and ROLLOUT_INDEX_KEY_NAME in row + ] + no_persist_dispositions = [row for row in dispositions if row[NG_NO_PERSIST_KEY]] + persisted_dispositions = [row for row in dispositions if not row[NG_NO_PERSIST_KEY]] + # A no-persist attempt is a tombstone for any older timeout sidecar. + # Journal it before routing so a crash can only cause a conservative + # retry, never acceptance based on the stale timeout. + if no_persist_dispositions: + append_journal_record( + journal_fpath, + { + "stage_index": index, + "status": "attempt_dispositions", + "attempts": no_persist_dispositions, + }, + fingerprint, + ) route_stage_rows(output_fpath, rows) + if persisted_dispositions: + append_journal_record( + journal_fpath, + { + "stage_index": index, + "status": "attempt_dispositions", + "attempts": persisted_dispositions, + }, + fingerprint, + ) + + def on_restart(index: int) -> None: + append_journal_record( + journal_fpath, + {"stage_index": index, "status": "restart_from_stage"}, + fingerprint, + ) + aggregate_metrics_path_for(output_fpath).unlink(missing_ok=True) + _prune_downstream_files(output_fpath, index) + append_journal_record( + journal_fpath, + {"stage_index": index, "status": "restart_cleanup_complete"}, + fingerprint, + ) return StageResume( plans=plans, @@ -792,6 +2214,12 @@ def on_rows(index: int, rows: List[Dict[str, Any]]) -> None: on_plan=on_plan, on_outcome=on_outcome, on_rows=on_rows, + on_restart=on_restart, + elapsed_by_stage=elapsed_by_stage, + reuse_cached_keys=reuse_cached_keys, + attempts_by_stage=attempts_by_stage, + latest_failures_by_stage=latest_failures_by_stage, + latest_attempt_dispositions_by_stage=latest_attempt_dispositions_by_stage, ) @@ -806,15 +2234,18 @@ async def run_rollout_collection( """Rollout-collection driver entrypoint (wired via ``rollout_collection_driver``). Runs the multi-stage adaptive ELO procedure when ``multistage.enabled=true``; - otherwise delegates to the standard single-pass collection so rubric and - non-staged comparison runs behave exactly as they would without a driver. + otherwise delegates to the standard single-pass collection, so rubric and + non-staged comparison runs behave as they would without a driver except that + ``GDPVAL_FAILURE_CONTRACT`` fills in any key the config leaves unset. """ if (global_config_dict.get("multistage") or {}).get("enabled"): return await run_e2e_multistage(rollout_collection_config, global_config_dict) from nemo_gym.rollout_collection import RolloutCollectionHelper - await RolloutCollectionHelper().run_from_config(rollout_collection_config) + await RolloutCollectionHelper().run_from_config( + rollout_collection_config.model_copy(update=_failure_contract_updates(rollout_collection_config)) + ) return None @@ -826,8 +2257,8 @@ async def run_e2e_multistage( Called by ``ng_e2e_collect_rollouts`` when ``multistage.enabled=true``. Brings nothing up itself (the caller's ``RunHelper`` has already started the servers); it preprocesses the prepared dataset into materialized rows, samples/judges - stage-by-stage via the helper's ``run_examples``, writes the merged rollouts, - and runs the standard stage-aware ``_call_aggregate_metrics``. + stage-by-stage via the helper's ``_run_examples_with_metadata``, writes the + merged rollouts, and runs the standard stage-aware ``_call_aggregate_metrics``. """ from contextlib import nullcontext @@ -835,6 +2266,7 @@ async def run_e2e_multistage( from nemo_gym.rollout_collection import RolloutCollectionHelper, _attach_ng_perf multistage_config = parse_multistage_config(global_config_dict.get("multistage") or {}) + _reject_failure_contract_opt_out(rollout_collection_config) observability_enabled = observability_enabled_from_config(global_config_dict) helper = RolloutCollectionHelper() @@ -855,6 +2287,20 @@ async def run_e2e_multistage( ) semaphore_size = getattr(rollout_collection_config, "num_samples_in_parallel", None) + dispatch_budget_s = getattr(rollout_collection_config, "dispatch_budget_s", None) + drain_margin_s = getattr(rollout_collection_config, "drain_margin_s", None) + # The multistage driver bypasses run_from_config, so the failure-routing + # knob must be forwarded explicitly or a failed /run silently reverts to + # raising and killing every in-flight rollout of the stage. + route_failures_to_sidecar = bool(getattr(rollout_collection_config, "route_failures_to_sidecar", False)) + if dispatch_budget_s is not None and (semaphore_size is None or semaphore_size <= 0): + raise ValueError( + "dispatch_budget_s requires a finite positive num_samples_in_parallel; " + "unbounded dispatch can POST the entire queue before the budget is re-checked" + ) + + latency_tracker = DispatchLatencyTracker() + dispatch_started_at = time.monotonic() async def run_rollouts(rows: List[Dict[str, Any]]) -> List[Tuple[Dict[str, Any], Dict[str, Any]]]: semaphore = None @@ -863,7 +2309,17 @@ async def run_rollouts(rows: List[Dict[str, Any]]) -> List[Tuple[Dict[str, Any], semaphore = Semaphore(semaphore_size) results: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] - for future in helper._run_examples_with_metadata(rows, semaphore=semaphore or nullcontext()): + remaining_budget_s = None + if dispatch_budget_s is not None: + remaining_budget_s = max(0.0, dispatch_budget_s - (time.monotonic() - dispatch_started_at)) + for future in helper._run_examples_with_metadata( + rows, + semaphore=semaphore or nullcontext(), + dispatch_budget_s=remaining_budget_s, + drain_margin_s=drain_margin_s, + latency_tracker=latency_tracker, + route_failures_to_sidecar=route_failures_to_sidecar, + ): completed = await future _attach_ng_perf( completed.result, @@ -875,7 +2331,26 @@ async def run_rollouts(rows: List[Dict[str, Any]]) -> List[Tuple[Dict[str, Any], output_fpath = Path(rollout_collection_config.output_jsonl_fpath) journal_fpath = journal_path_for(output_fpath) - fingerprint = compute_fingerprint(multistage_config, reference_elos, distribution) + fingerprint = compute_fingerprint( + multistage_config, + reference_elos, + distribution, + materialized_rows=materialized_rows, + rollout_collection_config=rollout_collection_config, + resolved_global_config=global_config_dict, + ) + for row in materialized_rows: + row["verify_cache_namespace"] = fingerprint + + # Validate the repair config before ``_prepare_resume`` can set prior state + # aside, so a misconfigured run fails without touching it. + assignment_repair = None + transport_config = (global_config_dict.get("multistage") or {}).get("transport_assignment_repair") + if isinstance(transport_config, Mapping) and transport_config.get("enabled"): + from resources_servers.gdpval.transport_assignment import make_assignment_repair + + assignment_repair = make_assignment_repair(global_config_dict, transport_config) + resume = _prepare_resume(rollout_collection_config, output_fpath, journal_fpath, fingerprint) all_results, stage_summaries = await run_multistage_stages( @@ -886,10 +2361,19 @@ async def run_rollouts(rows: List[Dict[str, Any]]) -> List[Tuple[Dict[str, Any], run_rollouts, on_event=_log_event, resume=resume, + dispatch_longest_first=bool(getattr(rollout_collection_config, "dispatch_longest_first", False)), + assignment_repair=assignment_repair, ) + print(latency_tracker.summary()) + write_rollouts(all_results, output_fpath) + # After the rollouts file is on disk (so the evidence survives) but before + # any aggregate metrics are emitted: a run that never reached its declared + # final stage has no result to report. + _require_final_stage_rows(all_results, len(multistage_config.stages), stage_summaries) + print("[multistage-elo] computing stage-aware aggregate metrics") aggregate_metrics_fpath = await helper._call_aggregate_metrics(all_results, all_results, output_fpath) print( @@ -912,8 +2396,9 @@ def _prepare_resume( Always returns a writing StageResume so even a fresh run persists the journal and rows incrementally, giving a later resume state to read. Prior state is reused only when ``resume_from_cache`` is set and both the rollouts file and a - fingerprint-matching journal exist; every other case clears stale files and - starts fresh (with the reason logged). + fingerprint-matching journal exist; every other case renames the prior + rollouts, failure sidecar, journal and aggregate metrics to + ``.stale.`` and starts fresh (with the reason logged). """ import sys @@ -929,9 +2414,19 @@ def _prepare_resume( if reason is not None: print(f"[multistage-elo] starting fresh: {reason}", file=sys.stderr, flush=True) - output_fpath.unlink(missing_ok=True) - failures_path_for(output_fpath).unlink(missing_ok=True) - journal_fpath.unlink(missing_ok=True) + # Completed judged rollouts are expensive provider evidence. A stale + # fingerprint must never destroy them: set them aside recoverably. + stale_suffix = f".stale.{time.time_ns()}" + for stale in ( + output_fpath, + failures_path_for(output_fpath), + journal_fpath, + aggregate_metrics_path_for(output_fpath), + ): + if stale.exists(): + quarantined = stale.with_name(stale.name + stale_suffix) + stale.replace(quarantined) + print(f"[multistage-elo] quarantined stale state: {quarantined}", file=sys.stderr, flush=True) else: print("[multistage-elo] resuming multi-stage run from cache (fingerprint match)", file=sys.stderr, flush=True) return build_file_resume(output_fpath, journal_fpath, fingerprint) @@ -972,9 +2467,39 @@ def _log_event(name: str, data: dict) -> None: # pragma: no cover elif name == "stage_cached": elo = data.get("eval_elo") elo_str = f"{elo:.1f}" if isinstance(elo, (int, float)) else "unset (no games)" + cache_kind = "partial result" if data.get("partial") else "complete result" print( - f"[multistage-elo] stage {data['index'] + 1}/{data['total_stages']} reused from cache: " + f"[multistage-elo] stage {data['index'] + 1}/{data['total_stages']} reused {cache_kind} from cache: " f"eval ELO = {elo_str} ({data.get('num_rollouts')} cached rollout(s))", file=sys.stderr, flush=True, ) + elif name == "stage_partial_complete": + print( + f"[multistage-elo] stage {data['index'] + 1}/{data['total_stages']} accepted partial calibration: " + f"success coverage {data.get('success_fraction', 0):.1%}, omitted " + f"{data.get('num_omitted')} rollout(s)", + file=sys.stderr, + flush=True, + ) + elif name == "stage_retry": + print( + f"[multistage-elo] stage {data['index'] + 1}/{data['total_stages']} in-process retry: " + f"{data.get('num_selected')} rollout(s) selected, {data.get('num_dispatched')} re-dispatched, " + f"{data.get('num_drained')} drained without dispatch, " + f"{data.get('num_recovered')} recovered", + file=sys.stderr, + flush=True, + ) + elif name == "stage_incomplete": + detail = ( + f"coverage policy rejected {data.get('num_omitted')} resolved omission(s)" + if data.get("coverage_blocked") + else f"{data.get('num_pending')} rollout(s) retryable or undispatched" + ) + print( + f"[multistage-elo] stage {data['index'] + 1}/{data['total_stages']} remains incomplete " + f"({detail}); stopping before downstream stages", + file=sys.stderr, + flush=True, + ) diff --git a/resources_servers/gdpval/tests/test_hardening_fixes.py b/resources_servers/gdpval/tests/test_hardening_fixes.py new file mode 100644 index 00000000000..711d6244075 --- /dev/null +++ b/resources_servers/gdpval/tests/test_hardening_fixes.py @@ -0,0 +1,226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Resume safety and transport-repair hardening.""" + +import zipfile +from pathlib import Path +from types import SimpleNamespace + +import orjson +import pytest +from omegaconf import OmegaConf + +from resources_servers.gdpval.multistage_orchestrator import ( + MultiStageRunConfig, + _prepare_resume, + _prune_downstream_files, + compute_fingerprint, + failures_path_for, + parse_multistage_config, +) +from resources_servers.gdpval.transport_assignment import PairCost, _footprint, _pair_cost + + +REF_ELOS = {"ref_a": 1000.0} +DIST = {"grp": {"task_ids": ["t0"], "percentage": 100.0}} + + +def _cfg() -> MultiStageRunConfig: + return MultiStageRunConfig( + enabled=True, + stages=parse_multistage_config({"enabled": True, "stages": [{"num_tasks": 1}]}).stages, + seed=0, + ) + + +class TestFingerprintConnectionFields: + def test_endpoint_and_credential_changes_do_not_invalidate(self) -> None: + def runtime(base_url: str, api_key: str) -> dict: + return { + "judge": { + "responses_api_models": { + "minimax": { + "openai_base_url": base_url, + "openai_api_key": api_key, + "openai_model": "minimax-m3", + } + } + } + } + + cfg = _cfg() + baseline = compute_fingerprint( + cfg, REF_ELOS, DIST, resolved_global_config=runtime("http://10.0.0.1:8000/v1", "sk-a") + ) + rotated = compute_fingerprint( + cfg, REF_ELOS, DIST, resolved_global_config=runtime("http://10.0.9.9:8000/v1", "sk-b") + ) + assert rotated == baseline + + changed_model = dict(runtime("http://10.0.0.1:8000/v1", "sk-a")) + changed_model["judge"]["responses_api_models"]["minimax"]["openai_model"] = "other-model" + assert compute_fingerprint(cfg, REF_ELOS, DIST, resolved_global_config=changed_model) != baseline + + def test_bare_connection_fields_and_headers_do_not_invalidate(self) -> None: + def runtime(**overrides) -> dict: + server = { + "base_url": ["http://10.0.0.1:8000/v1"], + "api_key": "sk-a", + "default_headers": {"Authorization": "Bearer a"}, + "openai_default_headers": {"X-Key": "a"}, + "model": "policy", + } + server.update(overrides) + return {"policy_model": {"responses_api_models": {"vllm_model": server}}} + + cfg = _cfg() + baseline = compute_fingerprint(cfg, REF_ELOS, DIST, resolved_global_config=runtime()) + rotated = runtime( + base_url=["http://10.0.9.9:8000/v1", "http://10.0.9.10:8000/v1"], + api_key="sk-b", + default_headers={"Authorization": "Bearer b"}, + openai_default_headers={"X-Key": "b"}, + ) + assert compute_fingerprint(cfg, REF_ELOS, DIST, resolved_global_config=rotated) == baseline + assert compute_fingerprint(cfg, REF_ELOS, DIST, resolved_global_config=runtime(model="other")) != baseline + + +class TestFingerprintBenchmarkConfig: + """Resume across allocations with the shipped GDPVal benchmark and policy configs.""" + + REPO_ROOT = Path(__file__).resolve().parents[3] + + def _fingerprint(self, monkeypatch: pytest.MonkeyPatch, *, rerun_incomplete: str = "false", **policy) -> str: + monkeypatch.setenv("RERUN_INCOMPLETE", rerun_incomplete) + policy_values = { + "policy_base_url": "http://10.0.0.1:8000/v1", + "policy_api_key": "key-a", + "policy_model_name": "policy-model", + **policy, + } + config = OmegaConf.merge( + OmegaConf.load(self.REPO_ROOT / "benchmarks/gdpval/config.yaml"), + OmegaConf.load(self.REPO_ROOT / "responses_api_models/vllm_model/configs/vllm_model.yaml"), + OmegaConf.create(policy_values), + ) + resolved = OmegaConf.to_container(config, resolve=True) + return compute_fingerprint(_cfg(), REF_ELOS, DIST, resolved_global_config=resolved) + + def test_new_allocation_and_resume_knobs_keep_the_fingerprint(self, monkeypatch: pytest.MonkeyPatch) -> None: + baseline = self._fingerprint(monkeypatch) + + assert self._fingerprint(monkeypatch, policy_base_url="http://10.0.9.9:8000/v1") == baseline + assert self._fingerprint(monkeypatch, policy_api_key="key-b") == baseline + assert self._fingerprint(monkeypatch, rerun_incomplete="true") == baseline + + def test_policy_model_change_invalidates(self, monkeypatch: pytest.MonkeyPatch) -> None: + baseline = self._fingerprint(monkeypatch) + + assert self._fingerprint(monkeypatch, policy_model_name="other-model") != baseline + + def test_agent_concurrency_is_operational_but_turn_cap_is_not(self) -> None: + def runtime(**agent) -> dict: + return { + "gdpval_stirrup_agent": { + "responses_api_agents": {"stirrup_agent": {"concurrency": 32, "agent_max_turns": 250, **agent}} + } + } + + cfg = _cfg() + baseline = compute_fingerprint(cfg, REF_ELOS, DIST, resolved_global_config=runtime()) + assert compute_fingerprint(cfg, REF_ELOS, DIST, resolved_global_config=runtime(concurrency=8)) == baseline + assert ( + compute_fingerprint(cfg, REF_ELOS, DIST, resolved_global_config=runtime(agent_max_turns=100)) != baseline + ) + + +class TestStaleResumeQuarantine: + def test_stale_state_is_quarantined_not_deleted(self, tmp_path: Path) -> None: + output = tmp_path / "rollouts.jsonl" + journal = tmp_path / "rollouts_multistage_state.jsonl" + payload = b'{"stage_index": 0, "reward": 1.0}\n' + output.write_bytes(payload) + journal.write_bytes(b'{"stage_index": 0, "status": "planned", "fingerprint": "OLD"}\n') + + config = SimpleNamespace(resume_from_cache=True) + _prepare_resume(config, output, journal, "NEW") + + assert not journal.exists() + quarantined = sorted(tmp_path.glob("rollouts.jsonl.stale.*")) + assert quarantined and quarantined[0].read_bytes() == payload + + +class TestPruneQuarantine: + def test_pruned_rows_are_preserved(self, tmp_path: Path) -> None: + output = tmp_path / "rollouts.jsonl" + rows = [ + {"stage_index": 0, "task_id": "keep"}, + {"stage_index": 1, "task_id": "pruned"}, + ] + output.write_bytes(b"".join(orjson.dumps(row) + b"\n" for row in rows)) + failures_path_for(output).write_bytes(orjson.dumps({"stage_index": 2, "task_id": "pruned_failure"}) + b"\n") + + _prune_downstream_files(output, restart_stage=0) + + assert orjson.loads(output.read_bytes())["task_id"] == "keep" + main_quarantine = sorted(tmp_path.glob("rollouts.jsonl.pruned.*")) + assert main_quarantine and orjson.loads(main_quarantine[0].read_bytes())["task_id"] == "pruned" + sidecar_quarantine = sorted(tmp_path.glob("*failures*.pruned.*")) + assert sidecar_quarantine + + +class TestTransportRepairHardening: + def _cost(self, candidate: object, reference: object) -> PairCost: + return _pair_cost( + candidate, + reference, + max_file_bytes=250 * 1024 * 1024, + max_raw_bytes=300 * 1024 * 1024, + max_wire_bytes=420 * 1024 * 1024, + max_section_raw_bytes=96 * 1024 * 1024, + framing_reserve_bytes=4 * 1024 * 1024, + ) + + def test_corrupt_zip_is_costed_incompatible_not_fatal(self, tmp_path: Path) -> None: + (tmp_path / "broken.zip").write_bytes(b"this is not a zip archive") + footprint = _footprint(tmp_path) + assert footprint.defective is True + + from resources_servers.gdpval.transport_assignment import Footprint + + clean = Footprint(raw_bytes=1, max_file_bytes=1, has_av=False, file_count=1) + cost = self._cost(footprint, clean) + assert cost.compatible is False and "corrupt_archive" in cost.reasons + + def test_valid_zip_members_still_enumerated(self, tmp_path: Path) -> None: + with zipfile.ZipFile(tmp_path / "ok.zip", "w") as archive: + archive.writestr("page.pdf", b"%PDF-fake") + footprint = _footprint(tmp_path) + assert footprint.defective is False and footprint.file_count == 1 + + def test_section_budget_marks_pair_incompatible(self) -> None: + from resources_servers.gdpval.transport_assignment import Footprint + + big_section = Footprint( + raw_bytes=150 * 1024 * 1024, max_file_bytes=150 * 1024 * 1024, has_av=False, file_count=1 + ) + small = Footprint(raw_bytes=1, max_file_bytes=1, has_av=False, file_count=1) + cost = self._cost(big_section, small) + assert cost.compatible is False and "section_raw_over_cap" in cost.reasons + + def test_new_av_extensions_are_recognized(self, tmp_path: Path) -> None: + (tmp_path / "voice.opus").write_bytes(b"0" * 10) + footprint = _footprint(tmp_path) + assert footprint.has_av is True and footprint.file_count == 1 diff --git a/resources_servers/gdpval/tests/test_multistage_elo.py b/resources_servers/gdpval/tests/test_multistage_elo.py index 9aea3beb4f8..20d14d49425 100644 --- a/resources_servers/gdpval/tests/test_multistage_elo.py +++ b/resources_servers/gdpval/tests/test_multistage_elo.py @@ -110,6 +110,38 @@ def test_deterministic_for_same_rng_seed(self) -> None: def test_empty_reference_set_yields_no_assignment(self) -> None: assert assign_task_references(["t0", "t1"], [], rng=random.Random(0)) == {} + def test_default_preserves_independent_uniform_draws(self) -> None: + task_ids = [f"t{i}" for i in range(20)] + refs = ["a", "b", "c"] + expected_rng = random.Random(11) + expected = {task_id: expected_rng.choice(refs) for task_id in task_ids} + + assert assign_task_references(task_ids, refs, rng=random.Random(11)) == expected + + def test_balanced_counts_differ_by_at_most_one(self) -> None: + task_ids = [f"t{i}" for i in range(10)] + refs = ["a", "b", "c"] + assignment = assign_task_references(task_ids, refs, rng=random.Random(7), balanced=True) + counts = {ref: list(assignment.values()).count(ref) for ref in refs} + + assert set(assignment) == set(task_ids) + assert max(counts.values()) - min(counts.values()) <= 1 + + def test_balanced_assignment_is_seeded_and_randomized(self) -> None: + task_ids = [f"t{i}" for i in range(20)] + refs = ["a", "b", "c"] + first = assign_task_references(task_ids, refs, rng=random.Random(3), balanced=True) + replay = assign_task_references(task_ids, refs, rng=random.Random(3), balanced=True) + other = assign_task_references(task_ids, refs, rng=random.Random(4), balanced=True) + + assert first == replay + assert first != other + + def test_balanced_with_fewer_tasks_than_references(self) -> None: + assignment = assign_task_references(["t0", "t1"], ["a", "b", "c", "d"], rng=random.Random(5), balanced=True) + + assert len(set(assignment.values())) == 2 + class TestStageAssignmentRng: def test_distinct_streams_per_stage_and_seed(self) -> None: diff --git a/resources_servers/gdpval/tests/test_multistage_orchestrator.py b/resources_servers/gdpval/tests/test_multistage_orchestrator.py index 7afb91316ab..11ab739530e 100644 --- a/resources_servers/gdpval/tests/test_multistage_orchestrator.py +++ b/resources_servers/gdpval/tests/test_multistage_orchestrator.py @@ -16,14 +16,21 @@ from __future__ import annotations import json -from contextlib import nullcontext +from copy import deepcopy from pathlib import Path from types import SimpleNamespace -from typing import Any, Dict, List, Tuple +from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple import pytest -from nemo_gym.global_config import AGENT_REF_KEY_NAME, ROLLOUT_INDEX_KEY_NAME, TASK_INDEX_KEY_NAME +import nemo_gym.rollout_collection as rollout_collection_module +import resources_servers.gdpval.multistage_orchestrator as multistage_module +from nemo_gym.global_config import ( + AGENT_REF_KEY_NAME, + ATTEMPT_INDEX_KEY_NAME, + ROLLOUT_INDEX_KEY_NAME, + TASK_INDEX_KEY_NAME, +) from nemo_gym.path_utils import failures_path_for from nemo_gym.rollout_collection import ( NG_FAILURE_CLASS_KEY, @@ -31,23 +38,38 @@ NG_TERMINAL_KEY, SharedRolloutCollectionConfig, ) +from resources_servers.gdpval.multistage_elo import PartialStagePolicy from resources_servers.gdpval.multistage_orchestrator import ( + GDPVAL_FAILURE_CONTRACT, MultiStageRunConfig, StageResume, + _cached_partial_snapshot_is_valid, + _is_in_process_retryable, + _partial_stage_outcome, _prepare_resume, + _require_final_stage_rows, + aggregate_metrics_path_for, + append_journal_record, build_file_resume, build_stage_rows, compute_fingerprint, find_gdpval_reference_elos, index_rows_by_task, journal_path_for, + load_failure_attempts, + load_failure_timings, load_gated_keys, + load_latest_attempt_dispositions, + load_latest_failures, load_persisted_rows, + load_reuse_cached_keys, parse_multistage_config, read_journal, route_stage_rows, row_task_id, + run_e2e_multistage, run_multistage_stages, + run_rollout_collection, tag_results, write_rollouts, ) @@ -56,6 +78,39 @@ REF_ELOS = {"a": 1000.0, "b": 1200.0, "c": 1400.0, "d": 1600.0} +def _runtime_components_with_bind_addresses(host_prefix: str, port_offset: int) -> Dict[str, Any]: + return { + "agent": { + "responses_api_agents": { + "stirrup": { + "host": f"{host_prefix}-agent", + "port": 8001 + port_offset, + "task": "gdpval", + "worker_target": {"host": "sandbox.internal", "port": 9000}, + } + } + }, + "policy": { + "responses_api_models": { + "vllm": { + "host": f"{host_prefix}-policy", + "port": 8002 + port_offset, + "model": "/checkpoints/policy-a", + } + } + }, + "resources": { + "resources_servers": { + "gdpval": { + "host": f"{host_prefix}-resources", + "port": 8003 + port_offset, + "num_comparison_trials": 2, + } + } + }, + } + + # --------------------------------------------------------------------------- # Config parsing # --------------------------------------------------------------------------- @@ -94,6 +149,66 @@ def test_empty_stages_raises(self) -> None: with pytest.raises(ValueError): parse_multistage_config({"enabled": True, "stages": []}) + def test_parses_partial_completion_policy(self) -> None: + cfg = parse_multistage_config( + { + "enabled": True, + "stages": [ + { + "num_tasks": 45, + "partial_completion": { + "min_success_fraction": 0.9, + "min_per_reference_success_fraction": 0.5, + "min_successful_rows_per_reference": 1, + }, + }, + {"num_tasks": 220, "num_models": 4}, + ], + } + ) + + policy = cfg.stages[0].partial_completion + assert policy is not None + assert policy.min_success_fraction == 0.9 + assert policy.min_per_reference_success_fraction == 0.5 + assert policy.min_successful_rows_per_reference == 1 + assert cfg.stages[1].partial_completion is None + + @pytest.mark.parametrize( + "partial_completion", + [ + {"min_success_fraction": 0}, + {"min_success_fraction": True}, + {"min_per_reference_success_fraction": 1.1}, + {"min_successful_rows_per_reference": 0}, + {"min_successful_rows_per_reference": True}, + {"allowed_failure_classes": ["transient"]}, + ], + ) + def test_rejects_unsafe_partial_completion_policy(self, partial_completion: Dict[str, Any]) -> None: + with pytest.raises(ValueError): + parse_multistage_config( + { + "enabled": True, + "stages": [ + {"num_tasks": 5, "partial_completion": partial_completion}, + {"num_tasks": 10, "num_models": 2}, + ], + } + ) + + 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}}, + ], + } + ) + class TestFindReferenceElos: def test_extracts_from_nel_style_config(self) -> None: @@ -197,10 +312,17 @@ def test_tag_results_stamps_identity(self) -> None: "task_id": "t3", } result = {"per_reference": {}, "reward": 1.0} - tagged = tag_results([(row, result)], stage_index=1) + tagged = tag_results( + [(row, result)], + stage_index=1, + expected_final_stage_index=2, + expected_stage_row_count=17, + ) assert tagged[0][TASK_INDEX_KEY_NAME] == 3 assert tagged[0][ROLLOUT_INDEX_KEY_NAME] == 7 assert tagged[0]["stage_index"] == 1 + assert tagged[0]["expected_final_stage_index"] == 2 + assert tagged[0]["expected_stage_row_count"] == 17 assert tagged[0]["task_id"] == "t3" @@ -240,6 +362,75 @@ async def fake_run_rollouts(rows: List[Dict[str, Any]]) -> List[Tuple[Dict[str, return fake_run_rollouts +class TestFinalStageIsTerminal: + """A coverage-rejected run that never reached its final stage must fail.""" + + @staticmethod + def _summary(stage_index: int, **extra: Any) -> Dict[str, Any]: + return {"stage_index": stage_index, **extra} + + def test_coverage_rejected_empty_final_stage_raises(self) -> None: + # Stage 0 was rejected on coverage with nothing left to retry, so stage 1 + # was never planned and the declared final stage has no rows. + with pytest.raises(RuntimeError, match="produced no rows for it"): + _require_final_stage_rows( + [{"stage_index": 0}], + 2, + [self._summary(0, incomplete=True, coverage_blocked=True)], + ) + + def test_retryable_stop_does_not_raise(self, capsys: pytest.CaptureFixture[str]) -> None: + """Rows still retryable: a resume continues the run, so this is not terminal.""" + _require_final_stage_rows( + [{"stage_index": 0}], + 2, + [self._summary(0, incomplete=True, coverage_blocked=False)], + ) + + message = capsys.readouterr().err + assert "stopped after stage 1/2" in message + assert "--resume" in message + + def test_final_stage_with_some_rows_does_not_raise(self) -> None: + """A partial final stage is left to aggregation (reported degraded), not fatal.""" + _require_final_stage_rows( + [{"stage_index": 0}, {"stage_index": 1}], + 2, + [self._summary(0), self._summary(1, incomplete=True, coverage_blocked=True)], + ) + + def test_complete_run_does_not_raise(self, capsys: pytest.CaptureFixture[str]) -> None: + _require_final_stage_rows( + [{"stage_index": 0}, {"stage_index": 1}], + 2, + [self._summary(0), self._summary(1)], + ) + + assert capsys.readouterr().err == "" + + async def test_coverage_rejection_marks_the_stage_summary(self) -> None: + """The loop itself still returns; it only records why it stopped.""" + task_ids = [f"t{i}" for i in range(4)] + rows = _materialized_rows(task_ids) + cfg = MultiStageRunConfig( + enabled=True, + stages=parse_multistage_config({"enabled": True, "stages": [{"num_models": 4}, {"num_models": 2}]}).stages, + seed=0, + ) + + async def no_battle_evidence(rows_in: List[Dict[str, Any]]): + return [(row, {"task_id": row["task_id"], "per_reference": {}}) for row in rows_in] + + all_results, summaries = await run_multistage_stages( + cfg, REF_ELOS, _distribution(task_ids), rows, no_battle_evidence + ) + + assert not any(row.get("stage_index") == 1 for row in all_results) + assert summaries[-1]["coverage_blocked"] is True + with pytest.raises(RuntimeError, match="produced no rows for it"): + _require_final_stage_rows(all_results, len(cfg.stages), summaries) + + class TestRunStages: async def test_threads_elo_and_shrinks_references(self) -> None: task_ids = [f"t{i}" for i in range(20)] @@ -258,6 +449,9 @@ async def test_threads_elo_and_shrinks_references(self) -> None: _fake_run_rollouts_factory(), ) + assert all(row["expected_final_stage_index"] == 1 for row in all_results) + assert all(row["expected_stage_row_count"] == 20 for row in all_results) + # Stage 0 uses all references; stage 1 shrinks to the 2 closest to the # stage-0 estimate (~1300 ⇒ b=1200, c=1400). assert summaries[0]["reference_ids"] == ["a", "b", "c", "d"] @@ -404,6 +598,67 @@ async def test_emits_lifecycle_events(self) -> None: assert events.count("stage_end") == 2 +class TestAssignmentRepairSeam: + """The planner dispatches the repair callback's mapping and records its receipt.""" + + @staticmethod + def _cfg() -> MultiStageRunConfig: + return MultiStageRunConfig( + enabled=True, + stages=parse_multistage_config({"enabled": True, "stages": [{"num_tasks": 4}]}).stages, + seed=0, + ) + + async def test_repaired_assignment_is_dispatched_and_recorded(self) -> None: + task_ids = [f"t{i}" for i in range(4)] + calls: List[Tuple[int, List[str], Dict[str, str]]] = [] + + def repair(stage_index, reference_ids, original): + calls.append((stage_index, list(reference_ids), dict(original))) + target = sorted(reference_ids)[0] + return {task_id: target for task_id in original}, {"schema": "test", "changes": []} + + resume = RecordingResume() + results, _ = await run_multistage_stages( + self._cfg(), + REF_ELOS, + _distribution(task_ids), + _materialized_rows(task_ids), + _fake_run_rollouts_factory(), + resume=resume, + assignment_repair=repair, + ) + + assert [call[0] for call in calls] == [0] + # The uniform draw spread the tasks; only the repair collapsed them onto "a". + assert set(calls[0][2].values()) != {"a"} + assert {reference_id for row in results for reference_id in row["per_reference"]} == {"a"} + plan = resume.planned[0][1] + assert set(plan["task_reference_ids"].values()) == {"a"} + assert plan["transport_assignment_repair"] == {"schema": "test", "changes": []} + + @pytest.mark.parametrize( + ("mutate", "message"), + [ + (lambda mapping: dict(list(mapping.items())[1:]), "changed the stage task set"), + (lambda mapping: {task_id: "not_a_reference" for task_id in mapping}, "outside the stage"), + ], + ids=["drops-a-task", "unknown-reference"], + ) + async def test_contract_violations_raise(self, mutate, message: str) -> None: + task_ids = [f"t{i}" for i in range(4)] + + with pytest.raises(ValueError, match=message): + await run_multistage_stages( + self._cfg(), + REF_ELOS, + _distribution(task_ids), + _materialized_rows(task_ids), + _fake_run_rollouts_factory(), + assignment_repair=lambda stage_index, reference_ids, original: (mutate(dict(original)), {}), + ) + + class TestWriteRollouts: def test_writes_sorted_jsonl(self, tmp_path: Path) -> None: results = [ @@ -414,6 +669,17 @@ def test_writes_sorted_jsonl(self, tmp_path: Path) -> None: lines = [json.loads(line) for line in out.read_text().splitlines()] assert [line["task_id"] for line in lines] == ["t0", "t1"] + def test_fresh_file_uses_normal_umask_permissions(self, tmp_path: Path) -> None: + control = tmp_path / "normal.jsonl" + control.write_bytes(b"normal\n") + + out = write_rollouts( + [{TASK_INDEX_KEY_NAME: 0, ROLLOUT_INDEX_KEY_NAME: 0}], + tmp_path / "rollouts.jsonl", + ) + + assert out.stat().st_mode & 0o777 == control.stat().st_mode & 0o777 + def test_dedupes_by_stage_task_rollout(self, tmp_path: Path) -> None: results = [ {"stage_index": 0, TASK_INDEX_KEY_NAME: 0, ROLLOUT_INDEX_KEY_NAME: 0, "task_id": "old"}, @@ -425,6 +691,32 @@ def test_dedupes_by_stage_task_rollout(self, tmp_path: Path) -> None: # Dedup keeps the last write per (stage, task, rollout); stage 1 is distinct. assert [line["task_id"] for line in lines] == ["new", "other"] + def test_failed_rewrite_preserves_previous_file(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + out = tmp_path / "rollouts.jsonl" + original = b'{"previous":true}\n' + out.write_bytes(original) + real_dumps = multistage_module.orjson.dumps + calls = 0 + + def fail_on_second_row(value, *args, **kwargs): + nonlocal calls + calls += 1 + if calls == 2: + raise RuntimeError("injected serialization failure") + return real_dumps(value, *args, **kwargs) + + monkeypatch.setattr(multistage_module.orjson, "dumps", fail_on_second_row) + rows = [ + {"stage_index": 0, TASK_INDEX_KEY_NAME: 0, ROLLOUT_INDEX_KEY_NAME: 0}, + {"stage_index": 0, TASK_INDEX_KEY_NAME: 1, ROLLOUT_INDEX_KEY_NAME: 0}, + ] + + with pytest.raises(RuntimeError, match="injected serialization failure"): + write_rollouts(rows, out) + + assert out.read_bytes() == original + assert list(tmp_path.glob(".rollouts.jsonl.merge-*")) == [] + # --------------------------------------------------------------------------- # Resume seam (pure, in-memory) @@ -438,10 +730,22 @@ class RecordingResume(StageResume): explicitly to model terminal / max-attempt gating from the sidecar. """ - def __init__(self, plans=None, outcomes=None, rows_by_stage=None, gated_keys=None) -> None: + def __init__( + self, + plans=None, + outcomes=None, + rows_by_stage=None, + gated_keys=None, + elapsed_by_stage=None, + reuse_cached_keys=None, + attempts_by_stage=None, + latest_failures_by_stage=None, + latest_attempt_dispositions_by_stage=None, + ) -> None: self.planned: List[Tuple[int, dict]] = [] self.completed: List[Tuple[int, dict]] = [] self.appended: Dict[int, List[Dict[str, Any]]] = {} + self.restarted: List[int] = [] rows_by_stage = dict(rows_by_stage or {}) if gated_keys is None: gated_keys = { @@ -456,6 +760,12 @@ def __init__(self, plans=None, outcomes=None, rows_by_stage=None, gated_keys=Non on_plan=lambda i, p: self.planned.append((i, p)), on_outcome=lambda i, o: self.completed.append((i, o)), on_rows=lambda i, r: self.appended.setdefault(i, []).extend(r), + on_restart=self.restarted.append, + elapsed_by_stage=dict(elapsed_by_stage or {}), + reuse_cached_keys=dict(reuse_cached_keys or {}), + attempts_by_stage=dict(attempts_by_stage or {}), + latest_failures_by_stage=dict(latest_failures_by_stage or {}), + latest_attempt_dispositions_by_stage=dict(latest_attempt_dispositions_by_stage or {}), ) @@ -497,7 +807,7 @@ async def test_complete_stage_skips_dispatch_and_threads_elo(self, incomplete) - stage0_plan = { "stage_index": 0, "reference_ids": base_summaries[0]["reference_ids"], - "task_ids": [f"t{i}" for i in range(3)], + "task_ids": list(dict.fromkeys(row["task_id"] for row in stage0_rows)), } resume = RecordingResume( plans={0: stage0_plan}, @@ -515,19 +825,22 @@ async def counting_run(rows_in: List[Dict[str, Any]]): cfg, REF_ELOS, _distribution(task_ids), rows, counting_run, resume=resume ) - # Stage 0 was not dispatched; only stage 1 ran. - assert len(dispatched) == 1 - # Stage 0 ELO was re-fit from cached rows and threaded into stage 1's - # reference selection (same as the original full run). - assert summaries[0]["cached"] is True - if not incomplete: + if incomplete: + # A stage recorded complete but missing a planned row is stale: it is + # re-dispatched rather than reused, so the run still ends complete. + assert len(dispatched) == 2 + assert not summaries[0].get("cached") + else: + # Stage 0 was not dispatched; only stage 1 ran. Its ELO was re-fit from + # the cached rows and threaded into stage 1's reference selection. + assert len(dispatched) == 1 + assert summaries[0]["cached"] is True assert summaries[1]["reference_ids"] == base_summaries[1]["reference_ids"] assert summaries[0]["num_rollouts"] == 6 - if incomplete: - results.append(results[0]) # A duplicate cannot replace the missing repeat. config = SharedRolloutCollectionConfig(output_jsonl_fpath="rollouts.jsonl", require_complete=True) - with pytest.raises(RuntimeError, match="15/16 samples completed") if incomplete else nullcontext(): - config.check_completion(expected=sum(s["num_rollouts"] for s in summaries), results=results) + config.check_completion(expected=sum(s["num_rollouts"] for s in summaries), results=results) + with pytest.raises(RuntimeError, match="15/16 samples completed"): + config.check_completion(expected=sum(s["num_rollouts"] for s in summaries), results=results[:-1]) async def test_interrupted_stage_redispatches_only_missing(self) -> None: task_ids = [f"t{i}" for i in range(10)] @@ -574,6 +887,203 @@ async def capturing_run(rows_in: List[Dict[str, Any]]): # Final stage-0 result count equals the full run (cached + re-dispatched). assert summaries[0]["num_rollouts"] == base_summaries[0]["num_rollouts"] + async def test_completed_cached_stage_stamps_full_expected_row_count(self) -> None: + task_ids = ["t0", "t1"] + rows = _materialized_rows(task_ids, repeats=2) + cfg = MultiStageRunConfig( + enabled=True, + stages=parse_multistage_config({"enabled": True, "stages": ["2"]}).stages, + seed=0, + ) + full_results, summaries = await run_multistage_stages( + cfg, REF_ELOS, _distribution(task_ids), rows, _fake_run_rollouts_factory() + ) + # Model a completed stage with one persisted success and the other rows + # accounted for by terminal/max-attempt sidecars. The success must still + # declare the full planned cardinality, not the cached success count. + resume = RecordingResume( + plans={ + 0: { + "stage_index": 0, + "reference_ids": summaries[0]["reference_ids"], + "task_ids": task_ids, + } + }, + outcomes={0: {"stage_index": 0, "status": "complete"}}, + rows_by_stage={0: [full_results[0]]}, + gated_keys={0: {(row[TASK_INDEX_KEY_NAME], row[ROLLOUT_INDEX_KEY_NAME]) for row in full_results}}, + ) + + async def no_dispatch(rows_in: List[Dict[str, Any]]): + raise AssertionError("completed cached stage must not dispatch") + + cached_results, _ = await run_multistage_stages( + cfg, REF_ELOS, _distribution(task_ids), rows, no_dispatch, resume=resume + ) + + assert len(cached_results) == 1 + assert cached_results[0]["expected_stage_row_count"] == 4 + assert cached_results[0]["expected_final_stage_index"] == 0 + + async def test_completed_stage_with_missing_cached_key_is_reopened(self) -> None: + task_ids = ["t0", "t1"] + rows = _materialized_rows(task_ids) + cfg = MultiStageRunConfig( + enabled=True, + stages=parse_multistage_config({"enabled": True, "stages": ["2"]}).stages, + seed=0, + ) + full_run = _fake_run_rollouts_factory() + full_results, summaries = await run_multistage_stages(cfg, REF_ELOS, _distribution(task_ids), rows, full_run) + missing = full_results[-1] + missing_key = (missing[TASK_INDEX_KEY_NAME], missing[ROLLOUT_INDEX_KEY_NAME]) + cached = full_results[:-1] + resume = RecordingResume( + plans={ + 0: { + "stage_index": 0, + "reference_ids": summaries[0]["reference_ids"], + "task_ids": task_ids, + } + }, + outcomes={0: {"stage_index": 0, "status": "complete"}}, + rows_by_stage={0: cached}, + ) + dispatched: List[Tuple[int, int]] = [] + + async def capture(rows_in: List[Dict[str, Any]]): + dispatched.extend((row[TASK_INDEX_KEY_NAME], row[ROLLOUT_INDEX_KEY_NAME]) for row in rows_in) + return await full_run(rows_in) + + results, result_summaries = await run_multistage_stages( + cfg, REF_ELOS, _distribution(task_ids), rows, capture, resume=resume + ) + + assert dispatched == [missing_key] + assert len(results) == len(full_results) + assert result_summaries[0].get("cached") is not True + + async def test_reopened_calibration_invalidates_dependent_stage(self) -> None: + task_ids = ["t0", "t1", "t2"] + rows = _materialized_rows(task_ids) + cfg = MultiStageRunConfig( + enabled=True, + stages=parse_multistage_config({"enabled": True, "stages": ["3", "3:2"]}).stages, + seed=0, + ) + baseline_resume = RecordingResume() + successful_run = _fake_run_rollouts_factory() + baseline_results, _ = await run_multistage_stages( + cfg, + REF_ELOS, + _distribution(task_ids), + rows, + successful_run, + resume=baseline_resume, + ) + rows_by_stage = { + stage_index: [row for row in baseline_results if row["stage_index"] == stage_index] + for stage_index in (0, 1) + } + missing = rows_by_stage[0].pop() + missing_key = (missing[TASK_INDEX_KEY_NAME], missing[ROLLOUT_INDEX_KEY_NAME]) + resumed = RecordingResume( + plans=dict(baseline_resume.planned), + outcomes=dict(baseline_resume.completed), + rows_by_stage=rows_by_stage, + ) + dispatched: List[Tuple[int, Tuple[int, int]]] = [] + + async def capture(rows_in: List[Dict[str, Any]]): + dispatched.extend( + ( + row["stage_index"], + (row[TASK_INDEX_KEY_NAME], row[ROLLOUT_INDEX_KEY_NAME]), + ) + for row in rows_in + ) + return await successful_run(rows_in) + + await run_multistage_stages( + cfg, + REF_ELOS, + _distribution(task_ids), + rows, + capture, + resume=resumed, + ) + + assert resumed.restarted == [0] + assert [key for stage, key in dispatched if stage == 0] == [missing_key] + assert len([key for stage, key in dispatched if stage == 1]) == 3 + + async def test_resumed_empty_fit_blocks_adaptive_next_stage(self) -> None: + task_ids = [f"t{i}" for i in range(6)] + rows = _materialized_rows(task_ids) + cfg = MultiStageRunConfig( + enabled=True, + stages=parse_multistage_config({"enabled": True, "stages": ["2:4", "2:2", "2:1"]}).stages, + seed=0, + reuse_cached_deliverables=False, + ) + baseline_resume = RecordingResume() + successful_run = _fake_run_rollouts_factory() + + async def terminal_middle_stage(rows_in: List[Dict[str, Any]]): + if rows_in and rows_in[0]["stage_index"] == 1: + return [ + ( + row, + { + NG_FAILURE_CLASS_KEY: "skipped", + NG_TERMINAL_KEY: True, + }, + ) + for row in rows_in + ] + return await successful_run(rows_in) + + baseline_results, baseline_summaries = await run_multistage_stages( + cfg, + REF_ELOS, + _distribution(task_ids), + rows, + terminal_middle_stage, + resume=baseline_resume, + ) + assert baseline_summaries[0]["eval_elo"] is not None + assert baseline_summaries[1]["eval_elo"] is None + + stage0_rows = [row for row in baseline_results if row["stage_index"] == 0] + stage0_keys = {(row[TASK_INDEX_KEY_NAME], row[ROLLOUT_INDEX_KEY_NAME]) for row in stage0_rows} + stage1_keys = {(row[TASK_INDEX_KEY_NAME], row[ROLLOUT_INDEX_KEY_NAME]) for row in baseline_resume.appended[1]} + resumed = RecordingResume( + plans=dict(baseline_resume.planned[:2]), + outcomes=dict(baseline_resume.completed[:2]), + rows_by_stage={0: stage0_rows}, + gated_keys={0: stage0_keys, 1: stage1_keys}, + ) + + dispatched_stages: List[int] = [] + + async def capture_stage(rows_in: List[Dict[str, Any]]): + dispatched_stages.extend(row["stage_index"] for row in rows_in) + return await successful_run(rows_in) + + _, resumed_summaries = await run_multistage_stages( + cfg, + REF_ELOS, + _distribution(task_ids), + _materialized_rows(task_ids), + capture_stage, + resume=resumed, + ) + + assert dispatched_stages == [] + assert len(resumed_summaries) == 2 + assert resumed_summaries[0]["cached"] is True + assert resumed_summaries[1]["eval_elo"] is None + async def test_plan_replay_is_deterministic_without_seed(self) -> None: task_ids = [f"t{i}" for i in range(10)] rows = _materialized_rows(task_ids) @@ -642,54 +1152,1555 @@ async def capturing_run(rows_in: List[Dict[str, Any]]): await run_multistage_stages(cfg, REF_ELOS, _distribution(task_ids), rows, capturing_run, resume=resume) assert dispatched_keys == [failed_key] + async def test_resume_longest_first_uses_stage_aware_failure_timings(self) -> None: + task_ids = [f"t{i}" for i in range(4)] + rows = _materialized_rows(task_ids) + cfg = MultiStageRunConfig( + enabled=True, + stages=parse_multistage_config({"enabled": True, "stages": ["4"]}).stages, + seed=0, + ) + resume = RecordingResume( + plans={ + 0: { + "stage_index": 0, + "reference_ids": ["a"], + "task_ids": task_ids, + "task_reference_ids": {task_id: "a" for task_id in task_ids}, + } + }, + elapsed_by_stage={0: {(0, 0): 60.0, (2, 0): 9000.0}}, + ) + dispatched: List[int] = [] + run = _fake_run_rollouts_factory() -class TestFingerprint: - def test_stable_and_config_sensitive(self) -> None: - dist = _distribution(["t0", "t1", "t2"]) - cfg = _two_stage_cfg() - fp1 = compute_fingerprint(cfg, REF_ELOS, dist) - fp2 = compute_fingerprint(cfg, REF_ELOS, dist) - assert fp1 == fp2 + async def capturing_run(rows_in: List[Dict[str, Any]]): + dispatched.extend(row[TASK_INDEX_KEY_NAME] for row in rows_in) + return await run(rows_in) - other_dist = _distribution(["t0", "t1", "t9"]) - assert compute_fingerprint(cfg, REF_ELOS, other_dist) != fp1 + await run_multistage_stages( + cfg, + REF_ELOS, + _distribution(task_ids), + rows, + capturing_run, + resume=resume, + dispatch_longest_first=True, + ) - other_elos = dict(REF_ELOS, a=999.0) - assert compute_fingerprint(cfg, other_elos, dist) != fp1 + assert dispatched == [2, 0, 1, 3] - cfg2 = MultiStageRunConfig( + async def test_resume_reuses_failed_judge_deliverable_for_matching_stage_only(self) -> None: + task_ids = ["t0"] + rows = _materialized_rows(task_ids) + cfg = MultiStageRunConfig( enabled=True, - stages=parse_multistage_config({"enabled": True, "stages": ["4", "5:2"]}).stages, + stages=parse_multistage_config({"enabled": True, "stages": ["1", "1"]}).stages, seed=0, + # Isolate sidecar propagation from normal cross-stage reuse. + reuse_cached_deliverables=False, ) - assert compute_fingerprint(cfg2, REF_ELOS, dist) != fp1 + resume = RecordingResume(reuse_cached_keys={1: {(0, 0)}}) + seen: List[Tuple[int, bool]] = [] + run = _fake_run_rollouts_factory() - def test_percentage_change_invalidates(self) -> None: - # Same task-id sets but different group weights ⇒ seeded sampling draws - # different tasks, so the fingerprint must differ. - cfg = _two_stage_cfg() - dist_a = { - "g0": {"percentage": 0.5, "task_ids": ["t0", "t1"]}, - "g1": {"percentage": 0.5, "task_ids": ["t2", "t3"]}, - } - dist_b = { - "g0": {"percentage": 0.9, "task_ids": ["t0", "t1"]}, - "g1": {"percentage": 0.1, "task_ids": ["t2", "t3"]}, - } - assert compute_fingerprint(cfg, REF_ELOS, dist_a) != compute_fingerprint(cfg, REF_ELOS, dist_b) + async def capturing_run(rows_in: List[Dict[str, Any]]): + seen.extend((row["stage_index"], bool(row.get("reuse_cached_deliverable"))) for row in rows_in) + return await run(rows_in) + await run_multistage_stages(cfg, REF_ELOS, _distribution(task_ids), rows, capturing_run, resume=resume) -class TestJournalIO: - def test_journal_round_trip_latest_wins(self, tmp_path: Path) -> None: - from resources_servers.gdpval.multistage_orchestrator import append_journal_record + assert seen == [(0, False), (1, True)] - journal = journal_path_for(tmp_path / "rollouts.jsonl") - assert journal.name == "rollouts_multistage_state.jsonl" - # Plan carries references/tasks; completion is just a marker (eval_elo is - # re-fit from rows on resume, so it is not stored). - append_journal_record(journal, {"stage_index": 0, "status": "planned", "reference_ids": ["a"]}, "FP") - append_journal_record(journal, {"stage_index": 0, "status": "complete"}, "FP") - append_journal_record(journal, {"stage_index": 1, "status": "planned", "reference_ids": ["b", "c"]}, "FP") + async def test_gated_sidecar_artifact_is_produced_for_later_stages(self) -> None: + task_ids = ["t0", "t1"] + rows = _materialized_rows(task_ids) + cfg = MultiStageRunConfig( + enabled=True, + stages=parse_multistage_config( + { + "enabled": True, + "stages": [ + { + "num_tasks": 2, + "partial_completion": { + "min_success_fraction": 0.5, + "min_per_reference_success_fraction": 0.5, + }, + }, + {"num_tasks": 2}, + ], + } + ).stages, + seed=0, + ) + cached_success = { + **rows[1], + "stage_index": 0, + "reference_ids": ["a"], + "per_reference": {"a": {"wins": 1, "losses": 0, "ties": 0, "reference_elo": REF_ELOS["a"]}}, + } + # Stage 0's judge-invalid row exhausted its attempts, but its sidecar + # flag proves that the reference-independent policy artifact exists. + resume = RecordingResume( + plans={ + 0: { + "stage_index": 0, + "status": "planned", + "reference_ids": ["a"], + "task_ids": task_ids, + "task_reference_ids": {task_id: "a" for task_id in task_ids}, + } + }, + rows_by_stage={0: [cached_success]}, + gated_keys={0: {(0, 0), (1, 0)}}, + reuse_cached_keys={0: {(0, 0)}}, + attempts_by_stage={0: {(0, 0): 3}}, + ) + seen: List[Tuple[int, int, bool]] = [] + run = _fake_run_rollouts_factory() + + async def capturing_run(rows_in: List[Dict[str, Any]]): + seen.extend( + (row["stage_index"], row[TASK_INDEX_KEY_NAME], bool(row.get("reuse_cached_deliverable"))) + for row in rows_in + ) + return await run(rows_in) + + await run_multistage_stages(cfg, REF_ELOS, _distribution(task_ids), rows, capturing_run, resume=resume) + + assert sorted(seen) == [(1, 0, True), (1, 1, True)] + + async def test_new_terminal_judge_failure_reuses_artifact_in_next_stage(self) -> None: + task_ids = ["t0", "t1"] + rows = _materialized_rows(task_ids) + cfg = MultiStageRunConfig( + enabled=True, + stages=parse_multistage_config( + { + "enabled": True, + "stages": [ + { + "num_tasks": 2, + "partial_completion": { + "min_success_fraction": 0.5, + "min_per_reference_success_fraction": 0.5, + }, + }, + {"num_tasks": 2}, + ], + } + ).stages, + seed=0, + ) + seen: List[Tuple[int, int, bool]] = [] + success_run = _fake_run_rollouts_factory() + + async def terminal_then_success(rows_in: List[Dict[str, Any]]): + seen.extend( + (row["stage_index"], row[TASK_INDEX_KEY_NAME], bool(row.get("reuse_cached_deliverable"))) + for row in rows_in + ) + if rows_in and rows_in[0]["stage_index"] == 0: + pairs = await success_run(rows_in) + row, _ = pairs[0] + pairs[0] = ( + row, + { + NG_FAILURE_CLASS_KEY: "permanent", + NG_TERMINAL_KEY: True, + "reuse_cached_deliverable": True, + "deliverables_dir": "/cached/task_t0/repeat_0", + }, + ) + return pairs + return await success_run(rows_in) + + await run_multistage_stages( + cfg, + {"a": REF_ELOS["a"]}, + _distribution(task_ids), + rows, + terminal_then_success, + ) + + assert sorted(seen) == [ + (0, 0, False), + (0, 1, False), + (1, 0, True), + (1, 1, True), + ] + + async def test_third_failure_stops_adaptive_stage_without_fourth_invocation(self) -> None: + task_ids = ["t0"] + rows = _materialized_rows(task_ids) + cfg = MultiStageRunConfig( + enabled=True, + stages=parse_multistage_config({"enabled": True, "stages": ["1", "1"]}).stages, + seed=0, + reuse_cached_deliverables=False, + ) + resume = RecordingResume(attempts_by_stage={0: {(0, 0): 2}}) + seen_attempts: List[Tuple[int, Any]] = [] + run = _fake_run_rollouts_factory() + + async def fail_third_then_succeed(rows_in: List[Dict[str, Any]]): + stage_index = rows_in[0]["stage_index"] + seen_attempts.append((stage_index, rows_in[0].get(ATTEMPT_INDEX_KEY_NAME))) + if stage_index == 0: + return [(rows_in[0], {NG_FAILURE_CLASS_KEY: "judge_invalid"})] + return await run(rows_in) + + await run_multistage_stages( + cfg, REF_ELOS, _distribution(task_ids), rows, fail_third_then_succeed, resume=resume + ) + + assert seen_attempts == [(0, 2)] + assert resume.completed == [] + assert resume.appended[0][0][ATTEMPT_INDEX_KEY_NAME] == 2 + + async def test_drained_stage_is_not_marked_complete(self) -> None: + task_ids = ["t0"] + rows = _materialized_rows(task_ids) + cfg = MultiStageRunConfig( + enabled=True, + stages=parse_multistage_config({"enabled": True, "stages": ["1", "1"]}).stages, + seed=0, + ) + resume = RecordingResume() + dispatches = 0 + + async def drain(rows_in: List[Dict[str, Any]]): + nonlocal dispatches + dispatches += 1 + return [ + ( + row, + {NG_FAILURE_CLASS_KEY: "kill_shaped", NG_NO_PERSIST_KEY: True}, + ) + for row in rows_in + ] + + results, _ = await run_multistage_stages(cfg, REF_ELOS, _distribution(task_ids), rows, drain, resume=resume) + + assert results == [] + assert resume.completed == [] + assert dispatches == 1 + + async def test_partial_final_stage_declares_full_expected_row_count(self) -> None: + task_ids = ["t0", "t1"] + rows = _materialized_rows(task_ids) + cfg = MultiStageRunConfig( + enabled=True, + stages=parse_multistage_config({"enabled": True, "stages": ["2", "2"]}).stages, + seed=0, + ) + resume = RecordingResume() + base_run = _fake_run_rollouts_factory() + calls = 0 + + async def partial_final(rows_in: List[Dict[str, Any]]): + nonlocal calls + calls += 1 + pairs = await base_run(rows_in) + if calls == 2: + row, _ = pairs[-1] + pairs[-1] = (row, {NG_FAILURE_CLASS_KEY: "kill_shaped", NG_NO_PERSIST_KEY: True}) + return pairs + + results, _ = await run_multistage_stages( + cfg, REF_ELOS, _distribution(task_ids), rows, partial_final, resume=resume + ) + + final_rows = [row for row in results if row["stage_index"] == 1] + assert len(final_rows) == 1 + assert final_rows[0]["expected_final_stage_index"] == 1 + assert final_rows[0]["expected_stage_row_count"] == 2 + assert [index for index, _ in resume.completed] == [0] + + +class TestPartialStageCompletion: + @staticmethod + def _config( + *, + enabled: bool, + min_success_fraction: float = 0.6, + min_per_reference_success_fraction: float = 0.6, + min_successful_rows_per_reference: int = 1, + waivable_failure_classes: Optional[Sequence[str]] = None, + ) -> MultiStageRunConfig: + first_stage: Dict[str, Any] = {"num_tasks": 10} + if enabled: + first_stage["partial_completion"] = { + "min_success_fraction": min_success_fraction, + "min_per_reference_success_fraction": min_per_reference_success_fraction, + "min_successful_rows_per_reference": min_successful_rows_per_reference, + } + if waivable_failure_classes is not None: + first_stage["partial_completion"]["waivable_failure_classes"] = list(waivable_failure_classes) + return MultiStageRunConfig( + enabled=True, + stages=parse_multistage_config( + { + "enabled": True, + "stages": [first_stage, {"num_tasks": 10, "num_models": 1}], + } + ).stages, + seed=0, + ) + + @staticmethod + def _balanced_resume(task_ids: List[str]) -> RecordingResume: + return RecordingResume( + plans={ + 0: { + "stage_index": 0, + "status": "planned", + "reference_ids": ["a", "b"], + "task_ids": task_ids, + "task_reference_ids": { + task_id: ("a" if index % 2 == 0 else "b") for index, task_id in enumerate(task_ids) + }, + "seed": None, + "prior_eval_elo": None, + } + } + ) + + @staticmethod + def _runner( + failed_indices: set[int], + *, + failure_class: str = "timeout_exceeded", + no_persist: bool = False, + terminal: bool = False, + empty_successes: bool = False, + dispatched_stages: Optional[List[int]] = None, + ): + successful_run = _fake_run_rollouts_factory(target_elo=1100.0) + + async def run(rows_in: List[Dict[str, Any]]) -> List[Tuple[Dict[str, Any], Dict[str, Any]]]: + if dispatched_stages is not None: + dispatched_stages.extend(row["stage_index"] for row in rows_in) + pairs = await successful_run(rows_in) + if not rows_in or rows_in[0]["stage_index"] != 0: + return pairs + + rewritten: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] + for row, result in pairs: + if row[TASK_INDEX_KEY_NAME] in failed_indices: + failure = {NG_FAILURE_CLASS_KEY: failure_class} + if no_persist: + failure[NG_NO_PERSIST_KEY] = True + if terminal: + failure[NG_TERMINAL_KEY] = True + rewritten.append((row, failure)) + elif empty_successes: + rewritten.append((row, {"task_id": row["task_id"], "per_reference": {}})) + else: + rewritten.append((row, result)) + return rewritten + + return run + + @pytest.mark.parametrize("fail_one", [False, True]) + async def test_full_calibration_floor_requires_45_judgments(self, fail_one: bool) -> None: + task_ids = [f"t{i}" for i in range(45)] + config = parse_multistage_config( + { + "enabled": True, + "stages": [ + { + "num_tasks": 45, + "partial_completion": { + "min_success_fraction": 1.0, + "min_per_reference_success_fraction": 0.8, + "min_successful_rows_per_reference": 1, + "tolerate_unresolved": True, + }, + }, + {"num_tasks": 45, "num_models": 1}, + ], + "seed": 0, + } + ) + dispatched: List[int] = [] + results, _ = await run_multistage_stages( + config, + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + self._runner( + {0} if fail_one else set(), + failure_class="transport_ineligible", + terminal=True, + dispatched_stages=dispatched, + ), + resume=self._balanced_resume(task_ids), + ) + assert sum(row["stage_index"] == 0 for row in results) == 45 - int(fail_one) + assert set(dispatched) == ({0} if fail_one else {0, 1}) + + async def test_four_retryable_timeouts_keep_default_retry_stage_open(self) -> None: + task_ids = [f"t{i}" for i in range(10)] + resume = self._balanced_resume(task_ids) + dispatched_stages: List[int] = [] + + results, summaries = await run_multistage_stages( + self._config(enabled=False), + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + self._runner({0, 1, 2, 3}, dispatched_stages=dispatched_stages), + resume=resume, + ) + + assert set(dispatched_stages) == {0} + assert len([row for row in results if row["stage_index"] == 0]) == 6 + assert len(summaries) == 1 + assert resume.completed == [] + + async def test_existing_41_of_45_timeout_stage_advances_without_redispatch(self, tmp_path: Path) -> None: + task_ids = [f"t{i}" for i in range(45)] + rows = _materialized_rows(task_ids) + reference_elos = {f"ref_{index}": 600.0 + 100.0 * index for index in range(9)} + reference_ids = list(reference_elos) + strict_cfg = parse_multistage_config( + { + "enabled": True, + "stages": [{"num_tasks": 45}, {"num_tasks": 45, "num_models": 4}], + "seed": 0, + } + ) + partial_cfg = parse_multistage_config( + { + "enabled": True, + "stages": [ + { + "num_tasks": 45, + "partial_completion": { + "min_success_fraction": 0.9, + "min_per_reference_success_fraction": 0.8, + "min_successful_rows_per_reference": 1, + }, + }, + {"num_tasks": 45, "num_models": 4}, + ], + "seed": 0, + } + ) + distribution = _distribution(task_ids) + strict_fingerprint = compute_fingerprint(strict_cfg, reference_elos, distribution) + partial_fingerprint = compute_fingerprint(partial_cfg, reference_elos, distribution) + assert partial_fingerprint == strict_fingerprint + + output = tmp_path / "rollouts.jsonl" + output.write_bytes(b"") + journal = journal_path_for(output) + task_reference_ids = {task_id: reference_ids[index // 5] for index, task_id in enumerate(task_ids)} + append_journal_record( + journal, + { + "stage_index": 0, + "status": "planned", + "reference_ids": reference_ids, + "task_ids": task_ids, + "task_reference_ids": task_reference_ids, + "seed": None, + "prior_eval_elo": None, + }, + strict_fingerprint, + ) + + async def successful(rows_in: List[Dict[str, Any]]): + pairs = [] + for row in rows_in: + reference_id = row["reference_ids"][0] + reference_elo = reference_elos[reference_id] + wins = 6 if reference_elo < 1000 else 4 + pairs.append( + ( + row, + { + "task_id": row["task_id"], + "per_reference": { + reference_id: { + "wins": wins, + "losses": 10 - wins, + "ties": 0, + "reference_elo": reference_elo, + } + }, + }, + ) + ) + return pairs + + timeout_indices = {0, 5, 10, 15} + + async def first_allocation(rows_in: List[Dict[str, Any]]): + assert {row["stage_index"] for row in rows_in} == {0} + pairs = await successful(rows_in) + return [ + (row, {NG_FAILURE_CLASS_KEY: "timeout_exceeded"}) + if row[TASK_INDEX_KEY_NAME] in timeout_indices + else (row, result) + for row, result in pairs + ] + + first_results, first_summaries = await run_multistage_stages( + strict_cfg, + reference_elos, + distribution, + rows, + first_allocation, + resume=build_file_resume(output, journal, strict_fingerprint), + ) + assert len(first_results) == 41 + assert len(first_summaries) == 1 + assert read_journal(journal)[1] == {} + assert set(load_latest_failures(output)[0]) == {(index, 0) for index in timeout_indices} + + dispatched_stages: List[int] = [] + + async def resumed_allocation(rows_in: List[Dict[str, Any]]): + dispatched_stages.extend(row["stage_index"] for row in rows_in) + if any(row["stage_index"] == 0 for row in rows_in): + pytest.fail("persisted Stage-0 timeouts were unexpectedly redispatched") + pairs = await successful(rows_in) + row, _ = pairs[-1] + pairs[-1] = (row, {NG_FAILURE_CLASS_KEY: "kill_shaped", NG_NO_PERSIST_KEY: True}) + return pairs + + resumed_results, resumed_summaries = await run_multistage_stages( + partial_cfg, + reference_elos, + distribution, + rows, + resumed_allocation, + resume=build_file_resume(output, journal, partial_fingerprint), + ) + + assert set(dispatched_stages) == {1} + assert len([row for row in resumed_results if row["stage_index"] == 0]) == 41 + assert len([row for row in resumed_results if row["stage_index"] == 1]) == 44 + assert resumed_summaries[0]["partial"] is True + assert resumed_summaries[0]["success_fraction"] == pytest.approx(41 / 45) + plans, outcomes, fingerprint = read_journal(journal) + assert fingerprint == partial_fingerprint + assert 1 in plans + assert outcomes[0]["status"] == "partial_complete" + assert len(outcomes[0]["omitted_keys"]) == 4 + assert 1 not in outcomes + assert all(record["success_fraction"] >= 0.8 for record in outcomes[0]["per_reference"].values()) + + final_dispatches: List[Tuple[int, int]] = [] + + async def finish_final_stage(rows_in: List[Dict[str, Any]]): + final_dispatches.extend((row["stage_index"], row[TASK_INDEX_KEY_NAME]) for row in rows_in) + return await successful(rows_in) + + final_results, final_summaries = await run_multistage_stages( + partial_cfg, + reference_elos, + distribution, + rows, + finish_final_stage, + resume=build_file_resume(output, journal, partial_fingerprint), + ) + + assert len(final_dispatches) == 1 + assert final_dispatches[0][0] == 1 + assert len([row for row in final_results if row["stage_index"] == 0]) == 41 + assert len([row for row in final_results if row["stage_index"] == 1]) == 45 + assert len(final_summaries) == 2 + assert read_journal(journal)[1][1]["status"] == "complete" + + async def test_no_persist_retry_overrides_older_timeout_on_resume(self, tmp_path: Path) -> None: + task_ids = [f"t{i}" for i in range(10)] + rows = _materialized_rows(task_ids) + cfg = parse_multistage_config( + { + "enabled": True, + "stages": [ + { + "num_tasks": 10, + "partial_completion": { + "min_success_fraction": 0.8, + "min_per_reference_success_fraction": 0.8, + }, + }, + {"num_tasks": 10, "num_models": 1}, + ], + "seed": 0, + } + ) + reference_elos = {"a": 1000.0} + distribution = _distribution(task_ids) + output = tmp_path / "rollouts.jsonl" + journal = journal_path_for(output) + fingerprint = compute_fingerprint(cfg, reference_elos, distribution) + successful_run = _fake_run_rollouts_factory(target_elo=1100.0) + + async def first_attempt(rows_in: List[Dict[str, Any]]): + pairs = await successful_run(rows_in) + return [ + (row, {NG_FAILURE_CLASS_KEY: "timeout_exceeded"}) + if row["stage_index"] == 0 and row[TASK_INDEX_KEY_NAME] >= 6 + else (row, result) + for row, result in pairs + ] + + _, first_summaries = await run_multistage_stages( + cfg, + reference_elos, + distribution, + rows, + first_attempt, + resume=build_file_resume(output, journal, fingerprint), + ) + assert len(first_summaries) == 1 + + async def drained_retry(rows_in: List[Dict[str, Any]]): + pairs = await successful_run(rows_in) + return [ + ( + row, + {NG_FAILURE_CLASS_KEY: "kill_shaped", NG_NO_PERSIST_KEY: True}, + ) + if row[TASK_INDEX_KEY_NAME] >= 8 + else (row, result) + for row, result in pairs + ] + + _, second_summaries = await run_multistage_stages( + cfg, + reference_elos, + distribution, + rows, + drained_retry, + resume=build_file_resume(output, journal, fingerprint), + ) + assert len(second_summaries) == 1 + dispositions = load_latest_attempt_dispositions(journal)[0] + assert dispositions[(8, 0)][NG_NO_PERSIST_KEY] is True + assert dispositions[(9, 0)][NG_NO_PERSIST_KEY] is True + + dispatched: List[Tuple[int, int]] = [] + + async def final_retry(rows_in: List[Dict[str, Any]]): + dispatched.extend((row["stage_index"], row[TASK_INDEX_KEY_NAME]) for row in rows_in) + return await successful_run(rows_in) + + _, final_summaries = await run_multistage_stages( + cfg, + reference_elos, + distribution, + rows, + final_retry, + resume=build_file_resume(output, journal, fingerprint), + ) + + assert sorted(index for stage, index in dispatched if stage == 0) == [8, 9] + assert {stage for stage, _ in dispatched} == {0, 1} + assert len(final_summaries) == 2 + + async def test_newer_sidecar_failure_overrides_older_timeout_disposition(self, tmp_path: Path) -> None: + task_ids = [f"t{i}" for i in range(5)] + rows = _materialized_rows(task_ids) + strict_cfg = parse_multistage_config( + { + "enabled": True, + "stages": [{"num_tasks": 5}, {"num_tasks": 5, "num_models": 1}], + "seed": 0, + } + ) + partial_cfg = parse_multistage_config( + { + "enabled": True, + "stages": [ + { + "num_tasks": 5, + "partial_completion": { + "min_success_fraction": 0.8, + "min_per_reference_success_fraction": 0.8, + }, + }, + {"num_tasks": 5, "num_models": 1}, + ], + "seed": 0, + } + ) + reference_elos = {"a": 1000.0} + distribution = _distribution(task_ids) + output = tmp_path / "rollouts.jsonl" + journal = journal_path_for(output) + fingerprint = compute_fingerprint(strict_cfg, reference_elos, distribution) + assert compute_fingerprint(partial_cfg, reference_elos, distribution) == fingerprint + successful_run = _fake_run_rollouts_factory(target_elo=1100.0) + + async def initial_timeout(rows_in: List[Dict[str, Any]]): + pairs = await successful_run(rows_in) + return [ + (row, {NG_FAILURE_CLASS_KEY: "timeout_exceeded"}) if row[TASK_INDEX_KEY_NAME] == 4 else (row, result) + for row, result in pairs + ] + + _, first_summaries = await run_multistage_stages( + strict_cfg, + reference_elos, + distribution, + rows, + initial_timeout, + resume=build_file_resume(output, journal, fingerprint), + ) + assert len(first_summaries) == 1 + timeout = load_latest_failures(output)[0][(4, 0)] + assert load_latest_attempt_dispositions(journal)[0][(4, 0)][NG_FAILURE_CLASS_KEY] == "timeout_exceeded" + + # Simulate a crash after the newer sidecar row is fsynced but before its + # disposition is appended to the journal. + route_stage_rows(output, [{**timeout, NG_FAILURE_CLASS_KEY: "transient"}]) + assert load_latest_failures(output)[0][(4, 0)][NG_FAILURE_CLASS_KEY] == "transient" + assert load_latest_attempt_dispositions(journal)[0][(4, 0)][NG_FAILURE_CLASS_KEY] == "timeout_exceeded" + + dispatched: List[Tuple[int, int]] = [] + + async def capture(rows_in: List[Dict[str, Any]]): + dispatched.extend((row["stage_index"], row[TASK_INDEX_KEY_NAME]) for row in rows_in) + return await successful_run(rows_in) + + _, summaries = await run_multistage_stages( + partial_cfg, + reference_elos, + distribution, + rows, + capture, + resume=build_file_resume(output, journal, fingerprint), + ) + + assert [task_index for stage, task_index in dispatched if stage == 0] == [4] + assert {stage for stage, _ in dispatched} == {0, 1} + assert len(summaries) == 2 + + async def test_explicit_timeout_only_policy_advances_with_balanced_coverage(self) -> None: + task_ids = [f"t{i}" for i in range(10)] + resume = self._balanced_resume(task_ids) + dispatched_stages: List[int] = [] + + results, summaries = await run_multistage_stages( + self._config(enabled=True), + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + self._runner({0, 1, 2, 3}, dispatched_stages=dispatched_stages), + resume=resume, + ) + + assert set(dispatched_stages) == {0, 1} + assert len(summaries) == 2 + assert summaries[0]["partial"] is True + assert summaries[0]["success_fraction"] == pytest.approx(0.6) + assert summaries[0]["num_omitted"] == 4 + assert len([row for row in results if row["stage_index"] == 1]) == 10 + assert [index for index, _ in resume.completed] == [0, 1] + outcome = resume.completed[0][1] + assert outcome["status"] == "partial_complete" + assert len(outcome["included_keys"]) == 6 + assert len(outcome["omitted_keys"]) == 4 + assert outcome["per_reference"]["a"]["success_fraction"] == pytest.approx(0.6) + assert outcome["per_reference"]["b"]["success_fraction"] == pytest.approx(0.6) + + @pytest.mark.parametrize( + ("failure_class", "no_persist"), + [("transient", False), ("timeout_exceeded", True)], + ) + async def test_policy_rejects_non_timeout_or_unpersisted_omissions( + self, failure_class: str, no_persist: bool + ) -> None: + task_ids = [f"t{i}" for i in range(10)] + resume = self._balanced_resume(task_ids) + dispatched_stages: List[int] = [] + + _, summaries = await run_multistage_stages( + self._config(enabled=True), + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + self._runner( + {0, 1, 2, 3}, + failure_class=failure_class, + no_persist=no_persist, + dispatched_stages=dispatched_stages, + ), + resume=resume, + ) + + assert set(dispatched_stages) == {0} + assert len(summaries) == 1 + assert resume.completed == [] + + async def test_transient_omission_advances_when_explicitly_waivable(self) -> None: + """A judge-failure `transient` row is waivable only when the policy says so. + + This is the `dc6c776f3af506df` case: one unresolved `transient` rollout + (an empty-PDF 400 from the judge panel) held the whole run at stage 1 + under the timeout-only default, costing 176 of 220 tasks. + """ + task_ids = [f"t{i}" for i in range(10)] + resume = self._balanced_resume(task_ids) + dispatched_stages: List[int] = [] + + _, summaries = await run_multistage_stages( + self._config(enabled=True, waivable_failure_classes=["timeout_exceeded", "transient"]), + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + self._runner( + {0, 1, 2, 3}, + failure_class="transient", + dispatched_stages=dispatched_stages, + ), + resume=resume, + ) + + assert set(dispatched_stages) == {0, 1} + assert len(summaries) == 2 + assert summaries[0]["partial"] is True + assert summaries[0]["num_omitted"] == 4 + outcome = resume.completed[0][1] + assert outcome["status"] == "partial_complete" + assert outcome["policy"]["newly_waivable_failure_classes"] == ["timeout_exceeded", "transient"] + + async def test_unknown_waivable_failure_class_is_rejected(self) -> None: + with pytest.raises(ValueError, match="waivable_failure_classes"): + parse_multistage_config( + { + "enabled": True, + "stages": [ + {"num_tasks": 10, "partial_completion": {"waivable_failure_classes": ["skipped"]}}, + {"num_tasks": 10}, + ], + } + ) + + async def test_empty_waivable_failure_classes_is_rejected(self) -> None: + with pytest.raises(ValueError, match="waivable_failure_classes"): + parse_multistage_config( + { + "enabled": True, + "stages": [ + {"num_tasks": 10, "partial_completion": {"waivable_failure_classes": []}}, + {"num_tasks": 10}, + ], + } + ) + + async def test_policy_rejects_missing_elo_fit(self) -> None: + task_ids = [f"t{i}" for i in range(10)] + resume = self._balanced_resume(task_ids) + + _, summaries = await run_multistage_stages( + self._config(enabled=True), + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + self._runner({0, 1, 2, 3}, empty_successes=True), + resume=resume, + ) + + assert summaries[0]["eval_elo"] is None + assert len(summaries) == 1 + assert resume.completed == [] + + async def test_policy_rejects_one_success_row_without_battle_evidence(self) -> None: + task_ids = [f"t{i}" for i in range(10)] + resume = self._balanced_resume(task_ids) + successful_run = _fake_run_rollouts_factory(target_elo=1100.0) + + async def one_empty_battle(rows_in: List[Dict[str, Any]]): + pairs = await successful_run(rows_in) + if rows_in and rows_in[0]["stage_index"] == 0: + row, result = pairs[0] + pairs[0] = (row, {**result, "per_reference": {}}) + return pairs + + _, summaries = await run_multistage_stages( + self._config(enabled=True, min_success_fraction=0.8), + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + one_empty_battle, + resume=resume, + ) + + assert len(summaries) == 1 + assert resume.completed == [] + + async def test_policy_requires_evidence_for_every_selected_reference(self) -> None: + task_ids = [f"t{i}" for i in range(10)] + resume = RecordingResume( + plans={ + 0: { + "stage_index": 0, + "status": "planned", + "reference_ids": ["a", "b"], + "task_ids": task_ids, + "task_reference_ids": {task_id: "a" for task_id in task_ids}, + "seed": None, + "prior_eval_elo": None, + } + } + ) + + _, summaries = await run_multistage_stages( + self._config(enabled=True), + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + _fake_run_rollouts_factory(target_elo=1100.0), + resume=resume, + ) + + assert len(summaries) == 1 + assert resume.completed == [] + + async def test_policy_rejects_inadequate_per_reference_coverage(self) -> None: + task_ids = [f"t{i}" for i in range(10)] + resume = self._balanced_resume(task_ids) + + _, summaries = await run_multistage_stages( + self._config(enabled=True), + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + self._runner({0, 2, 4, 6}), + resume=resume, + ) + + assert summaries[0]["eval_elo"] is not None + assert len(summaries) == 1 + assert resume.completed == [] + + @pytest.mark.parametrize( + ("min_success_fraction", "min_per_reference_success_fraction", "min_successful_rows_per_reference"), + [ + (0.8, 0.5, 1), + (0.6, 0.5, 4), + ], + ids=["overall-fraction", "minimum-rows-per-reference"], + ) + async def test_policy_enforces_independent_coverage_floors( + self, + min_success_fraction: float, + min_per_reference_success_fraction: float, + min_successful_rows_per_reference: int, + ) -> None: + task_ids = [f"t{i}" for i in range(10)] + resume = self._balanced_resume(task_ids) + cfg = self._config( + enabled=True, + min_success_fraction=min_success_fraction, + min_per_reference_success_fraction=min_per_reference_success_fraction, + min_successful_rows_per_reference=min_successful_rows_per_reference, + ) + + _, summaries = await run_multistage_stages( + cfg, + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + self._runner({0, 1, 2, 3}), + resume=resume, + ) + + assert len(summaries) == 1 + assert resume.completed == [] + + async def test_policy_separates_existing_terminal_and_new_timeout_omissions(self) -> None: + task_ids = [f"t{i}" for i in range(10)] + resume = self._balanced_resume(task_ids) + successful_run = _fake_run_rollouts_factory(target_elo=1100.0) + + async def mixed_failures(rows_in: List[Dict[str, Any]]): + pairs = await successful_run(rows_in) + if not rows_in or rows_in[0]["stage_index"] != 0: + return pairs + rewritten = [] + for row, result in pairs: + task_index = row[TASK_INDEX_KEY_NAME] + if task_index == 0: + rewritten.append((row, {NG_FAILURE_CLASS_KEY: "skipped", NG_TERMINAL_KEY: True})) + elif task_index in {1, 2, 3}: + rewritten.append((row, {NG_FAILURE_CLASS_KEY: "timeout_exceeded"})) + else: + rewritten.append((row, result)) + return rewritten + + _, summaries = await run_multistage_stages( + self._config(enabled=True), + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + mixed_failures, + resume=resume, + ) + + assert len(summaries) == 2 + outcome = resume.completed[0][1] + assert outcome["status"] == "partial_complete" + assert outcome["accepted_unresolved_keys"] == [[1, 0], [2, 0], [3, 0]] + assert outcome["already_resolved_omitted_keys"] == [[0, 0]] + + async def test_enabled_policy_is_a_coverage_floor_for_terminal_omissions(self) -> None: + task_ids = [f"t{i}" for i in range(10)] + resume = self._balanced_resume(task_ids) + + _, summaries = await run_multistage_stages( + self._config( + enabled=True, + min_success_fraction=0.9, + min_per_reference_success_fraction=0.5, + ), + {"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(summaries) == 1 + assert resume.completed == [] + + async def test_policy_rejects_undispatched_row(self) -> None: + task_ids = [f"t{i}" for i in range(10)] + resume = self._balanced_resume(task_ids) + successful_run = _fake_run_rollouts_factory(target_elo=1100.0) + + async def omit_one_result(rows_in: List[Dict[str, Any]]): + pairs = await successful_run(rows_in) + if rows_in and rows_in[0]["stage_index"] == 0: + return pairs[:-1] + return pairs + + _, summaries = await run_multistage_stages( + self._config(enabled=True), + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + omit_one_result, + resume=resume, + ) + + 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 + + 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_invalid_directly_constructed_policy_is_rejected(self) -> None: + task_ids = ["t0"] + cfg = self._config(enabled=True) + assert cfg.stages[0].partial_completion is not None + cfg.stages[0].partial_completion.min_success_fraction = 1.1 + + with pytest.raises(ValueError, match="min_success_fraction"): + await run_multistage_stages( + cfg, + {"a": 1000.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + _fake_run_rollouts_factory(), + ) + + async def test_cached_partial_snapshot_with_invalid_evidence_is_not_reused(self) -> None: + task_ids = [f"t{i}" for i in range(10)] + cfg = self._config(enabled=True) + initial_resume = self._balanced_resume(task_ids) + results, _ = await run_multistage_stages( + cfg, + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + self._runner({0, 1, 2, 3}), + resume=initial_resume, + ) + stage0_rows = [row for row in results if row["stage_index"] == 0] + corrupted_rows = [dict(row, per_reference={}) for row in stage0_rows] + outcome = initial_resume.completed[0][1] + all_stage0_keys = {(index, 0) for index in range(10)} + resumed = RecordingResume( + plans={0: initial_resume.plans[0]}, + outcomes={0: outcome}, + rows_by_stage={0: corrupted_rows}, + gated_keys={0: all_stage0_keys}, + ) + dispatched_stages: List[int] = [] + + with pytest.raises(RuntimeError, match="cached partial stage 0 snapshot is invalid"): + await run_multistage_stages( + cfg, + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + self._runner(set(), dispatched_stages=dispatched_stages), + resume=resumed, + ) + + assert dispatched_stages == [] + assert resumed.completed == [] + + async def test_cached_partial_snapshot_rejects_changed_valid_evidence(self) -> None: + task_ids = [f"t{i}" for i in range(10)] + cfg = self._config(enabled=True) + initial_resume = self._balanced_resume(task_ids) + results, _ = await run_multistage_stages( + cfg, + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + self._runner({0, 1, 2, 3}), + resume=initial_resume, + ) + stage0_rows = [deepcopy(row) for row in results if row["stage_index"] == 0] + reference_id = next(iter(stage0_rows[0]["per_reference"])) + counts = stage0_rows[0]["per_reference"][reference_id] + counts["wins"] += 1 + resumed = RecordingResume( + plans={0: initial_resume.plans[0]}, + outcomes={0: initial_resume.completed[0][1]}, + rows_by_stage={0: stage0_rows}, + gated_keys={0: {(index, 0) for index in range(10)}}, + ) + dispatched_stages: List[int] = [] + + with pytest.raises(RuntimeError, match="cached partial stage 0 snapshot is invalid"): + await run_multistage_stages( + cfg, + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + self._runner(set(), dispatched_stages=dispatched_stages), + resume=resumed, + ) + + assert dispatched_stages == [] + + async def test_changed_policy_does_not_reuse_frozen_partial_outcome(self) -> None: + task_ids = [f"t{i}" for i in range(10)] + initial_cfg = self._config(enabled=True) + initial_resume = self._balanced_resume(task_ids) + results, _ = await run_multistage_stages( + initial_cfg, + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + self._runner({0, 1, 2, 3}), + resume=initial_resume, + ) + outcome = initial_resume.completed[0][1] + stage0_rows = [row for row in results if row["stage_index"] == 0] + resumed = RecordingResume( + plans={0: initial_resume.plans[0]}, + outcomes={0: outcome}, + rows_by_stage={0: stage0_rows}, + gated_keys={0: {(index, 0) for index in range(10)}}, + ) + + with pytest.raises(RuntimeError, match="cached partial stage 0 snapshot is invalid"): + await run_multistage_stages( + self._config(enabled=True, min_success_fraction=0.5), + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + self._runner(set()), + resume=resumed, + ) + + async def test_file_resume_freezes_partial_rows_and_ignores_late_success(self, tmp_path: Path) -> None: + task_ids = [f"t{i}" for i in range(5)] + rows = _materialized_rows(task_ids) + cfg = MultiStageRunConfig( + enabled=True, + stages=parse_multistage_config( + { + "enabled": True, + "stages": [ + { + "num_tasks": 5, + "partial_completion": { + "min_success_fraction": 0.8, + "min_per_reference_success_fraction": 0.8, + }, + }, + {"num_tasks": 5, "num_models": 1}, + ], + } + ).stages, + seed=0, + ) + reference_elos = {"a": 1000.0} + distribution = _distribution(task_ids) + output = tmp_path / "rollouts.jsonl" + journal = journal_path_for(output) + fingerprint = compute_fingerprint(cfg, reference_elos, distribution) + + first_resume = build_file_resume(output, journal, fingerprint) + _, first_summaries = await run_multistage_stages( + cfg, + reference_elos, + distribution, + rows, + self._runner({4}), + resume=first_resume, + ) + plans, outcomes, _ = read_journal(journal) + assert outcomes[0]["status"] == "partial_complete" + assert first_summaries[0]["num_omitted"] == 1 + original_stage0_elo = first_summaries[0]["eval_elo"] + original_stage1_plan = deepcopy(plans[1]) + + omitted_key = tuple(outcomes[0]["omitted_keys"][0]) + late_base = next(row for row in rows if (row[TASK_INDEX_KEY_NAME], row[ROLLOUT_INDEX_KEY_NAME]) == omitted_key) + route_stage_rows( + output, + [ + { + TASK_INDEX_KEY_NAME: omitted_key[0], + ROLLOUT_INDEX_KEY_NAME: omitted_key[1], + AGENT_REF_KEY_NAME: late_base[AGENT_REF_KEY_NAME], + "stage_index": 0, + "task_id": late_base["task_id"], + "per_reference": {"a": {"wins": 1, "losses": 0, "ties": 0, "reference_elo": 1000.0}}, + } + ], + ) + + resumed = build_file_resume(output, journal, fingerprint) + + async def no_dispatch(rows_in: List[Dict[str, Any]]): + pytest.fail(f"completed partial run unexpectedly dispatched {len(rows_in)} row(s)") + + resumed_results, resumed_summaries = await run_multistage_stages( + cfg, + reference_elos, + distribution, + rows, + no_dispatch, + resume=resumed, + ) + + resumed_stage0 = [row for row in resumed_results if row["stage_index"] == 0] + assert len(resumed_stage0) == 4 + assert omitted_key not in {(row[TASK_INDEX_KEY_NAME], row[ROLLOUT_INDEX_KEY_NAME]) for row in resumed_stage0} + assert resumed_summaries[0]["partial"] is True + assert resumed_summaries[0]["num_rollouts"] == 5 + assert resumed_summaries[0]["num_successful"] == 4 + assert resumed_summaries[0]["eval_elo"] == pytest.approx(original_stage0_elo) + assert resumed.plans[1] == original_stage1_plan + + +class TestFingerprint: + def test_stable_and_config_sensitive(self) -> None: + dist = _distribution(["t0", "t1", "t2"]) + cfg = _two_stage_cfg() + fp1 = compute_fingerprint(cfg, REF_ELOS, dist) + fp2 = compute_fingerprint(cfg, REF_ELOS, dist) + assert fp1 == fp2 + + other_dist = _distribution(["t0", "t1", "t9"]) + assert compute_fingerprint(cfg, REF_ELOS, other_dist) != fp1 + + other_elos = dict(REF_ELOS, a=999.0) + assert compute_fingerprint(cfg, other_elos, dist) != fp1 + + cfg2 = MultiStageRunConfig( + enabled=True, + stages=parse_multistage_config({"enabled": True, "stages": ["4", "5:2"]}).stages, + seed=0, + ) + assert compute_fingerprint(cfg2, REF_ELOS, dist) != fp1 + + def test_percentage_change_invalidates(self) -> None: + # Same task-id sets but different group weights ⇒ seeded sampling draws + # different tasks, so the fingerprint must differ. + cfg = _two_stage_cfg() + dist_a = { + "g0": {"percentage": 0.5, "task_ids": ["t0", "t1"]}, + "g1": {"percentage": 0.5, "task_ids": ["t2", "t3"]}, + } + dist_b = { + "g0": {"percentage": 0.9, "task_ids": ["t0", "t1"]}, + "g1": {"percentage": 0.1, "task_ids": ["t2", "t3"]}, + } + assert compute_fingerprint(cfg, REF_ELOS, dist_a) != compute_fingerprint(cfg, REF_ELOS, dist_b) + + def test_partial_completion_policy_preserves_rollout_cache_fingerprint(self) -> None: + dist = _distribution(["t0", "t1"]) + strict = parse_multistage_config( + {"enabled": True, "stages": [{"num_tasks": 1}, {"num_tasks": 2, "num_models": 1}]} + ) + partial = parse_multistage_config( + { + "enabled": True, + "stages": [ + { + "num_tasks": 1, + "partial_completion": { + "min_success_fraction": 0.9, + "min_per_reference_success_fraction": 0.5, + }, + }, + {"num_tasks": 2, "num_models": 1}, + ], + } + ) + changed_threshold = deepcopy(partial) + assert changed_threshold.stages[0].partial_completion is not None + changed_threshold.stages[0].partial_completion.min_success_fraction = 0.8 + + strict_fingerprint = compute_fingerprint(strict, REF_ELOS, dist) + partial_fingerprint = compute_fingerprint(partial, REF_ELOS, dist) + assert partial_fingerprint == strict_fingerprint + assert compute_fingerprint(changed_threshold, REF_ELOS, dist) == partial_fingerprint + + def test_materialized_rows_and_result_affecting_run_config_invalidate(self) -> None: + cfg = _two_stage_cfg() + dist = _distribution(["t0"]) + rows = _materialized_rows(["t0"]) + run_config = SimpleNamespace( + agent_name="gdpval_stirrup_agent", + input_jsonl_fpath="tasks.jsonl", + limit=None, + num_repeats=1, + num_repeats_add_seed=False, + responses_create_params={"temperature": 0.6}, + prompt_config="prompt.yaml", + skills=None, + ) + baseline = compute_fingerprint( + cfg, + REF_ELOS, + dist, + materialized_rows=rows, + rollout_collection_config=run_config, + ) + + changed_prompt_row = [dict(rows[0], responses_create_params={"input": [{"role": "user", "content": "new"}]})] + assert ( + compute_fingerprint( + cfg, + REF_ELOS, + dist, + materialized_rows=changed_prompt_row, + rollout_collection_config=run_config, + ) + != baseline + ) + + changed_repeats = SimpleNamespace(**vars(run_config)) + changed_repeats.num_repeats = 2 + assert ( + compute_fingerprint( + cfg, + REF_ELOS, + dist, + materialized_rows=rows, + rollout_collection_config=changed_repeats, + ) + != baseline + ) + + remapped = [dict(rows[0], **{TASK_INDEX_KEY_NAME: 99})] + assert ( + compute_fingerprint( + cfg, + REF_ELOS, + dist, + materialized_rows=remapped, + rollout_collection_config=run_config, + ) + != baseline + ) + + def test_fixed_policy_and_judge_config_invalidate_fingerprint(self) -> None: + cfg = _two_stage_cfg() + dist = _distribution(["t0"]) + rows = _materialized_rows(["t0"]) + runtime = { + "policy_model": {"responses_api_models": {"vllm": {"model": "/checkpoints/policy-a"}}}, + "gdpval_resources_server": { + "resources_servers": { + "gdpval": { + "num_comparison_trials": 2, + "judge_responses_create_params_overrides": {"model": "judge-a"}, + } + } + }, + "operational_logging": {"level": "INFO"}, + } + baseline = compute_fingerprint( + cfg, + REF_ELOS, + dist, + materialized_rows=rows, + resolved_global_config=runtime, + ) + + changed_policy = deepcopy(runtime) + changed_policy["policy_model"]["responses_api_models"]["vllm"]["model"] = "/checkpoints/policy-b" + changed_judge = deepcopy(runtime) + changed_judge["gdpval_resources_server"]["resources_servers"]["gdpval"]["num_comparison_trials"] = 4 + changed_strict_trials = deepcopy(runtime) + changed_strict_trials["gdpval_resources_server"]["resources_servers"]["gdpval"]["strict_comparison_trials"] = ( + True + ) + changed_logging = deepcopy(runtime) + changed_logging["operational_logging"]["level"] = "DEBUG" + + assert ( + compute_fingerprint( + cfg, + REF_ELOS, + dist, + materialized_rows=rows, + resolved_global_config=changed_policy, + ) + != baseline + ) + assert ( + compute_fingerprint( + cfg, + REF_ELOS, + dist, + materialized_rows=rows, + resolved_global_config=changed_strict_trials, + ) + != baseline + ) + assert ( + compute_fingerprint( + cfg, + REF_ELOS, + dist, + materialized_rows=rows, + resolved_global_config=changed_judge, + ) + != baseline + ) + assert ( + compute_fingerprint( + cfg, + REF_ELOS, + dist, + materialized_rows=rows, + resolved_global_config=changed_logging, + ) + == baseline + ) + + def test_input_path_relocation_does_not_invalidate_fingerprint(self) -> None: + """The CLI's generated preprocessed-input path moves between releases. + + Input identity is the materialized rows' content, which is hashed + already; the path string must not stale a journal on its own. + """ + cfg = _two_stage_cfg() + dist = _distribution(["t0"]) + rows = _materialized_rows(["t0"]) + + def run_config(path: str) -> SimpleNamespace: + return SimpleNamespace( + agent_name="gdpval_stirrup_agent", + input_jsonl_fpath=path, + limit=None, + num_repeats=1, + num_repeats_add_seed=False, + responses_create_params=None, + prompt_config=None, + skills=None, + ) + + old_layout = "/runs/out/preprocessed_datasets/tasks.jsonl" + new_layout = "/runs/out/rollouts/preprocessed_datasets/tasks.jsonl" + assert compute_fingerprint( + cfg, REF_ELOS, dist, materialized_rows=rows, rollout_collection_config=run_config(old_layout) + ) == compute_fingerprint( + cfg, REF_ELOS, dist, materialized_rows=rows, rollout_collection_config=run_config(new_layout) + ) + + def test_endpoint_liveness_knobs_do_not_invalidate_fingerprint(self) -> None: + cfg = _two_stage_cfg() + dist = _distribution(["t0"]) + + def runtime(**endpoint_knobs: Any) -> Dict[str, Any]: + return { + "policy": { + "responses_api_models": { + "vllm": {"model": "/checkpoints/policy-a", "uses_reasoning_parser": True, **endpoint_knobs} + } + } + } + + baseline = compute_fingerprint(cfg, REF_ELOS, dist, resolved_global_config=runtime()) + with_knobs = runtime( + endpoint_file="/tmp/endpoint.json", + endpoint_stale_grace_s=300.0, + endpoint_connection_retries=8, + endpoint_check_interval_s=10.0, + ) + assert compute_fingerprint(cfg, REF_ELOS, dist, resolved_global_config=with_knobs) == baseline + + changed_parser = runtime() + changed_parser["policy"]["responses_api_models"]["vllm"]["uses_reasoning_parser"] = False + assert compute_fingerprint(cfg, REF_ELOS, dist, resolved_global_config=changed_parser) != baseline + + def test_runtime_bind_addresses_do_not_invalidate_fingerprint(self) -> None: + cfg = _two_stage_cfg() + dist = _distribution(["t0"]) + runtime = _runtime_components_with_bind_addresses("node-a", 0) + rebound_runtime = _runtime_components_with_bind_addresses("node-b", 100) + + baseline = compute_fingerprint(cfg, REF_ELOS, dist, resolved_global_config=runtime) + assert compute_fingerprint(cfg, REF_ELOS, dist, resolved_global_config=rebound_runtime) == baseline + + changed_nested_host = deepcopy(rebound_runtime) + changed_nested_host["agent"]["responses_api_agents"]["stirrup"]["worker_target"]["host"] = ( + "other-sandbox.internal" + ) + assert compute_fingerprint(cfg, REF_ELOS, dist, resolved_global_config=changed_nested_host) != baseline + + changed_nested_port = deepcopy(rebound_runtime) + changed_nested_port["agent"]["responses_api_agents"]["stirrup"]["worker_target"]["port"] = 9001 + assert compute_fingerprint(cfg, REF_ELOS, dist, resolved_global_config=changed_nested_port) != baseline + + changed_model = deepcopy(rebound_runtime) + changed_model["policy"]["responses_api_models"]["vllm"]["model"] = "/checkpoints/policy-b" + assert compute_fingerprint(cfg, REF_ELOS, dist, resolved_global_config=changed_model) != baseline + + +class TestJournalIO: + def test_journal_round_trip_latest_wins(self, tmp_path: Path) -> None: + from resources_servers.gdpval.multistage_orchestrator import append_journal_record + + journal = journal_path_for(tmp_path / "rollouts.jsonl") + assert journal.name == "rollouts_multistage_state.jsonl" + # Plan carries references/tasks; completion is just a marker (eval_elo is + # re-fit from rows on resume, so it is not stored). + append_journal_record(journal, {"stage_index": 0, "status": "planned", "reference_ids": ["a"]}, "FP") + append_journal_record(journal, {"stage_index": 0, "status": "complete"}, "FP") + append_journal_record(journal, {"stage_index": 1, "status": "planned", "reference_ids": ["b", "c"]}, "FP") plans, outcomes, fingerprint = read_journal(journal) assert fingerprint == "FP" @@ -780,50 +2791,253 @@ def test_route_stage_rows_splits_by_outcome(self, tmp_path: Path) -> None: ] route_stage_rows(out, rows) - main = [json.loads(line) for line in out.read_text().splitlines()] - sidecar = [json.loads(line) for line in failures_path_for(out).read_text().splitlines()] - # Success -> main; failure -> sidecar (with stage_index); kill_shaped -> nowhere. - assert [r["task_id"] for r in main] == ["t0"] - assert [r["task_id"] for r in sidecar] == ["t1"] - assert sidecar[0]["stage_index"] == 0 + main = [json.loads(line) for line in out.read_text().splitlines()] + sidecar = [json.loads(line) for line in failures_path_for(out).read_text().splitlines()] + # Success -> main; failure -> sidecar (with stage_index); kill_shaped -> nowhere. + assert [r["task_id"] for r in main] == ["t0"] + assert [r["task_id"] for r in sidecar] == ["t1"] + assert sidecar[0]["stage_index"] == 0 + + def test_load_gated_keys_terminal_and_max_attempts(self, tmp_path: Path) -> None: + out = tmp_path / "rollouts.jsonl" + # One success in main jsonl (stage 0, task 0). + with out.open("wb") as handle: + handle.write( + json.dumps( + {"stage_index": 0, TASK_INDEX_KEY_NAME: 0, ROLLOUT_INDEX_KEY_NAME: 0, "task_id": "t0"} + ).encode() + + b"\n" + ) + # Sidecar: task 1 terminal (never retried), task 2 hit 3 attempts (gated), + # task 3 has 1 attempt (still re-dispatchable). + sidecar = failures_path_for(out) + entries = [ + { + "stage_index": 0, + TASK_INDEX_KEY_NAME: 1, + ROLLOUT_INDEX_KEY_NAME: 0, + NG_FAILURE_CLASS_KEY: "x", + NG_TERMINAL_KEY: True, + }, + {"stage_index": 0, TASK_INDEX_KEY_NAME: 2, ROLLOUT_INDEX_KEY_NAME: 0, NG_FAILURE_CLASS_KEY: "x"}, + {"stage_index": 0, TASK_INDEX_KEY_NAME: 2, ROLLOUT_INDEX_KEY_NAME: 0, NG_FAILURE_CLASS_KEY: "x"}, + {"stage_index": 0, TASK_INDEX_KEY_NAME: 2, ROLLOUT_INDEX_KEY_NAME: 0, NG_FAILURE_CLASS_KEY: "x"}, + {"stage_index": 0, TASK_INDEX_KEY_NAME: 3, ROLLOUT_INDEX_KEY_NAME: 0, NG_FAILURE_CLASS_KEY: "x"}, + ] + with sidecar.open("wb") as handle: + for e in entries: + handle.write(json.dumps(e).encode() + b"\n") + + rows_by_stage = load_persisted_rows(out) + gated = load_gated_keys(out, rows_by_stage) + # Success + terminal + max-attempts are gated; the single-attempt one is not. + assert (0, 0) in gated[0] + assert (1, 0) in gated[0] + assert (2, 0) in gated[0] + assert (3, 0) not in gated[0] + + def test_reopening_stage_prunes_all_downstream_persisted_state(self, tmp_path: Path) -> None: + out = tmp_path / "rollouts.jsonl" + journal = journal_path_for(out) + from resources_servers.gdpval.multistage_orchestrator import append_journal_record + + main_rows = [ + {"stage_index": 0, TASK_INDEX_KEY_NAME: 9, ROLLOUT_INDEX_KEY_NAME: 0, "task_id": "keep"}, + {"stage_index": 1, TASK_INDEX_KEY_NAME: 1, ROLLOUT_INDEX_KEY_NAME: 0, "task_id": "stale-1"}, + {"stage_index": 2, TASK_INDEX_KEY_NAME: 2, ROLLOUT_INDEX_KEY_NAME: 0, "task_id": "stale-2"}, + ] + out.write_text("".join(json.dumps(row) + "\n" for row in main_rows)) + failures = [ + { + "stage_index": 0, + TASK_INDEX_KEY_NAME: 0, + ROLLOUT_INDEX_KEY_NAME: 0, + NG_FAILURE_CLASS_KEY: "judge_invalid", + "elapsed_seconds": 10, + "reuse_cached_deliverable": True, + }, + { + "stage_index": 1, + TASK_INDEX_KEY_NAME: 1, + ROLLOUT_INDEX_KEY_NAME: 0, + NG_FAILURE_CLASS_KEY: "judge_invalid", + "elapsed_seconds": 20, + "reuse_cached_deliverable": True, + }, + ] + failures_path_for(out).write_text("".join(json.dumps(row) + "\n" for row in failures)) + for stage_index in range(3): + append_journal_record( + journal, + { + "stage_index": stage_index, + "status": "planned", + "reference_ids": ["a"], + "task_ids": [f"t{stage_index}"], + }, + "FP", + ) + append_journal_record(journal, {"stage_index": stage_index, "status": "complete"}, "FP") + + resume = build_file_resume(out, journal, "FP") + + assert set(resume.plans) == {0} + assert resume.outcomes == {} + assert set(resume.rows_by_stage) == {0} + assert set(resume.gated_keys) <= {0} + assert set(resume.elapsed_by_stage) == {0} + assert set(resume.reuse_cached_keys) == {0} + assert set(resume.attempts_by_stage) == {0} + assert {json.loads(line)["stage_index"] for line in out.read_text().splitlines()} == {0} + assert {json.loads(line)["stage_index"] for line in failures_path_for(out).read_text().splitlines()} == {0} + journal_records = [json.loads(line) for line in journal.read_text().splitlines()] + assert any(row.get("status") == "restart_from_stage" and row["stage_index"] == 0 for row in journal_records) + assert journal_records[-1]["status"] == "restart_cleanup_complete" + + def test_historical_failure_does_not_reopen_stage_after_later_success(self, tmp_path: Path) -> None: + out = tmp_path / "rollouts.jsonl" + journal = journal_path_for(out) + from resources_servers.gdpval.multistage_orchestrator import append_journal_record + + main_rows = [ + {"stage_index": 0, TASK_INDEX_KEY_NAME: 0, ROLLOUT_INDEX_KEY_NAME: 0, "task_id": "resolved"}, + {"stage_index": 1, TASK_INDEX_KEY_NAME: 0, ROLLOUT_INDEX_KEY_NAME: 0, "task_id": "downstream"}, + ] + out.write_text("".join(json.dumps(row) + "\n" for row in main_rows)) + failures_path_for(out).write_text( + json.dumps( + { + "stage_index": 0, + TASK_INDEX_KEY_NAME: 0, + ROLLOUT_INDEX_KEY_NAME: 0, + NG_FAILURE_CLASS_KEY: "judge_invalid", + } + ) + + "\n" + ) + for stage_index in range(2): + append_journal_record( + journal, + {"stage_index": stage_index, "status": "planned", "reference_ids": ["a"], "task_ids": ["t0"]}, + "FP", + ) + append_journal_record(journal, {"stage_index": stage_index, "status": "complete"}, "FP") + + resume = build_file_resume(out, journal, "FP") + + assert set(resume.outcomes) == {0, 1} + assert set(resume.rows_by_stage) == {0, 1} + assert not any( + json.loads(line).get("status") == "restart_from_stage" for line in journal.read_text().splitlines() + ) + + def test_legacy_terminal_timeout_reopens_stage_and_preserves_timing(self, tmp_path: Path) -> None: + out = tmp_path / "rollouts.jsonl" + out.write_bytes(b"") + journal = journal_path_for(out) + from resources_servers.gdpval.multistage_orchestrator import append_journal_record + + append_journal_record( + journal, + { + "stage_index": 1, + "status": "planned", + "reference_ids": ["a"], + "task_ids": ["t0"], + }, + "FP", + ) + append_journal_record(journal, {"stage_index": 1, "status": "complete"}, "FP") + failures_path_for(out).write_text( + json.dumps( + { + "stage_index": 1, + TASK_INDEX_KEY_NAME: 7, + ROLLOUT_INDEX_KEY_NAME: 0, + NG_FAILURE_CLASS_KEY: "timeout_exceeded", + NG_TERMINAL_KEY: True, + "elapsed_seconds": 123.0, + "reuse_cached_deliverable": True, + } + ) + + "\n" + ) + + resume = build_file_resume(out, journal, "FP") + + assert 1 not in resume.outcomes + assert (7, 0) not in resume.gated_keys.get(1, set()) + assert load_failure_timings(out)[1][(7, 0)] == 123.0 + assert load_failure_attempts(out)[1][(7, 0)] == 1 + assert load_reuse_cached_keys(out)[1] == {(7, 0)} + assert resume.reuse_cached_keys[1] == {(7, 0)} + assert resume.attempts_by_stage[1][(7, 0)] == 1 + + def test_legacy_invalid_judge_main_row_migrates_before_resume(self, tmp_path: Path) -> None: + out = tmp_path / "rollouts.jsonl" + deliverables_dir = tmp_path / "deliverables" / "task-4" + deliverables_dir.mkdir(parents=True) + (deliverables_dir / "answer.txt").write_text("answer") + out.write_text( + json.dumps( + { + "stage_index": 0, + TASK_INDEX_KEY_NAME: 4, + ROLLOUT_INDEX_KEY_NAME: 0, + "task_id": "t4", + "invalid_judge_response": True, + "reward": 0.0, + "deliverables_dir": str(deliverables_dir), + } + ) + + "\n" + ) + journal = journal_path_for(out) + from resources_servers.gdpval.multistage_orchestrator import append_journal_record + + append_journal_record( + journal, + {"stage_index": 0, "status": "planned", "reference_ids": ["a"], "task_ids": ["t4"]}, + "FP", + ) + append_journal_record(journal, {"stage_index": 0, "status": "complete"}, "FP") + + resume = build_file_resume(out, journal, "FP") - def test_load_gated_keys_terminal_and_max_attempts(self, tmp_path: Path) -> None: + assert out.read_bytes() == b"" + assert resume.rows_by_stage == {} + assert resume.outcomes == {} + assert resume.reuse_cached_keys[0] == {(4, 0)} + assert resume.attempts_by_stage[0][(4, 0)] == 1 + migrated = [json.loads(line) for line in failures_path_for(out).read_text().splitlines()] + assert len(migrated) == 1 + assert migrated[0][NG_FAILURE_CLASS_KEY] == "judge_invalid" + assert migrated[0]["reuse_cached_deliverable"] is True + + def test_skipped_failure_keeps_stage_complete(self, tmp_path: Path) -> None: out = tmp_path / "rollouts.jsonl" - # One success in main jsonl (stage 0, task 0). - with out.open("wb") as handle: - handle.write( - json.dumps( - {"stage_index": 0, TASK_INDEX_KEY_NAME: 0, ROLLOUT_INDEX_KEY_NAME: 0, "task_id": "t0"} - ).encode() - + b"\n" + out.write_bytes(b"") + journal = journal_path_for(out) + from resources_servers.gdpval.multistage_orchestrator import append_journal_record + + append_journal_record(journal, {"stage_index": 0, "status": "complete"}, "FP") + failures_path_for(out).write_text( + json.dumps( + { + "stage_index": 0, + TASK_INDEX_KEY_NAME: 7, + ROLLOUT_INDEX_KEY_NAME: 0, + NG_FAILURE_CLASS_KEY: "skipped", + NG_TERMINAL_KEY: True, + } ) - # Sidecar: task 1 terminal (never retried), task 2 hit 3 attempts (gated), - # task 3 has 1 attempt (still re-dispatchable). - sidecar = failures_path_for(out) - entries = [ - { - "stage_index": 0, - TASK_INDEX_KEY_NAME: 1, - ROLLOUT_INDEX_KEY_NAME: 0, - NG_FAILURE_CLASS_KEY: "x", - NG_TERMINAL_KEY: True, - }, - {"stage_index": 0, TASK_INDEX_KEY_NAME: 2, ROLLOUT_INDEX_KEY_NAME: 0, NG_FAILURE_CLASS_KEY: "x"}, - {"stage_index": 0, TASK_INDEX_KEY_NAME: 2, ROLLOUT_INDEX_KEY_NAME: 0, NG_FAILURE_CLASS_KEY: "x"}, - {"stage_index": 0, TASK_INDEX_KEY_NAME: 2, ROLLOUT_INDEX_KEY_NAME: 0, NG_FAILURE_CLASS_KEY: "x"}, - {"stage_index": 0, TASK_INDEX_KEY_NAME: 3, ROLLOUT_INDEX_KEY_NAME: 0, NG_FAILURE_CLASS_KEY: "x"}, - ] - with sidecar.open("wb") as handle: - for e in entries: - handle.write(json.dumps(e).encode() + b"\n") + + "\n" + ) - rows_by_stage = load_persisted_rows(out) - gated = load_gated_keys(out, rows_by_stage) - # Success + terminal + max-attempts are gated; the single-attempt one is not. - assert (0, 0) in gated[0] - assert (1, 0) in gated[0] - assert (2, 0) in gated[0] - assert (3, 0) not in gated[0] + resume = build_file_resume(out, journal, "FP") + + assert 0 in resume.outcomes + assert (7, 0) in resume.gated_keys[0] async def test_failure_row_redispatched_but_terminal_not(self) -> None: # A failed (non-terminal, below max) row is re-dispatched; a terminal one is not. @@ -886,6 +3100,77 @@ async def mixed_run(rows_in: List[Dict[str, Any]]): assert all(NG_FAILURE_CLASS_KEY not in r for r in all_results) assert all(not r.get(NG_NO_PERSIST_KEY) for r in all_results) + async def test_runtime_stale_stage_prunes_file_backed_downstream(self, tmp_path: Path) -> None: + task_ids = ["t0", "t1", "t2"] + rows = _materialized_rows(task_ids) + cfg = MultiStageRunConfig( + enabled=True, + stages=parse_multistage_config({"enabled": True, "stages": ["3", "3:2"]}).stages, + seed=0, + ) + distribution = _distribution(task_ids) + output = tmp_path / "rollouts.jsonl" + journal = journal_path_for(output) + fingerprint = compute_fingerprint(cfg, REF_ELOS, distribution) + successful_run = _fake_run_rollouts_factory() + + await run_multistage_stages( + cfg, + REF_ELOS, + distribution, + rows, + successful_run, + resume=build_file_resume(output, journal, fingerprint), + ) + persisted = [json.loads(line) for line in output.read_text().splitlines()] + missing = next(row for row in persisted if row["stage_index"] == 0) + missing_key = (missing[TASK_INDEX_KEY_NAME], missing[ROLLOUT_INDEX_KEY_NAME]) + output.write_text( + "".join( + json.dumps(row) + "\n" + for row in persisted + if (row["stage_index"], row[TASK_INDEX_KEY_NAME], row[ROLLOUT_INDEX_KEY_NAME]) != (0, *missing_key) + ) + ) + aggregate_path = aggregate_metrics_path_for(output) + aggregate_path.write_text("stale") + + resume = build_file_resume(output, journal, fingerprint) + assert set(resume.outcomes) == {0, 1} + dispatched: List[Tuple[int, Tuple[int, int]]] = [] + + async def fresh_retry(rows_in: List[Dict[str, Any]]): + dispatched.extend( + ( + row["stage_index"], + (row[TASK_INDEX_KEY_NAME], row[ROLLOUT_INDEX_KEY_NAME]), + ) + for row in rows_in + ) + pairs = await successful_run(rows_in) + return [(row, {**result, "fresh_retry": True}) for row, result in pairs] + + await run_multistage_stages( + cfg, + REF_ELOS, + distribution, + rows, + fresh_retry, + resume=resume, + ) + + assert [key for stage, key in dispatched if stage == 0] == [missing_key] + assert len([key for stage, key in dispatched if stage == 1]) == 3 + assert not aggregate_path.exists() + final_rows = [json.loads(line) for line in output.read_text().splitlines()] + assert len([row for row in final_rows if row["stage_index"] == 1]) == 3 + assert all(row["fresh_retry"] is True for row in final_rows if row["stage_index"] == 1) + assert any( + json.loads(line).get("status") == "restart_from_stage" and json.loads(line)["stage_index"] == 0 + for line in journal.read_text().splitlines() + ) + assert set(read_journal(journal)[1]) == {0, 1} + async def test_file_backed_resume_end_to_end(self, tmp_path: Path) -> None: task_ids = [f"t{i}" for i in range(10)] rows = _materialized_rows(task_ids) @@ -951,12 +3236,14 @@ def test_fresh_returns_writing_resume_with_empty_state(self, tmp_path: Path) -> def test_resume_disabled_clears_existing_and_empties_state(self, tmp_path: Path) -> None: out = tmp_path / "rollouts.jsonl" journal = journal_path_for(out) + metrics = tmp_path / "rollouts_aggregate_metrics.json" out.write_text('{"x": 1}\n') journal.write_text('{"stage_index": 0, "status": "complete"}\n') + metrics.write_text('{"stale": true}\n') fp = compute_fingerprint(_two_stage_cfg(), REF_ELOS, _distribution(["t0"])) resume = _prepare_resume(self._cfg(False), out, journal, fp) assert isinstance(resume, StageResume) - assert not out.exists() and not journal.exists() + assert not out.exists() and not journal.exists() and not metrics.exists() assert resume.outcomes == {} def test_stale_fingerprint_clears_and_starts_fresh(self, tmp_path: Path) -> None: @@ -972,6 +3259,32 @@ def test_stale_fingerprint_clears_and_starts_fresh(self, tmp_path: Path) -> None assert not out.exists() and not journal.exists() assert resume.outcomes == {} + def test_runtime_bind_address_change_preserves_resume_state(self, tmp_path: Path) -> None: + from resources_servers.gdpval.multistage_orchestrator import append_journal_record + + out = tmp_path / "rollouts.jsonl" + journal = journal_path_for(out) + out.write_text("") + dist = _distribution(["t0"]) + old_fingerprint = compute_fingerprint( + _two_stage_cfg(), + REF_ELOS, + dist, + resolved_global_config=_runtime_components_with_bind_addresses("node-a", 0), + ) + append_journal_record(journal, {"stage_index": 0, "status": "complete"}, old_fingerprint) + rebound_fingerprint = compute_fingerprint( + _two_stage_cfg(), + REF_ELOS, + dist, + resolved_global_config=_runtime_components_with_bind_addresses("node-b", 100), + ) + + resume = _prepare_resume(self._cfg(True), out, journal, rebound_fingerprint) + + assert out.exists() and journal.exists() + assert set(resume.outcomes) == {0} + async def test_fresh_run_persists_journal_then_resume_reuses_all(self, tmp_path: Path) -> None: # Regression: a fresh run through _prepare_resume must write the journal + # rows, so a second _prepare_resume resumes without re-dispatching anything. @@ -1000,3 +3313,679 @@ async def no_dispatch(rows_in: List[Dict[str, Any]]): assert all(s["cached"] for s in again) assert [s["reference_ids"] for s in again] == [s["reference_ids"] for s in base] assert [s["eval_elo"] for s in again] == [s["eval_elo"] for s in base] + + +class TestIntegrationWiring: + async def test_multistage_forwards_one_budget_and_shared_tracker( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + calls: List[Dict[str, Any]] = [] + cache_namespaces: List[str] = [] + + class FakeHelper: + def _preprocess_rows_from_config(self, config): + return _materialized_rows(["t0"]) + + def _run_examples_with_metadata(self, rows, **kwargs): + calls.append(kwargs) + cache_namespaces.extend(row["verify_cache_namespace"] for row in rows) + + async def done(row): + reference_id = row["reference_ids"][0] + result = { + "task_id": row["task_id"], + "per_reference": { + reference_id: { + "wins": 1, + "losses": 0, + "ties": 0, + "reference_elo": REF_ELOS[reference_id], + } + }, + } + return rollout_collection_module._CompletedRollout(row=row, result=result, rollout_latency_ms=None) + + return [done(row) for row in rows] + + async def _call_aggregate_metrics(self, results, rows, output_fpath): + return tmp_path / "aggregate.json" + + monkeypatch.setattr(rollout_collection_module, "RolloutCollectionHelper", FakeHelper) + monkeypatch.setattr( + "resources_servers.gdpval.multistage_orchestrator.ensure_distribution", + lambda *args, **kwargs: (_distribution(["t0"]), None), + ) + config = SimpleNamespace( + input_jsonl_fpath=str(tmp_path / "input.jsonl"), + output_jsonl_fpath=str(tmp_path / "rollouts.jsonl"), + num_samples_in_parallel=2, + dispatch_budget_s=120.0, + drain_margin_s=15.0, + dispatch_longest_first=True, + resume_from_cache=False, + route_failures_to_sidecar=True, + check_completion=lambda **kwargs: None, + ) + global_config = { + "multistage": {"enabled": True, "stages": ["1", "1"], "seed": 0}, + "gdpval": { + "resources_servers": { + "gdpval": {"reference_models": {key: {"elo": elo} for key, elo in REF_ELOS.items()}} + } + }, + } + + await run_e2e_multistage(config, global_config) + + assert len(calls) == 2 + assert calls[0]["latency_tracker"] is calls[1]["latency_tracker"] + assert 0 <= calls[1]["dispatch_budget_s"] <= calls[0]["dispatch_budget_s"] <= 120.0 + assert [call["drain_margin_s"] for call in calls] == [15.0, 15.0] + # The multistage driver bypasses run_from_config, so the failure-routing + # knob must reach the dispatcher explicitly or it silently reverts to raise. + assert [call["route_failures_to_sidecar"] for call in calls] == [True, True] + persisted = [json.loads(line) for line in (tmp_path / "rollouts.jsonl").read_text().splitlines()] + assert {row["expected_final_stage_index"] for row in persisted} == {1} + assert {row["expected_stage_row_count"] for row in persisted} == {1} + assert len(set(cache_namespaces)) == 1 + assert len(cache_namespaces[0]) == 64 + + +class TestFailureContractOptIn: + """Shared collection keeps the failure contract off; GDPVal opts in unless a key is explicitly set.""" + + async def test_single_pass_delegate_opts_in_without_mutating_the_callers_config( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + seen: List[Any] = [] + + async def fake_run_from_config(self, config): + seen.append(config) + + monkeypatch.setattr(rollout_collection_module.RolloutCollectionHelper, "run_from_config", fake_run_from_config) + config = rollout_collection_module.RolloutCollectionConfig( + input_jsonl_fpath=str(tmp_path / "input.jsonl"), output_jsonl_fpath=str(tmp_path / "rollouts.jsonl") + ) + + assert await run_rollout_collection(config, {}) is None + + assert [(c.retry_terminal_timeouts, c.retry_invalid_judge_responses) for c in seen] == [(True, True)] + assert (config.retry_terminal_timeouts, config.retry_invalid_judge_responses) == (False, False) + + async def test_single_pass_delegate_keeps_an_explicit_false( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + seen: List[Any] = [] + + async def fake_run_from_config(self, config): + seen.append(config) + + monkeypatch.setattr(rollout_collection_module.RolloutCollectionHelper, "run_from_config", fake_run_from_config) + config = rollout_collection_module.RolloutCollectionConfig.model_validate( + { + "input_jsonl_fpath": str(tmp_path / "input.jsonl"), + "output_jsonl_fpath": str(tmp_path / "rollouts.jsonl"), + "retry_terminal_timeouts": False, + } + ) + + assert await run_rollout_collection(config, {}) is None + + assert [(c.retry_terminal_timeouts, c.retry_invalid_judge_responses) for c in seen] == [(False, True)] + + @pytest.mark.parametrize("opted_out", sorted(GDPVAL_FAILURE_CONTRACT)) + async def test_multistage_refuses_an_explicit_false(self, tmp_path: Path, opted_out: str) -> None: + config = rollout_collection_module.RolloutCollectionConfig.model_validate( + { + "input_jsonl_fpath": str(tmp_path / "input.jsonl"), + "output_jsonl_fpath": str(tmp_path / "rollouts.jsonl"), + opted_out: False, + } + ) + global_config = {"multistage": {"enabled": True, "stages": ["1"]}} + + with pytest.raises(ValueError, match=opted_out): + await run_rollout_collection(config, global_config) + + assert not (tmp_path / "rollouts.jsonl").exists() + + def test_benchmark_config_opts_in_where_the_driver_does_not_run(self) -> None: + from omegaconf import OmegaConf + + benchmark = OmegaConf.load(Path(__file__).resolve().parents[3] / "benchmarks/gdpval/config.yaml") + + assert {key: benchmark[key] for key in GDPVAL_FAILURE_CONTRACT} == GDPVAL_FAILURE_CONTRACT + + +class TestReferenceMissingCoverage: + """How a ``reference_missing`` row is counted by a partial-completion policy. + + Under GDPVal's failure contract (``retry_terminal_timeouts``) the row is not + terminal: it is retried across resumes, since the reference tree can be + repaired between runs, and it stays unresolved until it reaches + ``NEMO_GYM_MAX_ROLLOUT_ATTEMPTS``. The outcome-level tests below pass + ``unresolved_keys=set()``, i.e. the state at the max attempt, where the row + is an ordinary omission governed by the coverage floors. Before that, only + ``tolerate_unresolved`` accepts it (see the loop-level tests), or + ``transport_assignment_repair.reference_availability_only`` avoids it by + moving tasks off references without a completed deliverable. + """ + + @staticmethod + def _rows(n: int, ref: str): + return [{TASK_INDEX_KEY_NAME: i, ROLLOUT_INDEX_KEY_NAME: 0, "reference_ids": [ref]} for i in range(n)] + + def _outcome(self, *, judged: int, planned: int, ref: str = "r1"): + from resources_servers.gdpval.multistage_elo import PartialStagePolicy + from resources_servers.gdpval.multistage_orchestrator import _partial_stage_outcome + + stage_rows = self._rows(planned, ref) + successful = [] + for i in range(judged): + row = dict(stage_rows[i]) + row["per_reference"] = {ref: {"wins": 4, "losses": 0, "ties": 0}} + successful.append(row) + policy = PartialStagePolicy( + min_success_fraction=0.9, + min_per_reference_success_fraction=0.5, + min_successful_rows_per_reference=1, + waivable_failure_classes=("timeout_exceeded", "transient"), + ) + return _partial_stage_outcome(policy, stage_rows, successful, [], set(), [ref], 1200.0, 1) + + def test_terminal_omission_is_accepted_without_being_waivable(self) -> None: + # 19 of 20 judged: the missing one is the reference_missing row. It is + # NOT in successful_rows and NOT in unresolved_keys (terminal). + outcome = self._outcome(judged=19, planned=20) + assert outcome is not None + assert outcome["success_fraction"] == pytest.approx(0.95) + assert len(outcome["omitted_keys"]) == 1 + + def test_still_rejected_when_coverage_floor_is_breached(self) -> None: + # Forgiveness is bounded: 17/20 = 0.85 is below min_success_fraction 0.9. + assert self._outcome(judged=17, planned=20) is None + + def test_evidence_less_success_row_still_hard_rejects(self) -> None: + """The gate this change routes around is intentionally left intact: a + row that IS in the success set while carrying no battle stays fatal.""" + from resources_servers.gdpval.multistage_elo import PartialStagePolicy + from resources_servers.gdpval.multistage_orchestrator import _partial_stage_outcome + + ref = "r1" + stage_rows = self._rows(20, ref) + successful = [] + for i, row in enumerate(stage_rows): + r = dict(row) + r["per_reference"] = {} if i == 0 else {ref: {"wins": 4, "losses": 0, "ties": 0}} + successful.append(r) + policy = PartialStagePolicy( + min_success_fraction=0.9, + min_per_reference_success_fraction=0.5, + min_successful_rows_per_reference=1, + ) + assert _partial_stage_outcome(policy, stage_rows, successful, [], set(), [ref], 1200.0, 1) is None + + @pytest.mark.parametrize("tolerate_unresolved", [False, True]) + async def test_first_launch_advances_only_when_tolerating(self, tolerate_unresolved: bool) -> None: + task_ids = [f"t{i}" for i in range(20)] + partial_completion = { + "min_success_fraction": 0.9, + "min_per_reference_success_fraction": 0.5, + "min_successful_rows_per_reference": 1, + "tolerate_unresolved": tolerate_unresolved, + } + config = parse_multistage_config( + { + "enabled": True, + "stages": [{"num_tasks": 20, "partial_completion": partial_completion}, {"num_tasks": 20}], + "seed": 0, + } + ) + dispatched: List[int] = [] + + results, summaries = await run_multistage_stages( + config, + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + TestPartialStageCompletion._runner( + {0}, failure_class="reference_missing", terminal=True, dispatched_stages=dispatched + ), + ) + + if tolerate_unresolved: + assert set(dispatched) == {0, 1} + assert summaries[0]["partial"] is True + else: + # Retryable across resumes, so the stage stays open instead of failing. + assert set(dispatched) == {0} + assert summaries[-1]["incomplete"] is True + assert summaries[-1]["coverage_blocked"] is False + assert not any(row["stage_index"] == 1 for row in results) + + +class TestReferenceMissingIsTerminal: + """A missing reference deliverable can never be fixed by retrying. + + The verify response stamps the generic ``_ng_failure_terminal`` flag rather + than teaching ``is_terminal_failure`` a GDPVal-specific class name -- the + harness already honours that flag, so no change to nemo_gym is needed. + GDPVal runs the opt-in failure contract (``retry_terminal_timeouts``). + """ + + def test_generic_terminal_flag_is_honoured(self) -> None: + from nemo_gym.rollout_collection import NG_TERMINAL_KEY, is_terminal_failure + + assert is_terminal_failure({NG_TERMINAL_KEY: True}, retry_terminal_timeouts=True) is True + + def test_environment_faults_are_revalidated_on_resume(self) -> None: + # reference_missing/eval_missing are terminal within a run (retrying + # cannot make the file appear) but the tree can be repaired between + # runs, so resume re-dispatches them for a cheap /verify recheck. + from nemo_gym.rollout_collection import ( + NG_FAILURE_CLASS_KEY, + NG_TERMINAL_KEY, + is_terminal_failure, + ) + from resources_servers.gdpval.app import ( + EVAL_MISSING_FAILURE_CLASS, + REFERENCE_MISSING_FAILURE_CLASS, + ) + + for failure_class in (REFERENCE_MISSING_FAILURE_CLASS, EVAL_MISSING_FAILURE_CLASS): + row = {NG_FAILURE_CLASS_KEY: failure_class, NG_TERMINAL_KEY: True} + assert is_terminal_failure(row, retry_terminal_timeouts=True) is False + + def test_timeout_stays_retryable(self) -> None: + from nemo_gym.rollout_collection import NG_FAILURE_CLASS_KEY, is_terminal_failure + + assert is_terminal_failure({NG_FAILURE_CLASS_KEY: "timeout_exceeded"}, retry_terminal_timeouts=True) is False + + def test_reference_missing_row_is_not_a_success(self) -> None: + """The point of the change: it must not reach the success set, where the + non-final-stage coverage gate would reject it before any threshold.""" + from nemo_gym.rollout_collection import NG_FAILURE_CLASS_KEY + from resources_servers.gdpval.app import REFERENCE_MISSING_FAILURE_CLASS + from resources_servers.gdpval.multistage_orchestrator import _is_success_row + + row = { + "task_id": "t1", + "reward": 0.0, + "judge_response": {"error": "reference_missing"}, + NG_FAILURE_CLASS_KEY: REFERENCE_MISSING_FAILURE_CLASS, + } + assert _is_success_row(row) is False + + +# --------------------------------------------------------------------------- +# In-process retry gate + tolerant partial gate +# --------------------------------------------------------------------------- + + +class TestInProcessRetryGate: + """``_is_in_process_retryable`` must apply the gate its caller claims.""" + + def test_accepts_persisted_retryable_classes(self) -> None: + assert _is_in_process_retryable({NG_FAILURE_CLASS_KEY: "timeout_exceeded"}) is True + assert _is_in_process_retryable({NG_FAILURE_CLASS_KEY: "transient"}) is True + + def test_rejects_success_and_drained_rows(self) -> None: + assert _is_in_process_retryable({}) is False + assert _is_in_process_retryable({"per_reference": {}, "reward": 1.0}) is False + assert _is_in_process_retryable({NG_FAILURE_CLASS_KEY: "transient", NG_NO_PERSIST_KEY: True}) is False + + def test_rejects_terminal_rows(self) -> None: + # A resume gates a terminal row; an in-process retry must gate it too. + assert _is_in_process_retryable({NG_FAILURE_CLASS_KEY: "transient", NG_TERMINAL_KEY: True}) is False + # A timeout is never terminal, even when an old build stamped the flag. + assert _is_in_process_retryable({NG_FAILURE_CLASS_KEY: "timeout_exceeded", NG_TERMINAL_KEY: True}) is True + + @pytest.mark.parametrize( + "failure_class", + [ + "reference_missing", + "eval_missing", + "transport_ineligible", + "skipped", + "incomplete", + "judge_invalid", + "legitimate", + ], + ) + def test_rejects_classes_nothing_repairs_mid_process(self, failure_class: str) -> None: + assert _is_in_process_retryable({NG_FAILURE_CLASS_KEY: failure_class}) is False + + +class TestTolerateUnresolved: + """``tolerate_unresolved`` widens the class test and nothing else.""" + + REF = "r1" + + def _stage_rows(self) -> List[Dict[str, Any]]: + return [{TASK_INDEX_KEY_NAME: i, ROLLOUT_INDEX_KEY_NAME: 0, "reference_ids": [self.REF]} for i in range(4)] + + def _successful(self, stage_rows: Sequence[Dict[str, Any]]) -> List[Dict[str, Any]]: + rows: List[Dict[str, Any]] = [] + for row in stage_rows[:3]: + success = dict(row) + success["per_reference"] = {self.REF: {"wins": 4, "losses": 3, "ties": 1}} + rows.append(success) + return rows + + def _outcome(self, result: Optional[Dict[str, Any]], *, tolerate: bool) -> Optional[Dict[str, Any]]: + stage_rows = self._stage_rows() + policy = PartialStagePolicy( + min_success_fraction=0.7, + min_per_reference_success_fraction=0.7, + min_successful_rows_per_reference=1, + waivable_failure_classes=("timeout_exceeded",), + tolerate_unresolved=tolerate, + ) + new_results: List[Dict[str, Any]] = [] + if result is not None: + new_results.append(dict(result, **{TASK_INDEX_KEY_NAME: 3, ROLLOUT_INDEX_KEY_NAME: 0})) + return _partial_stage_outcome( + policy, + stage_rows, + self._successful(stage_rows), + new_results, + {(3, 0)}, + [self.REF], + 1200.0, + 1, + ) + + def test_waivable_class_is_accepted_either_way(self) -> None: + result = {NG_FAILURE_CLASS_KEY: "timeout_exceeded"} + assert self._outcome(result, tolerate=False) is not None + assert self._outcome(result, tolerate=True) is not None + + def test_non_waivable_class_is_accepted_only_when_tolerating(self) -> None: + result = {NG_FAILURE_CLASS_KEY: "agent_run_error"} + assert self._outcome(result, tolerate=False) is None + accepted = self._outcome(result, tolerate=True) + assert accepted is not None + assert accepted["accepted_unresolved_keys"] == [[3, 0]] + + def test_missing_failure_class_is_never_waived(self) -> None: + # The docstring promises "any persisted failure class"; a row carrying + # no class at all is not one. + assert self._outcome({"reward": 0.0}, tolerate=False) is None + assert self._outcome({"reward": 0.0}, tolerate=True) is None + + def test_drained_and_missing_rows_still_hold_the_stage_open(self) -> None: + drained = {NG_FAILURE_CLASS_KEY: "timeout_exceeded", NG_NO_PERSIST_KEY: True} + assert self._outcome(drained, tolerate=True) is None + assert self._outcome(None, tolerate=True) is None + + +class TestCachedPartialSnapshotPolicyRecord: + """A frozen snapshot survives a record without, or a flip of, ``tolerate_unresolved``.""" + + REF = "r1" + REF_ELOS = {"r1": 1200.0} + + def _fixture(self) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], PartialStagePolicy, Dict[str, Any]]: + stage_rows = [ + {TASK_INDEX_KEY_NAME: i, ROLLOUT_INDEX_KEY_NAME: 0, "reference_ids": [self.REF]} for i in range(20) + ] + persisted: List[Dict[str, Any]] = [] + for row in stage_rows[:19]: + success = dict(row) + success["per_reference"] = {self.REF: {"wins": 4, "losses": 3, "ties": 1}} + persisted.append(success) + policy = PartialStagePolicy( + min_success_fraction=0.9, + min_per_reference_success_fraction=0.5, + min_successful_rows_per_reference=1, + ) + outcome = _partial_stage_outcome(policy, stage_rows, persisted, [], set(), [self.REF], 1200.0, 1) + assert outcome is not None + return stage_rows, persisted, policy, outcome + + def _is_valid(self, outcome: Mapping[str, Any], policy: PartialStagePolicy) -> bool: + stage_rows, persisted, _, _ = self._fixture() + return _cached_partial_snapshot_is_valid(outcome, policy, stage_rows, persisted, [self.REF], self.REF_ELOS) + + def test_round_trips_its_own_record(self) -> None: + _, _, policy, outcome = self._fixture() + assert self._is_valid(outcome, policy) is True + + def test_policy_record_without_tolerate_unresolved_is_accepted(self) -> None: + _, _, policy, outcome = self._fixture() + record = dict(outcome) + record["policy"] = {k: v for k, v in outcome["policy"].items() if k != "tolerate_unresolved"} + assert "tolerate_unresolved" not in record["policy"] + assert self._is_valid(record, policy) is True + + def test_tolerate_unresolved_flip_does_not_invalidate(self) -> None: + # Flipping the flag mid-run is its intended use (breaking a hang); it + # changes no evidence, so an already-frozen snapshot stays valid. + _, _, policy, outcome = self._fixture() + flipped = PartialStagePolicy( + min_success_fraction=policy.min_success_fraction, + min_per_reference_success_fraction=policy.min_per_reference_success_fraction, + min_successful_rows_per_reference=policy.min_successful_rows_per_reference, + waivable_failure_classes=policy.waivable_failure_classes, + tolerate_unresolved=True, + ) + assert self._is_valid(outcome, flipped) is True + recorded_on = dict(outcome, policy=dict(outcome["policy"], tolerate_unresolved=True)) + assert self._is_valid(recorded_on, policy) is True + + def test_material_policy_change_still_invalidates(self) -> None: + _, _, _, outcome = self._fixture() + stricter = PartialStagePolicy( + min_success_fraction=0.95, + min_per_reference_success_fraction=0.5, + min_successful_rows_per_reference=1, + ) + assert self._is_valid(outcome, stricter) is False + + +class TestRetryInProcess: + """The opt-in retry loop must terminate and must not void persisted evidence.""" + + @staticmethod + def _tolerant_cfg(*, retry_inprocess: bool) -> MultiStageRunConfig: + return parse_multistage_config( + { + "enabled": True, + "stages": [ + { + "num_tasks": 10, + "partial_completion": { + "min_success_fraction": 0.7, + "min_per_reference_success_fraction": 0.7, + "min_successful_rows_per_reference": 1, + # Deliberately excludes the class the rows come back + # with, so only `tolerate_unresolved` can accept them. + "waivable_failure_classes": ["transient"], + "tolerate_unresolved": True, + }, + }, + {"num_tasks": 10, "num_models": 1}, + ], + "seed": 0, + "retry_inprocess": retry_inprocess, + } + ) + + @staticmethod + def _drained_retry_runner(failed_indices: set[int], dispatched_stages: List[int]): + """Stage-0 rows fail persisted once, then every retry comes back drained.""" + successful_run = _fake_run_rollouts_factory(target_elo=1100.0) + + async def run(rows_in: List[Dict[str, Any]]) -> List[Tuple[Dict[str, Any], Dict[str, Any]]]: + dispatched_stages.extend(row["stage_index"] for row in rows_in) + pairs = await successful_run(rows_in) + if not rows_in or rows_in[0]["stage_index"] != 0: + return pairs + rewritten: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] + for row, result in pairs: + if row[TASK_INDEX_KEY_NAME] not in failed_indices: + rewritten.append((row, result)) + elif row.get(ATTEMPT_INDEX_KEY_NAME): + rewritten.append( + (row, {NG_FAILURE_CLASS_KEY: "timeout_exceeded", NG_NO_PERSIST_KEY: True}), + ) + else: + rewritten.append((row, {NG_FAILURE_CLASS_KEY: "timeout_exceeded"})) + return rewritten + + return run + + async def test_tolerated_stage_advances_with_retry_off(self) -> None: + task_ids = [f"t{i}" for i in range(10)] + resume = TestPartialStageCompletion._balanced_resume(task_ids) + dispatched_stages: List[int] = [] + + _, summaries = await run_multistage_stages( + self._tolerant_cfg(retry_inprocess=False), + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + self._drained_retry_runner({0, 1}, dispatched_stages), + resume=resume, + ) + + assert set(dispatched_stages) == {0, 1} + assert len(summaries) == 2 + + async def test_drained_retry_does_not_stop_a_tolerated_stage(self, monkeypatch) -> None: + # Enabling the retry flag must not turn a completable stage into a + # stopped one once the dispatch budget is spent. + monkeypatch.setenv("NEMO_GYM_MAX_ROLLOUT_ATTEMPTS", "3") + task_ids = [f"t{i}" for i in range(10)] + resume = TestPartialStageCompletion._balanced_resume(task_ids) + dispatched_stages: List[int] = [] + + _, summaries = await run_multistage_stages( + self._tolerant_cfg(retry_inprocess=True), + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + self._drained_retry_runner({0, 1}, dispatched_stages), + resume=resume, + ) + + assert set(dispatched_stages) == {0, 1} + assert len(summaries) == 2 + + async def test_drained_retry_journals_no_tombstone(self, monkeypatch) -> None: + # A durable no-persist disposition over a key whose sidecar already holds + # a real failure would disable the pre-dispatch shortcut forever. + monkeypatch.setenv("NEMO_GYM_MAX_ROLLOUT_ATTEMPTS", "3") + task_ids = [f"t{i}" for i in range(10)] + resume = TestPartialStageCompletion._balanced_resume(task_ids) + + await run_multistage_stages( + self._tolerant_cfg(retry_inprocess=True), + {"a": 1000.0, "b": 1200.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + self._drained_retry_runner({0, 1}, []), + resume=resume, + ) + + assert resume.appended[0] + assert all(not row.get(NG_NO_PERSIST_KEY) for row in resume.appended[0]) + + @staticmethod + def _single_stage_cfg() -> MultiStageRunConfig: + return parse_multistage_config( + {"enabled": True, "stages": [{"num_tasks": 1}], "seed": 0, "retry_inprocess": True} + ) + + async def test_always_failing_row_makes_exactly_max_attempts_dispatches(self, monkeypatch) -> None: + monkeypatch.setenv("NEMO_GYM_MAX_ROLLOUT_ATTEMPTS", "3") + task_ids = ["t0"] + resume = RecordingResume() + dispatched: List[Dict[str, Any]] = [] + + async def run(rows_in: List[Dict[str, Any]]) -> List[Tuple[Dict[str, Any], Dict[str, Any]]]: + dispatched.extend(dict(row) for row in rows_in) + return [(row, {NG_FAILURE_CLASS_KEY: "timeout_exceeded"}) for row in rows_in] + + await run_multistage_stages( + self._single_stage_cfg(), + {"a": 1000.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + run, + resume=resume, + ) + + # One initial dispatch plus max_attempts - 1 retries, then the loop stops. + assert len(dispatched) == 3 + assert [row.get(ATTEMPT_INDEX_KEY_NAME) for row in dispatched] == [None, 1, 2] + + async def test_reuse_cached_deliverable_is_sticky_across_retries(self, monkeypatch) -> None: + # Attempt 0 reports a reusable deliverable; attempt 1 comes back stamped + # on the rollout path without the flag. Attempt 2 must still reuse it + # rather than re-running the whole policy rollout. + monkeypatch.setenv("NEMO_GYM_MAX_ROLLOUT_ATTEMPTS", "3") + task_ids = ["t0"] + resume = RecordingResume() + dispatched: List[Dict[str, Any]] = [] + + async def run(rows_in: List[Dict[str, Any]]) -> List[Tuple[Dict[str, Any], Dict[str, Any]]]: + pairs: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] + for row in rows_in: + dispatched.append(dict(row)) + if row.get(ATTEMPT_INDEX_KEY_NAME): + pairs.append((row, {NG_FAILURE_CLASS_KEY: "timeout_exceeded"})) + else: + pairs.append( + (row, {NG_FAILURE_CLASS_KEY: "transient", "reuse_cached_deliverable": True}), + ) + return pairs + + await run_multistage_stages( + self._single_stage_cfg(), + {"a": 1000.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + run, + resume=resume, + ) + + assert len(dispatched) == 3 + assert dispatched[0].get("reuse_cached_deliverable") is None + assert dispatched[1].get("reuse_cached_deliverable") is True + assert dispatched[2].get("reuse_cached_deliverable") is True + + async def test_all_drained_retry_still_emits_explicit_counts(self, monkeypatch) -> None: + # Suppressing the event when nothing was dispatched leaves an operator + # with no evidence that the loop ran and consumed selections. + monkeypatch.setenv("NEMO_GYM_MAX_ROLLOUT_ATTEMPTS", "3") + task_ids = ["t0"] + resume = RecordingResume() + events: List[Tuple[str, dict]] = [] + + async def run(rows_in: List[Dict[str, Any]]) -> List[Tuple[Dict[str, Any], Dict[str, Any]]]: + pairs: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] + for row in rows_in: + if row.get(ATTEMPT_INDEX_KEY_NAME): + pairs.append((row, {NG_FAILURE_CLASS_KEY: "timeout_exceeded", NG_NO_PERSIST_KEY: True})) + else: + pairs.append((row, {NG_FAILURE_CLASS_KEY: "timeout_exceeded"})) + return pairs + + await run_multistage_stages( + self._single_stage_cfg(), + {"a": 1000.0}, + _distribution(task_ids), + _materialized_rows(task_ids), + run, + resume=resume, + on_event=lambda name, data: events.append((name, data)), + ) + + retries = [data for name, data in events if name == "stage_retry"] + assert len(retries) == 2 + for data in retries: + assert data["num_selected"] == 1 + assert data["num_dispatched"] == 0 + assert data["num_drained"] == 1 + assert data["num_recovered"] == 0 diff --git a/resources_servers/gdpval/tests/test_transport_assignment.py b/resources_servers/gdpval/tests/test_transport_assignment.py new file mode 100644 index 00000000000..11c8ad1e2ab --- /dev/null +++ b/resources_servers/gdpval/tests/test_transport_assignment.py @@ -0,0 +1,240 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Transport-aware reference assignment for multistage stage plans.""" + +import pytest + +from resources_servers.gdpval.multistage_orchestrator import ( + MultiStageRunConfig, + compute_fingerprint, + parse_multistage_config, +) +from resources_servers.gdpval.transport_assignment import PairCost, _solve_capacity_assignment, make_assignment_repair + + +def _global_config(candidate_root, references, *, agent_extra=None) -> dict: + """A resolved config shaped like benchmarks/gdpval/config.yaml.""" + agent = {"resources_server": {"type": "resources_servers", "name": "gdpval_resources_server"}} + if candidate_root is not None: + agent["persist_deliverables_dir"] = str(candidate_root) + agent.update(agent_extra or {}) + return { + "gdpval_resources_server": { + "resources_servers": { + "gdpval": { + "reference_models": {name: {"deliverables_dir": str(root)} for name, root in references.items()}, + } + } + }, + "gdpval_stirrup_agent": {"responses_api_agents": {"stirrup_agent": agent}}, + } + + +def _write_task(root, task_id: str, *, payload: bytes = b"x", finished: bool = True) -> None: + repeat = root / f"task_{task_id}" / "repeat_0" + repeat.mkdir(parents=True) + if finished: + (repeat / "finish_params.json").write_text("{}") + (repeat / "evidence.mp4").write_bytes(payload) + + +class TestSolveCapacityAssignment: + def test_minimum_change_reassignment_respects_capacities(self) -> None: + tasks = ["t0", "t1", "t2"] + references = ["ra", "rb"] + original = {"t0": "ra", "t1": "ra", "t2": "rb"} + + def cost(compatible: bool, wire: int = 1) -> PairCost: + return PairCost(compatible=compatible, wire_bytes=wire, raw_bytes=wire, max_file_bytes=wire, reasons=()) + + # t1/ra is incompatible; the only capacity-preserving repair swaps t1 and t2. + costs = { + ("t0", "ra"): cost(True), + ("t0", "rb"): cost(True), + ("t1", "ra"): cost(False), + ("t1", "rb"): cost(True), + ("t2", "ra"): cost(True), + ("t2", "rb"): cost(True), + } + + result = _solve_capacity_assignment(tasks, references, original, costs) + + assert result == {"t0": "ra", "t1": "rb", "t2": "ra"} + assert sorted(result.values()) == sorted(original.values()) + + @pytest.mark.parametrize("candidate_present", [True, False]) + @pytest.mark.parametrize( + ("marker", "available"), + [(None, False), ("", False), ("{", False), ("[]", False), ("null", True), ("{}", True)], + ids=["missing-task", "missing-marker", "corrupt-marker", "wrong-shape", "null-finish", "explicit-finish"], + ) + def test_availability_only_preserves_counts_and_completion_contract( + self, tmp_path, marker, available, candidate_present + ) -> None: + roots = {name: tmp_path / name for name in ("candidate", "ra", "rb")} + for name, root in roots.items(): + for task in ("t0", "t1", "t2"): + if name == "candidate" and task == "t0" and not candidate_present: + continue + if name == "ra" and task == "t0" and marker is None: + continue + repeat = root / f"task_{task}" / "repeat_0" + repeat.mkdir(parents=True) + value = marker if name == "ra" and task == "t0" else "{}" + if value: + (repeat / "finish_params.json").write_text(value) + (repeat / "evidence.mp4").write_bytes(b"over-cap") + global_config = _global_config(roots["candidate"], {name: roots[name] for name in ("ra", "rb")}) + limits = {"max_file_bytes": 1, "max_raw_bytes": 1, "max_wire_bytes": 1, "max_section_raw_bytes": 1} + original = {"t0": "ra", "t1": "rb", "t2": "ra"} + # Default repair still considers transport; availability-only is opt-in. + # Every pair is over the 1-byte caps, so no task has a route: each keeps + # its original draw instead of aborting the stage. + transport_result, transport_receipt = make_assignment_repair(global_config, limits)(0, ["ra", "rb"], original) + assert transport_result == original + assert transport_receipt["changes"] == [] + assert transport_receipt["unrepairable_tasks"] == ([] if candidate_present else ["t0"]) + assert transport_receipt["unroutable_tasks"] == (["t0", "t1", "t2"] if candidate_present else ["t1", "t2"]) + assert transport_receipt["infeasible"] is None + repair = make_assignment_repair(global_config, {**limits, "reference_availability_only": True}) + + result, receipt = repair(0, ["ra", "rb"], original) + + assert result == (original if available else {"t0": "rb", "t1": "ra", "t2": "ra"}) + assert sorted(result.values()) == sorted(original.values()) + assert receipt["reference_availability_only"] is True + assert len(receipt["changes"]) == (0 if available else 2) + assert receipt["initially_incompatible"] == ( + [] if available else [{"task_id": "t0", "reference_id": "ra", "reasons": ["reference_incomplete"]}] + ) + + +class TestFingerprintTransportRepair: + REF_ELOS = {"ref_a": 1000.0} + DIST = {"grp": {"task_ids": ["t0"], "percentage": 100.0}} + + def _cfg(self) -> MultiStageRunConfig: + return MultiStageRunConfig( + enabled=True, + stages=parse_multistage_config({"enabled": True, "stages": [{"num_tasks": 1}]}).stages, + seed=0, + ) + + def test_repair_settings_invalidate_fingerprint(self) -> None: + cfg = self._cfg() + base = compute_fingerprint(cfg, self.REF_ELOS, self.DIST, resolved_global_config={"multistage": {}}) + repair_a = compute_fingerprint( + cfg, + self.REF_ELOS, + self.DIST, + resolved_global_config={"multistage": {"transport_assignment_repair": {"max_file_bytes": 1}}}, + ) + repair_b = compute_fingerprint( + cfg, + self.REF_ELOS, + self.DIST, + resolved_global_config={"multistage": {"transport_assignment_repair": {"max_file_bytes": 2}}}, + ) + assert repair_a != base + assert repair_b != base + assert repair_a != repair_b + + def test_absent_repair_config_adds_nothing(self) -> None: + cfg = self._cfg() + without_block = compute_fingerprint(cfg, self.REF_ELOS, self.DIST, resolved_global_config={}) + empty_block = compute_fingerprint(cfg, self.REF_ELOS, self.DIST, resolved_global_config={"multistage": {}}) + assert without_block == empty_block + + +class TestAssignmentRepairConfig: + def test_candidate_root_is_the_linked_agents_persist_dir(self, tmp_path) -> None: + candidate, ra, rb = (tmp_path / name for name in ("candidate", "ra", "rb")) + for root in (candidate, ra, rb): + _write_task(root, "t0") + repair = make_assignment_repair(_global_config(candidate, {"ra": ra, "rb": rb}), {}) + + _, receipt = repair(0, ["ra", "rb"], {"t0": "ra"}) + + # The candidate deliverable was found, so the task was costed, not pinned. + assert receipt["unrepairable_tasks"] == [] + + def test_resources_server_persist_dir_is_not_a_candidate_root(self, tmp_path) -> None: + references = {"ra": tmp_path / "ra"} + references["ra"].mkdir() + config = _global_config(None, references) + config["gdpval_resources_server"]["resources_servers"]["gdpval"]["persist_deliverables_dir"] = str(tmp_path) + + with pytest.raises(ValueError, match="exactly one persist_deliverables_dir"): + make_assignment_repair(config, {}) + + def test_agent_of_another_resources_server_is_ignored(self, tmp_path) -> None: + references = {"ra": tmp_path / "ra"} + references["ra"].mkdir() + config = _global_config( + tmp_path / "candidate", + references, + agent_extra={"resources_server": {"type": "resources_servers", "name": "other_server"}}, + ) + + with pytest.raises(ValueError, match="exactly one persist_deliverables_dir"): + make_assignment_repair(config, {}) + + def test_relative_candidate_root_is_rejected(self, tmp_path) -> None: + references = {"ra": tmp_path / "ra"} + references["ra"].mkdir() + + with pytest.raises(ValueError, match="absolute path"): + make_assignment_repair(_global_config("output/gdpval/deliverables", references), {}) + + def test_candidate_root_that_is_a_file_is_rejected(self, tmp_path) -> None: + references = {"ra": tmp_path / "ra"} + references["ra"].mkdir() + (tmp_path / "candidate").write_text("") + + with pytest.raises(ValueError, match="not a directory"): + make_assignment_repair(_global_config(tmp_path / "candidate", references), {}) + + def test_fresh_run_pins_every_task(self, tmp_path) -> None: + references = {"ra": tmp_path / "ra", "rb": tmp_path / "rb"} + for root in references.values(): + _write_task(root, "t0") + _write_task(root, "t1") + # Nothing has been persisted yet, so the candidate root does not exist. + repair = make_assignment_repair(_global_config(tmp_path / "not_yet_created", references), {}) + + original = {"t0": "ra", "t1": "rb"} + result, receipt = repair(0, ["ra", "rb"], original) + + assert result == original + assert receipt["unrepairable_tasks"] == ["t0", "t1"] + assert receipt["changes"] == [] + + def test_capacity_infeasible_stage_keeps_its_original_draw(self, tmp_path) -> None: + # Both tasks can only be judged against rb, which holds one slot. + references = {"ra": tmp_path / "ra", "rb": tmp_path / "rb"} + references["ra"].mkdir() + _write_task(references["rb"], "t0") + _write_task(references["rb"], "t1") + repair = make_assignment_repair( + _global_config(tmp_path / "candidate", references), {"reference_availability_only": True} + ) + + original = {"t0": "ra", "t1": "rb"} + result, receipt = repair(0, ["ra", "rb"], original) + + assert result == original + assert receipt["unroutable_tasks"] == [] + assert "no count-preserving" in receipt["infeasible"] + assert receipt["changes"] == [] diff --git a/resources_servers/gdpval/transport_assignment.py b/resources_servers/gdpval/transport_assignment.py new file mode 100644 index 00000000000..b7f5cbbd3ac --- /dev/null +++ b/resources_servers/gdpval/transport_assignment.py @@ -0,0 +1,640 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Deterministic provider-free repair for transport-incompatible GDPVal pairs. + +Runs when a multistage stage is planned. Only a task whose candidate +deliverable already exists at planning time can be reassigned on transport +grounds; on a fresh run the policy has not produced anything yet, so those +tasks keep their original reference (``unrepairable_tasks`` in the receipt). +``reference_availability_only`` needs no candidate deliverable: it moves tasks +away from references without a completed deliverable. A task with no +compatible reference, or a stage with no count-preserving assignment, keeps its +original draw and is recorded in the receipt; dispatch then fails that row as +it would without repair, instead of aborting the run. +""" + +from __future__ import annotations + +import hashlib +import heapq +import json +import math +import stat +import zipfile +from collections import Counter +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Dict, Mapping, Sequence, Tuple + + +_BINARY_EXTENSIONS = { + ".pdf", + ".png", + ".jpg", + ".jpeg", + ".gif", + ".webp", + ".bmp", + ".tif", + ".tiff", + ".wav", + ".wave", + ".flac", + ".mp3", + ".m4a", + ".aac", + ".ogg", + ".aif", + ".aiff", + ".mp4", + ".m4v", + ".mov", + ".mkv", + ".webm", + ".avi", + ".wmv", + ".flv", + ".mpeg", + ".mpg", + ".3gp", + ".oga", + ".opus", + ".wma", +} +_AV_EXTENSIONS = { + ".wav", + ".wave", + ".flac", + ".mp3", + ".m4a", + ".aac", + ".ogg", + ".oga", + ".opus", + ".wma", + ".aif", + ".aiff", + ".mp4", + ".m4v", + ".mov", + ".mkv", + ".webm", + ".avi", + ".wmv", + ".flv", + ".mpeg", + ".mpg", + ".3gp", +} +_IGNORED_NAMES = { + "finish_params.json", + "last_responses.jsonl", + "last_responses.jsonl.lock", + "last_responses.jsonl.offset", +} + + +@dataclass(frozen=True) +class Footprint: + raw_bytes: int + max_file_bytes: int + has_av: bool + file_count: int + # A directory whose archive contents could not be enumerated. Its true + # footprint is unknown, so every pair using it is costed incompatible + # instead of aborting the whole planning pass. + defective: bool = False + + +@dataclass(frozen=True) +class PairCost: + compatible: bool + wire_bytes: int + raw_bytes: int + max_file_bytes: int + reasons: Tuple[str, ...] + + +@dataclass(frozen=True) +class _BinaryArtifact: + path: Path + member: str | None + suffix: str + size: int + + +def _artifact_sha256(artifact: _BinaryArtifact) -> str: + digest = hashlib.sha256() + if artifact.member is None: + with artifact.path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + else: + with zipfile.ZipFile(artifact.path, "r") as archive, archive.open(artifact.member, "r") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _find_gdpval_config(global_config: Mapping[str, Any]) -> tuple[str, Mapping[str, Any]]: + """Return ``(server_name, config)`` of the one GDPVal server with references.""" + matches = [] + for name, value in global_config.items(): + if not isinstance(value, Mapping): + continue + resources = value.get("resources_servers") + if not isinstance(resources, Mapping): + continue + gdpval = resources.get("gdpval") + if isinstance(gdpval, Mapping) and gdpval.get("reference_models"): + matches.append((str(name), gdpval)) + if len(matches) != 1: + raise ValueError(f"expected one GDPVal reference config, found {len(matches)}") + return matches[0] + + +def _find_candidate_root(global_config: Mapping[str, Any], gdpval_server_name: str) -> Path: + """The ``persist_deliverables_dir`` of the agent that serves the GDPVal server. + + Candidate deliverables are persisted by the agent, not the resources server, + so the root is read from every agent whose ``resources_server`` is the GDPVal + server. Like the Stirrup agent, a relative path is refused. A root that does + not exist yet is accepted: nothing has been persisted on a fresh run. + """ + roots = set() + for value in global_config.values(): + agents = value.get("responses_api_agents") if isinstance(value, Mapping) else None + if not isinstance(agents, Mapping): + continue + for agent in agents.values(): + if not isinstance(agent, Mapping): + continue + resources_server = agent.get("resources_server") + if not isinstance(resources_server, Mapping) or resources_server.get("name") != gdpval_server_name: + continue + root = agent.get("persist_deliverables_dir") + if root: + roots.add(str(root)) + if len(roots) != 1: + raise ValueError( + "multistage.transport_assignment_repair needs exactly one persist_deliverables_dir on the agent " + f"serving {gdpval_server_name!r} (found {sorted(roots)})" + ) + candidate_root = Path(roots.pop()) + if not candidate_root.is_absolute(): + raise ValueError(f"persist_deliverables_dir must be an absolute path (got {str(candidate_root)!r})") + if candidate_root.exists() and not candidate_root.is_dir(): + raise ValueError(f"candidate transport view is not a directory: {candidate_root}") + return candidate_root + + +def _repeat_dirs(root: Path, task_id: str) -> list[Path]: + task = root / f"task_{task_id}" + if not task.is_dir(): + return [] + repeats = sorted(path for path in task.iterdir() if path.is_dir() and path.name.startswith("repeat_")) + return repeats or [task] + + +def _has_valid_finish_marker(repeat_dirs: Sequence[Path]) -> bool: + """Whether a reference task has at least one persisted completion marker.""" + for repeat in repeat_dirs: + marker = repeat / "finish_params.json" + try: + marker_stat = marker.stat() # Follow the immutable transport-view symlink. + if not stat.S_ISREG(marker_stat.st_mode) or marker_stat.st_size > 16 * 1024 * 1024: + continue + document = json.loads(marker.read_text()) + except (OSError, UnicodeDecodeError, json.JSONDecodeError): + continue + # Stirrup writes null for a normal max-turn completion without explicit + # finish arguments. Atomic persistence is the completion signal; reject + # other JSON shapes and corrupt files. + if document is None or isinstance(document, dict): + return True + return False + + +def _footprint(directory: Path, *, include_reference_files: bool = True) -> Footprint: + artifacts: list[_BinaryArtifact] = [] + defective = False + # Match app.py exactly: the submission directory and its optional + # reference_files directory are built as two semantic sections. Reference + # assets recurse below reference_files; candidate run-state subdirectories + # are not evidence and must not affect routing. + paths = list(directory.iterdir()) + reference_files = directory / "reference_files" + if include_reference_files and reference_files.is_dir(): + paths.extend(reference_files.rglob("*")) + for path in sorted(paths, key=lambda item: item.as_posix()): + if not path.is_file() or path.name in _IGNORED_NAMES: + continue + suffix = path.suffix.lower() + if suffix == ".zip": + try: + archive_context = zipfile.ZipFile(path, "r") + except (zipfile.BadZipFile, OSError): + defective = True + continue + with archive_context as archive: + for info in archive.infolist(): + member = Path(info.filename.replace("\\", "/")) + if ( + info.is_dir() + or member.is_absolute() + or ".." in member.parts + or stat.S_ISLNK(info.external_attr >> 16) + ): + continue + member_suffix = member.suffix.lower() + if member_suffix not in _BINARY_EXTENSIONS: + continue + artifacts.append( + _BinaryArtifact( + path=path, + member=info.filename, + suffix=member_suffix, + size=info.file_size, + ) + ) + continue + if suffix not in _BINARY_EXTENSIONS: + continue + artifacts.append(_BinaryArtifact(path=path, member=None, suffix=suffix, size=path.stat().st_size)) + + # A common deliverable layout includes both loose stems and a convenience + # ZIP containing those exact stems. The judge renderer retains the bytes + # once and labels byte-identical duplicates, so the provider-cap planner + # must use the same lossless semantic footprint. Only exact AV payloads are + # deduplicated; same-size but different files remain distinct evidence. + av_size_counts = Counter(artifact.size for artifact in artifacts if artifact.suffix in _AV_EXTENSIONS) + seen_av: set[tuple[int, str]] = set() + retained: list[_BinaryArtifact] = [] + for artifact in artifacts: + if artifact.suffix in _AV_EXTENSIONS and av_size_counts[artifact.size] > 1: + identity = (artifact.size, _artifact_sha256(artifact)) + if identity in seen_av: + continue + seen_av.add(identity) + retained.append(artifact) + + raw = sum(artifact.size for artifact in retained) + maximum = max((artifact.size for artifact in retained), default=0) + count = len(retained) + has_av = any(artifact.suffix in _AV_EXTENSIONS for artifact in retained) + return Footprint(raw_bytes=raw, max_file_bytes=maximum, has_av=has_av, file_count=count, defective=defective) + + +def _pair_cost( + candidate: Footprint, + reference: Footprint, + *, + max_file_bytes: int, + max_raw_bytes: int, + max_wire_bytes: int, + max_section_raw_bytes: int, + framing_reserve_bytes: int, +) -> PairCost: + if candidate.defective or reference.defective: + return PairCost(False, 0, 0, 0, ("corrupt_archive",)) + raw = candidate.raw_bytes + reference.raw_bytes + maximum = max(candidate.max_file_bytes, reference.max_file_bytes) + # Each judge section is built under its own raw budget; a directory over + # that share emits an omission marker at dispatch regardless of media type, + # which preflight then rejects for every judge. + section_reasons = [] + if max(candidate.raw_bytes, reference.raw_bytes) > max_section_raw_bytes: + section_reasons.append("section_raw_over_cap") + # Audio/video is routed only to Gemini, so its native base64 request is the + # binding case. PDF/image-only tasks retain GPT raster and native-PDF paths; + # their exact provider request is gated later by the resource server. + if not (candidate.has_av or reference.has_av): + reasons = tuple((["file_over_cap"] if maximum > max_file_bytes else []) + section_reasons) + return PairCost(not reasons, raw, raw, maximum, reasons) + wire = 4 * math.ceil(raw / 3) + framing_reserve_bytes + reasons_list = list(section_reasons) + if maximum > max_file_bytes: + reasons_list.append("file_over_cap") + if raw > max_raw_bytes: + reasons_list.append("aggregate_raw_over_cap") + if wire >= max_wire_bytes: + reasons_list.append("serialized_request_over_cap") + return PairCost(not reasons_list, wire, raw, maximum, tuple(reasons_list)) + + +class _Edge: + __slots__ = ("to", "rev", "capacity", "cost") + + def __init__(self, to: int, rev: int, capacity: int, cost: int) -> None: + self.to = to + self.rev = rev + self.capacity = capacity + self.cost = cost + + +def _add_edge(graph: list[list[_Edge]], source: int, target: int, capacity: int, cost: int) -> None: + graph[source].append(_Edge(target, len(graph[target]), capacity, cost)) + graph[target].append(_Edge(source, len(graph[source]) - 1, 0, -cost)) + + +def _solve_capacity_assignment( + task_ids: Sequence[str], + reference_ids: Sequence[str], + original: Mapping[str, str], + costs: Mapping[tuple[str, str], PairCost], +) -> Dict[str, str]: + """Minimum-change min-cost flow with exact original reference capacities.""" + tasks = sorted(task_ids) + references = sorted(reference_ids) + capacities = Counter(original.values()) + maximum_pair_cost = max((pair.wire_bytes for pair in costs.values() if pair.compatible), default=0) + change_penalty = (maximum_pair_cost + len(references) + 1) * (len(tasks) + 1) + + source = 0 + task_base = 1 + reference_base = task_base + len(tasks) + sink = reference_base + len(references) + graph: list[list[_Edge]] = [[] for _ in range(sink + 1)] + for task_index, task_id in enumerate(tasks): + _add_edge(graph, source, task_base + task_index, 1, 0) + for reference_index, reference_id in enumerate(references): + pair = costs[(task_id, reference_id)] + if not pair.compatible: + continue + changed = reference_id != original[task_id] + cost = (change_penalty if changed else 0) + pair.wire_bytes + reference_index + _add_edge(graph, task_base + task_index, reference_base + reference_index, 1, cost) + for reference_index, reference_id in enumerate(references): + _add_edge(graph, reference_base + reference_index, sink, capacities[reference_id], 0) + + potential = [0] * len(graph) + flow = 0 + while flow < len(tasks): + infinity = 10**40 + distance = [infinity] * len(graph) + previous_node = [-1] * len(graph) + previous_edge = [-1] * len(graph) + distance[source] = 0 + queue: list[tuple[int, int]] = [(0, source)] + while queue: + current_distance, node = heapq.heappop(queue) + if current_distance != distance[node]: + continue + for edge_index, edge in enumerate(graph[node]): + if edge.capacity <= 0: + continue + candidate = current_distance + edge.cost + potential[node] - potential[edge.to] + if candidate < distance[edge.to]: + distance[edge.to] = candidate + previous_node[edge.to] = node + previous_edge[edge.to] = edge_index + heapq.heappush(queue, (candidate, edge.to)) + if previous_node[sink] < 0: + blocked = sorted( + task_id + for task_id in tasks + if not any(costs[(task_id, reference_id)].compatible for reference_id in references) + ) + raise ValueError(f"no count-preserving transport-compatible assignment; tasks_without_any_route={blocked}") + for node, value in enumerate(distance): + if value < infinity: + potential[node] += value + node = sink + while node != source: + parent = previous_node[node] + edge = graph[parent][previous_edge[node]] + edge.capacity -= 1 + graph[node][edge.rev].capacity += 1 + node = parent + flow += 1 + + result: Dict[str, str] = {} + for task_index, task_id in enumerate(tasks): + node = task_base + task_index + for edge in graph[node]: + if reference_base <= edge.to < sink and edge.capacity == 0: + result[task_id] = references[edge.to - reference_base] + break + if len(result) != len(tasks) or Counter(result.values()) != capacities: + raise RuntimeError("transport assignment solver violated its capacity contract") + return result + + +def make_assignment_repair( + global_config: Mapping[str, Any], + config: Mapping[str, Any], +) -> Callable[[int, Sequence[str], Mapping[str, str]], tuple[Dict[str, str], Dict[str, Any]]]: + """Bind filesystem roots and return the multistage plan-repair callback. + + Config errors (no candidate root, unreadable reference directories, bad + knob types) raise here, at startup. The returned callback does not raise + for tasks it cannot repair; see the module docstring. + """ + gdpval_server_name, gdpval = _find_gdpval_config(global_config) + candidate_root = _find_candidate_root(global_config, gdpval_server_name) + all_references = { + str(reference_id): Path(str(reference_config["deliverables_dir"])) + for reference_id, reference_config in (gdpval.get("reference_models") or {}).items() + if isinstance(reference_config, Mapping) + } + if not all_references or any(not path.is_dir() for path in all_references.values()): + raise ValueError("one or more reference transport views are unreadable") + + # Defaults come from the judge-side enforcement constants so the plan-time + # model cannot silently drift from what dispatch actually rejects. Run + # configs may still override every knob. + from resources_servers.gdpval.comparison import ( + MAX_FILE_BYTES_FOR_JUDGE, + MAX_SECTION_RAW_ATTACHMENT_BYTES_FOR_JUDGE, + MAX_TOTAL_RAW_ATTACHMENT_BYTES_FOR_JUDGE, + MAX_TOTAL_SERIALIZED_REQUEST_BYTES_FOR_JUDGE, + ) + + max_file_bytes = int(config.get("max_file_bytes", MAX_FILE_BYTES_FOR_JUDGE)) + max_raw_bytes = int(config.get("max_raw_bytes", MAX_TOTAL_RAW_ATTACHMENT_BYTES_FOR_JUDGE)) + max_wire_bytes = int(config.get("max_wire_bytes", MAX_TOTAL_SERIALIZED_REQUEST_BYTES_FOR_JUDGE)) + max_section_raw_bytes = int(config.get("max_section_raw_bytes", MAX_SECTION_RAW_ATTACHMENT_BYTES_FOR_JUDGE)) + framing_reserve_bytes = int(config.get("framing_reserve_bytes", 4 * 1024 * 1024)) + reference_availability_only = config.get("reference_availability_only", False) + if not isinstance(reference_availability_only, bool): + raise ValueError("reference_availability_only must be a boolean") + footprint_cache: dict[tuple[str, str, bool, bool], list[Footprint]] = {} + + def footprints( + root: Path, + task_id: str, + *, + require_completed_reference: bool = False, + include_reference_files: bool = False, + ) -> list[Footprint]: + key = (str(root), task_id, require_completed_reference, include_reference_files) + if key not in footprint_cache: + directories = _repeat_dirs(root, task_id) + if not directories: + if require_completed_reference: + return [] + raise ValueError(f"missing transport deliverable: {root}/task_{task_id}") + if require_completed_reference and not _has_valid_finish_marker(directories): + return [] + footprint_cache[key] = [ + _footprint(directory, include_reference_files=include_reference_files) for directory in directories + ] + return footprint_cache[key] + + def repair( + stage_index: int, + reference_ids: Sequence[str], + original: Mapping[str, str], + ) -> tuple[Dict[str, str], Dict[str, Any]]: + references = sorted(str(reference_id) for reference_id in reference_ids) + missing = sorted(set(references) - set(all_references)) + if missing: + raise ValueError(f"stage selected unknown reference transport views: {missing}") + costs: dict[tuple[str, str], PairCost] = {} + unrepairable_tasks: list[str] = [] + for task_id in sorted(original): + try: + candidates = [] if reference_availability_only else footprints(candidate_root, task_id) + except ValueError: + # No candidate deliverable yet (fresh run: repair fires at + # stage planning, before any rollout exists). The repair cannot + # reason about this task; pin it to its original reference. + unrepairable_tasks.append(task_id) + for reference_id in references: + pinned = reference_id == original[task_id] + costs[(task_id, reference_id)] = PairCost( + compatible=pinned, + wire_bytes=0, + raw_bytes=0, + max_file_bytes=0, + reasons=() if pinned else ("candidate_deliverable_missing",), + ) + continue + for reference_id in references: + references_for_task = footprints( + all_references[reference_id], + task_id, + require_completed_reference=True, + include_reference_files=True, + ) + if not references_for_task: + costs[(task_id, reference_id)] = PairCost( + compatible=False, + wire_bytes=0, + raw_bytes=0, + max_file_bytes=0, + reasons=("reference_incomplete",), + ) + continue + if reference_availability_only: + # Repair missing reference completions; dispatch still + # enforces the judge's transport and attachment limits. + costs[(task_id, reference_id)] = PairCost( + compatible=True, + wire_bytes=0, + raw_bytes=0, + max_file_bytes=0, + reasons=(), + ) + continue + combinations = [ + _pair_cost( + candidate, + reference, + max_file_bytes=max_file_bytes, + max_raw_bytes=max_raw_bytes, + max_wire_bytes=max_wire_bytes, + max_section_raw_bytes=max_section_raw_bytes, + framing_reserve_bytes=framing_reserve_bytes, + ) + for candidate in candidates + for reference in references_for_task + ] + # One assignment must work for every candidate/reference repeat. + incompatible = [pair for pair in combinations if not pair.compatible] + if incompatible: + worst = max(incompatible, key=lambda pair: (pair.wire_bytes, pair.raw_bytes)) + costs[(task_id, reference_id)] = worst + else: + costs[(task_id, reference_id)] = max( + combinations, + key=lambda pair: (pair.wire_bytes, pair.raw_bytes), + ) + + # A task with no compatible reference cannot be repaired. Pin it to its + # original draw so the rest of the stage can still be repaired; its row + # then fails at dispatch exactly as it would without repair. + unroutable_tasks = [ + task_id + for task_id in sorted(original) + if not any(costs[(task_id, reference_id)].compatible for reference_id in references) + ] + solver_costs = dict(costs) + for task_id in unroutable_tasks: + solver_costs[(task_id, original[task_id])] = PairCost( + compatible=True, wire_bytes=0, raw_bytes=0, max_file_bytes=0, reasons=() + ) + # The exact per-reference counts can still leave no feasible assignment + # (several tasks compatible only with the same reference). Keep the + # original draw for the whole stage rather than aborting the run. + infeasible_reason = None + try: + repaired = _solve_capacity_assignment(list(original), references, original, solver_costs) + except ValueError as error: + infeasible_reason = str(error) + repaired = dict(original) + changes = [] + for task_id in sorted(original): + before = original[task_id] + after = repaired[task_id] + if before != after: + before_cost = costs[(task_id, before)] + after_cost = costs[(task_id, after)] + changes.append( + { + "task_id": task_id, + "before": before, + "after": after, + "before_compatible": before_cost.compatible, + "before_reasons": list(before_cost.reasons), + "before_wire_bytes": before_cost.wire_bytes, + "after_wire_bytes": after_cost.wire_bytes, + } + ) + initially_incompatible = [ + { + "task_id": task_id, + "reference_id": original[task_id], + "reasons": list(costs[(task_id, original[task_id])].reasons), + } + for task_id in sorted(original) + if not costs[(task_id, original[task_id])].compatible + ] + receipt = { + "schema": "gdpval.transport-assignment-repair.v1", + "stage_index": stage_index, + "policy": ( + "minimum changed tasks; exact selected-reference counts; " + + ("completed-reference availability only" if reference_availability_only else "provider-free") + ), + "reference_availability_only": reference_availability_only, + "limits": { + "max_file_bytes": max_file_bytes, + "max_raw_bytes": max_raw_bytes, + "max_wire_bytes": max_wire_bytes, + "max_section_raw_bytes": max_section_raw_bytes, + "framing_reserve_bytes": framing_reserve_bytes, + }, + "reference_counts": dict(sorted(Counter(original.values()).items())), + "initially_incompatible": initially_incompatible, + # No candidate deliverable at planning time: kept on the original reference. + "unrepairable_tasks": unrepairable_tasks, + # No compatible reference at all: kept on the original reference. + "unroutable_tasks": unroutable_tasks, + # Set when no count-preserving assignment exists: the whole stage keeps its original draw. + "infeasible": infeasible_reason, + "changes": changes, + } + return repaired, receipt + + return repair