diff --git a/resources_servers/genrm_compare/README.md b/resources_servers/genrm_compare/README.md index cf7804491c2..19fdf624f80 100644 --- a/resources_servers/genrm_compare/README.md +++ b/resources_servers/genrm_compare/README.md @@ -460,11 +460,39 @@ attempt. An individual rollout's retry counter does not advance this shared grou Run the resources server with one HTTP worker. State is process-local. Completed explicit-ID groups and failed groups retain compact response digests, rewards when completed, and failure/attempt information; -full answer bodies and waiters are released. Retention is bounded by `cohort_result_ttl_s` and `max_terminal_cohorts`. +full answer bodies and waiters are released. Retention is bounded by `cohort_result_ttl_s` and `max_terminal_cohorts`, +except for groups an open session still needs during a partial-rollout checkpoint (see below). The count cap can evict state before the TTL. After eviction or restart, replay and stale-attempt detection are no longer guaranteed. The caller must enforce accepted attempt identity across those boundaries and -use fresh group IDs for unrelated work. This cache does not provide durable recovery. +use fresh group IDs for unrelated work. Outside a partial-rollout checkpoint, this cache does not provide +durable recovery. + +### Partial-rollout checkpoints + +With checkpointing on, GenRM is an `exported` resources server whose `/verify` replays after a crash. Its +outcomes reach a group's members one reply at a time, so a checkpoint can land after one member recorded its +reward but before a sibling did. That sibling replays `/verify` after a restore; the member that already +recorded its reward never does. So each resources session still open on a finished explicit-ID group exports +its member's outcome, the reward or the judge failure that callers record as a masked row, and a restore +reinstalls that finished group. The replay then gets the recorded outcome without another judge call. A +group still collecting or evaluating exports nothing, and its members rebuild it by replaying, as before. Each saved +response digest names how the response was serialized; a restore refuses digests saved under a scheme it +does not know rather than refusing every replay. + +While a checkpoint is open, a finished evaluation holds its rewards or judge failure until resume, so no +member records an outcome after this server exported. If a restored member's session is retired before it +replays, for example after a failed restore, the member is dropped; a restored group left with no member is +dropped with the attempt watermark its restore installed, so the rollouts that start over from their inputs +re-form the group under the same attempt. + +A finished group that a member with an open session may still need is kept past `cohort_result_ttl_s` and +beyond `max_terminal_cohorts`, so pruning cannot evict an outcome a later checkpoint must export or +a replay must return. It becomes an ordinary tombstone once those members close their sessions or are retired. + +Only groups with an explicit `_ng_group_id` are exported. A member that is started over from its input while +its restored group still holds continued members gets a conflicting response for its slot; the controller +must replace the whole group by advancing `_ng_group_attempt`. ### Using the file collector @@ -499,7 +527,8 @@ when only some completed rewards reached disk. The collector's individual rollou not coordinate a group attempt. A failed legacy run requires a fresh explicit group ID and a complete replacement submission with a fresh output path. Plain `--resume` against the retained failed legacy key returns its failure again. Partial resume is also unsafe: if only some rewards reached disk, the missing -members cannot form a full group by themselves. Use caller-coordinated complete attempts for recovery. Durable answer reuse and training checkpoint recovery require caller-side coordination. +members cannot form a full group by themselves. Use caller-coordinated complete attempts for recovery. The file collector takes no checkpoints, so durable +answer reuse requires caller-side coordination; for training, see [Partial-rollout checkpoints](#partial-rollout-checkpoints). ### Reverification limits diff --git a/resources_servers/genrm_compare/app.py b/resources_servers/genrm_compare/app.py index 6176489edc8..5d584253032 100644 --- a/resources_servers/genrm_compare/app.py +++ b/resources_servers/genrm_compare/app.py @@ -42,8 +42,8 @@ from typing import Any, ClassVar, Dict, List, Literal, Optional, Tuple from aiohttp import ClientConnectionError, ClientPayloadError, ClientResponseError -from fastapi import FastAPI, HTTPException -from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, field_validator, model_validator +from fastapi import FastAPI, HTTPException, Request +from pydantic import BaseModel, ConfigDict, Field, JsonValue, PrivateAttr, field_validator, model_validator from nemo_gym.base_resources_server import ( BaseResourcesServerConfig, @@ -59,7 +59,7 @@ NeMoGymEasyInputMessage, NeMoGymResponseCreateParamsNonStreaming, ) -from nemo_gym.server_utils import get_response_json, raise_for_status +from nemo_gym.server_utils import SESSION_ID_KEY, get_response_json, raise_for_status from resources_servers.genrm_compare.utils import ( GenRMOutputParseError, aggregate_scores, @@ -145,6 +145,38 @@ class _CohortState: lock: asyncio.Lock = field(default_factory=asyncio.Lock) collection_timeout_task: Optional[asyncio.Task[None]] = None evaluation_task: Optional[asyncio.Task[None]] = None + # With checkpointing on, the resources sessions recorded as this cohort's members, and of those a restore + # installed, the ones whose member has not replayed yet. Both end with the cohort. + member_sessions: set[str] = field(default_factory=set) + unreplayed_sessions: set[str] = field(default_factory=set) + + +# How _response_digest serializes a response before hashing it. A checkpoint saves digests, so a change to that +# serialization must also change this name; a restore then refuses saved state it cannot compare. +RESPONSE_DIGEST_SCHEME = "json-ascii-sha256" + + +class _CompletedMemberState(BaseModel): + """One session's checkpointed share of a finished cohort: the reward or judge failure its member was given.""" + + model_config = ConfigDict(extra="forbid") + + cohort_key: str + prompt_digest: str + group_id: str + group_attempt: int = Field(ge=0) + rollout_index: int = Field(ge=0) + response_digest: str + response_digest_scheme: Literal["json-ascii-sha256"] + # Exactly one of these: a completed cohort's reward, or a judge failure, which the caller records as a masked row. + result: Optional[Dict[str, Any]] = None + judge_failure: Optional[str] = None + + @model_validator(mode="after") + def _one_outcome(self) -> "_CompletedMemberState": + if (self.result is None) == (self.judge_failure is None): + raise ValueError("a saved GenRM member needs exactly one of result and judge_failure") + return self @dataclass @@ -364,19 +396,27 @@ class GenRMCompareResourcesServer(SimpleResourcesServer): - Batch /compare: Direct comparison of N response_objs (e.g. for rollout_collection or tests). """ - # Cohorts are keyed by prompt, not by session, so there is no session state to export. # Verification must replay: a member waits for its siblings, # so waiting on it would deadlock a checkpoint while siblings are still generating. - # After a crash, every member that had not recorded its reward re-verifies and rebuilds the cohort. - # A crash between a cohort's result and all of its members recording it leaves the rest waiting for siblings - # that will not re-verify; they fail after cohort_collection_timeout_s and are retried from input. - checkpoint_mode = "stateless" + # After a crash, every member that had not recorded its reward re-verifies. + # A cohort's outcome reaches its members one reply at a time, so a checkpoint can land after one member recorded + # its reward, or its masked row after a judge failure, but before the others did. Those others replay + # verification, and their sibling never will: rebuilding the cohort from them alone would wait forever. + # So each session still open on a finished cohort exports its member's outcome, and a restore installs the + # finished cohort that answers those replays. While a checkpoint is open, outcomes are held until resume, + # so no member records one after this export. + checkpoint_mode = "exported" checkpoint_verify = "replay" config: GenRMCompareConfig _verify_cohorts: Dict[str, _CohortState] = PrivateAttr(default_factory=dict) _latest_group_attempts: Dict[str, _GroupAttemptWatermark] = PrivateAttr(default_factory=dict) _cohort_registry_lock: asyncio.Lock = PrivateAttr(default_factory=asyncio.Lock) + # With checkpointing on, the member of a cohort with a group ID that each resources session verified as. + # A restore installs the same session IDs, so this also serves a session that has not replayed yet. + _session_members: Dict[str, Tuple[str, int]] = PrivateAttr(default_factory=dict) + # Groups whose attempt watermark a restore installed and no verification has used since. + _restored_watermarks: set[str] = PrivateAttr(default_factory=set) _cohort_tasks: set[asyncio.Task] = PrivateAttr(default_factory=set) _closed: bool = PrivateAttr(default=False) @@ -387,8 +427,9 @@ def _own_task(self, coro, *, name: str) -> asyncio.Task: task.add_done_callback(self._cohort_tasks.discard) return task - async def verify(self, body: GenRMCompareVerifyRequest) -> GenRMCompareVerifyResponse: + async def verify(self, request: Request, body: GenRMCompareVerifyRequest) -> GenRMCompareVerifyResponse: """Verify one logical rollout slot as part of a prompt cohort.""" + session_id = request.session.get(SESSION_ID_KEY) if self._closed: raise HTTPException(status_code=503, detail="GenRM server is shutting down") cfg = self.config @@ -463,6 +504,7 @@ async def verify(self, body: GenRMCompareVerifyRequest) -> GenRMCompareVerifyRes f"rollout_index={rollout_index}" ), ) + self._remember_member(cohort, rollout_index, session_id) if cohort.phase == "completed": return self._verify_response(body, cohort.results[rollout_index]) member.waiters.append(future) @@ -480,6 +522,7 @@ async def verify(self, body: GenRMCompareVerifyRequest) -> GenRMCompareVerifyRes response_digest=response_digest, waiters=[future], ) + self._remember_member(cohort, rollout_index, session_id) if cohort.collection_timeout_task is None: cohort.collection_timeout_task = self._own_task( self._expire_collecting_cohort( @@ -592,6 +635,7 @@ async def _resolve_verify_cohort( ), ) + self._restored_watermarks.discard(body.group_id) if watermark is None or body.group_attempt > watermark.latest_attempt: self._latest_group_attempts[body.group_id] = _GroupAttemptWatermark( latest_attempt=body.group_attempt, @@ -676,7 +720,7 @@ async def _remove_waiter( @staticmethod def _response_digest(response: Any) -> str: - """Hash the exact response payload whose tokens will receive the reward.""" + """Hash the exact response payload whose tokens will receive the reward; see RESPONSE_DIGEST_SCHEME.""" payload = response.model_dump(mode="json") if hasattr(response, "model_dump") else response canonical = json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=True) return hashlib.sha256(canonical.encode("utf-8")).hexdigest() @@ -774,13 +818,16 @@ async def _evaluate_verify_cohort( ) raise except JudgeError as error: - await self._fail_verify_cohort(cohort, str(error), expected_phase="evaluating", failure_kind="judge") + await self._fail_verify_cohort( + cohort, str(error), expected_phase="evaluating", failure_kind="judge", hold_for_checkpoint=True + ) except TimeoutError: await self._fail_verify_cohort( cohort, f"GenRM cohort evaluation deadline exceeded after {self.config.cohort_evaluation_timeout_s}s", expected_phase="evaluating", failure_kind="judge", + hold_for_checkpoint=True, ) except Exception as error: logger.exception("GenRM cohort evaluation failed for %s", prompt_key) @@ -797,32 +844,51 @@ async def _publish_verify_cohort( results_by_index: Dict[int, VerifyResult], ) -> None: """Publish rewards, retiring legacy cohorts and compacting explicit-ID tombstones.""" - async with cohort.lock: - if cohort.phase == "failed": - logger.debug("Discarding late GenRM completion for key=%r", cohort.key[:160]) - return - if cohort.phase != "evaluating": - raise RuntimeError(f"cannot publish GenRM rewards while cohort is {cohort.phase}") - cohort.results = results_by_index - cohort.phase = "completed" - cohort.terminal_at = time.monotonic() - cohort.evaluation_task = None - for index, member in cohort.members.items(): - for waiter in member.waiters: - if not waiter.done(): - waiter.set_result(results_by_index[index]) - # A tombstone only needs the digest and result. Do not retain - # full response payloads for every completed training cohort. - member.body = None - member.waiters.clear() - if cohort.group_id is None and self._verify_cohorts.get(prompt_key) is cohort: - self._verify_cohorts.pop(prompt_key, None) - logger.info( - "GenRM cohort disposition=completed key=%r attempt=%s members=%s", - cohort.key[:160], - cohort.group_attempt, - len(cohort.members), - ) + participant = self.checkpoint_participant + # While a checkpoint is open, hold the rewards until resume. A member that recorded one after this server + # exported would leave a cut where it has its reward and this server does not, so its siblings would replay + # verification against a cohort that never completes. Checked again under the lock, where publication is + # atomic, since a new checkpoint can close admission while this waits for it. + while True: + while participant is not None and not participant.accepting: + await participant.wait_open() + async with cohort.lock: + if participant is None or participant.accepting: + self._publish_verify_cohort_locked(prompt_key, cohort, results_by_index) + return + + def _publish_verify_cohort_locked( + self, + prompt_key: str, + cohort: _CohortState, + results_by_index: Dict[int, VerifyResult], + ) -> None: + """Publish under ``cohort.lock``, with no await, so every member gets its reward at once.""" + if cohort.phase == "failed": + logger.debug("Discarding late GenRM completion for key=%r", cohort.key[:160]) + return + if cohort.phase != "evaluating": + raise RuntimeError(f"cannot publish GenRM rewards while cohort is {cohort.phase}") + cohort.results = results_by_index + cohort.phase = "completed" + cohort.terminal_at = time.monotonic() + cohort.evaluation_task = None + for index, member in cohort.members.items(): + for waiter in member.waiters: + if not waiter.done(): + waiter.set_result(results_by_index[index]) + # A tombstone only needs the digest and result. Do not retain + # full response payloads for every completed training cohort. + member.body = None + member.waiters.clear() + if cohort.group_id is None and self._verify_cohorts.get(prompt_key) is cohort: + self._verify_cohorts.pop(prompt_key, None) + logger.info( + "GenRM cohort disposition=completed key=%r attempt=%s members=%s", + cohort.key[:160], + cohort.group_attempt, + len(cohort.members), + ) async def _fail_verify_cohort( self, @@ -831,60 +897,104 @@ async def _fail_verify_cohort( *, expected_phase: Literal["collecting", "evaluating"], failure_kind: Literal["cohort", "judge"] = "cohort", + hold_for_checkpoint: bool = False, ) -> bool: - async with cohort.lock: - if cohort.phase != expected_phase: - return False - cohort.phase = "failed" - cohort.failure = message - cohort.failure_kind = failure_kind - cohort.terminal_at = time.monotonic() - timeout_task = cohort.collection_timeout_task - cohort.collection_timeout_task = None - current_task = asyncio.current_task() - if timeout_task is not None and timeout_task is not current_task: - timeout_task.cancel() - cohort.evaluation_task = None - for member in cohort.members.values(): - for waiter in member.waiters: - if not waiter.done(): - waiter.set_exception(CohortEvaluationError(message)) - member.body = None - member.waiters.clear() - # Keep failed legacy groups fenced too: delayed old members must not - # join a replacement under the same key. Recovery needs an explicit ID. - logger.warning( - "GenRM cohort disposition=failed key=%r attempt=%s members=%s kind=%s reason=%r", - cohort.key[:160], - cohort.group_attempt, - len(cohort.members), - failure_kind, - message[:500], - ) - return True + # A judge failure reaches callers as a masked row they record like a reward, so with + # ``hold_for_checkpoint`` it waits out an open checkpoint the same way; see _publish_verify_cohort. + participant = self.checkpoint_participant if hold_for_checkpoint else None + while True: + while participant is not None and not participant.accepting: + await participant.wait_open() + async with cohort.lock: + if participant is None or participant.accepting: + return self._fail_verify_cohort_locked(cohort, message, expected_phase, failure_kind) + + def _fail_verify_cohort_locked( + self, + cohort: _CohortState, + message: str, + expected_phase: Literal["collecting", "evaluating"], + failure_kind: Literal["cohort", "judge"], + ) -> bool: + """Fail under ``cohort.lock``, with no await, so every member gets the failure at once.""" + if cohort.phase != expected_phase: + return False + cohort.phase = "failed" + cohort.failure = message + cohort.failure_kind = failure_kind + cohort.terminal_at = time.monotonic() + timeout_task = cohort.collection_timeout_task + cohort.collection_timeout_task = None + current_task = asyncio.current_task() + if timeout_task is not None and timeout_task is not current_task: + timeout_task.cancel() + cohort.evaluation_task = None + for member in cohort.members.values(): + for waiter in member.waiters: + if not waiter.done(): + waiter.set_exception(CohortEvaluationError(message)) + member.body = None + member.waiters.clear() + # Keep failed legacy groups fenced too: delayed old members must not + # join a replacement under the same key. Recovery needs an explicit ID. + logger.warning( + "GenRM cohort disposition=failed key=%r attempt=%s members=%s kind=%s reason=%r", + cohort.key[:160], + cohort.group_attempt, + len(cohort.members), + failure_kind, + message[:500], + ) + return True + + def _needed_by_open_session(self, cohort: _CohortState) -> bool: + """Whether a member whose session is still open may need this cohort's outcome. + + A checkpoint exports the outcome for that session, and its replay after a restore is answered from it. + A member that recorded its outcome has closed its session. + """ + participant = self.checkpoint_participant + return participant is not None and any(participant.is_open(session) for session in cohort.member_sessions) def _prune_terminal_cohorts(self) -> None: - """Bound process-local cohort tombstones and attempt watermarks.""" + """Bound process-local cohort tombstones and attempt watermarks. + + A finished cohort an open member session may still need is never dropped: it outlives + ``cohort_result_ttl_s``, and the count pass skips it. Only eviction candidates are checked for open + sessions, so the cost follows what is evicted or kept rather than everything retained. + """ now = time.monotonic() - terminal = [(key, cohort) for key, cohort in self._verify_cohorts.items() if cohort.terminal_at is not None] + # Once checked, a needed cohort stays needed for the rest of this pass. + needed: set[str] = set() terminal_ttl_s = self.config.cohort_result_ttl_s if terminal_ttl_s is not None: - for key, cohort in terminal: - if now - cohort.terminal_at >= terminal_ttl_s: - self._verify_cohorts.pop(key, None) + for key, cohort in list(self._verify_cohorts.items()): + if cohort.terminal_at is None or now - cohort.terminal_at < terminal_ttl_s: + continue + if self._needed_by_open_session(cohort): + needed.add(key) + else: + self._drop_cohort(key) terminal = sorted( - ((key, cohort) for key, cohort in self._verify_cohorts.items() if cohort.terminal_at is not None), + ( + (key, cohort) + for key, cohort in self._verify_cohorts.items() + if cohort.terminal_at is not None and key not in needed + ), key=lambda item: item[1].terminal_at or 0.0, ) - for key, _ in terminal[: -self.config.max_terminal_cohorts]: - self._verify_cohorts.pop(key, None) + excess = len(terminal) - self.config.max_terminal_cohorts + for key, cohort in terminal: + if excess <= 0: + break + if self._needed_by_open_session(cohort): + continue + self._drop_cohort(key) + excess -= 1 - active_group_ids = { - cohort.group_id - for cohort in self._verify_cohorts.values() - if cohort.group_id is not None and cohort.terminal_at is None - } + # A kept cohort keeps its group's watermark, so the group still refuses older attempts. + active_group_ids = {cohort.group_id for cohort in self._verify_cohorts.values() if cohort.group_id is not None} if terminal_ttl_s is not None: for group_id, watermark in list(self._latest_group_attempts.items()): if group_id not in active_group_ids and now - watermark.updated_at >= terminal_ttl_s: @@ -902,6 +1012,149 @@ def _prune_terminal_cohorts(self) -> None: for group_id, _ in prunable_watermarks[: max(0, excess)]: self._latest_group_attempts.pop(group_id, None) + def _drop_cohort(self, key: str) -> None: + """Remove a cohort and the member records that name it.""" + cohort = self._verify_cohorts.pop(key, None) + if cohort is None: + return + for session_id in cohort.member_sessions: + if self._session_members.get(session_id, (None,))[0] == key: + del self._session_members[session_id] + + def _remember_member(self, cohort: _CohortState, rollout_index: int, session_id: Optional[str]) -> None: + """Record which member of a cohort with a group ID a session verified as; only checkpointing needs it.""" + if cohort.group_id is None or session_id is None or self.checkpoint_participant is None: + return + previous = self._session_members.get(session_id) + if previous is not None and previous[0] != cohort.key and (other := self._verify_cohorts.get(previous[0])): + other.member_sessions.discard(session_id) + self._session_members[session_id] = (cohort.key, rollout_index) + cohort.member_sessions.add(session_id) + # This member continued from the restored cohort, so a later retire of its session keeps it. + cohort.unreplayed_sessions.discard(session_id) + + def _finished_member_state(self, session_id: str) -> Dict[str, Any]: + """The outcome a session's member was given, or an empty state while its cohort has none to give.""" + member = self._session_members.get(session_id) + if member is None: + return {} + cohort_key, rollout_index = member + cohort = self._verify_cohorts.get(cohort_key) + if cohort is None or cohort.group_id is None or rollout_index not in cohort.members: + return {} + if cohort.phase == "completed": + outcome: Dict[str, Any] = {"result": cohort.results[rollout_index]} + elif cohort.phase == "failed" and cohort.failure_kind == "judge": + outcome = {"judge_failure": cohort.failure or "GenRM cohort evaluation failed"} + else: + # Still collecting or evaluating, or failed without a recorded row: every member replays, as before. + return {} + return _CompletedMemberState( + cohort_key=cohort_key, + prompt_digest=cohort.prompt_digest, + group_id=cohort.group_id, + group_attempt=cohort.group_attempt, + rollout_index=rollout_index, + response_digest=cohort.members[rollout_index].response_digest, + response_digest_scheme=RESPONSE_DIGEST_SCHEME, + **outcome, + ).model_dump(mode="json", exclude_none=True) + + async def export_session_states(self, session_ids: list[str]) -> dict[str, JsonValue]: + """Export the outcome of each session's member whose cohort has finished; an empty state otherwise.""" + return {session_id: self._finished_member_state(session_id) for session_id in session_ids} + + async def restore_session_states(self, states: dict[str, JsonValue]) -> None: + """Rebuild every finished cohort the sessions name, so a replayed verification gets its saved outcome.""" + cohorts: Dict[str, _CohortState] = {} + session_members: Dict[str, Tuple[str, int]] = {} + for session_id, raw in states.items(): + if not raw: + continue + state = _CompletedMemberState.model_validate(raw) + session_members[session_id] = (state.cohort_key, state.rollout_index) + phase: Literal["completed", "failed"] = "completed" if state.result is not None else "failed" + cohort = cohorts.get(state.cohort_key) + if cohort is None: + cohort = _CohortState( + prompt_digest=state.prompt_digest, + key=state.cohort_key, + group_id=state.group_id, + group_attempt=state.group_attempt, + phase=phase, + failure=state.judge_failure, + failure_kind="judge", + ) + cohorts[state.cohort_key] = cohort + elif (cohort.prompt_digest, cohort.group_id, cohort.group_attempt, cohort.phase, cohort.failure) != ( + state.prompt_digest, + state.group_id, + state.group_attempt, + phase, + state.judge_failure, + ): + raise ValueError(f"session {session_id!r} names cohort {state.cohort_key!r} inconsistently") + existing = cohort.members.get(state.rollout_index) + if existing is not None and ( + existing.response_digest != state.response_digest + or cohort.results.get(state.rollout_index) != state.result + ): + raise ValueError( + f"sessions disagree on rollout_index={state.rollout_index} of cohort {state.cohort_key!r}" + ) + cohort.members[state.rollout_index] = _CohortMember(body=None, response_digest=state.response_digest) + if state.result is not None: + cohort.results[state.rollout_index] = state.result + cohort.member_sessions.add(session_id) + cohort.unreplayed_sessions.add(session_id) + # Validated in full before anything is installed. + self._session_members.update(session_members) + now = time.monotonic() + for key, cohort in cohorts.items(): + cohort.terminal_at = now + self._verify_cohorts[key] = cohort + assert cohort.group_id is not None + watermark = self._latest_group_attempts.get(cohort.group_id) + if watermark is None or watermark.latest_attempt < cohort.group_attempt: + self._latest_group_attempts[cohort.group_id] = _GroupAttemptWatermark( + latest_attempt=cohort.group_attempt, + prompt_digest=cohort.prompt_digest, + updated_at=now, + ) + self._restored_watermarks.add(cohort.group_id) + + async def retire_session_state(self, session_id: str) -> None: + """Forget a retired or unclaimed session's member, and the restored cohort once nothing continues from it. + + A failed restore, or a commit that leaves a restored episode out of scope, starts the rollout over from its + input under the same group attempt. A restored cohort left in place would refuse its new response, so a + member whose restored session goes before it replays is dropped, and so is a restored cohort left with + none. A member that replayed stays, since it continued from the cohort. + """ + member = self._session_members.pop(session_id, None) + if member is None: + return + cohort_key, rollout_index = member + cohort = self._verify_cohorts.get(cohort_key) + if cohort is None: + return + cohort.member_sessions.discard(session_id) + if session_id not in cohort.unreplayed_sessions: + return + cohort.unreplayed_sessions.discard(session_id) + if not any(self._session_members.get(other) == member for other in cohort.member_sessions): + cohort.members.pop(rollout_index, None) + cohort.results.pop(rollout_index, None) + if cohort.members: + return + self._drop_cohort(cohort_key) + group_id = cohort.group_id + if group_id in self._restored_watermarks and not any( + other.group_id == group_id for other in self._verify_cohorts.values() + ): + self._restored_watermarks.discard(group_id) + self._latest_group_attempts.pop(group_id, None) + def setup_webserver(self) -> FastAPI: app = super().setup_webserver() app.post("/compare")(self.compare) @@ -933,6 +1186,8 @@ async def aclose(self) -> None: await asyncio.gather(*tasks, return_exceptions=True) self._verify_cohorts.clear() self._latest_group_attempts.clear() + self._session_members.clear() + self._restored_watermarks.clear() def _get_verify_cohort_key( self, diff --git a/resources_servers/genrm_compare/tests/support.py b/resources_servers/genrm_compare/tests/support.py new file mode 100644 index 00000000000..ef40a5bde3b --- /dev/null +++ b/resources_servers/genrm_compare/tests/support.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Helpers shared by the GenRM tests.""" + +from typing import Optional + +from starlette.requests import Request + +from nemo_gym.server_utils import SESSION_ID_KEY + + +def request_with_session(session_id: Optional[str] = None) -> Request: + """A request as the session middleware leaves it, for calling ``verify`` directly.""" + return Request({"type": "http", "session": {} if session_id is None else {SESSION_ID_KEY: session_id}}) diff --git a/resources_servers/genrm_compare/tests/test_app.py b/resources_servers/genrm_compare/tests/test_app.py index 10d708db01e..fd9a1bf3a21 100644 --- a/resources_servers/genrm_compare/tests/test_app.py +++ b/resources_servers/genrm_compare/tests/test_app.py @@ -48,6 +48,7 @@ _input_to_conversation_history, _output_budget_exhausted, ) +from resources_servers.genrm_compare.tests.support import request_with_session from resources_servers.genrm_compare.tests.test_cohort_lifecycle import collected_comparisons from resources_servers.genrm_compare.utils import get_prompt_key_from_input @@ -356,12 +357,13 @@ async def run_single_comparison_mock(*args, **kwargs): for rollout_index in range(16): tasks.append( server.verify( + request_with_session(), request.model_copy( update={ "task_index": 7, "rollout_index": rollout_index, } - ) + ), ) ) @@ -484,7 +486,7 @@ def request(rollout_index: int) -> GenRMCompareVerifyRequest: run_compare = AsyncMock(return_value=collected_comparisons([10.0, 20.0, 30.0])) monkeypatch.setattr(server, "_collect_comparisons", run_compare) - results = await asyncio.gather(*(server.verify(request(index)) for index in (2, 0, 1))) + results = await asyncio.gather(*(server.verify(request_with_session(), request(index)) for index in (2, 0, 1))) assert [result.reward for result in results] == [30.0, 10.0, 20.0] run_compare.assert_awaited_once() @@ -518,7 +520,11 @@ def request(rollout_index: int) -> GenRMCompareVerifyRequest: monkeypatch.setattr(server, "_collect_comparisons", run_compare) first = request(0) - results = await asyncio.gather(server.verify(first), server.verify(first), server.verify(request(1))) + results = await asyncio.gather( + server.verify(request_with_session(), first), + server.verify(request_with_session(), first), + server.verify(request_with_session(), request(1)), + ) assert [result.reward for result in results] == [1.0, 1.0, 2.0] run_compare.assert_awaited_once() @@ -549,15 +555,15 @@ def request(rollout_index: int, response_id: str) -> GenRMCompareVerifyRequest: run_compare = AsyncMock(return_value=collected_comparisons([1.0, 2.0])) monkeypatch.setattr(server, "_collect_comparisons", run_compare) - original = asyncio.create_task(server.verify(request(0, "original"))) + original = asyncio.create_task(server.verify(request_with_session(), request(0, "original"))) await asyncio.sleep(0) with pytest.raises(HTTPException) as error: - await server.verify(request(0, "replacement")) + await server.verify(request_with_session(), request(0, "replacement")) assert error.value.status_code == 409 assert len(next(iter(server._verify_cohorts.values())).members) == 1 - second = asyncio.create_task(server.verify(request(1, "sibling"))) + second = asyncio.create_task(server.verify(request_with_session(), request(1, "sibling"))) results = await asyncio.gather(original, second) assert [result.reward for result in results] == [1.0, 2.0] assert server._verify_cohorts == {} @@ -569,8 +575,8 @@ async def test_verify_supports_legacy_prompt_only_cohort(self, config, monkeypat monkeypatch.setattr(server, "_collect_comparisons", run_compare) results = await asyncio.gather( - server.verify(self._verify_request(0, task_index=None)), - server.verify(self._verify_request(1, task_index=None)), + server.verify(request_with_session(), self._verify_request(0, task_index=None)), + server.verify(request_with_session(), self._verify_request(1, task_index=None)), ) assert [result.reward for result in results] == [1.0, 2.0] @@ -593,12 +599,12 @@ async def test_completed_legacy_cohort_does_not_block_sequential_rerun( monkeypatch.setattr(server, "_collect_comparisons", run_compare) first = await asyncio.gather( - server.verify(self._verify_request(0, response_id="first-0")), - server.verify(self._verify_request(1, response_id="first-1")), + server.verify(request_with_session(), self._verify_request(0, response_id="first-0")), + server.verify(request_with_session(), self._verify_request(1, response_id="first-1")), ) second = await asyncio.gather( - server.verify(self._verify_request(0, response_id="second-0")), - server.verify(self._verify_request(1, response_id="second-1")), + server.verify(request_with_session(), self._verify_request(0, response_id="second-0")), + server.verify(request_with_session(), self._verify_request(1, response_id="second-1")), ) assert [result.reward for result in first] == [1.0, 2.0] @@ -611,9 +617,9 @@ async def test_verify_rejects_invalid_logical_coordinates(self, config): server = GenRMCompareResourcesServer.model_construct(config=config, server_client=MagicMock()) with pytest.raises(HTTPException) as missing_rollout_index: - await server.verify(self._verify_request(None)) + await server.verify(request_with_session(), self._verify_request(None)) with pytest.raises(HTTPException) as invalid_rollout_index: - await server.verify(self._verify_request(2)) + await server.verify(request_with_session(), self._verify_request(2)) assert missing_rollout_index.value.status_code == 422 assert invalid_rollout_index.value.status_code == 422 @@ -627,10 +633,10 @@ async def test_late_identical_duplicate_receives_cached_reward(self, config, mon first = self._verify_request(0, task_index=None, group_id="legacy-21") await asyncio.gather( - server.verify(first), - server.verify(self._verify_request(1, task_index=None, group_id="legacy-21")), + server.verify(request_with_session(), first), + server.verify(request_with_session(), self._verify_request(1, task_index=None, group_id="legacy-21")), ) - duplicate = await server.verify(first) + duplicate = await server.verify(request_with_session(), first) assert duplicate.reward == 1.0 run_compare.assert_awaited_once() @@ -649,27 +655,35 @@ async def test_new_group_attempt_is_isolated_from_completed_cohort(self, config, monkeypatch.setattr(server, "_collect_comparisons", run_compare) first_attempt = await asyncio.gather( - server.verify(self._verify_request(0, task_index=None, group_id="completed-group", response_id="old-0")), - server.verify(self._verify_request(1, task_index=None, group_id="completed-group", response_id="old-1")), + server.verify( + request_with_session(), + self._verify_request(0, task_index=None, group_id="completed-group", response_id="old-0"), + ), + server.verify( + request_with_session(), + self._verify_request(1, task_index=None, group_id="completed-group", response_id="old-1"), + ), ) replacement_attempt = await asyncio.gather( server.verify( + request_with_session(), self._verify_request( 0, task_index=None, group_id="completed-group", group_attempt=1, response_id="replacement-0", - ) + ), ), server.verify( + request_with_session(), self._verify_request( 1, task_index=None, group_id="completed-group", group_attempt=1, response_id="replacement-1", - ) + ), ), ) @@ -689,34 +703,37 @@ async def test_partial_old_attempt_does_not_mix_with_completed_replacement(self, old_waiter = asyncio.create_task( server.verify( + request_with_session(), self._verify_request( 0, task_index=None, group_id="replacement-group", group_attempt=0, response_id="old-0", - ) + ), ) ) await asyncio.sleep(0) replacement = await asyncio.gather( server.verify( + request_with_session(), self._verify_request( 0, task_index=None, group_id="replacement-group", group_attempt=1, response_id="new-0", - ) + ), ), server.verify( + request_with_session(), self._verify_request( 1, task_index=None, group_id="replacement-group", group_attempt=1, response_id="new-1", - ) + ), ), ) @@ -744,34 +761,37 @@ async def test_late_older_group_attempt_is_rejected(self, config, monkeypatch: M replacement = await asyncio.gather( server.verify( + request_with_session(), self._verify_request( 0, task_index=None, group_id="stale-group", group_attempt=1, response_id="new-0", - ) + ), ), server.verify( + request_with_session(), self._verify_request( 1, task_index=None, group_id="stale-group", group_attempt=1, response_id="new-1", - ) + ), ), ) with pytest.raises(HTTPException) as error: await server.verify( + request_with_session(), self._verify_request( 0, task_index=None, group_id="stale-group", group_attempt=0, response_id="old-0", - ) + ), ) assert [result.reward for result in replacement] == [3.0, 4.0] @@ -784,12 +804,13 @@ async def test_prompt_mismatch_cannot_supersede_active_group_attempt(self, confi server = GenRMCompareResourcesServer.model_construct(config=config, server_client=MagicMock()) first_waiter = asyncio.create_task( server.verify( + request_with_session(), self._verify_request( 0, task_index=None, group_id="stable-group", group_attempt=0, - ) + ), ) ) await asyncio.sleep(0) @@ -813,7 +834,7 @@ async def test_prompt_mismatch_cannot_supersede_active_group_attempt(self, confi ) with pytest.raises(HTTPException) as error: - await server.verify(mismatched) + await server.verify(request_with_session(), mismatched) assert error.value.status_code == 409 assert "inconsistent prompt" in str(error.value.detail) @@ -827,7 +848,7 @@ async def test_prompt_mismatch_in_group_attempt_is_rejected(self, config): config = config.model_copy(update={"num_rollouts_per_prompt": 2}) server = GenRMCompareResourcesServer.model_construct(config=config, server_client=MagicMock()) first = self._verify_request(0, task_index=None, group_id="prompt-mismatch", group_attempt=0) - first_waiter = asyncio.create_task(server.verify(first)) + first_waiter = asyncio.create_task(server.verify(request_with_session(), first)) await asyncio.sleep(0) mismatched = self._verify_request(1, task_index=None, group_id="prompt-mismatch", group_attempt=0).model_copy( update={ @@ -838,7 +859,7 @@ async def test_prompt_mismatch_in_group_attempt_is_rejected(self, config): ) with pytest.raises(HTTPException) as error: - await server.verify(mismatched) + await server.verify(request_with_session(), mismatched) assert error.value.status_code == 409 assert "inconsistent prompt" in error.value.detail @@ -854,7 +875,7 @@ async def test_incomplete_cohort_times_out_and_releases_waiters(self, config): ) server = GenRMCompareResourcesServer.model_construct(config=config, server_client=MagicMock()) - waiter = asyncio.create_task(server.verify(self._verify_request(0, task_index=22))) + waiter = asyncio.create_task(server.verify(request_with_session(), self._verify_request(0, task_index=22))) await asyncio.sleep(0) cohort = next(iter(server._verify_cohorts.values())) with pytest.raises(HTTPException, match="did not collect 2 unique rollout indices") as error: @@ -874,7 +895,7 @@ async def test_disconnected_waiter_does_not_retire_logical_cohort(self, config): } ) server = GenRMCompareResourcesServer.model_construct(config=config, server_client=MagicMock()) - waiter = asyncio.create_task(server.verify(self._verify_request(0, task_index=27))) + waiter = asyncio.create_task(server.verify(request_with_session(), self._verify_request(0, task_index=27))) await asyncio.sleep(0) cohort = next(iter(server._verify_cohorts.values())) @@ -892,8 +913,8 @@ async def test_evaluation_failure_releases_every_waiter(self, config, monkeypatc monkeypatch.setattr(server, "_collect_comparisons", AsyncMock(side_effect=ValueError("judge failed"))) results = await asyncio.gather( - server.verify(self._verify_request(0, task_index=23)), - server.verify(self._verify_request(1, task_index=23)), + server.verify(request_with_session(), self._verify_request(0, task_index=23)), + server.verify(request_with_session(), self._verify_request(1, task_index=23)), return_exceptions=True, ) @@ -918,8 +939,8 @@ def fail_model_dump(*args, **kwargs): monkeypatch.setattr(NeMoGymResponse, "model_dump", fail_model_dump) results = await asyncio.gather( - server.verify(self._verify_request(0, task_index=26)), - server.verify(self._verify_request(1, task_index=26)), + server.verify(request_with_session(), self._verify_request(0, task_index=26)), + server.verify(request_with_session(), self._verify_request(1, task_index=26)), return_exceptions=True, ) @@ -946,8 +967,8 @@ async def blocked_compare(*args, **kwargs): monkeypatch.setattr(server, "_collect_comparisons", blocked_compare) waiters = [ - asyncio.create_task(server.verify(self._verify_request(0, task_index=24))), - asyncio.create_task(server.verify(self._verify_request(1, task_index=24))), + asyncio.create_task(server.verify(request_with_session(), self._verify_request(0, task_index=24))), + asyncio.create_task(server.verify(request_with_session(), self._verify_request(1, task_index=24))), ] await asyncio.wait_for(comparison_started.wait(), timeout=1.0) cohort = next(iter(server._verify_cohorts.values())) @@ -970,7 +991,9 @@ async def test_server_instances_do_not_share_cohorts(self, config): first = GenRMCompareResourcesServer.model_construct(config=config, server_client=MagicMock()) second = GenRMCompareResourcesServer.model_construct(config=config, server_client=MagicMock()) - first_waiter = asyncio.create_task(first.verify(self._verify_request(0, task_index=25))) + first_waiter = asyncio.create_task( + first.verify(request_with_session(), self._verify_request(0, task_index=25)) + ) await asyncio.sleep(0) assert len(first._verify_cohorts) == 1 @@ -994,8 +1017,12 @@ async def test_terminal_tombstones_are_bounded_and_expire(self, config, monkeypa for task_index in (30, 31, 32): await asyncio.gather( - server.verify(self._verify_request(0, task_index=None, group_id=f"group-{task_index}")), - server.verify(self._verify_request(1, task_index=None, group_id=f"group-{task_index}")), + server.verify( + request_with_session(), self._verify_request(0, task_index=None, group_id=f"group-{task_index}") + ), + server.verify( + request_with_session(), self._verify_request(1, task_index=None, group_id=f"group-{task_index}") + ), ) server._prune_terminal_cohorts() @@ -1051,7 +1078,7 @@ async def test_singleton_verify_calls_judge(self) -> None: rollout_index=0, ) - result = await server.verify(request) + result = await server.verify(request_with_session(), request) assert result.reward == 5.0 server._run_single_comparison.assert_awaited_once() diff --git a/resources_servers/genrm_compare/tests/test_checkpoint.py b/resources_servers/genrm_compare/tests/test_checkpoint.py new file mode 100644 index 00000000000..7f89466786b --- /dev/null +++ b/resources_servers/genrm_compare/tests/test_checkpoint.py @@ -0,0 +1,545 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A checkpoint must not lose a finished cohort that only some of its members have recorded.""" + +import asyncio +import tempfile +import time +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest +from omegaconf import DictConfig + +import resources_servers.genrm_compare.app as genrm +from nemo_gym._checkpoint.control import CheckpointRequest, next_attempt +from nemo_gym._checkpoint.resources import ResourcesParticipant +from nemo_gym._checkpoint.steps import CHECKPOINT_CLOSE_FOLLOWS_HEADER +from nemo_gym.episode_types import EpisodeId +from nemo_gym.judge import JudgeError +from nemo_gym.openai_utils import NeMoGymResponse, NeMoGymResponseCreateParamsNonStreaming +from nemo_gym.server_utils import ServerClient +from resources_servers.genrm_compare.tests.support import request_with_session + + +def comparison_result(score_1=4.0, score_2=2.0, ranking=1.0): + scores = (score_1, score_2, ranking) + return (*scores, *scores, -1.0, -1.0, -1.0, 0.0, 0.0) + + +def member_0_wins(*_args, pair_idx, **_kwargs): + """Member 0 scores higher in either order, so the members finish with different rewards and a mix-up shows.""" + return comparison_result(5.0, 1.0) if pair_idx[0] == 0 else comparison_result(1.0, 5.0) + + +def member(index, *, group="group", attempt=0, answer=None): + return genrm.GenRMCompareVerifyRequest( + responses_create_params=NeMoGymResponseCreateParamsNonStreaming(input=[{"role": "user", "content": "2+2?"}]), + response=NeMoGymResponse( + id=answer or f"answer-{index}", + created_at=0.0, + model="test", + tools=[], + parallel_tool_calls=True, + tool_choice="auto", + output=[], + object="response", + ), + group_id=group, + group_attempt=attempt, + rollout_index=index, + ) + + +EPISODES = [EpisodeId(rollout_id=f"group_g{index}") for index in range(2)] +SESSIONS = [f"session-{index}" for index in range(2)] +CHECKPOINT = CheckpointRequest(checkpoint_id="c1", deadline_ts=time.time() + 3600) + + +def participant_server(config, *, judge=None) -> tuple[genrm.GenRMCompareResourcesServer, ResourcesParticipant]: + """A server behind the real resources participant, so sessions bind and restore as they do in Gym.""" + server = genrm.GenRMCompareResourcesServer.model_construct(config=config, server_client=MagicMock()) + participant = ResourcesParticipant(server, "exported", verify_step_type="replay") + server._checkpoint = participant + server._run_single_comparison = judge or AsyncMock(return_value=comparison_result()) + return server, participant + + +def seeded(config, *, judge=None) -> tuple[genrm.GenRMCompareResourcesServer, ResourcesParticipant]: + server, participant = participant_server(config, judge=judge) + for session_id, episode_id in zip(SESSIONS, EPISODES): + participant.seeded(session_id, episode_id) + return server, participant + + +async def verify_as(server, session_id: str, body): + return await server.verify(request_with_session(session_id), body) + + +async def until(condition, timeout: float = 2) -> None: + async with asyncio.timeout(timeout): + while not condition(): + await asyncio.sleep(0) + + +def count_wait_open(participant) -> dict[str, int]: + """Count the waits for admission to reopen, so a test sees a held outcome actually reach the hold.""" + counts = {"entered": 0} + wait_open = participant.wait_open + + async def counting_wait_open() -> None: + counts["entered"] += 1 + await wait_open() + + participant.wait_open = counting_wait_open + return counts + + +async def finish_cohort(server) -> list: + """Verify both members; member 0 then records its outcome and closes, member 1 has not recorded its own.""" + outcomes = await asyncio.gather( + *(verify_as(server, SESSIONS[index], member(index)) for index in range(2)), return_exceptions=True + ) + server._checkpoint.ended(SESSIONS[0]) + return outcomes + + +async def restore(config, records, *, judge=None): + server, participant = participant_server(config, judge=judge) + await participant.install(records, [record.episode_id for record in records]) + return server, participant + + +async def test_a_member_that_has_not_recorded_its_reward_gets_it_back_after_a_restore(config): + server, participant = seeded(config, judge=AsyncMock(side_effect=member_0_wins)) + outcomes = await finish_cohort(server) + assert outcomes[0].reward != outcomes[1].reward + records = await participant.export(None) + assert [record.session_id for record in records] == [SESSIONS[1]] + await server.aclose() + + restored, _ = await restore(config, records) + replayed = await asyncio.wait_for(verify_as(restored, SESSIONS[1], member(1)), timeout=0.5) + + assert replayed.reward == outcomes[1].reward + # Answered from the restored cohort: no sibling had to verify again, and the judge did not run. + restored._run_single_comparison.assert_not_awaited() + await restored.aclose() + + +@pytest.mark.parametrize("recorded_before_checkpoint", [[0], []], ids=["member_1_open", "both_open"]) +async def test_a_restored_cohort_survives_another_checkpoint_before_its_members_replay( + config, recorded_before_checkpoint +): + """Each member gets back its own outcome, not a sibling's, after two checkpoint-crash-restore cycles.""" + server, participant = seeded(config, judge=AsyncMock(side_effect=member_0_wins)) + outcomes = await asyncio.gather(*(verify_as(server, SESSIONS[index], member(index)) for index in range(2))) + assert outcomes[0].reward != outcomes[1].reward + for index in recorded_before_checkpoint: + participant.ended(SESSIONS[index]) + checkpoint_a = await participant.export(None) + await server.aclose() + + restored, restored_participant = await restore(config, checkpoint_a) + # Checkpoint B lands before the open members replay their verification. + checkpoint_b = await restored_participant.export(None) + assert [record.state for record in checkpoint_b] == [record.state for record in checkpoint_a] + await restored.aclose() + + again, _ = await restore(config, checkpoint_b) + open_members = [index for index in range(2) if index not in recorded_before_checkpoint] + for index in open_members: + replayed = await asyncio.wait_for(verify_as(again, SESSIONS[index], member(index)), timeout=0.5) + assert replayed.reward == outcomes[index].reward + again._run_single_comparison.assert_not_awaited() + await again.aclose() + + +async def test_a_member_that_has_not_recorded_its_judge_failure_gets_it_back_after_a_restore(config): + """A judge failure reaches callers as a masked row they record, so it must survive a crash like a reward.""" + server, participant = seeded(config, judge=AsyncMock(side_effect=JudgeError("judge is down"))) + outcomes = await finish_cohort(server) + assert all(isinstance(outcome, JudgeError) for outcome in outcomes) + records = await participant.export(None) + assert records[0].state["judge_failure"] == str(outcomes[1]) + await server.aclose() + + restored, _ = await restore(config, records) + with pytest.raises(JudgeError, match="judge is down"): + await asyncio.wait_for(verify_as(restored, SESSIONS[1], member(1)), timeout=0.5) + restored._run_single_comparison.assert_not_awaited() + await restored.aclose() + + +@pytest.mark.parametrize("judge_fails", [False, True], ids=["reward", "judge_failure"]) +async def test_outcomes_are_held_while_a_checkpoint_is_open(config, judge_fails): + """An outcome published after this server exported would let one member record it while its siblings could not.""" + judge = AsyncMock(side_effect=JudgeError("judge is down")) if judge_fails else None + server, participant = seeded(config, judge=judge) + counts = count_wait_open(participant) + await participant.close_admission(CHECKPOINT) + members = [asyncio.create_task(verify_as(server, SESSIONS[index], member(index))) for index in range(2)] + await until(lambda: counts["entered"] == 1) + + # The judge ran, but no member has its outcome, and the export holds none. + assert not any(task.done() for task in members) + assert [record.state for record in await participant.export(None)] == [{}, {}] + + await participant.open_admission() + outcomes = await asyncio.wait_for(asyncio.gather(*members, return_exceptions=True), timeout=0.5) + if judge_fails: + assert all(isinstance(outcome, JudgeError) for outcome in outcomes) + else: + assert all(outcome.reward == 3.0 for outcome in outcomes) + await server.aclose() + + +async def test_a_session_whose_cohort_has_no_outcome_yet_exports_nothing(config): + server, participant = seeded(config) + waiting = asyncio.create_task(verify_as(server, SESSIONS[0], member(0))) + await until(lambda: server._session_members) + + assert [record.state for record in await participant.export(None)] == [{}, {}] + waiting.cancel() + await asyncio.gather(waiting, return_exceptions=True) + await server.aclose() + + +@pytest.mark.parametrize("how", ["failed_restore", "left_out_of_commit"]) +async def test_a_restored_cohort_nothing_continues_from_is_dropped(config, how): + """The rollout starts over from its input under the same group attempt; a leftover cohort would refuse it.""" + server, participant = seeded(config) + await finish_cohort(server) + records = await participant.export(None) + await server.aclose() + + restored, restored_participant = await restore(config, records) + replacement = next_attempt(EPISODES[1]) + if how == "failed_restore": + await restored_participant.mark_retired([replacement]) + await restored_participant.retire(replacement) + else: + await restored_participant.delete_restored(replacement) + assert restored._verify_cohorts == {} + assert restored._latest_group_attempts == {} + assert restored._session_members == {} + + # Both members regenerate and re-form the group, as they would on a process with no restored cohort. + rewards = await asyncio.wait_for( + asyncio.gather( + *(verify_as(restored, f"fresh-{index}", member(index, answer=f"new-{index}")) for index in range(2)) + ), + timeout=0.5, + ) + assert all(reward.reward == 3.0 for reward in rewards) + await restored.aclose() + + +async def test_a_restored_cohort_stays_once_a_member_has_replayed(config): + sessions = ["s0", "s1", "s2"] + config3 = config.model_copy(update={"num_rollouts_per_prompt": 3}) + server, participant = participant_server(config3, judge=AsyncMock(return_value=comparison_result())) + episodes = [EpisodeId(rollout_id=f"group_g{index}") for index in range(3)] + for session_id, episode_id in zip(sessions, episodes): + participant.seeded(session_id, episode_id) + await asyncio.gather(*(verify_as(server, sessions[index], member(index)) for index in range(3))) + participant.ended("s0") + records = await participant.export(None) + await server.aclose() + + restored, restored_participant = await restore(config3, records) + # Member 1 replays and continues; member 2's restored session is then left out of a commit. + await verify_as(restored, "s1", member(1)) + await restored_participant.delete_restored(next_attempt(episodes[2])) + [cohort] = restored._verify_cohorts.values() + assert set(cohort.members) == {1} + + # Retiring the member that replayed keeps it: it continued from the cohort. + await restored_participant.mark_retired([next_attempt(episodes[1])]) + await restored_participant.retire(next_attempt(episodes[1])) + assert list(restored._verify_cohorts.values()) == [cohort] + assert set(cohort.members) == {1} + await restored.aclose() + + +async def test_a_restore_that_names_a_cohort_inconsistently_installs_nothing(config): + server, participant = seeded(config) + await asyncio.gather(*(verify_as(server, SESSIONS[index], member(index)) for index in range(2))) + records = await participant.export(None) + await server.aclose() + records[1] = records[1].model_copy(update={"state": {**records[1].state, "group_attempt": 1}}) + + restored, restored_participant = participant_server(config) + with pytest.raises(ValueError, match="inconsistently"): + await restored_participant.install(records, [record.episode_id for record in records]) + assert restored._verify_cohorts == {} + assert restored._latest_group_attempts == {} + assert restored._session_members == {} + await restored.aclose() + + +async def test_a_restore_whose_sessions_disagree_on_a_member_installs_nothing(config): + server, participant = seeded(config) + await asyncio.gather(*(verify_as(server, SESSIONS[index], member(index)) for index in range(2))) + records = await participant.export(None) + await server.aclose() + # A second session claims member 1 with another response. + impostor = records[1].model_copy( + update={"session_id": "impostor", "state": {**records[1].state, "response_digest": "0" * 64}} + ) + + restored, restored_participant = participant_server(config) + with pytest.raises(ValueError, match="disagree"): + await restored_participant.install([*records, impostor], [record.episode_id for record in records]) + assert restored._verify_cohorts == {} + assert restored._session_members == {} + await restored.aclose() + + +async def test_a_restore_refuses_digests_saved_under_another_scheme(config): + """A digest from another serialization would refuse every exact replay; refuse the saved state instead.""" + server, participant = seeded(config) + await finish_cohort(server) + records = await participant.export(None) + assert records[0].state["response_digest_scheme"] == genrm.RESPONSE_DIGEST_SCHEME + await server.aclose() + records[0] = records[0].model_copy( + update={"state": {**records[0].state, "response_digest_scheme": "orjson-sha256"}} + ) + + restored, restored_participant = participant_server(config) + with pytest.raises(ValueError, match="response_digest_scheme"): + await restored_participant.install(records, [record.episode_id for record in records]) + assert restored._verify_cohorts == {} + await restored.aclose() + + +async def test_a_restored_cohort_refuses_an_older_attempt_of_its_group(config): + server, participant = seeded(config) + await asyncio.gather(*(verify_as(server, SESSIONS[index], member(index, attempt=1)) for index in range(2))) + participant.ended(SESSIONS[0]) + records = await participant.export(None) + await server.aclose() + + restored, _ = await restore(config, records) + with pytest.raises(genrm.HTTPException, match="superseded"): + await verify_as(restored, "other", member(0, attempt=0)) + await restored.aclose() + + +async def test_member_records_end_with_their_sessions_and_the_server(config): + server, participant = seeded(config) + await asyncio.gather(*(verify_as(server, SESSIONS[index], member(index)) for index in range(2))) + assert set(server._session_members) == set(SESSIONS) + + await participant.mark_retired([EPISODES[0]]) + await participant.retire(EPISODES[0]) + assert set(server._session_members) == {SESSIONS[1]} + + await server.aclose() + assert server._session_members == {} + + +async def test_a_cohort_without_a_group_id_records_no_members(config): + server, participant = seeded(config) + await asyncio.gather(*(verify_as(server, SESSIONS[index], member(index, group=None)) for index in range(2))) + assert server._session_members == {} + await server.aclose() + + +async def test_member_records_end_with_their_pruned_cohort(config): + config = config.model_copy(update={"cohort_result_ttl_s": 0.01}) + server, participant = seeded(config) + await asyncio.gather(*(verify_as(server, SESSIONS[index], member(index)) for index in range(2))) + assert set(server._session_members) == set(SESSIONS) + # Both members recorded their outcomes and closed their sessions, so nothing needs the cohort any more. + for session_id in SESSIONS: + participant.ended(session_id) + + # The next cohort's verification prunes the expired one, and its member records go with it. + await asyncio.sleep(0.02) + pending = asyncio.create_task(verify_as(server, "later", member(0, group="later"))) + await until(lambda: "later" in server._session_members) + assert set(server._session_members) == {"later"} + pending.cancel() + await asyncio.gather(pending, return_exceptions=True) + await server.aclose() + + +@pytest.mark.parametrize("judge_fails", [False, True], ids=["reward", "judge_failure"]) +@pytest.mark.parametrize("pruned_by", ["count", "age"]) +async def test_pruning_keeps_an_outcome_an_open_session_still_needs(config, pruned_by, judge_fails): + """A member that has not recorded its outcome replays after a restore; an evicted cohort would leave it waiting + for a sibling that already finished, until the collection timeout.""" + limit = {"max_terminal_cohorts": 1} if pruned_by == "count" else {"cohort_result_ttl_s": 0.01} + config = config.model_copy(update=limit) + judge = AsyncMock(side_effect=JudgeError("judge is down") if judge_fails else member_0_wins) + groups = ["a", "b"] + server, participant = participant_server(config, judge=judge) + for group in groups: + for index in range(2): + participant.seeded(f"{group}{index}", EpisodeId(rollout_id=f"{group}_g{index}")) + outcomes = {} + for group in groups: + results = await asyncio.gather( + *(verify_as(server, f"{group}{index}", member(index, group=group)) for index in range(2)), + return_exceptions=True, + ) + outcomes[group] = results[1] + # Member 0 recorded its outcome; member 1 has not. + participant.ended(f"{group}0") + if pruned_by == "age": + await asyncio.sleep(0.02) + records = await participant.export(None) + assert sorted(record.session_id for record in records) == ["a1", "b1"] + assert all(record.state for record in records) + await server.aclose() + + restored, restored_participant = await restore(config, records) + if pruned_by == "age": + await asyncio.sleep(0.02) + for group in groups: + # Each replay prunes first; the other group's cohort, and its own, are past the limit. + replay = asyncio.wait_for(verify_as(restored, f"{group}1", member(1, group=group)), timeout=0.5) + if judge_fails: + with pytest.raises(JudgeError, match="judge is down"): + await replay + else: + assert (await replay).reward == outcomes[group].reward + restored._run_single_comparison.assert_not_awaited() + # A kept cohort keeps its group's attempt watermark, so the group still refuses older attempts. + restored._prune_terminal_cohorts() + assert set(restored._latest_group_attempts) == set(groups) + + # Once the members close their sessions, the cohorts are ordinary tombstones again. + for group in groups: + restored_participant.ended(f"{group}1") + if pruned_by == "age": + await asyncio.sleep(0.02) + restored._prune_terminal_cohorts() + assert len(restored._verify_cohorts) == (1 if pruned_by == "count" else 0) + await restored.aclose() + + +@pytest.mark.parametrize("judge_fails", [False, True], ids=["reward", "judge_failure"]) +async def test_an_outcome_stays_held_when_a_checkpoint_closes_admission_again_before_publication(config, judge_fails): + """Admission can reopen and close again while the outcome waits for the cohort lock; the check under it holds.""" + judge = AsyncMock(side_effect=JudgeError("judge is down")) if judge_fails else None + server, participant = seeded(config, judge=judge) + counts = count_wait_open(participant) + await participant.close_admission(CHECKPOINT) + members = [asyncio.create_task(verify_as(server, SESSIONS[index], member(index))) for index in range(2)] + await until(lambda: counts["entered"] == 1) + [cohort] = server._verify_cohorts.values() + + async with cohort.lock: + await participant.open_admission() + # The held outcome wakes and waits for the lock; a new checkpoint closes admission meanwhile. + for _ in range(5): + await asyncio.sleep(0) + await participant.close_admission(CHECKPOINT) + await until(lambda: counts["entered"] == 2) + assert not any(task.done() for task in members) + + await participant.open_admission() + outcomes = await asyncio.wait_for(asyncio.gather(*members, return_exceptions=True), timeout=0.5) + assert all(isinstance(outcome, JudgeError) for outcome in outcomes) == judge_fails + await server.aclose() + + +AUTH = {"authorization": "Bearer t"} + + +def http_server(config, *, judge_fails: bool): + """A server behind its real app, with checkpointing on, so requests pass the checkpoint middleware.""" + server_client = MagicMock(spec=ServerClient) + server_client.global_config_dict = DictConfig({"checkpoint": {"enabled": True, "control_auth_token": "t"}}) + server = genrm.GenRMCompareResourcesServer(config=config, server_client=server_client) + if judge_fails: + server._run_single_comparison = AsyncMock(side_effect=JudgeError("judge is down")) + else: + server._run_single_comparison = AsyncMock(side_effect=member_0_wins) + return server, httpx.ASGITransport(app=server.setup_webserver()) + + +def control(checkpoint_id: str, **extra) -> dict: + return {"checkpoint_id": checkpoint_id, "deadline_ts": time.time() + 10, **extra} + + +def body_json(index: int) -> dict: + return member(index).model_dump(mode="json", by_alias=True) + + +def outcome(response: httpx.Response) -> object: + assert response.status_code == 200, response.text + row = response.json() + return ("masked", row["failure_kind"]) if row.get("mask_sample") else row["reward"] + + +async def commit(client: httpx.AsyncClient, checkpoint_id: str) -> str: + directory = tempfile.mkdtemp() + for route, body in ( + ("prepare", control(checkpoint_id)), + ("commit", control(checkpoint_id, checkpoint_dir=directory)), + ): + (await client.post(f"/ng-control/v1/checkpoint/{route}", json=body, headers=AUTH)).raise_for_status() + return directory + + +async def restore_and_resume(client: httpx.AsyncClient, directory: str, scope: list[dict]) -> None: + restored = await client.post( + "/ng-control/v1/checkpoint/restore", + json=control("restore", checkpoint_dir=directory, episode_ids=scope), + headers=AUTH, + ) + restored.raise_for_status() + (await client.post("/ng-control/v1/checkpoint/resume", json=control("restore"), headers=AUTH)).raise_for_status() + + +@pytest.mark.parametrize("checkpoints", [1, 2]) +@pytest.mark.parametrize("judge_fails", [False, True], ids=["reward", "judge_failure"]) +async def test_a_replay_through_the_checkpoint_control_plane_gets_the_recorded_outcome( + config, judge_fails, checkpoints +): + """Member 0 recorded its outcome; member 1, whose caller closes its session later, replays after the crash.""" + server, transport = http_server(config, judge_fails=judge_fails) + clients = [httpx.AsyncClient(transport=transport, base_url="http://genrm") for _ in range(3)] + async with clients[0], clients[1], clients[2]: + for index in range(2): + seed = await clients[index].post( + f"/ng-rollout/g{index}/seed_session", json={"responses_create_params": {"input": "hi"}} + ) + seed.raise_for_status() + replies = await asyncio.gather( + clients[0].post("/ng-rollout/g0/verify", json=body_json(0)), + clients[1].post( + "/ng-rollout/g1/verify", json=body_json(1), headers={CHECKPOINT_CLOSE_FOLLOWS_HEADER: "1"} + ), + ) + recorded = outcome(replies[1]) + assert recorded != outcome(replies[0]) or judge_fails + directory = await commit(clients[2], "save-a") + cookies = dict(clients[1].cookies) + await server.aclose() + + attempt = 0 + for _ in range(checkpoints): + server, transport = http_server(config, judge_fails=judge_fails) + control_client = httpx.AsyncClient(transport=transport, base_url="http://genrm") + async with control_client: + await restore_and_resume(control_client, directory, [{"rollout_id": "g1", "attempt": attempt}]) + attempt += 1 + if attempt < checkpoints: + # Another checkpoint and crash before member 1 replays. + directory = await commit(control_client, f"save-{attempt}") + await server.aclose() + + replay = httpx.AsyncClient(transport=transport, base_url="http://genrm", cookies=cookies) + async with replay: + # Without the saved outcome, this would wait for member 0 until the collection timeout. + response = await asyncio.wait_for( + replay.post(f"/ng-rollout/g1-a{attempt}/verify", json=body_json(1)), timeout=0.5 + ) + assert outcome(response) == recorded + server._run_single_comparison.assert_not_awaited() + await server.aclose() diff --git a/resources_servers/genrm_compare/tests/test_cohort_lifecycle.py b/resources_servers/genrm_compare/tests/test_cohort_lifecycle.py index 864f2e324a9..9b98951d1d8 100644 --- a/resources_servers/genrm_compare/tests/test_cohort_lifecycle.py +++ b/resources_servers/genrm_compare/tests/test_cohort_lifecycle.py @@ -24,6 +24,7 @@ import resources_servers.genrm_compare.app as genrm from nemo_gym.openai_utils import NeMoGymResponse, NeMoGymResponseCreateParamsNonStreaming +from resources_servers.genrm_compare.tests.support import request_with_session def comparison_result(score_1=4.0, score_2=2.0, ranking=1.0): @@ -70,12 +71,12 @@ def test_deadline_must_be_finite_positive(config, field, value): async def test_judging_waits_for_full_group_and_exact_retry_replays(server): judge = AsyncMock(return_value=comparison_result()) server._run_single_comparison = judge - first = asyncio.create_task(server.verify(member(0))) + first = asyncio.create_task(server.verify(request_with_session(), member(0))) await asyncio.sleep(0) assert judge.await_count == 0 - result = await server.verify(member(1)) + result = await server.verify(request_with_session(), member(1)) assert result.reward == (await first).reward == 3.0 - assert (await server.verify(member(0))).reward == 3.0 + assert (await server.verify(request_with_session(), member(0))).reward == 3.0 assert judge.await_count == 2 @@ -89,9 +90,9 @@ async def compare(*args, **kwargs): return collected_comparisons([1.0, 2.0]) server._collect_comparisons = AsyncMock(side_effect=compare) - first = asyncio.create_task(server.verify(member(0))) + first = asyncio.create_task(server.verify(request_with_session(), member(0))) await asyncio.sleep(0) - second = asyncio.create_task(server.verify(member(1))) if during_judging else None + second = asyncio.create_task(server.verify(request_with_session(), member(1))) if during_judging else None if second: await started.wait() first.cancel() @@ -99,11 +100,11 @@ async def compare(*args, **kwargs): cohort = next(iter(server._verify_cohorts.values())) assert cohort.members[0].body is not None and not cohort.members[0].waiters with pytest.raises(HTTPException) as error: - await server.verify(member(0, response_id="different-answer")) + await server.verify(request_with_session(), member(0, response_id="different-answer")) assert error.value.status_code == 409 - retry = asyncio.create_task(server.verify(member(0))) + retry = asyncio.create_task(server.verify(request_with_session(), member(0))) if second is None: - second = asyncio.create_task(server.verify(member(1))) + second = asyncio.create_task(server.verify(request_with_session(), member(1))) release.set() assert [r.reward for r in await asyncio.gather(retry, second)] == [1.0, 2.0] server._collect_comparisons.assert_awaited_once() @@ -122,7 +123,7 @@ async def judge(*args, pair_idx, **kwargs): cancelled.add(pair_idx) server._run_single_comparison = judge - tasks = [asyncio.create_task(server.verify(member(i))) for i in range(2)] + tasks = [asyncio.create_task(server.verify(request_with_session(), member(i))) for i in range(2)] await asyncio.sleep(0.08) assert all(not task.done() for task in tasks) results = await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), 1) @@ -133,13 +134,13 @@ async def judge(*args, pair_idx, **kwargs): async def test_simultaneous_groups_do_not_mix(server): server._run_single_comparison = AsyncMock(return_value=comparison_result()) - a = asyncio.create_task(server.verify(member(0, group="run-a"))) - b = asyncio.create_task(server.verify(member(1, group="run-b"))) + a = asyncio.create_task(server.verify(request_with_session(), member(0, group="run-a"))) + b = asyncio.create_task(server.verify(request_with_session(), member(1, group="run-b"))) await asyncio.sleep(0) assert not a.done() and not b.done() - await server.verify(member(1, group="run-a")) + await server.verify(request_with_session(), member(1, group="run-a")) assert (await a).reward == 3.0 and not b.done() - await server.verify(member(0, group="run-b")) + await server.verify(request_with_session(), member(0, group="run-b")) assert (await b).reward == 3.0 @@ -147,9 +148,11 @@ async def test_simultaneous_groups_do_not_mix(server): async def test_final_arrival_deadline_race_never_publishes_partial_reward(server, delay): server.config.cohort_collection_timeout_s = 0.01 server._run_single_comparison = AsyncMock(return_value=comparison_result()) - first = asyncio.create_task(server.verify(member(0))) + first = asyncio.create_task(server.verify(request_with_session(), member(0))) await asyncio.sleep(delay) - results = await asyncio.wait_for(asyncio.gather(first, server.verify(member(1)), return_exceptions=True), 1) + results = await asyncio.wait_for( + asyncio.gather(first, server.verify(request_with_session(), member(1)), return_exceptions=True), 1 + ) cohort = next(iter(server._verify_cohorts.values())) if cohort.phase == "completed": assert [r.reward for r in results] == [3.0, 3.0] @@ -160,19 +163,19 @@ async def test_final_arrival_deadline_race_never_publishes_partial_reward(server async def test_shutdown_fails_waiters_and_prevents_new_work(server): - first = asyncio.create_task(server.verify(member(0))) + first = asyncio.create_task(server.verify(request_with_session(), member(0))) await asyncio.sleep(0) await server.aclose() with pytest.raises(HTTPException, match="shutting down"): await first with pytest.raises(HTTPException, match="shutting down"): - await server.verify(member(1)) + await server.verify(request_with_session(), member(1)) assert not server._cohort_tasks and not server._verify_cohorts async def test_abandoned_group_expires_without_another_request(server): server.config.cohort_collection_timeout_s = 0.02 - first = asyncio.create_task(server.verify(member(0))) + first = asyncio.create_task(server.verify(request_with_session(), member(0))) await asyncio.sleep(0) cohort = next(iter(server._verify_cohorts.values())) timer = cohort.collection_timeout_task @@ -201,9 +204,9 @@ async def compare(*args, **kwargs): return collected_comparisons([1.0, 2.0]) server._collect_comparisons = compare - old = [asyncio.create_task(server.verify(member(i))) for i in range(2)] + old = [asyncio.create_task(server.verify(request_with_session(), member(i))) for i in range(2)] await asyncio.wait_for(started.wait(), 1) - new = await asyncio.gather(*(server.verify(member(i, attempt=1)) for i in range(2))) + new = await asyncio.gather(*(server.verify(request_with_session(), member(i, attempt=1)) for i in range(2))) old_results = await asyncio.gather(*old, return_exceptions=True) assert all(isinstance(r, HTTPException) and r.status_code == 503 for r in old_results) assert [r.reward for r in new] == [1.0, 2.0] @@ -213,7 +216,7 @@ async def compare(*args, **kwargs): async def test_failed_legacy_group_cannot_mix_replacement_with_delayed_old_member(server): server.config.cohort_collection_timeout_s = 0.02 - first = asyncio.create_task(server.verify(member(0, group=None))) + first = asyncio.create_task(server.verify(request_with_session(), member(0, group=None))) await asyncio.sleep(0) old = next(iter(server._verify_cohorts.values())) with pytest.raises(HTTPException, match="did not collect"): @@ -223,8 +226,8 @@ async def test_failed_legacy_group_cannot_mix_replacement_with_delayed_old_membe server._run_single_comparison = AsyncMock(return_value=comparison_result()) results = await asyncio.gather( - server.verify(member(0, group=None, response_id="new-0")), - server.verify(member(1, group=None, response_id="old-1")), + server.verify(request_with_session(), member(0, group=None, response_id="new-0")), + server.verify(request_with_session(), member(1, group=None, response_id="old-1")), return_exceptions=True, ) assert all(isinstance(r, HTTPException) and r.status_code == 503 for r in results) @@ -233,7 +236,7 @@ async def test_failed_legacy_group_cannot_mix_replacement_with_delayed_old_membe await server._publish_verify_cohort(old.key, old, {0: 99.0, 1: 99.0}) assert old.phase == "failed" and not old.results replacement = await asyncio.gather( - *(server.verify(member(i, group="fresh", response_id=f"new-{i}")) for i in range(2)) + *(server.verify(request_with_session(), member(i, group="fresh", response_id=f"new-{i}")) for i in range(2)) ) assert [r.reward for r in replacement] == [3.0, 3.0] assert [r.response.id for r in replacement] == ["new-0", "new-1"] diff --git a/resources_servers/genrm_compare/tests/test_review_fixes.py b/resources_servers/genrm_compare/tests/test_review_fixes.py index fdb3ad4c64e..5613eb66173 100644 --- a/resources_servers/genrm_compare/tests/test_review_fixes.py +++ b/resources_servers/genrm_compare/tests/test_review_fixes.py @@ -36,6 +36,7 @@ from nemo_gym.reward_profile import RewardProfiler from nemo_gym.rollout_correlation import current_rollout_id, rollout_context from nemo_gym.task_data import TaskDataValidator, load_task_data_schema +from resources_servers.genrm_compare.tests.support import request_with_session from resources_servers.genrm_compare.tests.test_cohort_lifecycle import collected_comparisons, member @@ -59,18 +60,20 @@ def failing_judge(status): async def test_judge_http_failure_never_completes_cohort(server, status): response = failing_judge(status) server.server_client.post = AsyncMock(return_value=response) - results = await asyncio.gather(*(server.verify(member(i)) for i in range(2)), return_exceptions=True) + results = await asyncio.gather( + *(server.verify(request_with_session(), member(i)) for i in range(2)), return_exceptions=True + ) assert all(isinstance(result, genrm.JudgeError) and "judge offline" in str(result) for result in results) cohort = next(iter(server._verify_cohorts.values())) assert cohort.phase == "failed" and not cohort.results response.read.assert_not_awaited() with pytest.raises(genrm.JudgeError, match="judge offline"): - await server.verify(member(0, response_id="regenerated-answer")) + await server.verify(request_with_session(), member(0, response_id="regenerated-answer")) calls_after_failure = server.server_client.post.await_count max_calls = 2 * (1 if status == 401 else server.config.genrm_parse_retries + 1) assert 1 <= calls_after_failure <= max_calls server._collect_comparisons = AsyncMock(return_value=collected_comparisons([1.0, 2.0])) - recovered = await asyncio.gather(*(server.verify(member(i, attempt=1)) for i in range(2))) + recovered = await asyncio.gather(*(server.verify(request_with_session(), member(i, attempt=1)) for i in range(2))) assert [result.reward for result in recovered] == [1.0, 2.0] server._collect_comparisons.assert_awaited_once() @@ -78,7 +81,9 @@ async def test_judge_http_failure_never_completes_cohort(server, status): @pytest.mark.parametrize("error", [ConnectionError("offline"), TimeoutError("judge timeout")]) async def test_judge_transport_failure_never_defaults(server, error): server.server_client.post = AsyncMock(side_effect=error) - results = await asyncio.gather(*(server.verify(member(i)) for i in range(2)), return_exceptions=True) + results = await asyncio.gather( + *(server.verify(request_with_session(), member(i)) for i in range(2)), return_exceptions=True + ) assert all(isinstance(result, genrm.JudgeError) for result in results) assert not next(iter(server._verify_cohorts.values())).results @@ -90,7 +95,9 @@ async def test_unsuccessful_http_200_judge_is_failure(server, payload): response = MagicMock(ok=True) response.read = AsyncMock(return_value=json.dumps(payload).encode()) server.server_client.post = AsyncMock(return_value=response) - results = await asyncio.gather(*(server.verify(member(i)) for i in range(2)), return_exceptions=True) + results = await asyncio.gather( + *(server.verify(request_with_session(), member(i)) for i in range(2)), return_exceptions=True + ) assert all(isinstance(r, genrm.JudgeError) for r in results) with pytest.raises(HTTPException, match="503"): await server.compare(genrm.GenRMCompareRequest(conversation_history=[], response_objs=[{}, {}])) @@ -122,7 +129,7 @@ def output(text): async def test_group_attempt_metadata_is_not_a_reward_metric(server): server._run_single_comparison = AsyncMock(return_value=comparison_result()) rows = [member(i, attempt=2) for i in range(2)] - results = await asyncio.gather(*(server.verify(row) for row in reversed(rows))) + results = await asyncio.gather(*(server.verify(request_with_session(), row) for row in reversed(rows))) assert [r.rollout_index for r in results] == [1, 0] assert [r.reward for r in results] == [3.0, 3.0] for result in results: @@ -139,7 +146,10 @@ async def test_actual_judge_receives_normalized_prompt(server, input_value): server._run_single_comparison = AsyncMock(return_value=comparison_result()) params = NeMoGymResponseCreateParamsNonStreaming(input=input_value) await asyncio.gather( - *(server.verify(member(i).model_copy(update={"responses_create_params": params})) for i in range(2)) + *( + server.verify(request_with_session(), member(i).model_copy(update={"responses_create_params": params})) + for i in range(2) + ) ) assert server._run_single_comparison.await_args.args[0] == [{"role": "user", "content": "Hello"}] @@ -153,10 +163,10 @@ async def judge(*args, **kwargs): server._run_single_comparison = AsyncMock(side_effect=judge) with rollout_context("member-a"): - first = asyncio.create_task(server.verify(member(0))) + first = asyncio.create_task(server.verify(request_with_session(), member(0))) await asyncio.sleep(0) with rollout_context("member-b"): - second = asyncio.create_task(server.verify(member(1))) + second = asyncio.create_task(server.verify(request_with_session(), member(1))) await asyncio.gather(first, second) assert seen == [None, None] @@ -292,9 +302,9 @@ async def test_late_member_receives_original_timeout_failure(server, index): server.config.cohort_collection_timeout_s = 0.01 server._run_single_comparison = AsyncMock() with pytest.raises(HTTPException) as first: - await server.verify(member(0)) + await server.verify(request_with_session(), member(0)) with pytest.raises(HTTPException) as late: - await server.verify(member(index, response_id="new-answer")) + await server.verify(request_with_session(), member(index, response_id="new-answer")) assert first.value.status_code == late.value.status_code == 503 assert late.value.detail == first.value.detail assert "did not collect" in late.value.detail @@ -305,7 +315,7 @@ async def test_judge_failure_preserves_existing_instance_config_and_masks_all_me server._collect_comparisons = AsyncMock(side_effect=genrm.JudgeError("judge offline")) bodies = [member(i) for i in range(2)] bodies[0].instance_config = {"task": "keep", "mask_sample": False} - results = await asyncio.gather(*(judge_failsafe(server.verify)(body) for body in bodies)) + results = await asyncio.gather(*(judge_failsafe(server.verify)(request_with_session(), body) for body in bodies)) rows = [json.loads(result.body) for result in results] assert rows[0]["instance_config"] == {"task": "keep", "mask_sample": True} assert rows[1]["instance_config"] == {"mask_sample": True} @@ -461,7 +471,7 @@ async def test_verify_maps_selected_and_overall_diagnostics(server: genrm.GenRMC # The shortest answer wins on rubrics but loses on overall scores. verdict = (5.0, 2.0, 1.0, 1.0, 4.0, 6.0, -1.0, -1.0, -1.0, 0.0, 0.0) server._collect_comparisons = AsyncMock(return_value=([verdict], [(0, 1, 0)])) - results = await asyncio.gather(*(server.verify(row) for row in rows)) + results = await asyncio.gather(*(server.verify(request_with_session(), row) for row in rows)) assert [r.reward for r in results] == pytest.approx([6.2, 1.8]) assert [r.reward_score_raw for r in results] == [5.0, 2.0] assert [r.reward_rubric_aggregate_valid for r in results] == [5.0, 2.0] @@ -480,14 +490,14 @@ async def test_reordered_rubric_ids_share_cohort_and_completed_replay( genrm.GenRMCompareVerifyRequest.model_validate(member(i).model_dump() | {"expected_rubric_ids": ids}) for i, ids in enumerate([[2, 1], [1, 2]]) ] - results = await asyncio.gather(*(server.verify(row) for row in rows)) + results = await asyncio.gather(*(server.verify(request_with_session(), row) for row in rows)) assert [result.reward for result in results] == [3.0, 3.0] replay = genrm.GenRMCompareVerifyRequest.model_validate(rows[0].model_dump() | {"expected_rubric_ids": [1, 2]}) - assert (await server.verify(replay)).model_dump() == results[0].model_dump() + assert (await server.verify(request_with_session(), replay)).model_dump() == results[0].model_dump() assert server._run_single_comparison.await_count == 2 different = genrm.GenRMCompareVerifyRequest.model_validate(rows[0].model_dump() | {"expected_rubric_ids": [1, 3]}) with pytest.raises(HTTPException) as error: - await server.verify(different) + await server.verify(request_with_session(), different) assert error.value.status_code == 409 diff --git a/tests/unit_tests/test_checkpoint_server_declarations.py b/tests/unit_tests/test_checkpoint_server_declarations.py index e00b15036b5..592f523670a 100644 --- a/tests/unit_tests/test_checkpoint_server_declarations.py +++ b/tests/unit_tests/test_checkpoint_server_declarations.py @@ -20,7 +20,6 @@ "competitive_coding_challenges": "CompetitiveCodingChallengesResourcesServer", "equivalence_llm_judge": "LLMJudgeResourcesServer", "math_with_judge": "LibraryJudgeMathResourcesServer", - "genrm_compare": "GenRMCompareResourcesServer", "abstention": "AbstentionServer", "calendar": "CalendarResourcesServer", "equivalence_rule": "EquivalenceRuleResourcesServer", @@ -58,3 +57,11 @@ def test_server_is_stateless_with_a_replayable_verification(server: str) -> None "checkpoint_mode": "stateless", "checkpoint_verify": "replay", } + + +def test_genrm_exports_completed_cohorts_and_replays_verification() -> None: + """GenRM's cohorts span sessions, so it exports each session's share of a completed cohort.""" + assert class_attributes("genrm_compare", "GenRMCompareResourcesServer") == { + "checkpoint_mode": "exported", + "checkpoint_verify": "replay", + }