diff --git a/src/prime_rl/orchestrator/envs.py b/src/prime_rl/orchestrator/envs.py index a0c86e1908..44910bbd37 100644 --- a/src/prime_rl/orchestrator/envs.py +++ b/src/prime_rl/orchestrator/envs.py @@ -130,8 +130,10 @@ async def run( detail = f"{error.type}: {error.message}" if error is not None else "no traces and no error recorded" raise RuntimeError(f"episode failed before any trace was produced — {detail}") rollouts = [ROLLOUT_TYPE.model_construct(**dict(wire)) for wire in episode.traces] + native_episode = episode.model_copy(update={"traces": rollouts}) for rollout in rollouts: rollout.episode_id = episode.id + rollout.native_episode = native_episode if not episode.ok and rollout.ok: error = episode.last_error or vf.Error( type="EpisodeFailed", message="A sibling trace in this episode failed" diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 65c56d3181..4d45404221 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -94,6 +94,9 @@ class Rollout(vf.Trace[DataT], Generic[DataT]): # Links the traces of one episode; stamped into ``info`` on arrival so # saved records keep their grouping. episode_id: str = Field(default="", exclude=True) + # The original v1 envelope, retained only in memory so monitors can upload the + # complete multi-trace Episode without making Episode the orchestrator's unit. + native_episode: vf.WireEpisode | None = Field(default=None, exclude=True, repr=False) policy_version: int = Field(default=0, exclude=True) off_policy_steps: int = Field(default=0, exclude=True) samples: list[TrainingSample] = Field(default_factory=list, exclude=True) diff --git a/src/prime_rl/utils/monitor/prime.py b/src/prime_rl/utils/monitor/prime.py index eee454a9aa..82c5f254ec 100644 --- a/src/prime_rl/utils/monitor/prime.py +++ b/src/prime_rl/utils/monitor/prime.py @@ -8,14 +8,16 @@ from datetime import datetime, timezone from pathlib import Path from threading import Thread -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import httpx import pyarrow as pa import pyarrow.parquet as pq +import verifiers.v1 as vf 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.episode import EnvInfo +from verifiers.v1.utils.platform import build_samples from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.configs.shared import PrimeMonitorConfig @@ -262,21 +264,25 @@ def log_samples(self, rollouts: list[Rollout], step: int) -> None: ): return - rollouts = sample_items_for_logging( - rollouts, + episodes = sample_items_for_logging( + self._rollouts_to_episodes(rollouts), self.config.log_extras.sample_ratio, ) - if not rollouts: + if not episodes: return assert self.last_log_samples_step <= step, "Step must be greater than last logged step" assert step not in self._pending_sample_steps, f"Step {step} upload already in progress" assert self.logger is not None, "Logger is required for sample logging" - self.logger.info(f"Logging {len(rollouts)} samples to Prime Intellect API at step {step}") + 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(rollouts, 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,53 +297,86 @@ def log_samples(self, rollouts: list[Rollout], 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, rollouts: list[Rollout], step: int) -> bytes | None: - """Convert rollouts to Parquet bytes for upload. One row per rollout. 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.""" + @staticmethod + def _rollouts_to_episodes(rollouts: list[Rollout]) -> list[vf.WireEpisode]: + """Recover each rollout's original v1 envelope without changing orchestration. + + Legacy/group paths do not return an Episode envelope, so those remain compatible as + one-trace Episodes. Multiple effective traces from one native Episode share the same + in-memory envelope and are emitted only once. + """ + episodes: list[vf.WireEpisode] = [] + seen_native_envelopes: set[int] = set() + for rollout in rollouts: + episode = rollout.native_episode + if episode is None: + episodes.append( + vf.WireEpisode.model_construct( + id=rollout.episode_id or rollout.id, + env=EnvInfo(id=rollout.env_name), + ok=rollout.ok, + errors=list(rollout.errors), + traces=[rollout], + ) + ) + continue + envelope_identity = id(episode) + if envelope_identity not in seen_native_envelopes: + seen_native_envelopes.add(envelope_identity) + episodes.append(episode) + return episodes + + def _episodes_to_parquet_bytes(self, episodes: list[vf.WireEpisode], step: int) -> bytes | None: + """Convert native Episodes to the existing training sample Parquet schema.""" now = datetime.now(timezone.utc) rows = [] - for sample_id, rollout in enumerate(rollouts): - sample = trace_to_sample(rollout, rollout_number=sample_id + 1, episode_id=rollout.episode_id or None) - 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] + 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 = cast("Rollout", episode.traces[summary_trace_index]) + else: + rollout = cast( + "Rollout", + next(trace for trace in episode.traces if trace.id == sample["sample_id"]), + ) - 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": rollout.env_name, - "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, - } - ) + 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": rollout.env_name or episode.env.id, + "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 798e041885..269c00998f 100644 --- a/tests/unit/utils/test_prime_monitor.py +++ b/tests/unit/utils/test_prime_monitor.py @@ -4,6 +4,7 @@ import pyarrow.parquet as pq import verifiers.v1 as vf +from verifiers.v1.episode import EnvInfo from prime_rl.orchestrator.types import Rollout from prime_rl.utils.monitor.prime import PrimeMonitor @@ -15,11 +16,10 @@ def _new_monitor() -> PrimeMonitor: return monitor -def _build_rollout(*, example_id: int, reward: float, task: str) -> 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 - branches (its ``completion`` column is the last branch's messages, ``trajectory`` is one - message list per branch).""" +def _build_rollout( + *, example_id: int, reward: float, task: str, agent_name: str = "agent", trainable: bool = True +) -> Rollout: + """Build a v1 ``Rollout`` with one user/assistant message branch.""" nodes = [ vf.MessageNode( message=vf.UserMessage(content=f"prompt-{example_id}"), @@ -37,7 +37,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=vf.AgentConfig()), + agent=vf.AgentInfo(config=vf.AgentConfig(), name=agent_name, trainable=trainable), nodes=nodes, rewards={"reward": vf.Reward(score=reward)}, ) @@ -48,17 +48,37 @@ def _build_rollout(*, example_id: int, reward: float, task: str) -> Rollout: return rollout -def test_rollouts_to_parquet_bytes_preserves_all_rollouts_and_ids(): +def _attach_episode( + *rollouts: Rollout, + episode_id: str, + env_id: str = "task-a-v1", + ok: bool = True, + errors: list[vf.Error] | None = None, +) -> vf.WireEpisode: + episode = vf.WireEpisode.model_construct( + id=episode_id, + env=EnvInfo(id=env_id), + ok=ok, + errors=errors or [], + traces=list(rollouts), + ) + for rollout in rollouts: + rollout.episode_id = episode_id + rollout.native_episode = episode + return episode + + +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( - [ - _build_rollout(example_id=101, reward=1.0, task="task-a"), - _build_rollout(example_id=202, reward=0.0, task="task-b"), - ], - step=7, - ) + first = _build_rollout(example_id=101, reward=1.0, task="task-a") + second = _build_rollout(example_id=202, reward=0.0, task="task-b") + _attach_episode(first, episode_id="episode-101") + _attach_episode(second, episode_id="episode-202", env_id="task-b-v1") + + episodes = monitor._rollouts_to_episodes([first, second]) + parquet_bytes = monitor._episodes_to_parquet_bytes(episodes, step=7) assert parquet_bytes is not None @@ -74,23 +94,45 @@ 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) -def test_rollouts_to_parquet_bytes_skips_rollouts_without_trajectory(): +def test_effective_rollout_retains_fixed_and_errored_sibling_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 = Rollout[vf.TaskData]( task=vf.TraceTask(type="Task", data=vf.TaskData(idx=2, prompt="missing-trajectory")), - agent=vf.AgentInfo(config=vf.AgentConfig()), + agent=vf.AgentInfo(config=vf.AgentConfig(), name="judge", trainable=False), ) - assert rollout_without_branches.branches == [] - - parquet_bytes = monitor._rollouts_to_parquet_bytes( - [rollout_with_branches, rollout_without_branches], - step=3, + trainable_trace = _build_rollout( + example_id=3, + reward=1.0, + task="task-a", + agent_name="solver", + ) + error = vf.Error(type="JudgeError", message="judge failed after scoring") + failed_trace = Rollout[vf.TaskData]( + task=vf.TraceTask(type="Task", data=vf.TaskData(idx=3, prompt="prompt-3")), + agent=vf.AgentInfo(config=vf.AgentConfig(), name="critic", trainable=False), + errors=[error], ) + _attach_episode( + fixed_trace, + trainable_trace, + failed_trace, + episode_id="multi-trace-episode", + ok=False, + errors=[error], + ) + + # The orchestrator still passes only effective rollouts. Its retained native envelope + # carries the fixed and errored siblings without changing the rollout-based pipeline. + episodes = monitor._rollouts_to_episodes([trainable_trace]) + parquet_bytes = monitor._episodes_to_parquet_bytes(episodes, step=3) assert parquet_bytes is not None @@ -98,8 +140,44 @@ 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 + info = json.loads(rows[0]["info"]) + assert info["native_trace_index"] == 1 + assert info["native_wrapper"]["ok"] is False + assert info["native_wrapper"]["errors"] == [{"type": "JudgeError", "message": "judge failed after scoring"}] + assert [trace["agent"]["name"] for trace in info["native_wrapper"]["traces"]] == [ + "judge", + "solver", + "critic", + ] + + +def test_rollout_without_native_envelope_uses_single_trace_fallback(): + monitor = _new_monitor() + monitor.run_id = "run-legacy" + rollout = _build_rollout(example_id=4, reward=0.25, task="legacy-env") + + episodes = monitor._rollouts_to_episodes([rollout]) + + assert len(episodes) == 1 + assert episodes[0].id == rollout.id + assert episodes[0].env.id == "legacy-env" + assert episodes[0].traces == [rollout] + assert "native_episode" not in rollout.model_dump() + + +def test_legacy_rollouts_with_same_episode_id_are_not_deduplicated(): + first = _build_rollout(example_id=4, reward=0.25, task="legacy-env") + second = _build_rollout(example_id=4, reward=0.5, task="legacy-env") + first.episode_id = second.episode_id = "legacy-group" + + episodes = PrimeMonitor._rollouts_to_episodes([first, second]) + + assert len(episodes) == 2 + assert [episode.traces for episode in episodes] == [[first], [second]] def test_sanitize_json_payload_drops_non_finite_values_and_logs_paths(): @@ -117,3 +195,26 @@ 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() + rollout = _build_rollout(example_id=5, reward=1.0, task="task-a") + + monitor.log_samples([rollout], 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()