From cfe7cc7fb2b142c7d7764218997b36f9bec8ad16 Mon Sep 17 00:00:00 2001 From: William Brown Date: Tue, 27 Jan 2026 17:15:52 -0800 Subject: [PATCH 1/6] serializable output type; generateoutputs builder --- docs/reference.md | 27 +++- tests/conftest.py | 21 +++ tests/test_env_group.py | 2 +- tests/test_environment.py | 31 ++-- tests/test_environment_extra.py | 12 +- tests/test_eval_cli.py | 4 +- tests/test_eval_utils.py | 12 +- tests/test_gym_env.py | 32 ++-- tests/test_save_utils.py | 10 +- tests/test_singleturn_env.py | 8 +- verifiers/envs/environment.py | 86 ++++------- verifiers/gepa/adapter.py | 44 +++--- verifiers/rl/trainer/orchestrator.py | 21 +-- verifiers/types.py | 38 ++++- verifiers/utils/eval_display.py | 14 +- verifiers/utils/eval_utils.py | 80 +++++----- verifiers/utils/save_utils.py | 219 ++++++++++++++++++++------- 17 files changed, 422 insertions(+), 239 deletions(-) diff --git a/docs/reference.md b/docs/reference.md index 7913e98f62..0d43d7aca3 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -116,6 +116,29 @@ class RolloutInput(TypedDict): info: Info # Optional ``` +### RolloutOutput + +```python +class RolloutOutput(TypedDict): + # Required fields + example_id: int + task: str + prompt: Messages | None + completion: Messages | None + reward: float + timing: RolloutTiming + is_completed: bool + is_truncated: bool + metrics: dict[str, float] + # Optional fields + answer: str + info: Info + error: str | None + stop_condition: str | None +``` + +Serialized output from a rollout. This is the JSON-serializable representation of a `State` object, used in `GenerateOutputs` and for saving results to disk. Additional fields from `state_columns` may also be present. + ### TrajectoryStep ```python @@ -162,11 +185,11 @@ class RolloutTiming(TypedDict, total=False): ```python class GenerateOutputs(TypedDict): - states: list[State] + outputs: list[RolloutOutput] metadata: GenerateMetadata ``` -Output from `Environment.generate()`. Contains a list of `State` objects (one per rollout) and generation metadata. Each `State` contains the rollout's prompt, completion, answer, reward, metrics, and other per-rollout data. +Output from `Environment.generate()`. Contains a list of `RolloutOutput` objects (one per rollout) and generation metadata. Each `RolloutOutput` is a serialized, JSON-compatible dict containing the rollout's prompt, completion, answer, reward, metrics, timing, and other per-rollout data. ### GenerateMetadata diff --git a/tests/conftest.py b/tests/conftest.py index 453cdf97a1..c08869546f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -27,10 +27,12 @@ GenerateMetadata, Info, RolloutInput, + RolloutOutput, RolloutTiming, SamplingArgs, TrajectoryStep, ) +from verifiers.utils.save_utils import state_to_output @pytest.fixture @@ -492,6 +494,25 @@ def _make_state( return _make_state +@pytest.fixture +def make_output(make_state) -> Callable[..., RolloutOutput]: + """Fixture to make RolloutOutput objects for testing. + + This creates a State first, then converts it to a RolloutOutput using + state_to_output(). This ensures the output matches the serialized format + used in GenerateOutputs. + """ + + def _make_output( + state_columns: list[str] = ["foo"], + **kwargs, + ) -> RolloutOutput: + state = make_state(**kwargs) + return state_to_output(state, state_columns) + + return _make_output + + @pytest.fixture def make_metadata() -> Callable[..., GenerateMetadata]: """Fixture to make GenerateMetadata objects for testing.""" diff --git a/tests/test_env_group.py b/tests/test_env_group.py index 0eea1d2c92..63ecf206e4 100644 --- a/tests/test_env_group.py +++ b/tests/test_env_group.py @@ -406,7 +406,7 @@ async def mock_score_group(states, score_sem=None): inputs, client=mock_openai_client, model="test-model" ) - states = outputs["states"] + states = outputs["outputs"] assert len(states) == 2 for state in states: assert "completion" in state diff --git a/tests/test_environment.py b/tests/test_environment.py index ed2fc30ef9..349501c0f7 100644 --- a/tests/test_environment.py +++ b/tests/test_environment.py @@ -249,7 +249,7 @@ async def mock_score_group(states, score_sem=None): model="test-model", ) - states = outputs["states"] + states = outputs["outputs"] assert len(states) == 1 assert "completion" in states[0] assert "reward" in states[0] @@ -280,26 +280,25 @@ async def mock_score_group(states, score_sem=None): model="test-model", ) - states = outputs["states"] + states = outputs["outputs"] assert len(states) == 1 assert "completion" in states[0] assert "reward" in states[0] assert states[0]["reward"] == 1.0 - def test_make_dataset(self, make_metadata, make_state): + def test_make_dataset(self, make_metadata, make_output): """Test creating a dataset from evaluation results.""" - outputs = GenerateOutputs(states=[make_state()], metadata=make_metadata()) - dataset = build_dataset(outputs) + results = GenerateOutputs(outputs=[make_output()], metadata=make_metadata()) + dataset = build_dataset(results) assert len(dataset) == 1 assert "prompt" in dataset.column_names assert "completion" in dataset.column_names - assert "answer" in dataset.column_names assert "reward" in dataset.column_names assert "task" in dataset.column_names assert "example_id" in dataset.column_names - assert "foo" in dataset.column_names # custom field from make_state fixture + assert "foo" in dataset.column_names # custom field from make_output fixture @pytest.mark.asyncio async def test_generate_updates_metadata(self, mock_openai_client): @@ -484,7 +483,7 @@ async def test_retry_after_retryable_error(self, mock_openai_client, make_input) outputs = await env.generate( inputs, client=mock_openai_client, model="test-model", max_retries=3 ) - states = outputs["states"] + states = outputs["outputs"] assert states[0].get("error") is None assert env.call_counts[0] == 3 @@ -508,10 +507,12 @@ async def test_no_retry_after_non_retryable_error( inputs, client=mock_openai_client, model="test-model", max_retries=3 ) - states = outputs["states"] + rollout_outputs = outputs["outputs"] assert env.call_counts[0] == 1 # No retries for non-retryable error - assert states[0].get("error") is not None - assert isinstance(states[0]["error"], vf.ToolError) + assert rollout_outputs[0].get("error") is not None + assert ( + "ToolError" in rollout_outputs[0]["error"] + ) # Error is serialized as repr string @pytest.mark.asyncio async def test_error_in_state_after_max_retries_exhausted( @@ -528,10 +529,12 @@ async def test_error_in_state_after_max_retries_exhausted( inputs, client=mock_openai_client, model="test-model", max_retries=2 ) - states = outputs["states"] + rollout_outputs = outputs["outputs"] assert env.call_counts[0] == 3 # 1 initial + 2 retries - assert states[0].get("error") is not None - assert isinstance(states[0]["error"], vf.InfraError) + assert rollout_outputs[0].get("error") is not None + assert ( + "InfraError" in rollout_outputs[0]["error"] + ) # Error is serialized as repr string class TestEmptyModelResponseErrors: diff --git a/tests/test_environment_extra.py b/tests/test_environment_extra.py index 63833dad5b..e268185014 100644 --- a/tests/test_environment_extra.py +++ b/tests/test_environment_extra.py @@ -159,7 +159,7 @@ def test_run_rollouts_with_max_concurrent( max_concurrent=2, ) ) - states = outputs["states"] + states = outputs["outputs"] assert len(states) == 3 @@ -178,7 +178,7 @@ def test_evaluate_fallback_and_repeat(mock_openai_client, make_dummy_env, make_i ) ) # Expect n * r rollouts in outputs - states = outputs["states"] + states = outputs["outputs"] assert len(states) == 2 * 2 @@ -190,7 +190,7 @@ async def test_generate_inside_running_loop( inputs = [make_input(example_id=0)] # Call the async API directly inside a running event loop to avoid nested sync wrapper issues outputs = await env.generate(inputs, client=mock_openai_client, model="test-model") - states = outputs["states"] + states = outputs["outputs"] assert len(states) == 1 assert states[0].get("completion") is not None @@ -218,7 +218,7 @@ def model_dump(self): assert isinstance(sanitized[0]["tool_calls"][0], str) -def test_make_dataset_basic_without_tools(make_metadata, make_state): - outputs = GenerateOutputs(states=[make_state()], metadata=make_metadata()) - ds = build_dataset(outputs) +def test_make_dataset_basic_without_tools(make_metadata, make_output): + results = GenerateOutputs(outputs=[make_output()], metadata=make_metadata()) + ds = build_dataset(results) assert len(ds) == 1 and "foo" in ds.column_names diff --git a/tests/test_eval_cli.py b/tests/test_eval_cli.py index c6d9cd9c60..c62e33e43a 100644 --- a/tests/test_eval_cli.py +++ b/tests/test_eval_cli.py @@ -9,6 +9,7 @@ import verifiers.utils.eval_utils from verifiers.types import GenerateOutputs from verifiers.utils.eval_utils import load_toml_config +from verifiers.utils.save_utils import states_to_outputs @pytest.fixture @@ -74,6 +75,7 @@ async def fake_run_evaluation(config, **kwargs): r = config.rollouts_per_example inputs = [_make_input(example_id=i // r) for i in range(n * r)] states = [_make_state(**inputs[i]) for i in range(n * r)] + rollout_outputs = states_to_outputs(states) metadata = _make_metadata( env_id=config.env_id, model=config.model, @@ -81,7 +83,7 @@ async def fake_run_evaluation(config, **kwargs): num_examples=n, rollouts_per_example=r, ) - return GenerateOutputs(states=states, metadata=metadata) + return GenerateOutputs(outputs=rollout_outputs, metadata=metadata) monkeypatch.setattr( verifiers.utils.eval_utils, "run_evaluation", fake_run_evaluation diff --git a/tests/test_eval_utils.py b/tests/test_eval_utils.py index dfdae4116c..430afe38a2 100644 --- a/tests/test_eval_utils.py +++ b/tests/test_eval_utils.py @@ -5,6 +5,7 @@ """ from verifiers.types import GenerateOutputs +from verifiers.utils.save_utils import states_to_outputs def test_print_results_rollout_indexing(capsys, make_metadata, make_state, make_input): @@ -39,9 +40,10 @@ def test_print_results_rollout_indexing(capsys, make_metadata, make_state, make_ make_state(**input, reward=reward, metrics={"test_metric": metric_value}) for input, reward, metric_value in zip(inputs, rewards, metric_values) ] + rollout_outputs = states_to_outputs(states) - outputs = GenerateOutputs(states=states, metadata=metadata) - print_results(outputs) + results = GenerateOutputs(outputs=rollout_outputs, metadata=metadata) + print_results(results) captured = capsys.readouterr() # Verify rollout groupings are correct @@ -72,8 +74,9 @@ def test_print_results_single_rollout(capsys, make_metadata, make_state, make_in make_state(**make_input(example_id=example_id), reward=reward) for example_id, reward in zip(example_ids, rewards) ] + rollout_outputs = states_to_outputs(states) - results = GenerateOutputs(states=states, metadata=metadata) + results = GenerateOutputs(outputs=rollout_outputs, metadata=metadata) print_results(results) captured = capsys.readouterr() @@ -97,11 +100,12 @@ def test_print_results_three_rollouts(capsys, make_metadata, make_state, make_in states = [ make_state(**input, reward=reward) for input, reward in zip(inputs, rewards) ] + rollout_outputs = states_to_outputs(states) metadata = make_metadata( num_examples=num_examples, rollouts_per_example=rollouts_per_example ) - results = GenerateOutputs(states=states, metadata=metadata) + results = GenerateOutputs(outputs=rollout_outputs, metadata=metadata) print_results(results) captured = capsys.readouterr() diff --git a/tests/test_gym_env.py b/tests/test_gym_env.py index a05b6ce532..d71e069db7 100644 --- a/tests/test_gym_env.py +++ b/tests/test_gym_env.py @@ -147,8 +147,10 @@ def test_basic_rollout_and_reward_sum(toy_env_class, client): num_eval_episodes=1, ) - outputs = env.evaluate_sync(client=client, model="mock") - st = outputs["states"][0] + outputs = env.evaluate_sync( + client=client, model="mock", state_columns=["trajectory", "gym_done"] + ) + st = outputs["outputs"][0] steps = st.get("trajectory", []) assert len(steps) > 0 @@ -173,8 +175,10 @@ def bad_parser(_txt: str) -> int: num_eval_episodes=1, ) - res = env.evaluate_sync(client=client, model="mock") - st = res["states"][0] + res = env.evaluate_sync( + client=client, model="mock", state_columns=["trajectory", "gym_done"] + ) + st = res["outputs"][0] steps = st.get("trajectory", []) assert st.get("gym_done") is True @@ -200,8 +204,10 @@ def step(self, action: int): num_train_episodes=0, num_eval_episodes=1, ) - res = env.evaluate_sync(client=client, model="mock") - st = res["states"][0] + res = env.evaluate_sync( + client=client, model="mock", state_columns=["trajectory", "gym_done"] + ) + st = res["outputs"][0] steps = st.get("trajectory", []) assert len(steps) == 3 @@ -229,7 +235,7 @@ def test_system_prompt_and_few_shot(toy_env_class, client): num_eval_episodes=1, ) res = env.evaluate_sync(client=client, model="mock") - st = res["states"][0] + st = res["outputs"][0] first_prompt = st["prompt"] roles = [m["role"] for m in first_prompt] @@ -256,8 +262,10 @@ def step(self, action: int): num_train_episodes=0, num_eval_episodes=1, ) - res = env.evaluate_sync(client=client, model="mock") - st = res["states"][0] + res = env.evaluate_sync( + client=client, model="mock", state_columns=["trajectory", "gym_done"] + ) + st = res["outputs"][0] steps = st.get("trajectory", []) assert steps[0]["extras"]["gym_info"] == {"info": "done"} @@ -275,7 +283,7 @@ def test_env_kwargs_passed_to_env(toy_env_class, client): num_eval_episodes=1, ) res = env.evaluate_sync(client=client, model="mock") - st = res["states"][0] + st = res["outputs"][0] first_obs_msg = st["prompt"][-1]["content"] assert first_obs_msg == "x=5" @@ -303,7 +311,7 @@ def obs_to_text(self, obs: Any) -> str: num_eval_episodes=1, ) res = env.evaluate_sync(client=client, model="mock") - st = res["states"][0] + st = res["outputs"][0] assert st["prompt"][-1]["content"] == "obs_is_0" @@ -324,7 +332,7 @@ def test_completion_mode(toy_env_class, client): ) res = env.evaluate_sync(client=client, model="mock") - st = res["states"][0] + st = res["outputs"][0] assert isinstance(st["prompt"], str) assert st["prompt"] == "x=0" diff --git a/tests/test_save_utils.py b/tests/test_save_utils.py index c22d3fd39e..9111ca3944 100644 --- a/tests/test_save_utils.py +++ b/tests/test_save_utils.py @@ -2,7 +2,7 @@ Covers: - make_serializable: JSON serialization for non-standard types -- sanitize_states: state sanitization before saving +- states_to_outputs: state to output conversion before saving - sanitize_metadata: metadata sanitization before saving - save_to_disk: disk saving with proper serialization """ @@ -16,7 +16,7 @@ from verifiers.utils.save_utils import ( make_serializable, - sanitize_states, + states_to_outputs, ) @@ -127,7 +127,7 @@ def test_serialize_metadata(self, make_metadata): class TestSavingResults: - def test_serialize_states(self, make_state): + def test_states_to_outputs(self, make_state): states = [ make_state( prompt=[{"role": "user", "content": "What is 2+2?"}], @@ -138,8 +138,8 @@ def test_serialize_states(self, make_state): client=OpenAI(api_key="EMPTY"), ), ] - sanitized_states = sanitize_states(states, state_columns=["client"]) - result = json.loads(json.dumps(sanitized_states, default=make_serializable)) + outputs = states_to_outputs(states, state_columns=["client"]) + result = json.loads(json.dumps(outputs, default=make_serializable)) assert isinstance(result, list) assert len(result) == 1 assert result[0]["example_id"] == 0 diff --git a/tests/test_singleturn_env.py b/tests/test_singleturn_env.py index e7f38b7d29..fb0775c063 100644 --- a/tests/test_singleturn_env.py +++ b/tests/test_singleturn_env.py @@ -260,7 +260,7 @@ async def mock_score_group(states, score_sem=None): model="test-model", ) - states = outputs["states"] + states = outputs["outputs"] assert len(states) == 2 for state in states: assert "completion" in state @@ -287,7 +287,7 @@ async def mock_score_group(states, score_sem=None): model="test-model", ) - states = outputs["states"] + states = outputs["outputs"] assert len(states) == 2 for state in states: assert "completion" in state @@ -304,7 +304,7 @@ async def test_a_generate_no_scoring(self, mock_singleturn_env, make_input): model="test-model", ) - states = outputs["states"] + states = outputs["outputs"] assert len(states) == 1 assert "completion" in states[0] assert "reward" in states[0] @@ -327,7 +327,7 @@ async def mock_score_group(states, score_sem=None): model="test-model", ) - states = outputs["states"] + states = outputs["outputs"] assert len(states) == 1 assert "completion" in states[0] assert "reward" in states[0] diff --git a/verifiers/envs/environment.py b/verifiers/envs/environment.py index 372c5f0e92..968796a80b 100644 --- a/verifiers/envs/environment.py +++ b/verifiers/envs/environment.py @@ -54,9 +54,9 @@ strip_nones_from_content, ) from verifiers.utils.save_utils import ( + GenerateOutputsBuilder, make_dataset, save_generate_outputs, - states_to_generate_metadata, ) from verifiers.utils.token_utils import ( get_prompt_ids, @@ -795,30 +795,6 @@ async def run_group( await self.rubric.dummy_score_group(group_states) return list(group_states) - def _build_generate_outputs( - self, - states: list[State], - model: str, - client: AsyncOpenAI, - state_columns: list[str] | None, - results_path: Path | None, - sampling_args: SamplingArgs, - start_time: float, - ) -> GenerateOutputs: - """Prepare GenerateOutputs from a list of completed states.""" - metadata = states_to_generate_metadata( - self.env_id, - self.env_args, - model, - client, - states, - state_columns, - sampling_args, - start_time, - results_path, - ) - return GenerateOutputs(states=states, metadata=metadata) - async def generate( self, inputs: Dataset | List[RolloutInput], @@ -870,7 +846,16 @@ async def generate( if sampling_args is not None: gen_sampling_args.update(sampling_args) - start_time = time.time() + # Initialize builder for incremental serialization + builder = GenerateOutputsBuilder( + env_id=self.env_id, + env_args=self.env_args, + model=model, + client=client, + state_columns=state_columns, + sampling_args=gen_sampling_args, + results_path=results_path, + ) # create tasks based on mode tasks: dict[asyncio.Task, int] = {} @@ -924,18 +909,19 @@ async def generate( # process tasks as they complete reward_sum, reward_count = 0, 0 groups_or_rollouts_completed = 0 - all_states: list[State] = [] try: for coro in asyncio.as_completed(tasks.keys()): result = await coro # normalize: independent_scoring returns State, group returns list[State] states = [result] if independent_scoring else result - all_states.extend(states) + + # Serialize states to outputs immediately (serialization happens once here) + new_outputs = builder.add_states(states) groups_or_rollouts_completed += 1 - # track reward for rolling average - for s in states: - r = s.get("reward") + # track reward for rolling average (from outputs) + for o in new_outputs: + r = o.get("reward") if r is not None: reward_sum += r reward_count += 1 @@ -946,51 +932,33 @@ async def generate( if reward_count > 0: pbar.set_postfix(reward=f"{reward_sum / reward_count:.3f}") elif on_progress is not None: - on_progress(all_states, states) + on_progress(builder.outputs, new_outputs) - # save intermediate results + # save intermediate results (outputs already serialized, no redundant work) if ( save_results and save_every > 0 and groups_or_rollouts_completed % save_every == 0 ): - intermediate_outputs = self._build_generate_outputs( - all_states, - model, - client, - state_columns, - results_path, - gen_sampling_args, - start_time, - ) + intermediate_results = builder.build() self.logger.debug( - f"Saving intermediate results to {intermediate_outputs['metadata']['path_to_save']}" + f"Saving intermediate results to {intermediate_results['metadata']['path_to_save']}" ) - save_generate_outputs(intermediate_outputs) + save_generate_outputs(intermediate_results) finally: if pbar is not None: pbar.close() - # sort by example_id to ensure deterministic ordering regardless of completion order - all_states.sort(key=lambda s: s.get("example_id", 0)) - - outputs = self._build_generate_outputs( - all_states, - model, - client, - state_columns, - results_path, - gen_sampling_args, - start_time, - ) + # Build final results (sorted by example_id for deterministic ordering) + results = builder.build(sort_by_example_id=True) # save if requested if save_results: - save_generate_outputs(outputs, push_to_hf_hub, hf_hub_dataset_name) + save_generate_outputs(results, push_to_hf_hub, hf_hub_dataset_name) if on_log is not None: - on_log(f"Saved final outputs to {outputs['metadata']['path_to_save']}") + on_log(f"Saved final outputs to {results['metadata']['path_to_save']}") - return outputs + return results def generate_sync( self, diff --git a/verifiers/gepa/adapter.py b/verifiers/gepa/adapter.py index b750cf5f66..492f8d0db5 100644 --- a/verifiers/gepa/adapter.py +++ b/verifiers/gepa/adapter.py @@ -7,7 +7,13 @@ from openai import AsyncOpenAI, OpenAI from verifiers.envs.environment import Environment -from verifiers.types import ClientConfig, Messages, RolloutInput, SamplingArgs, State +from verifiers.types import ( + ClientConfig, + Messages, + RolloutInput, + RolloutOutput, + SamplingArgs, +) from verifiers.utils.message_utils import message_to_printable, messages_to_printable from verifiers.utils.save_utils import make_serializable @@ -76,7 +82,7 @@ def evaluate( batch: list[RolloutInput], candidate: dict[str, str], capture_traces: bool = False, - ) -> EvaluationBatch[State, State]: + ) -> EvaluationBatch[RolloutOutput, RolloutOutput]: """ Run verifiers evaluation with the candidate system prompt. """ @@ -93,9 +99,9 @@ def evaluate( ) ) - states = results["states"] - example_ids = [s["example_id"] for s in states] - rewards = [s["reward"] for s in states] + outputs = results["outputs"] + example_ids = [o["example_id"] for o in outputs] + rewards = [o["reward"] for o in outputs] # Update display if configured if self.display is not None: @@ -112,41 +118,41 @@ def evaluate( ) return EvaluationBatch( - outputs=states, + outputs=outputs, scores=rewards, - trajectories=states if capture_traces else None, + trajectories=outputs if capture_traces else None, ) def make_reflective_dataset( self, candidate: dict[str, str], # noqa: ARG002 - required by GEPA adapter protocol - eval_batch: EvaluationBatch[State, dict[str, Any]], + eval_batch: EvaluationBatch[RolloutOutput, dict[str, Any]], components_to_update: list[str], ) -> Mapping[str, Sequence[Mapping[str, Any]]]: """Build reflective dataset for GEPA teacher LLM.""" - outputs: list[dict[str, Any]] = eval_batch.outputs - states: list[State] = eval_batch.trajectories or [] + outputs: list[RolloutOutput] = eval_batch.outputs + trajectories: list[RolloutOutput] = eval_batch.trajectories or [] scores = eval_batch.scores records = [] - # outputs, states, and scores should be the same length - for output, state, score in zip(outputs, states, scores): + # outputs, trajectories, and scores should be the same length + for output, trajectory, score in zip(outputs, trajectories, scores): record: dict[str, Any] = { "query": _extract_user_query(output["prompt"]), "completion": messages_to_printable(output["completion"]), - "expected_answer": output["answer"], + "expected_answer": output.get("answer", ""), "reward": score, } - if state.get("error"): - record["error"] = repr(state["error"]) + if trajectory.get("error"): + record["error"] = repr(trajectory["error"]) - if state.get("stop_condition"): - record["stop_condition"] = state["stop_condition"] + if trajectory.get("stop_condition"): + record["stop_condition"] = trajectory["stop_condition"] for col in self.state_columns: - if col in state: - record[col] = make_serializable(state[col]) + if col in trajectory: + record[col] = make_serializable(trajectory[col]) records.append(record) diff --git a/verifiers/rl/trainer/orchestrator.py b/verifiers/rl/trainer/orchestrator.py index cf52a8bb15..cdfd9ed4b2 100644 --- a/verifiers/rl/trainer/orchestrator.py +++ b/verifiers/rl/trainer/orchestrator.py @@ -225,6 +225,7 @@ async def generate_batch(self, batch_id: int) -> Batch: model=self.model_name, sampling_args=self.sampling_args, max_concurrent=self.max_concurrent, + state_columns=["trajectory"], # Include trajectory for RL training ) self.is_generating = False wall_clock_s = time.time() - start_time @@ -241,9 +242,9 @@ async def generate_batch(self, batch_id: int) -> Batch: metrics: dict[str, list[float]] = {} advantages: list[float] = [] - states = env_results["states"] - for state in states: - trajectory = state["trajectory"] + outputs = env_results["outputs"] + for output in outputs: + trajectory = output["trajectory"] for step in trajectory: tokens = step["tokens"] if tokens is None: @@ -254,10 +255,10 @@ async def generate_batch(self, batch_id: int) -> Batch: completion_mask.append(tokens["completion_mask"]) completion_logprobs.append(tokens["completion_logprobs"]) advantages.append(step["advantage"]) - prompts.append(state["prompt"]) - completions.append(state["completion"]) - rewards.append(state["reward"]) - for k, v in state["metrics"].items(): + prompts.append(output["prompt"]) + completions.append(output["completion"]) + rewards.append(output["reward"]) + for k, v in output["metrics"].items(): if k not in metrics: metrics[k] = [] metrics[k].append(v) @@ -303,8 +304,8 @@ async def generate_batch(self, batch_id: int) -> Batch: generation_ms: list[float] = [] scoring_ms: list[float] = [] total_ms: list[float] = [] - for state in states: - timing = state.get("timing", {}) + for output in outputs: + timing = output.get("timing", {}) if "generation_ms" in timing: generation_ms.append(float(timing["generation_ms"])) if "scoring_ms" in timing: @@ -320,7 +321,7 @@ async def generate_batch(self, batch_id: int) -> Batch: metrics_dict["timing/total_ms"] = float(np.mean(total_ms)) metrics_dict["wall_clock/generate_s"] = float(wall_clock_s) - errors = [state.get("error") for state in states] + errors = [output.get("error") for output in outputs] # build per-process microbatches N = len(advantages) diff --git a/verifiers/types.py b/verifiers/types.py index 8aeb415621..6e1e8a8f20 100644 --- a/verifiers/types.py +++ b/verifiers/types.py @@ -95,6 +95,36 @@ class RolloutTiming(TypedDict, total=False): total_ms: float +class BaseRolloutOutput(TypedDict): + """Required fields for serialized rollout output.""" + + example_id: int + task: str + prompt: Messages | None + completion: Messages | None + reward: float + timing: RolloutTiming + is_completed: bool + is_truncated: bool + metrics: dict[str, float] + + +class RolloutOutput(BaseRolloutOutput, total=False): + """Serialized output from a rollout (mirrors RolloutInput). + + Required fields: example_id, task, prompt, completion, reward, timing, + is_completed, is_truncated, metrics + Optional fields: answer, info, error, stop_condition + + Note: Additional fields may be present from state_columns. + """ + + answer: str + info: Info + error: str | None + stop_condition: str | None + + class State(dict): INPUT_FIELDS = ["prompt", "answer", "task", "info", "example_id"] # rollout inputs @@ -144,7 +174,9 @@ def get(self, key: str, default: Any = None) -> Any: # callbacks StartCallback = Callable[[int], None] # total rollouts -ProgressCallback = Callable[[list[State], list[State]], None] # all_states, new_states +ProgressCallback = Callable[ + [list[RolloutOutput], list[RolloutOutput]], None +] # all_outputs, new_outputs LogCallback = Callable[[str], None] # log messages @@ -168,9 +200,9 @@ class GenerateMetadata(TypedDict): class GenerateOutputs(TypedDict): - """TypedDict for generation outputs.""" + """TypedDict for generation outputs (results).""" - states: list[State] + outputs: list[RolloutOutput] metadata: GenerateMetadata diff --git a/verifiers/utils/eval_display.py b/verifiers/utils/eval_display.py index f144b136a7..b0dc232003 100644 --- a/verifiers/utils/eval_display.py +++ b/verifiers/utils/eval_display.py @@ -569,12 +569,12 @@ def _make_env_detail( items: list[Panel] = [] # Example 0 prompt/completion - states = results["states"] - if states and states[0]["prompt"] and states[0]["completion"]: - prompt = messages_to_printable(states[0]["prompt"]) - completion = messages_to_printable(states[0]["completion"]) - reward_0 = states[0]["reward"] if states[0]["reward"] else 0.0 - error_0 = states[0].get("error") if states[0] else None + outputs = results["outputs"] + if outputs and outputs[0]["prompt"] and outputs[0]["completion"]: + prompt = messages_to_printable(outputs[0]["prompt"]) + completion = messages_to_printable(outputs[0]["completion"]) + reward_0 = outputs[0]["reward"] if outputs[0]["reward"] else 0.0 + error_0 = outputs[0].get("error") if outputs[0] else None # Prompt panel items.append( @@ -602,7 +602,7 @@ def _make_env_detail( ) # Reward distribution - rewards = [s["reward"] for s in states] + rewards = [o["reward"] for o in outputs] if rewards: # All rollouts histogram all_rollouts_content = Group( diff --git a/verifiers/utils/eval_utils.py b/verifiers/utils/eval_utils.py index 9a68e52270..5ffef2726b 100644 --- a/verifiers/utils/eval_utils.py +++ b/verifiers/utils/eval_utils.py @@ -25,8 +25,8 @@ GenerateOutputs, LogCallback, ProgressCallback, + RolloutOutput, StartCallback, - State, ) from verifiers.utils.async_utils import EventLoopLagMonitor from verifiers.utils.client_utils import setup_client @@ -189,20 +189,20 @@ def to_col_order(list_of_dicts: list[dict[str, Any]]) -> dict[str, list[float]]: def get_task_outputs(results: GenerateOutputs, task: str) -> GenerateOutputs: """Get only the rollouts for a given task.""" - states = [s for s in results["states"] if s["task"] == task] + outputs = [o for o in results["outputs"] if o["task"] == task] return GenerateOutputs( - states=states, + outputs=outputs, metadata=results["metadata"], # duplicate metadata ) -def print_rewards(outputs: GenerateOutputs): - rewards = [s["reward"] for s in outputs["states"]] +def print_rewards(results: GenerateOutputs): + rewards = [o["reward"] for o in results["outputs"]] print("Rewards:") print( f"reward: avg - {sum(rewards) / len(rewards):.3f}, std - {np.std(rewards):.3f}" ) - r = outputs["metadata"]["rollouts_per_example"] + r = results["metadata"]["rollouts_per_example"] n = len(rewards) // r # results are sorted by example_id, so rollout i is at indices [i, i+r, i+2r, ...] for i in range(r): @@ -210,7 +210,7 @@ def print_rewards(outputs: GenerateOutputs): out = f"r{i + 1}: {trials}" print(out) - metrics = [s["metrics"] for s in outputs["states"]] + metrics = [o["metrics"] for o in results["outputs"]] metrics_col = to_col_order(metrics) for k in metrics_col.keys(): v = metrics_col[k] @@ -221,18 +221,18 @@ def print_rewards(outputs: GenerateOutputs): print(out) -def print_info(outputs: GenerateOutputs): - is_truncated = [s["is_truncated"] for s in outputs["states"]] +def print_info(results: GenerateOutputs): + is_truncated = [o["is_truncated"] for o in results["outputs"]] print("Info:") print( f"is_truncated: avg - {np.mean(is_truncated):.3f}, std - {np.std(is_truncated):.3f}" ) - stop_conditions = [s["stop_condition"] for s in outputs["states"]] + stop_conditions = [o["stop_condition"] for o in results["outputs"]] counter = Counter(stop_conditions) print( f"stop_conditions: {', '.join([f'{k}: {v / counter.total():.3f}' for k, v in counter.items()])}" ) - errors = [s.get("error") for s in outputs["states"]] + errors = [o.get("error") for o in results["outputs"]] has_errors = [e is not None for e in errors] if any(has_errors): print( @@ -245,9 +245,9 @@ def print_info(outputs: GenerateOutputs): print(f" - {repr(error_chain)}: {count / counter.total():.3f}") -def print_timing(outputs: GenerateOutputs): +def print_timing(results: GenerateOutputs): print("Timing:") - timing = [s["timing"] for s in outputs["states"]] + timing = [o["timing"] for o in results["outputs"]] timing_col = to_col_order(cast(list[dict], timing)) generation_ms_arr = np.array(timing_col["generation_ms"]) scoring_ms_arr = np.array(timing_col["scoring_ms"]) @@ -267,22 +267,22 @@ def print_timing(outputs: GenerateOutputs): ) -def print_results(outputs: GenerateOutputs, num_samples: int = 1): - assert outputs["metadata"] is not None +def print_results(results: GenerateOutputs, num_samples: int = 1): + assert results["metadata"] is not None print("--- Evaluation ---") - print(f"Environment: {outputs['metadata']['env_id']}") - print(f"Model: {outputs['metadata']['model']}") - print(f"Provider: {outputs['metadata']['base_url']}") - print(f"Examples: {outputs['metadata']['num_examples']}") - print(f"Rollouts per example: {outputs['metadata']['rollouts_per_example']}") + print(f"Environment: {results['metadata']['env_id']}") + print(f"Model: {results['metadata']['model']}") + print(f"Provider: {results['metadata']['base_url']}") + print(f"Examples: {results['metadata']['num_examples']}") + print(f"Rollouts per example: {results['metadata']['rollouts_per_example']}") print("--- Example ---") - printable_prompts = [messages_to_printable(s["prompt"]) for s in outputs["states"]] + printable_prompts = [messages_to_printable(o["prompt"]) for o in results["outputs"]] printable_completions = [ - messages_to_printable(s["completion"]) for s in outputs["states"] + messages_to_printable(o["completion"]) for o in results["outputs"] ] - rewards = [s["reward"] for s in outputs["states"]] - errors = [s.get("error") for s in outputs["states"]] + rewards = [o["reward"] for o in results["outputs"]] + errors = [o.get("error") for o in results["outputs"]] print_prompt_completions_sample( printable_prompts, printable_completions, @@ -292,18 +292,18 @@ def print_results(outputs: GenerateOutputs, num_samples: int = 1): num_samples=num_samples, ) print("--- All ---") - print_rewards(outputs) - print_info(outputs) - print_timing(outputs) + print_rewards(results) + print_info(results) + print_timing(results) - tasks = set([s["task"] for s in outputs["states"]]) + tasks = set([o["task"] for o in results["outputs"]]) if len(tasks) > 1: for task in tasks: - task_outputs = get_task_outputs(outputs, task) + task_results = get_task_outputs(results, task) print(f"\n--- {task} ---") - print_rewards(task_outputs) - print_info(task_outputs) - print_timing(task_outputs) + print_rewards(task_results) + print_info(task_results) + print_timing(task_results) @contextmanager @@ -432,20 +432,22 @@ def on_start(total: int) -> None: num_examples = total // env_config.rollouts_per_example display.update_env_state(env_idx, total=total, num_examples=num_examples) - def on_progress(all_states: list[State], new_states: list[State]) -> None: + def on_progress( + all_outputs: list[RolloutOutput], new_outputs: list[RolloutOutput] + ) -> None: nonlocal error_accum, reward_accum, metrics_accum # Progress is always rollout-based - completed = len(all_states) + completed = len(all_outputs) - for s in new_states: - if s.get("error") is not None: + for o in new_outputs: + if o.get("error") is not None: error_accum += 1 - reward = s.get("reward") + reward = o.get("reward") if reward is not None: reward_accum += reward - state_metrics = s.get("metrics") or {} - for name, value in state_metrics.items(): + output_metrics = o.get("metrics") or {} + for name, value in output_metrics.items(): if value is not None: metrics_accum[name] += value diff --git a/verifiers/utils/save_utils.py b/verifiers/utils/save_utils.py index a717d711e4..1eae24a6b3 100644 --- a/verifiers/utils/save_utils.py +++ b/verifiers/utils/save_utils.py @@ -11,8 +11,10 @@ from pydantic import BaseModel from verifiers.types import ( + ChatCompletionToolParam, GenerateMetadata, GenerateOutputs, + RolloutOutput, SamplingArgs, State, ) @@ -111,49 +113,161 @@ def get_hf_hub_dataset_name(outputs: GenerateOutputs) -> str: return dataset_name -def sanitize_states(states: list[State], state_columns: list[str] = []) -> list[dict]: - """Sanitizes a list of rollouts before saving to disk.""" - - def sanitize_state(state: State) -> dict: - sanitized_state = { - "example_id": state.get("example_id", 0), - "prompt": state.get("prompt"), - "completion": state.get("completion"), - "answer": state.get("answer", ""), - "task": state.get("task", "default"), - "info": state.get("info", {}), - "reward": state.get("reward", 0.0), - "error": state.get("error", None), - "total_ms": state.get("timing", {}).get("total_ms", 0.0), - "generation_ms": state.get("timing", {}).get("generation_ms", 0.0), - "scoring_ms": state.get("timing", {}).get("scoring_ms", 0.0), - } - # sanitize messages - sanitized_state["prompt"] = sanitize_tool_calls( - messages_to_printable(state["prompt"]) +def state_to_output(state: State, state_columns: list[str] = []) -> RolloutOutput: + """Convert a State to a serializable RolloutOutput.""" + output: RolloutOutput = { + "example_id": state.get("example_id", 0), + "prompt": state.get("prompt"), + "completion": state.get("completion"), + "answer": state.get("answer", ""), + "task": state.get("task", "default"), + "info": state.get("info", {}), + "reward": state.get("reward", 0.0), + "error": state.get("error", None), + # Include timing as dict for compatibility with print functions + "timing": state.get("timing", {}), + # Include these for print_info compatibility + "is_completed": state.get("is_completed", False), + "is_truncated": state.get("is_truncated", False), + "stop_condition": state.get("stop_condition", None), + # Include metrics as dict for print_rewards compatibility + "metrics": state.get("metrics", {}), + } + # sanitize messages (handle None for error cases) + prompt = state.get("prompt") + if prompt is not None: + output["prompt"] = sanitize_tool_calls(messages_to_printable(prompt)) + completion = state.get("completion") + if completion is not None: + output["completion"] = sanitize_tool_calls(messages_to_printable(completion)) + # use repr for error + if state.get("error") is not None: + output["error"] = repr(state.get("error")) + # only include optional fields if non-empty + if "answer" in output and not output["answer"]: + output.pop("answer") + if "info" in output and not output["info"]: + output.pop("info") + # flatten metrics to top-level keys (backwards compatibility) + state_metrics = state.get("metrics", {}) + for k, v in state_metrics.items(): + output[k] = v + # add state columns + for col in state_columns: + output[col] = state.get(col) + + return output + + +def states_to_outputs( + states: list[State], state_columns: list[str] = [] +) -> list[RolloutOutput]: + """Convert a list of States to serializable RolloutOutputs.""" + return [state_to_output(state, state_columns) for state in states] + + +class GenerateOutputsBuilder: + """Incrementally builds GenerateOutputs, serializing states exactly once. + + This builder accumulates outputs as rollouts complete, avoiding redundant + serialization for intermediate saves. Each state is serialized once via + add_states(), and the accumulated outputs can be retrieved at any time + via build(). + """ + + def __init__( + self, + env_id: str, + env_args: dict, + model: str, + client: AsyncOpenAI, + state_columns: list[str] | None, + sampling_args: SamplingArgs, + results_path: Path | None, + ): + self.env_id = env_id + self.env_args = env_args + self.model = model + self.client = client + self.state_columns = state_columns or [] + self.sampling_args = sampling_args + self.results_path = results_path or get_results_path(env_id, model) + self.start_time = time.time() + + # Accumulated outputs + self.outputs: list[RolloutOutput] = [] + # Track tools from states for metadata + self._tools_list: list[list[ChatCompletionToolParam] | None] = [] + + def add_states(self, new_states: list[State]) -> list[RolloutOutput]: + """Convert new states to outputs and accumulate. Returns the new outputs.""" + new_outputs = states_to_outputs(new_states, self.state_columns) + self.outputs.extend(new_outputs) + # Track tools for metadata computation + for state in new_states: + self._tools_list.append(state.get("oai_tools")) + return new_outputs + + def build(self, sort_by_example_id: bool = False) -> GenerateOutputs: + """Build GenerateOutputs from accumulated outputs.""" + outputs = self.outputs + if sort_by_example_id: + outputs = sorted(outputs, key=lambda o: o.get("example_id", 0)) + metadata = self._compute_metadata(outputs) + return GenerateOutputs(outputs=outputs, metadata=metadata) + + def _compute_metadata(self, outputs: list[RolloutOutput]) -> GenerateMetadata: + """Compute metadata from accumulated outputs.""" + base_url = str(self.client.base_url) if hasattr(self.client, "base_url") else "" + + # Compute reward stats from outputs + rewards = [o.get("reward", 0.0) for o in outputs] + avg_reward = sum(rewards) / len(rewards) if rewards else 0.0 + + # Compute metrics stats from outputs + metrics: dict[str, list[float]] = defaultdict(list) + for output in outputs: + output_metrics = output.get("metrics", {}) + if output_metrics: + for metric_name, metric_value in output_metrics.items(): + if isinstance(metric_value, (int, float)): + metrics[metric_name].append(metric_value) + avg_metrics = {k: sum(v) / len(v) if v else 0.0 for k, v in metrics.items()} + + # Compute example counts + example_ids = [o.get("example_id", 0) for o in outputs] + num_examples = len(set(example_ids)) if example_ids else 0 + rollouts_per_example = len(outputs) // num_examples if num_examples > 0 else 1 + + # Determine tools (use first non-None if all same) + def tools_key(tools: list | None) -> str: + if not tools: + return "" + return str(sorted([t.get("function", {}).get("name", "") for t in tools])) + + unique_tools = set(tools_key(t) for t in self._tools_list) + tools = ( + next((t for t in self._tools_list if t), None) + if len(unique_tools) == 1 + else None ) - sanitized_state["completion"] = sanitize_tool_calls( - messages_to_printable(state["completion"]) + + return GenerateMetadata( + env_id=self.env_id, + env_args=self.env_args, + model=self.model, + base_url=base_url, + num_examples=num_examples, + rollouts_per_example=rollouts_per_example, + sampling_args=self.sampling_args, + date=datetime.now().isoformat(), + time_ms=(time.time() - self.start_time) * 1000.0, + avg_reward=avg_reward, + avg_metrics=avg_metrics, + state_columns=self.state_columns, + path_to_save=self.results_path, + tools=tools, ) - # use repr for error - if state.get("error") is not None: - sanitized_state["error"] = repr(state.get("error")) - # only include optional fields if non-empty - if "answer" in sanitized_state and not sanitized_state["answer"]: - sanitized_state.pop("answer") - if "info" in sanitized_state and not sanitized_state["info"]: - sanitized_state.pop("info") - # flatten metrics - state_metrics = state.get("metrics", {}) - for k, v in state_metrics.items(): - sanitized_state[k] = v - # add state columns - for col in state_columns: - sanitized_state[col] = state.get(col) - - return sanitized_state - - return [sanitize_state(state) for state in states] def sanitize_metadata(metadata: GenerateMetadata) -> dict: @@ -193,27 +307,26 @@ def save_metadata(metadata: dict, metadata_path: Path): save_results(results, path / "results.jsonl") -def make_dataset(outputs: GenerateOutputs) -> Dataset: - state_columns = outputs["metadata"]["state_columns"] - sanitized_states = sanitize_states(outputs["states"], state_columns) - return Dataset.from_list(sanitized_states) +def make_dataset(results: GenerateOutputs) -> Dataset: + """Create a Dataset from GenerateOutputs (outputs are already serialized).""" + return Dataset.from_list(results["outputs"]) def save_generate_outputs( - outputs: GenerateOutputs, + results: GenerateOutputs, push_to_hf_hub: bool = False, hf_hub_dataset_name: str | None = None, ): - path_to_save = outputs["metadata"]["path_to_save"] - state_columns = outputs["metadata"]["state_columns"] - sanitized_states = sanitize_states(outputs["states"], state_columns) - sanitized_metadata = sanitize_metadata(outputs["metadata"]) - save_to_disk(sanitized_states, sanitized_metadata, path_to_save) + """Save GenerateOutputs to disk. Outputs are already serialized.""" + path_to_save = results["metadata"]["path_to_save"] + sanitized_metadata = sanitize_metadata(results["metadata"]) + # Outputs are already serialized RolloutOutputs + save_to_disk(results["outputs"], sanitized_metadata, path_to_save) logger.info(f"Results saved to {path_to_save}") if push_to_hf_hub: - dataset_name = hf_hub_dataset_name or get_hf_hub_dataset_name(outputs) + dataset_name = hf_hub_dataset_name or get_hf_hub_dataset_name(results) try: - dataset = make_dataset(outputs) + dataset = make_dataset(results) dataset.push_to_hub(dataset_name) logger.info(f"Dataset saved to Hugging Face Hub: {dataset_name}") except Exception as e: From 3364d6a12d13d37aa6c142ea8a7b8eeba19f4c83 Mon Sep 17 00:00:00 2001 From: William Brown Date: Tue, 27 Jan 2026 17:33:23 -0800 Subject: [PATCH 2/6] serializable output type; generateoutputs builder --- tests/test_save_utils.py | 21 +++++-- verifiers/gepa/adapter.py | 6 +- verifiers/types.py | 30 +++++----- verifiers/utils/eval_utils.py | 28 +++++---- verifiers/utils/logging_utils.py | 10 ++-- verifiers/utils/save_utils.py | 100 +++++++++++++++++++++---------- 6 files changed, 125 insertions(+), 70 deletions(-) diff --git a/tests/test_save_utils.py b/tests/test_save_utils.py index 9111ca3944..ffe94a1683 100644 --- a/tests/test_save_utils.py +++ b/tests/test_save_utils.py @@ -135,10 +135,9 @@ def test_states_to_outputs(self, make_state): answer="", info={}, reward=1.0, - client=OpenAI(api_key="EMPTY"), ), ] - outputs = states_to_outputs(states, state_columns=["client"]) + outputs = states_to_outputs(states, state_columns=["foo"]) result = json.loads(json.dumps(outputs, default=make_serializable)) assert isinstance(result, list) assert len(result) == 1 @@ -149,7 +148,19 @@ def test_states_to_outputs(self, make_state): ] assert result[0].get("answer") is None # empty answer not included assert result[0].get("info") is None # empty info not included - assert ( - result[0].get("client") is not None - ) # client included because of state_columns + assert result[0].get("foo") == "bar" # custom field from make_state fixture assert result[0]["reward"] == 1.0 + + def test_non_serializable_state_column_raises(self, make_state): + """Non-serializable state_columns should raise ValueError.""" + import pytest + + states = [ + make_state( + prompt=[{"role": "user", "content": "test"}], + completion=[{"role": "assistant", "content": "test"}], + client=OpenAI(api_key="EMPTY"), + ), + ] + with pytest.raises(ValueError, match="not JSON-serializable"): + states_to_outputs(states, state_columns=["client"]) diff --git a/verifiers/gepa/adapter.py b/verifiers/gepa/adapter.py index 492f8d0db5..90e4c11a2f 100644 --- a/verifiers/gepa/adapter.py +++ b/verifiers/gepa/adapter.py @@ -126,12 +126,12 @@ def evaluate( def make_reflective_dataset( self, candidate: dict[str, str], # noqa: ARG002 - required by GEPA adapter protocol - eval_batch: EvaluationBatch[RolloutOutput, dict[str, Any]], + eval_batch: EvaluationBatch[RolloutOutput, RolloutOutput], components_to_update: list[str], ) -> Mapping[str, Sequence[Mapping[str, Any]]]: """Build reflective dataset for GEPA teacher LLM.""" - outputs: list[RolloutOutput] = eval_batch.outputs - trajectories: list[RolloutOutput] = eval_batch.trajectories or [] + outputs = eval_batch.outputs + trajectories = eval_batch.trajectories or [] scores = eval_batch.scores records = [] diff --git a/verifiers/types.py b/verifiers/types.py index 6e1e8a8f20..5a7fc6fccf 100644 --- a/verifiers/types.py +++ b/verifiers/types.py @@ -95,9 +95,20 @@ class RolloutTiming(TypedDict, total=False): total_ms: float -class BaseRolloutOutput(TypedDict): - """Required fields for serialized rollout output.""" +class RolloutOutput(dict): + """Serialized output from a rollout (mirrors RolloutInput). + + A dict subclass that allows typed access to known fields while supporting + arbitrary additional fields from state_columns. All values must be + JSON-serializable. + + Required fields: example_id, task, prompt, completion, reward, timing, + is_completed, is_truncated, metrics + Optional fields: answer, info, error, stop_condition, trajectory, oai_tools + Additional fields: arbitrary serializable state_columns + """ + # Required fields example_id: int task: str prompt: Messages | None @@ -107,22 +118,13 @@ class BaseRolloutOutput(TypedDict): is_completed: bool is_truncated: bool metrics: dict[str, float] - - -class RolloutOutput(BaseRolloutOutput, total=False): - """Serialized output from a rollout (mirrors RolloutInput). - - Required fields: example_id, task, prompt, completion, reward, timing, - is_completed, is_truncated, metrics - Optional fields: answer, info, error, stop_condition - - Note: Additional fields may be present from state_columns. - """ - + # Optional fields answer: str info: Info error: str | None stop_condition: str | None + trajectory: list["TrajectoryStep"] + oai_tools: list["ChatCompletionToolParam"] class State(dict): diff --git a/verifiers/utils/eval_utils.py b/verifiers/utils/eval_utils.py index 5ffef2726b..9234b63a97 100644 --- a/verifiers/utils/eval_utils.py +++ b/verifiers/utils/eval_utils.py @@ -5,7 +5,8 @@ from collections import Counter, defaultdict from contextlib import contextmanager from pathlib import Path -from typing import Any, cast +from collections.abc import Mapping +from typing import cast from datasets import disable_progress_bar, enable_progress_bar from datasets.utils import logging as ds_logging @@ -30,7 +31,6 @@ ) from verifiers.utils.async_utils import EventLoopLagMonitor from verifiers.utils.client_utils import setup_client -from verifiers.utils.error_utils import ErrorChain from verifiers.utils.logging_utils import print_prompt_completions_sample, print_time from verifiers.utils.message_utils import messages_to_printable from verifiers.utils.path_utils import get_eval_results_path @@ -180,8 +180,8 @@ def load_toml_config(path: Path) -> list[dict]: return merged_eval_list -def to_col_order(list_of_dicts: list[dict[str, Any]]) -> dict[str, list[float]]: - """Convert a list of dictionaries to a dictionary of lists, ordered by the keys of the first dictionary.""" +def to_col_order(list_of_dicts: list[Mapping[str, float]]) -> dict[str, list[float]]: + """Convert a list of mappings to a dictionary of lists, ordered by the keys of the first mapping.""" if not list_of_dicts: return {} return {k: [m[k] for m in list_of_dicts] for k in list_of_dicts[0].keys()} @@ -238,17 +238,17 @@ def print_info(results: GenerateOutputs): print( f"errors: avg - {np.mean(has_errors):.3f}, std - {np.std(has_errors):.3f}" ) - errors = [e for e in errors if e is not None] - error_chains = [ErrorChain(e) for e in errors] - counter = Counter(error_chains) - for error_chain, count in counter.items(): - print(f" - {repr(error_chain)}: {count / counter.total():.3f}") + error_strs = [e for e in errors if e is not None] + # Errors are serialized as strings, count unique error types + counter = Counter(error_strs) + for error_str, count in counter.items(): + print(f" - {error_str}: {count / counter.total():.3f}") def print_timing(results: GenerateOutputs): print("Timing:") timing = [o["timing"] for o in results["outputs"]] - timing_col = to_col_order(cast(list[dict], timing)) + timing_col = to_col_order(timing) generation_ms_arr = np.array(timing_col["generation_ms"]) scoring_ms_arr = np.array(timing_col["scoring_ms"]) total_ms_arr = np.array(timing_col["total_ms"]) @@ -277,9 +277,13 @@ def print_results(results: GenerateOutputs, num_samples: int = 1): print(f"Rollouts per example: {results['metadata']['rollouts_per_example']}") print("--- Example ---") - printable_prompts = [messages_to_printable(o["prompt"]) for o in results["outputs"]] + printable_prompts = [ + messages_to_printable(o["prompt"]) if o["prompt"] else [] + for o in results["outputs"] + ] printable_completions = [ - messages_to_printable(o["completion"]) for o in results["outputs"] + messages_to_printable(o["completion"]) if o["completion"] else [] + for o in results["outputs"] ] rewards = [o["reward"] for o in results["outputs"]] errors = [o.get("error") for o in results["outputs"]] diff --git a/verifiers/utils/logging_utils.py b/verifiers/utils/logging_utils.py index 993a61de9f..7732a5c21f 100644 --- a/verifiers/utils/logging_utils.py +++ b/verifiers/utils/logging_utils.py @@ -76,7 +76,7 @@ def quiet_verifiers(): def print_prompt_completions_sample( prompts: list[Messages], completions: list[Messages], - errors: list[Error | None], + errors: list[str | Error | None], rewards: list[float], step: int, num_samples: int = 1, @@ -132,10 +132,12 @@ def _format_messages(messages) -> Text: return out - def _format_error(error: BaseException) -> Text: + def _format_error(error: str | BaseException) -> Text: out = Text() - out.append(f"error: {ErrorChain(error)}", style="bold red") - + if isinstance(error, str): + out.append(f"error: {error}", style="bold red") + else: + out.append(f"error: {ErrorChain(error)}", style="bold red") return out console = Console() diff --git a/verifiers/utils/save_utils.py b/verifiers/utils/save_utils.py index 1eae24a6b3..cafa374f9a 100644 --- a/verifiers/utils/save_utils.py +++ b/verifiers/utils/save_utils.py @@ -4,7 +4,6 @@ from collections import defaultdict from datetime import date, datetime from pathlib import Path -from typing import Any from datasets import Dataset from openai import AsyncOpenAI @@ -24,7 +23,29 @@ logger = logging.getLogger(__name__) -def make_serializable(value: Any) -> Any: +def is_json_serializable(value: object) -> bool: + """Check if a value is JSON-serializable without conversion. + + Returns True for JSON primitives, lists/dicts of primitives, + Pydantic models, datetime/date, Path, and exceptions. + """ + if value is None: + return True + if isinstance(value, (str, int, float, bool)): + return True + if isinstance(value, (list, tuple)): + return all(is_json_serializable(item) for item in value) + if isinstance(value, dict): + return all( + isinstance(k, str) and is_json_serializable(v) for k, v in value.items() + ) + # Types that make_serializable can handle + if isinstance(value, (BaseModel, datetime, date, Path, BaseException)): + return True + return False + + +def make_serializable(value: object) -> str | int | float | bool | list | dict | None: """Convert value to JSON-serializable types for non-standard types. Example: @@ -114,25 +135,34 @@ def get_hf_hub_dataset_name(outputs: GenerateOutputs) -> str: def state_to_output(state: State, state_columns: list[str] = []) -> RolloutOutput: - """Convert a State to a serializable RolloutOutput.""" - output: RolloutOutput = { - "example_id": state.get("example_id", 0), - "prompt": state.get("prompt"), - "completion": state.get("completion"), - "answer": state.get("answer", ""), - "task": state.get("task", "default"), - "info": state.get("info", {}), - "reward": state.get("reward", 0.0), - "error": state.get("error", None), - # Include timing as dict for compatibility with print functions - "timing": state.get("timing", {}), - # Include these for print_info compatibility - "is_completed": state.get("is_completed", False), - "is_truncated": state.get("is_truncated", False), - "stop_condition": state.get("stop_condition", None), - # Include metrics as dict for print_rewards compatibility - "metrics": state.get("metrics", {}), - } + """Convert a State to a serializable RolloutOutput. + + Args: + state: The State object to convert. + state_columns: Additional State fields to include. Values must be + JSON-serializable or an error will be raised. + + Returns: + A RolloutOutput dict with all standard fields plus state_columns. + + Raises: + ValueError: If a state_columns value is not JSON-serializable. + """ + output = RolloutOutput( + example_id=state.get("example_id", 0), + prompt=state.get("prompt"), + completion=state.get("completion"), + answer=state.get("answer", ""), + task=state.get("task", "default"), + info=state.get("info", {}), + reward=state.get("reward", 0.0), + error=state.get("error", None), + timing=state.get("timing", {}), + is_completed=state.get("is_completed", False), + is_truncated=state.get("is_truncated", False), + stop_condition=state.get("stop_condition", None), + metrics=state.get("metrics", {}), + ) # sanitize messages (handle None for error cases) prompt = state.get("prompt") if prompt is not None: @@ -152,9 +182,15 @@ def state_to_output(state: State, state_columns: list[str] = []) -> RolloutOutpu state_metrics = state.get("metrics", {}) for k, v in state_metrics.items(): output[k] = v - # add state columns + # add state columns (must be serializable) for col in state_columns: - output[col] = state.get(col) + value = state.get(col) + if not is_json_serializable(value): + raise ValueError( + f"state_columns value for '{col}' is not JSON-serializable: " + f"{type(value).__name__}. Only JSON-serializable types are allowed." + ) + output[col] = value return output @@ -280,16 +316,16 @@ def sanitize_metadata(metadata: GenerateMetadata) -> dict: return metadata_dict -def save_to_disk(results: list[dict], metadata: dict, path: Path): - """Saves (sanitized) rollouts and metadata to disk.""" +def save_to_disk(outputs: list[RolloutOutput], metadata: dict, path: Path): + """Saves rollout outputs and metadata to disk.""" path.mkdir(parents=True, exist_ok=True) - def save_results(results: list[dict], results_path: Path): + def save_outputs(outputs: list[RolloutOutput], results_path: Path): with open(results_path, "w") as f: - for idx, result in enumerate(results): - example_id = result.get("example_id") or "unknown" + for idx, output in enumerate(outputs): + example_id = output.get("example_id") or "unknown" try: - json.dump(result, f, default=make_serializable) + json.dump(output, f, default=make_serializable) f.write("\n") except Exception as e: logger.error( @@ -304,12 +340,13 @@ def save_metadata(metadata: dict, metadata_path: Path): logger.error(f"Failed to save metadata: {e}") save_metadata(metadata, path / "metadata.json") - save_results(results, path / "results.jsonl") + save_outputs(outputs, path / "results.jsonl") def make_dataset(results: GenerateOutputs) -> Dataset: """Create a Dataset from GenerateOutputs (outputs are already serialized).""" - return Dataset.from_list(results["outputs"]) + # RolloutOutput is a dict subclass, so this is type-safe at runtime + return Dataset.from_list(list(results["outputs"])) def save_generate_outputs( @@ -320,7 +357,6 @@ def save_generate_outputs( """Save GenerateOutputs to disk. Outputs are already serialized.""" path_to_save = results["metadata"]["path_to_save"] sanitized_metadata = sanitize_metadata(results["metadata"]) - # Outputs are already serialized RolloutOutputs save_to_disk(results["outputs"], sanitized_metadata, path_to_save) logger.info(f"Results saved to {path_to_save}") if push_to_hf_hub: From 8a9f81b33ce6e3a798e171cdc1e01b89ac3f776e Mon Sep 17 00:00:00 2001 From: William Brown Date: Tue, 27 Jan 2026 17:47:39 -0800 Subject: [PATCH 3/6] bugbot fixes --- verifiers/gepa/adapter.py | 2 +- verifiers/utils/save_utils.py | 56 ----------------------------------- 2 files changed, 1 insertion(+), 57 deletions(-) diff --git a/verifiers/gepa/adapter.py b/verifiers/gepa/adapter.py index 90e4c11a2f..1520728470 100644 --- a/verifiers/gepa/adapter.py +++ b/verifiers/gepa/adapter.py @@ -145,7 +145,7 @@ def make_reflective_dataset( } if trajectory.get("error"): - record["error"] = repr(trajectory["error"]) + record["error"] = trajectory["error"] if trajectory.get("stop_condition"): record["stop_condition"] = trajectory["stop_condition"] diff --git a/verifiers/utils/save_utils.py b/verifiers/utils/save_utils.py index cafa374f9a..13e00d27eb 100644 --- a/verifiers/utils/save_utils.py +++ b/verifiers/utils/save_utils.py @@ -63,62 +63,6 @@ def make_serializable(value: object) -> str | int | float | bool | list | dict | return str(value) -def states_to_generate_metadata( - env_id: str, - env_args: dict, - model: str, - client: AsyncOpenAI, - states: list[State], - state_columns: list[str] | None, - sampling_args: SamplingArgs, - start_time: float, - results_path: Path | None, -) -> GenerateMetadata: - """Converts a list of states to generate metadata.""" - base_url = str(client.base_url) if hasattr(client, "base_url") else "" - rewards = [s.get("reward", 0.0) for s in states] - avg_reward = sum(rewards) / len(rewards) if rewards else 0.0 - - metrics: dict[str, list[float]] = defaultdict(list) - for state in states: - if state.get("metrics"): - for metric_name, metric_value in state["metrics"].items(): - metrics[metric_name].append(metric_value) - avg_metrics = {k: sum(v) / len(v) if v else 0.0 for k, v in metrics.items()} - - example_ids = [s.get("example_id", 0) for s in states] - num_examples = len(set(example_ids)) if example_ids else 0 - rollouts_per_example = len(states) // num_examples if num_examples > 0 else 1 - - path_to_save = results_path or get_results_path(env_id, model) - - def tools_key(tools: list | None) -> str: - if not tools: - return "" - return str(sorted([t.get("function", {}).get("name", "") for t in tools])) - - all_tools = [s.get("oai_tools") for s in states] - unique_tools = set(tools_key(t) for t in all_tools) - tools = next((t for t in all_tools if t), None) if len(unique_tools) == 1 else None - - return GenerateMetadata( - env_id=env_id, - env_args=env_args, - model=model, - base_url=base_url, - num_examples=num_examples, - rollouts_per_example=rollouts_per_example, - sampling_args=sampling_args, - date=datetime.now().isoformat(), - time_ms=(time.time() - start_time) * 1000.0, - avg_reward=avg_reward, - avg_metrics=avg_metrics, - state_columns=state_columns or [], - path_to_save=path_to_save, - tools=tools, - ) - - def get_hf_hub_dataset_name(outputs: GenerateOutputs) -> str: """Auto-generates a dataset name.""" metadata = outputs["metadata"] From 32a792b64d51bbdbcf8755c3ab60a8ab6fd38ae9 Mon Sep 17 00:00:00 2001 From: William Brown Date: Tue, 27 Jan 2026 18:06:02 -0800 Subject: [PATCH 4/6] bugbot fixes --- verifiers/utils/eval_display.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/verifiers/utils/eval_display.py b/verifiers/utils/eval_display.py index b0dc232003..3be0d95481 100644 --- a/verifiers/utils/eval_display.py +++ b/verifiers/utils/eval_display.py @@ -22,7 +22,6 @@ from verifiers.types import EvalConfig, GenerateOutputs from verifiers.utils.display_utils import BaseDisplay, make_aligned_row -from verifiers.utils.error_utils import ErrorChain from verifiers.utils.message_utils import messages_to_printable @@ -589,7 +588,7 @@ def _make_env_detail( completion_text = _format_messages(completion) if error_0 is not None: completion_text.append("\n\nerror: ", style="bold red") - completion_text.append(str(ErrorChain(error_0)), style="bold red") + completion_text.append(error_0, style="bold red") completion_text.append("\n\nreward: ", style="bold cyan") completion_text.append(f"{reward_0:.3f}", style="bold cyan") From 8de7e47aacd3a90798c420948850162a08629f4f Mon Sep 17 00:00:00 2001 From: William Brown Date: Tue, 27 Jan 2026 18:24:05 -0800 Subject: [PATCH 5/6] bugbot fixes --- verifiers/gepa/adapter.py | 6 ++++-- verifiers/utils/eval_display.py | 7 +++---- verifiers/utils/eval_utils.py | 10 +++------- 3 files changed, 10 insertions(+), 13 deletions(-) diff --git a/verifiers/gepa/adapter.py b/verifiers/gepa/adapter.py index 1520728470..1c2cd97c4f 100644 --- a/verifiers/gepa/adapter.py +++ b/verifiers/gepa/adapter.py @@ -14,7 +14,7 @@ RolloutOutput, SamplingArgs, ) -from verifiers.utils.message_utils import message_to_printable, messages_to_printable +from verifiers.utils.message_utils import message_to_printable from verifiers.utils.save_utils import make_serializable if TYPE_CHECKING: @@ -96,6 +96,7 @@ def evaluate( sampling_args=self.sampling_args, max_concurrent=self.max_concurrent, use_tqdm=self.use_tqdm, + state_columns=self.state_columns, ) ) @@ -136,10 +137,11 @@ def make_reflective_dataset( records = [] # outputs, trajectories, and scores should be the same length + # Note: prompt/completion are already in printable format from state_to_output for output, trajectory, score in zip(outputs, trajectories, scores): record: dict[str, Any] = { "query": _extract_user_query(output["prompt"]), - "completion": messages_to_printable(output["completion"]), + "completion": output["completion"], "expected_answer": output.get("answer", ""), "reward": score, } diff --git a/verifiers/utils/eval_display.py b/verifiers/utils/eval_display.py index 3be0d95481..2342fd08cc 100644 --- a/verifiers/utils/eval_display.py +++ b/verifiers/utils/eval_display.py @@ -22,7 +22,6 @@ from verifiers.types import EvalConfig, GenerateOutputs from verifiers.utils.display_utils import BaseDisplay, make_aligned_row -from verifiers.utils.message_utils import messages_to_printable @dataclass @@ -567,11 +566,11 @@ def _make_env_detail( """Create detailed content for a single environment's summary.""" items: list[Panel] = [] - # Example 0 prompt/completion + # Example 0 prompt/completion (already in printable format from state_to_output) outputs = results["outputs"] if outputs and outputs[0]["prompt"] and outputs[0]["completion"]: - prompt = messages_to_printable(outputs[0]["prompt"]) - completion = messages_to_printable(outputs[0]["completion"]) + prompt = outputs[0]["prompt"] + completion = outputs[0]["completion"] reward_0 = outputs[0]["reward"] if outputs[0]["reward"] else 0.0 error_0 = outputs[0].get("error") if outputs[0] else None diff --git a/verifiers/utils/eval_utils.py b/verifiers/utils/eval_utils.py index 9234b63a97..03da565ae6 100644 --- a/verifiers/utils/eval_utils.py +++ b/verifiers/utils/eval_utils.py @@ -32,7 +32,6 @@ from verifiers.utils.async_utils import EventLoopLagMonitor from verifiers.utils.client_utils import setup_client from verifiers.utils.logging_utils import print_prompt_completions_sample, print_time -from verifiers.utils.message_utils import messages_to_printable from verifiers.utils.path_utils import get_eval_results_path logger = logging.getLogger(__name__) @@ -277,13 +276,10 @@ def print_results(results: GenerateOutputs, num_samples: int = 1): print(f"Rollouts per example: {results['metadata']['rollouts_per_example']}") print("--- Example ---") - printable_prompts = [ - messages_to_printable(o["prompt"]) if o["prompt"] else [] - for o in results["outputs"] - ] + # prompt/completion are already in printable format from state_to_output + printable_prompts = [o["prompt"] if o["prompt"] else [] for o in results["outputs"]] printable_completions = [ - messages_to_printable(o["completion"]) if o["completion"] else [] - for o in results["outputs"] + o["completion"] if o["completion"] else [] for o in results["outputs"] ] rewards = [o["reward"] for o in results["outputs"]] errors = [o.get("error") for o in results["outputs"]] From a17a1d17cf7822bd0d3f260d51203729e5b6a7d7 Mon Sep 17 00:00:00 2001 From: William Brown Date: Tue, 27 Jan 2026 18:42:35 -0800 Subject: [PATCH 6/6] bugbot fixes --- docs/reference.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/docs/reference.md b/docs/reference.md index 0d43d7aca3..88c7ab2757 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -119,7 +119,7 @@ class RolloutInput(TypedDict): ### RolloutOutput ```python -class RolloutOutput(TypedDict): +class RolloutOutput(dict): # Required fields example_id: int task: str @@ -135,9 +135,11 @@ class RolloutOutput(TypedDict): info: Info error: str | None stop_condition: str | None + trajectory: list[TrajectoryStep] + oai_tools: list[ChatCompletionToolParam] ``` -Serialized output from a rollout. This is the JSON-serializable representation of a `State` object, used in `GenerateOutputs` and for saving results to disk. Additional fields from `state_columns` may also be present. +Serialized output from a rollout. This is a `dict` subclass that provides typed access to known fields while supporting arbitrary additional fields from `state_columns`. All values must be JSON-serializable. Used in `GenerateOutputs` and for saving results to disk. ### TrajectoryStep