diff --git a/resources_servers/genrm_compare/app.py b/resources_servers/genrm_compare/app.py index 065cb11b8b..b857e06bc4 100644 --- a/resources_servers/genrm_compare/app.py +++ b/resources_servers/genrm_compare/app.py @@ -95,7 +95,11 @@ class GenRMCompareConfig(BaseResourcesServerConfig): genrm_model_server: ModelServerRef # Default: genrm_model (see config) genrm_responses_create_params: NeMoGymResponseCreateParamsNonStreaming + # Release a cohort that never fills to prevent deadlock. + cohort_timeout_s: float = 1800.0 + # Cohort-based verify: number of rollouts per prompt before running comparison (Difference 1) + # When > 1, verify() buffers by prompt and runs comparison when cohort is full; rewards are relative to cohort. # When <= 1, verify() returns default_score (no comparison). num_rollouts_per_prompt: int = 1 @@ -155,6 +159,15 @@ class GenRMCompareVerifyRequest(BaseVerifyRequest): prompt_id: Optional[str] = None # Optional stable prompt identifier from the caller +class GenRMCompareVerifyResponse(BaseVerifyResponse): + # None represents no failure. + # "cohort_timeout" represents a partial failure, where the cohort timed out but there were + # sufficient responses for a valid comparison. + # "no_comparisons" represents a failure where no comparisons could be made. + # "aggregation_failed" represents a failure where the scoring raised an exception. + failure_reason: Optional[str] = None + + class GenRMCompareRequest(BaseModel): """Request payload for GenRM pairwise comparison.""" @@ -209,7 +222,7 @@ async def verify(self, body: GenRMCompareVerifyRequest) -> BaseVerifyResponse: cfg = self.config principle = body.principle if cfg.num_rollouts_per_prompt <= 1: - return BaseVerifyResponse( + return GenRMCompareVerifyResponse( responses_create_params=body.responses_create_params, response=body.response, reward=cfg.default_score, @@ -221,7 +234,7 @@ async def verify(self, body: GenRMCompareVerifyRequest) -> BaseVerifyResponse: input_messages if isinstance(input_messages, list) else list(input_messages), principle, ) - future: asyncio.Future[float] = asyncio.get_running_loop().create_future() + future: asyncio.Future[Tuple[float, Optional[str]]] = asyncio.get_running_loop().create_future() _cohort_buffers[prompt_key].append((body, future)) @@ -246,41 +259,78 @@ async def verify(self, body: GenRMCompareVerifyRequest) -> BaseVerifyResponse: # Only run for the final response if cohort_ready: - existing_results, existing_metadata = _cohort_jit_buffers.pop(prompt_key) + async with _cohort_lock: + full_buf = _cohort_buffers.pop(prompt_key, None) + full_jit = _cohort_jit_buffers.pop(prompt_key, None) + # A waiter's timeout may claim the cohort while this final arrival was in flight. + if full_buf is not None: + self._resolve_cohort(full_buf, *(full_jit or ([], []))) - # Sort to match the ordering of the original `_run_compare` logic - existing_results, existing_metadata = zip( - *sorted( - zip(existing_results, existing_metadata), key=lambda pair: (pair[1][2], pair[1][0], pair[1][1]) + try: + reward, failure_reason = await asyncio.wait_for(asyncio.shield(future), timeout=cfg.cohort_timeout_s) + except asyncio.TimeoutError: + # The cohort never filled: a peer sub-request died upstream of verify(). + # The first timed-out waiter claims whatever arrived and scores it. + # Later waiters collect the resolved future. + async with _cohort_lock: + stale_buf = _cohort_buffers.pop(prompt_key, None) + stale_jit = _cohort_jit_buffers.pop(prompt_key, None) + if stale_buf: + logger.warning( + "[GenRM] Cohort for prompt_key=%s timed out with %d/%d rollouts; scoring the partial cohort.", + prompt_key, + len(stale_buf), + cfg.num_rollouts_per_prompt, ) - ) - - rewards, _, _, _ = aggregate_scores( - comparison_results=existing_results, - comparison_metadata=existing_metadata, - response_objs=response_objs, - aggregator_method=cfg.aggregator_method, - default_score=cfg.default_score, - reasoning_bonus=cfg.reasoning_bonus, - answer_bonus=cfg.answer_bonus, - top_percentile=cfg.top_percentile, - group_reasoning_length_penalty_coeff=cfg.group_reasoning_length_penalty_coeff, - group_answer_length_penalty_coeff=cfg.group_answer_length_penalty_coeff, - group_style_penalty_coeff=cfg.group_style_penalty_coeff, - ) - - cohort_buf = _cohort_buffers.pop(prompt_key) - for i, (_, f) in enumerate(cohort_buf): - if not f.done(): - f.set_result(rewards[i]) - - reward = await future - return BaseVerifyResponse( + self._resolve_cohort(stale_buf, *(stale_jit or ([], [])), cause="cohort_timeout") + reward, failure_reason = await future + return GenRMCompareVerifyResponse( responses_create_params=body.responses_create_params, response=body.response, reward=reward, + failure_reason=failure_reason, ) + def _resolve_cohort(self, cohort_buf, existing_results, existing_metadata, cause: Optional[str] = None) -> None: + """Aggregate a (possibly partial) cohort's comparisons and resolve every waiter future.""" + cfg = self.config + try: + response_objs = [ + (b.response.model_dump() if hasattr(b.response, "model_dump") else b.response) for b, _ in cohort_buf + ] + if len(response_objs) >= 2 and existing_results: + # Sort to match the ordering of the original `_run_compare` logic + existing_results, existing_metadata = zip( + *sorted( + zip(existing_results, existing_metadata), + key=lambda pair: (pair[1][2], pair[1][0], pair[1][1]), + ) + ) + rewards, _, _, _ = aggregate_scores( + comparison_results=existing_results, + comparison_metadata=existing_metadata, + response_objs=response_objs, + aggregator_method=cfg.aggregator_method, + default_score=cfg.default_score, + reasoning_bonus=cfg.reasoning_bonus, + answer_bonus=cfg.answer_bonus, + top_percentile=cfg.top_percentile, + group_reasoning_length_penalty_coeff=cfg.group_reasoning_length_penalty_coeff, + group_answer_length_penalty_coeff=cfg.group_answer_length_penalty_coeff, + group_style_penalty_coeff=cfg.group_style_penalty_coeff, + ) + for i, (_, f) in enumerate(cohort_buf): + if not f.done(): + f.set_result((rewards[i], cause)) + return + fallback_cause = cause or "no_comparisons" + except Exception: + logger.exception("[GenRM] Cohort scoring failed; scoring cohort with the default score.") + fallback_cause = cause or "aggregation_failed" + for _, f in cohort_buf: + if not f.done(): + f.set_result((cfg.default_score, fallback_cause)) + def setup_webserver(self) -> FastAPI: app = super().setup_webserver() app.post("/compare")(self.compare) diff --git a/resources_servers/genrm_compare/tests/test_app.py b/resources_servers/genrm_compare/tests/test_app.py index a4e99646a0..72dd26a342 100644 --- a/resources_servers/genrm_compare/tests/test_app.py +++ b/resources_servers/genrm_compare/tests/test_app.py @@ -366,6 +366,95 @@ async def run_single_comparison_mock(*args, **kwargs): actual_rewards = [r.reward for r in results] assert expected_rewards == actual_rewards + def _make_cohort_server(self, cohort_timeout_s: float) -> GenRMCompareResourcesServer: + config = GenRMCompareConfig( + host="localhost", + port=8000, + entrypoint="app.py", + domain="rlhf", + name="genrm_compare", + genrm_model_server=ModelServerRef(type="responses_api_models", name="genrm_model"), + genrm_responses_create_params=NeMoGymResponseCreateParamsNonStreaming(input=[], max_output_tokens=1024), + num_rollouts_per_prompt=2, + cohort_timeout_s=cohort_timeout_s, + ) + return GenRMCompareResourcesServer.model_construct(config=config, server_client=MagicMock()) + + def _make_cohort_body(self, task_index: int) -> GenRMCompareVerifyRequest: + response = NeMoGymResponse( + id="resp_test", + created_at=0.0, + model="dummy", + object="response", + output=[ + { + "id": "msg_1", + "role": "assistant", + "type": "message", + "status": "completed", + "content": [{"type": "output_text", "text": "an answer", "annotations": []}], + } + ], + parallel_tool_calls=False, + tool_choice="none", + tools=[], + ) + return GenRMCompareVerifyRequest( + responses_create_params=NeMoGymResponseCreateParamsNonStreaming( + input=[{"role": "user", "content": "same prompt"}] + ), + response=response, + task_index=task_index, + ) + + @pytest.mark.parametrize( + ("task_index", "arrivals", "cohort_timeout_s", "jit_mode", "reason"), + [ + # 1 of 2 ever arrives: the first timed-out waiter claims and scores what came. + (41, 1, 0.05, "empty", "cohort_timeout"), + # Full cohort resolves promptly; the JIT stub produced nothing to compare. + (42, 2, 30.0, "empty", "no_comparisons"), + # The cohort-completing arrival's JIT compare straddles the deadline: the timeout + # claim wins the race and the final arrival must not crash on the claimed buffer. + (43, 2, 0.05, "slow_final", "cohort_timeout"), + # Malformed metadata poisons the sort: scoring must still resolve every waiter + # instead of hanging them with their timeouts already spent. + (44, 2, 30.0, "poisoned", "aggregation_failed"), + ], + ) + async def test_cohorts_always_release( + self, + monkeypatch: MonkeyPatch, + task_index: int, + arrivals: int, + cohort_timeout_s: float, + jit_mode: str, + reason: str, + ) -> None: + """Every cohort resolves every waiter — full, timed out, raced, or poisoned.""" + server = self._make_cohort_server(cohort_timeout_s) + calls = {"n": 0} + + async def jit_compare(*args, **kwargs): + calls["n"] += 1 + if jit_mode == "slow_final" and calls["n"] == 2: + await asyncio.sleep(0.2) + if jit_mode == "poisoned": + return ([(1.0, 2.0, 3.0)], [None]) + return ([], []) + + monkeypatch.setattr( + GenRMCompareResourcesServer, "_run_jit_compare_using_most_recent_response_obj", jit_compare + ) + + responses = await asyncio.wait_for( + asyncio.gather(*(server.verify(self._make_cohort_body(task_index=task_index)) for _ in range(arrivals))), + timeout=5, + ) + + expected = (server.config.default_score, reason) + assert [(r.reward, r.failure_reason) for r in responses] == [expected] * arrivals + class TestRunSingleComparison: """Tests for GenRMCompareResourcesServer._run_single_comparison."""