From 3a6022b80fdafeae7ba34c09a17a1bdcffe109c1 Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Fri, 7 Aug 2026 11:04:38 +0000 Subject: [PATCH] Fix custom preprocess and postprocess in judger pools --- tests/rl/test_judger.py | 96 ++++++++++++++++++++++++++++++++--- xtuner/v1/rl/judger/native.py | 29 ++++++++--- 2 files changed, 113 insertions(+), 12 deletions(-) diff --git a/tests/rl/test_judger.py b/tests/rl/test_judger.py index c9de6ec86..067443372 100644 --- a/tests/rl/test_judger.py +++ b/tests/rl/test_judger.py @@ -140,20 +140,27 @@ def test_native_judger_batch_judge_not_supported(self): with self.assertRaisesRegex(NotImplementedError, "does not support batch_judge"): asyncio.run(judger.batch_judge([self._make_rollout_state("correctness")])) - def test_remote_judger_uses_driver_side_preprocess_judger(self): + def test_remote_judger_preserves_driver_side_judger_contract(self): from xtuner.v1.rl.judger import Judger, RemoteJudger - class CustomPreprocessJudger(Judger): + class CustomJudger(Judger): def preprocess(self, rollout_state): return {"custom_value": rollout_state.extra_fields["custom_value"]} + def postprocess(self, rollout_state, output): + rollout_state.reward = { + "score": output["raw_score"], + "postprocessed": True, + } + return rollout_state + class RemoteMethod: def __init__(self): self.payload = None async def remote(self, payload): self.payload = payload - return {"score": payload["custom_value"]} + return {"raw_score": payload["custom_value"]} class FakeActor: def __init__(self): @@ -162,8 +169,8 @@ def __init__(self): actor = FakeActor() judger = RemoteJudger( actor=actor, - judger_name="remote_custom_preprocess", - preprocess_judger=CustomPreprocessJudger(), + judger_name="remote_custom_contract", + preprocess_judger=CustomJudger(), ) rollout_state = self._make_rollout_state("correctness") rollout_state.extra_fields["custom_value"] = 7 @@ -171,7 +178,84 @@ def __init__(self): judged_state = asyncio.run(judger.judge(rollout_state)) self.assertEqual(actor.judge_payload.payload, {"custom_value": 7}) - self.assertEqual(judged_state.reward, {"score": 7}) + self.assertEqual(judged_state.reward, {"score": 7, "postprocessed": True}) + + def test_composed_judger_pool_preserves_branch_contract_and_routing(self): + from xtuner.v1.rl.judger import ComposedJudger, Judger, JudgerPool, RemoteJudger + + class CustomJudger(Judger): + def preprocess(self, rollout_state): + return { + "response": rollout_state.response, + "reward_model": rollout_state.reward_model, + "finish_reason": rollout_state.finish_reason, + "extra_fields": rollout_state.extra_fields, + } + + def postprocess(self, rollout_state, output): + rollout_state.reward = { + "score": output["raw_score"], + "postprocessed": True, + } + return rollout_state + + class RemoteMethod: + def __init__(self): + self.payload = None + + async def remote(self, payload): + self.payload = payload + return {"raw_score": payload["extra_fields"]["expected_score"]} + + class FakeActor: + def __init__(self): + self.judge_payload = RemoteMethod() + + branch_judger = CustomJudger() + actors = [FakeActor(), FakeActor()] + replicas = [ + RemoteJudger( + actor=actor, + judger_name="remote_custom_contract", + preprocess_judger=branch_judger, + ) + for actor in actors + ] + judger = ComposedJudger( + branches={ + "biology": JudgerPool( + replicas=replicas, + judger_name="biology", + ) + } + ) + rollout_states = [self._make_rollout_state("biology"), self._make_rollout_state("biology")] + for expected_score, rollout_state in enumerate(rollout_states, start=1): + rollout_state.finish_reason = "stop" + rollout_state.extra_fields = { + "expected_score": expected_score, + "task_name": "MCC", + } + + async def judge_all(): + return [await judger.judge(state) for state in rollout_states] + + judged_states = asyncio.run(judge_all()) + + for expected_score, (actor, state) in enumerate(zip(actors, judged_states), start=1): + self.assertEqual( + actor.judge_payload.payload, + { + "response": state.response, + "reward_model": state.reward_model, + "finish_reason": "stop", + "extra_fields": { + "expected_score": expected_score, + "task_name": "MCC", + }, + }, + ) + self.assertEqual(state.reward, {"score": expected_score, "postprocessed": True}) def test_composed_judger_config(self): def merge_fn(original, judged): diff --git a/xtuner/v1/rl/judger/native.py b/xtuner/v1/rl/judger/native.py index 9804d720f..18cae91fe 100644 --- a/xtuner/v1/rl/judger/native.py +++ b/xtuner/v1/rl/judger/native.py @@ -229,11 +229,11 @@ class RemoteJudger(Judger): ``RemoteJudger`` keeps the same ``Judger`` interface as local judgers, so callers still pass ``RolloutState`` to ``judge``. This proxy runs the same - ``preprocess`` implementation as the actor-side judger on the driver, then - sends only the lightweight payload to ``JudgerActor``. ``JudgerActor`` lives - in the Ray worker process and owns the real local judger instance that - executes ``judge_payload``. Batch support is determined by that actor-side - judger. + ``preprocess`` and ``postprocess`` implementations as the actor-side judger + on the driver, and sends only the lightweight payload to ``JudgerActor``. + ``JudgerActor`` lives in the Ray worker process and owns the real local + judger instance that executes ``judge_payload``. Batch support is determined + by that actor-side judger. """ def __init__(self, actor: RayJudgerProxy, judger_name: str, preprocess_judger: Judger | None = None): @@ -244,17 +244,28 @@ def __init__(self, actor: RayJudgerProxy, judger_name: str, preprocess_judger: J # Preprocess must run on the driver before the Ray call. Otherwise the full # RolloutState would be serialized to the actor, and remote branches with # custom preprocess logic would lose the payload fields they require. + # Postprocess also runs on the driver because it writes the output back to + # the original RolloutState retained by the caller. def preprocess(self, rollout_state: RolloutState) -> JudgerPayload: if self.preprocess_judger is None: return super().preprocess(rollout_state) return self.preprocess_judger.preprocess(rollout_state) + def postprocess(self, rollout_state: RolloutState, output: JudgerOutput) -> RolloutState: + if self.preprocess_judger is None: + return super().postprocess(rollout_state, output) + return self.preprocess_judger.postprocess(rollout_state, output) + async def judge_payload(self, payload: JudgerPayloadBatch) -> JudgerOutputBatch: return await self.actor.judge_payload.remote(payload) class JudgerPool(Judger): - """Round-robin dispatch across replicas of the same judger type.""" + """Round-robin dispatch across replicas of the same judger type. + + Replicas are homogeneous and therefore share the same payload contract. Preprocessing and postprocessing are + delegated to a representative replica, while payload execution is dispatched round-robin across all replicas. + """ def __init__(self, replicas: list[Judger], judger_name: str): super().__init__(judger_name=judger_name) @@ -265,6 +276,12 @@ def __init__(self, replicas: list[Judger], judger_name: str): self._lock = asyncio.Lock() self._worker_loads = dict.fromkeys(range(len(replicas)), 0) + def preprocess(self, rollout_state: RolloutState) -> JudgerPayload: + return self.replicas[0].preprocess(rollout_state) + + def postprocess(self, rollout_state: RolloutState, output: JudgerOutput) -> RolloutState: + return self.replicas[0].postprocess(rollout_state, output) + async def _pick_replica(self) -> tuple[int, Judger]: async with self._lock: replica_idx = self._rr_index % len(self.replicas)