feat(monitor): upload native v1 episodes to Platform

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],

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

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()