diff --git a/deps/verifiers b/deps/verifiers index 306a95fc0e..b43c355ac9 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 306a95fc0ede876f17889a28c36755d2a922d271 +Subproject commit b43c355ac9c0747c12c97298d699ea43084c541c diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index d79de8e6cd..0123241fe3 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -670,7 +670,7 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: metrics[f"pre_filters/all/{name}/rate"] = count / self.train_sink.pre_filter_seen self.monitor.log(metrics, step=step) self.wait_for_policy_time = 0.0 - self.monitor.log_samples(effective.rollouts, step=step) + self.monitor.log_samples(effective.episodes, step=step) self.monitor.log_distributions( distributions={ "rewards": [r.reward for r in effective], diff --git a/src/prime_rl/utils/monitor/prime.py b/src/prime_rl/utils/monitor/prime.py index 468a63653d..6be0bbdc3b 100644 --- a/src/prime_rl/utils/monitor/prime.py +++ b/src/prime_rl/utils/monitor/prime.py @@ -15,7 +15,7 @@ import pyarrow.parquet as pq from prime_cli.core.config import Config as PrimeConfig from transformers.tokenization_utils import PreTrainedTokenizer -from verifiers.v1.utils.platform import trace_to_sample +from verifiers.v1.utils.platform import build_samples from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.configs.shared import PrimeMonitorConfig @@ -249,7 +249,7 @@ def log(self, metrics: dict[str, Any], step: int) -> None: ) def log_samples(self, episodes: list[Episode], step: int) -> None: - """Logs rollouts to Prime Intellect API using presigned URLs for direct R2 upload.""" + """Log episodes to Prime Intellect API using presigned URLs for direct R2 upload.""" if not self.is_master: return if not self.enabled: @@ -276,7 +276,11 @@ def log_samples(self, episodes: list[Episode], step: int) -> None: self.logger.info(f"Logging {len(episodes)} episodes to Prime Intellect API at step {step}") start_time = time.perf_counter() - parquet_bytes = self._rollouts_to_parquet_bytes(episodes, step) + try: + parquet_bytes = self._episodes_to_parquet_bytes(episodes, step) + except Exception as e: + self.logger.warning(f"Failed to build Prime monitor samples at step {step}: {type(e).__name__}: {e}") + return if not parquet_bytes: self.logger.warning(f"No samples to log at step {step}") @@ -291,56 +295,60 @@ def log_samples(self, episodes: list[Episode], step: int) -> None: f"Initiated samples upload at step {step} to Prime Intellect API in {time.perf_counter() - start_time:.2f}s" ) - def _rollouts_to_parquet_bytes(self, episodes: list[Episode], step: int) -> bytes | None: - """Convert episodes to Parquet bytes for upload. One row per trace, carrying the episode - it belongs to. The conversation - is the unit (no prompt/completion split — meaningless mid-branch): `completion` is the - last branch's messages and `trajectory` is one message list per branch. Shares - `verifiers.v1.utils.platform.trace_to_sample` with verifiers' eval `--push`, so a training-run - sample and an eval sample land on the platform identically; the RFT-only columns - (run/step/advantage/problem_id/env_name) are layered on here.""" + def _episodes_to_parquet_bytes(self, episodes: list[Episode], step: int) -> bytes | None: + """Convert episodes to Parquet using verifiers' canonical Platform samples. + + A normal Episode produces one row whose ``info.native_wrapper`` is authoritative; the + selected trace only supplies legacy flat columns. Oversized Episodes retain verifiers' + per-trace projection fallback. Prime-RL layers its run, step, advantage, problem, and env + columns on either representation. + """ now = datetime.now(timezone.utc) rows = [] - for sample_id, (episode, rollout) in enumerate( - (episode, trace) for episode in episodes for trace in episode.traces - ): - sample = trace_to_sample(rollout, rollout_number=sample_id + 1, episode_id=episode.id) - trajectory = sample["trajectory"] - if not trajectory: # no branches (e.g. a rollout that errored before any message) - continue - advantage = rollout.scalar_advantage() - trajectory = [{**branch, "advantage": advantage} for branch in trajectory] - - example_id = sample["example_id"] - try: - problem_id = int(example_id) if example_id is not None else sample_id - except (TypeError, ValueError): - problem_id = sample_id - - rows.append( - { - "run_id": self.run_id, - "step": step, - "tag": "", - "problem_id": problem_id, - "sample_id": sample_id, - "prompt": "", - "completion": json.dumps(sample["completion"]), - "trajectory": json.dumps(trajectory), - "answer": "", - "env_name": episode.env.name or "", - "task": json.dumps(sample["task"]), - "info": json.dumps(rollout.info), - "reward": sample["reward"], - "advantage": advantage, - "metrics": json.dumps(sample["metrics"]), - "timing": json.dumps(sample["timing"]), - "num_input_tokens": trajectory[-1]["num_input_tokens"], - "num_output_tokens": trajectory[-1]["num_output_tokens"], - "created_at": now, - } - ) + for episode in episodes: + for sample in build_samples([episode]): + info = sample["info"] or {} + summary_trace_index = info.get("native_trace_index") + if isinstance(summary_trace_index, int): + rollout = episode.rollouts[summary_trace_index] + else: + rollout = next(trace for trace in episode.rollouts if trace.id == sample["sample_id"]) + + sample_id = len(rows) + trajectory = sample["trajectory"] + advantage = rollout.scalar_advantage() + trajectory = [{**branch, "advantage": advantage} for branch in trajectory] + + example_id = sample["example_id"] + try: + problem_id = int(example_id) if example_id is not None else sample_id + except (TypeError, ValueError): + problem_id = sample_id + + rows.append( + { + "run_id": self.run_id, + "step": step, + "tag": "", + "problem_id": problem_id, + "sample_id": sample_id, + "prompt": "", + "completion": json.dumps(sample["completion"]), + "trajectory": json.dumps(trajectory), + "answer": "", + "env_name": episode.env.name or "", + "task": json.dumps(sample["task"]), + "info": json.dumps(info), + "reward": sample["reward"], + "advantage": advantage, + "metrics": json.dumps(sample["metrics"]), + "timing": json.dumps(sample["timing"]), + "num_input_tokens": trajectory[-1]["num_input_tokens"] if trajectory else 0, + "num_output_tokens": trajectory[-1]["num_output_tokens"] if trajectory else 0, + "created_at": now, + } + ) if not rows: return None diff --git a/tests/unit/utils/test_prime_monitor.py b/tests/unit/utils/test_prime_monitor.py index 3456d0f947..6c838a2b3a 100644 --- a/tests/unit/utils/test_prime_monitor.py +++ b/tests/unit/utils/test_prime_monitor.py @@ -6,7 +6,7 @@ import verifiers.v1 as vf from verifiers.v1.configs.agent import WireAgentConfig -from prime_rl.orchestrator.types import Rollout +from prime_rl.orchestrator.types import Episode, Rollout from prime_rl.utils.monitor.prime import PrimeMonitor @@ -16,9 +16,11 @@ def _new_monitor() -> PrimeMonitor: return monitor -def _build_rollout(*, example_id: int, reward: float, task: str) -> Rollout: +def _build_rollout( + *, example_id: int, reward: float, task: str, agent_name: str = "agent", trainable: bool = True +) -> Rollout: """Build a v1 ``Rollout`` (message-graph trace). The user node carries the prompt and the - assistant node the completion; ``_rollouts_to_parquet_bytes`` reads the conversation off the + assistant node the completion; ``_episodes_to_parquet_bytes`` reads the conversation off the branches (its ``completion`` column is the last branch's messages, ``trajectory`` is one message list per branch).""" nodes = [ @@ -38,7 +40,7 @@ def _build_rollout(*, example_id: int, reward: float, task: str) -> Rollout: ] rollout = Rollout[vf.TaskData]( task=vf.TraceTask(type="Task", data=vf.TaskData(idx=example_id, prompt=f"prompt-{example_id}")), - agent=vf.AgentInfo(config=WireAgentConfig()), + agent=vf.AgentInfo(config=WireAgentConfig(), name=agent_name, trainable=trainable), nodes=nodes, rewards={"reward": vf.Reward(score=reward)}, ) @@ -47,19 +49,32 @@ def _build_rollout(*, example_id: int, reward: float, task: str) -> Rollout: return rollout -def _episode(*rollouts: Rollout, env_name: str = "task-a") -> vf.WireEpisode: - """The unit the monitor uploads: one episode, whose traces become the rows.""" - return vf.WireEpisode.model_construct(traces=list(rollouts), env=vf.EnvInfo(name=env_name)) +def _episode(*rollouts: Rollout, episode_id: str, env_name: str = "task-a") -> Episode: + """The unit the monitor uploads: one episode, represented by one native row.""" + return Episode.model_construct( + id=episode_id, + traces=list(rollouts), + env=vf.EnvInfo(id=f"{env_name}-v1", name=env_name), + run=vf.TrainRunInfo(id="training-run", metadata=vf.TrainMetadata(step=7)), + ok=True, + ) -def test_rollouts_to_parquet_bytes_preserves_all_rollouts_and_ids(): +def test_episodes_to_parquet_bytes_preserves_episode_rows_and_ids(): monitor = _new_monitor() monitor.run_id = "run-123" - parquet_bytes = monitor._rollouts_to_parquet_bytes( + parquet_bytes = monitor._episodes_to_parquet_bytes( [ - _episode(_build_rollout(example_id=101, reward=1.0, task="task-a")), - _episode(_build_rollout(example_id=202, reward=0.0, task="task-b"), env_name="task-b"), + _episode( + _build_rollout(example_id=101, reward=1.0, task="task-a"), + episode_id="episode-101", + ), + _episode( + _build_rollout(example_id=202, reward=0.0, task="task-b"), + episode_id="episode-202", + env_name="task-b", + ), ], step=7, ) @@ -78,21 +93,37 @@ def test_rollouts_to_parquet_bytes_preserves_all_rollouts_and_ids(): assert json.loads(rows[1]["completion"])[0]["content"] == "completion-202" trajectory = json.loads(rows[0]["trajectory"]) assert trajectory[0]["messages"][0]["content"] == "prompt-101" + infos = [json.loads(row["info"]) for row in rows] + assert [info["native_wrapper"]["id"] for info in infos] == ["episode-101", "episode-202"] + assert all(info["native_trace_index"] == 0 for info in infos) + assert all(len(info["native_wrapper"]["traces"]) == 1 for info in infos) + assert infos[0]["native_wrapper"]["run"]["id"] == "training-run" -def test_rollouts_to_parquet_bytes_skips_rollouts_without_trajectory(): +def test_episodes_to_parquet_bytes_uses_trainable_summary_and_preserves_all_traces(): monitor = _new_monitor() monitor.run_id = "run-456" - rollout_with_branches = _build_rollout(example_id=1, reward=1.0, task="task-a") - rollout_without_branches = Rollout[vf.TaskData]( + fixed_trace_without_branches = Rollout[vf.TaskData]( task=vf.TraceTask(type="Task", data=vf.TaskData(idx=2, prompt="missing-trajectory")), - agent=vf.AgentInfo(config=WireAgentConfig()), + agent=vf.AgentInfo(config=WireAgentConfig(), name="judge", trainable=False), + ) + trainable_trace = _build_rollout( + example_id=3, + reward=1.0, + task="task-a", + agent_name="solver", ) - assert rollout_without_branches.branches == [] + assert fixed_trace_without_branches.branches == [] - parquet_bytes = monitor._rollouts_to_parquet_bytes( - [_episode(rollout_with_branches, rollout_without_branches)], + parquet_bytes = monitor._episodes_to_parquet_bytes( + [ + _episode( + fixed_trace_without_branches, + trainable_trace, + episode_id="multi-trace-episode", + ) + ], step=3, ) @@ -102,8 +133,18 @@ def test_rollouts_to_parquet_bytes_skips_rollouts_without_trajectory(): rows = table.to_pylist() assert len(rows) == 1 - assert rows[0]["problem_id"] == 1 + assert rows[0]["problem_id"] == 3 assert rows[0]["sample_id"] == 0 + assert rows[0]["reward"] == 1.0 + assert rows[0]["advantage"] == 0.5 + assert json.loads(rows[0]["completion"])[0]["content"] == "completion-3" + info = json.loads(rows[0]["info"]) + assert info["native_trace_index"] == 1 + assert [trace["agent"]["name"] for trace in info["native_wrapper"]["traces"]] == [ + "judge", + "solver", + ] + assert info["native_wrapper"]["traces"][0]["nodes"] == [] def test_sanitize_json_payload_drops_non_finite_values_and_logs_paths(): @@ -121,3 +162,25 @@ def test_sanitize_json_payload_drops_non_finite_values_and_logs_paths(): monitor.logger.warning.assert_called_once_with( "Dropping 2 non-finite value(s) from Prime monitor metrics payload: metrics.nan, distributions[1]" ) + + +def test_log_samples_warns_and_skips_upload_when_serialization_fails(): + monitor = _new_monitor() + monitor.is_master = True + monitor.enabled = True + monitor.config = Mock() + monitor.config.log_extras.samples = True + monitor.config.log_extras.interval = 1 + monitor.config.log_extras.sample_ratio = 1.0 + monitor.last_log_samples_step = -1 + monitor._pending_sample_steps = set() + monitor.logger = Mock() + monitor._episodes_to_parquet_bytes = Mock(side_effect=ValueError("invalid episode")) + monitor._upload_samples_via_presigned_url = Mock() + + monitor.log_samples([Mock()], step=1) + + monitor.logger.warning.assert_called_once_with( + "Failed to build Prime monitor samples at step 1: ValueError: invalid episode" + ) + monitor._upload_samples_via_presigned_url.assert_not_called()