diff --git a/tests/unit/test_trace.py b/tests/unit/test_trace.py index 7796ffdbfb..13b96ef094 100644 --- a/tests/unit/test_trace.py +++ b/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 = { diff --git a/verifiers/v1/cli/eval/main.py b/verifiers/v1/cli/eval/main.py index e66101a8bb..bfde1c3c6d 100644 --- a/verifiers/v1/cli/eval/main.py +++ b/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 + ) + ) diff --git a/verifiers/v1/utils/trace_store.py b/verifiers/v1/utils/trace_store.py index febb7a2c5b..a74a071431 100644 --- a/verifiers/v1/utils/trace_store.py +++ b/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,):