fix(trace_store): exclude raw training arrays from saved episodes

tests/unit/test_trace.py

@@ -102,6 +102,54 @@ def test_custom_task_state_round_trip(tmp_path):
     assert tr.state.artifacts["/artifact.bin"] == b"\xff\x00\xfe"
 
 
+@pytest.mark.parametrize(
+    "arrays",
+    [
+        {"routed_experts": {"data": "//8=", "shape": [2, 1, 1]}},  # bytes ff ff
+        {"sampling_mask": [[128]]},
+        {"completion_top_ids": [[20]], "completion_top_logprobs": [[-0.5]]},
+        {"sampling_mask": [[1]]},  # ASCII-safe bytes: writes, then fails to reload
+    ],
+    ids=["routed_experts", "sampling_mask", "top_logprobs", "ascii_sampling_mask"],
+)
+def test_saved_episode_omits_raw_training_arrays(tmp_path, arrays):
+    """Raw training arrays stay on the msgpack wire; the saved episode drops them like
+    `Episode.to_record` does, so a train-client eval can save and reload its traces."""
+    from verifiers.v1.clients.train import response_from_generate
+    from verifiers.v1.utils.trace_store import read_episodes
+
+    tr = vf.Trace(
+        agent=vf.AgentInfo(config=vf.AgentConfig()),
+        task=vf.TraceTask(type="MyTask", data=MyTask(idx=0, prompt="q", answer="gold")),
+    )
+    prepare_turn(tr, [UserMessage(content="q")]).commit(
+        response_from_generate(
+            {
+                "content": "a",
+                "finish_reason": "stop",
+                "prompt_ids": [10],
+                "completion_ids": [20],
+                "completion_logprobs": [-0.5],
+                **arrays,
+            },
+            model="test",
+        )
+    )
+
+    write_episode(tmp_path, vf.Episode(task=tr.task, traces=[tr], ok=True))
+    [episode] = read_episodes(tmp_path, vf.Trace[MyTask, vf.State])
+
+    [saved] = episode.traces
+    assert saved.task.data.answer == "gold"
+    assert [m.content for m in saved.branches[0].messages] == ["q", "a"]
+    assert saved.nodes[-1].token_ids.tolist() == [10, 20]
+    assert saved.nodes[-1].logprobs.tolist() == [-0.5]
+    for node in saved.nodes:
+        assert node.routed_experts is None
+        assert node.sampling_mask is None
+        assert node.top_logprobs is None
+
+
 @pytest.mark.parametrize("history_type", ["additional_tools", "tool_search_output"])
 def test_wire_trace_round_trip(history_type):
     mcp = {

verifiers/v1/cli/eval/main.py

@@ -23,6 +23,7 @@
     with_positional_taskset,
 )
 from verifiers.v1.configs.cli.eval import EvalConfig
+from verifiers.v1.trace import EXCLUDE_FIELDS
 from verifiers.v1.utils.interrupt import install_interrupt
 from verifiers.v1.utils.logging import setup_logging
 from verifiers.v1.utils.trace_store import TRACES_FILE
@@ -137,4 +138,8 @@ def main(argv: list[str] | None = None) -> None:
     ):  # --rich is the whole output; otherwise dump each trace as JSON
         for episode in episodes:
             for trace in episode.traces:
-                print(trace.model_dump_json(indent=2, exclude_none=True))
+                print(
+                    trace.model_dump_json(
+                        indent=2, exclude=EXCLUDE_FIELDS, exclude_none=True
+                    )
+                )

verifiers/v1/utils/trace_store.py

@@ -12,7 +12,7 @@
 from verifiers.v1.episode import EnvInfo, Episode, WireEpisode
 from verifiers.v1.state import StateT
 from verifiers.v1.task import DataT
-from verifiers.v1.trace import AgentConfigT, Trace
+from verifiers.v1.trace import EXCLUDE_FIELDS, AgentConfigT, Trace
 from verifiers.v1.utils.aio import run_shielded
 from verifiers.v1.utils.artifacts import save_artifacts
 
@@ -24,8 +24,13 @@ def write_episode(
     results_dir: Path, episode: Episode[DataT, StateT, AgentConfigT]
 ) -> None:
     """Serialize and append one rollout episode in the worker thread."""
-    # Preserve fields declared by typed Trace subclasses nested in the episode.
-    data = type_adapter(type(episode)).dump_json(episode, exclude_none=True)
+    # Preserve fields declared by typed Trace subclasses nested in the episode. Raw tensors
+    # stay on the msgpack wire only, as in `Episode.to_record`: they don't survive JSON.
+    data = type_adapter(type(episode)).dump_json(
+        episode,
+        exclude={"traces": {"__all__": EXCLUDE_FIELDS}},
+        exclude_none=True,
+    )
     for trace in episode.traces:
         if trace.state.artifacts:
             if trace.id == ".." or Path(trace.id).parts != (trace.id,):