diff --git a/tests/test_client_multimodal_types.py b/tests/test_client_multimodal_types.py index 934cccb2f0..0e82a98bc5 100644 --- a/tests/test_client_multimodal_types.py +++ b/tests/test_client_multimodal_types.py @@ -1,7 +1,6 @@ from types import SimpleNamespace import pytest - from verifiers.clients.openai_chat_completions_client import OpenAIChatCompletionsClient from verifiers.types import ( AssistantMessage, @@ -214,7 +213,6 @@ async def test_anthropic_merges_consecutive_tool_results_into_single_user_messag async def test_anthropic_from_native_response_extracts_usage(): anthropic = pytest.importorskip("anthropic") from anthropic.types import Message as AnthropicMessage - from verifiers.clients.anthropic_messages_client import AnthropicMessagesClient client = AnthropicMessagesClient(object()) @@ -267,7 +265,6 @@ async def test_anthropic_tool_call_round_trips_thinking_blocks(): pytest.importorskip("anthropic") from anthropic.types import Message as AnthropicMessage from anthropic.types import Usage as AnthropicUsage - from verifiers.clients.anthropic_messages_client import AnthropicMessagesClient client = AnthropicMessagesClient(object()) @@ -296,3 +293,4 @@ async def test_anthropic_tool_call_round_trips_thinking_blocks(): {"type": "thinking", "thinking": "hidden chain", "signature": "sig_1"}, {"type": "tool_use", "id": "call_1", "name": "lookup", "input": {"q": "x"}}, ] + diff --git a/verifiers/types.py b/verifiers/types.py index 35eeefdeb6..5f9909cfed 100644 --- a/verifiers/types.py +++ b/verifiers/types.py @@ -214,7 +214,7 @@ class ResponseTokens(CustomBaseModel): completion_logprobs: list[float] routed_experts: RoutedExpertsPayload | None = None # Renderer-emitted multimodal sidecar (renderers.base.MultiModalData) - # carrying processed pixel_values / placeholder ranges per modality. + # carrying raw image descriptors / placeholder ranges per modality. # Populated by the renderer client when the rollout went through a # multimodal-aware renderer; ``None`` otherwise. Stored as ``Any`` to # avoid a hard import dependency on ``renderers`` at this layer. @@ -261,7 +261,7 @@ class TrajectoryStepTokens(TypedDict): is_truncated: bool routed_experts: RoutedExpertsPayload | None # Renderer-emitted multimodal sidecar (renderers.base.MultiModalData) - # carrying processed pixel_values / placeholder ranges per modality. + # carrying raw image descriptors / placeholder ranges per modality. # ``NotRequired`` because text-only rollouts (and non-renderer client # types) never populate it. multi_modal_data: NotRequired[Any] diff --git a/verifiers/v1/clients/train.py b/verifiers/v1/clients/train.py index 6b7e7fad8e..12eded4207 100644 --- a/verifiers/v1/clients/train.py +++ b/verifiers/v1/clients/train.py @@ -13,9 +13,9 @@ from typing import Any from openai import AsyncOpenAI, OpenAIError -from renderers import RenderedTokens from renderers import OverlongPromptError as RendererOverlongPromptError -from renderers import RendererConfig +from renderers import RenderedTokens, RendererConfig +from renderers.base import is_multimodal from verifiers.v1.clients.client import SESSION_ID_HEADER, Client from verifiers.v1.dialects import FINISH_REASONS, ChatDialect, Dialect, parse_tools @@ -170,16 +170,6 @@ def _is_valid_incremental_tail(messages: list[dict[str, Any]]) -> bool: return all(role == "tool" for role in roles) -def _has_multimodal_content(messages) -> bool: - for message in messages: - content = getattr(message, "content", None) - if not isinstance(content, list): - continue - if any(getattr(part, "type", None) == "image_url" for part in content): - return True - return False - - class TrainClient(Client): """Renders prompts to token ids and calls a vLLM `/inference/v1/generate` engine.""" @@ -266,23 +256,24 @@ async def get_response( ) bridged_turn: PendingTurn | None = None - # Only build the (O(context)) previous-turn token ids once the cheap guards pass — a - # multimodal prompt or a tail that isn't a clean `[tool*, user?]` extension can't bridge. - can_bridge = ( - turn is not None - and not _has_multimodal_content(prompt) - and _is_valid_incremental_tail(wire_messages) - ) + # Only build the (O(context)) previous-turn token ids once the cheap guards pass: a + # tail that isn't a clean `[tool*, user?]` extension can't bridge. + can_bridge = turn is not None and _is_valid_incremental_tail(wire_messages) previous_ids = turn.previous_token_ids() if can_bridge else None if previous_ids is not None: previous_prompt_ids, previous_completion_ids = previous_ids def bridge(): + kwargs: dict[str, Any] = {"tools": wire_tools} + if is_multimodal(renderer): + kwargs["previous_multi_modal_data"] = ( + turn.previous_multi_modal_data() + ) return renderer.bridge_to_next_turn( previous_prompt_ids, previous_completion_ids, wire_messages, - tools=wire_tools, + **kwargs, ) bridged = await _maybe_offload(renderer, bridge) diff --git a/verifiers/v1/graph.py b/verifiers/v1/graph.py index 8c9c17ec97..151e9c8a2c 100644 --- a/verifiers/v1/graph.py +++ b/verifiers/v1/graph.py @@ -31,6 +31,7 @@ from verifiers.v1.types import ( AssistantMessage, + FinishReason, KeptTokens, Message, Response, @@ -38,6 +39,7 @@ TextContentPart, Tool, ToolMessage, + Usage, ) if TYPE_CHECKING: @@ -62,6 +64,46 @@ def _decode_ndarray(d: dict) -> np.ndarray: return np.frombuffer(d["data"], dtype=np.dtype(d["dtype"])).reshape(d["shape"]) +_PROCESSED_MM_KEYS = frozenset({"pixel_values", "image_embeds", "image_features"}) + + +def _contains_processed_mm_key(value: Any) -> bool: + if isinstance(value, dict): + return bool(_PROCESSED_MM_KEYS.intersection(value)) or any( + _contains_processed_mm_key(v) for v in value.values() + ) + if isinstance(value, (list, tuple)): + return any(_contains_processed_mm_key(v) for v in value) + return False + + +def _validate_raw_mm_item(item: Any) -> dict[str, Any]: + if not isinstance(item, dict): + raise TypeError( + "v1 multimodal sidecars must be raw image descriptor dicts, " + f"got {type(item).__name__}" + ) + if _contains_processed_mm_key(item): + raise TypeError( + "v1 multimodal sidecars must be raw image descriptors, " + "not processed multimodal payloads" + ) + if not isinstance(item.get("raw_image_data"), str) or not item["raw_image_data"]: + raise ValueError("v1 multimodal sidecars require inline raw_image_data") + return dict(item) + + +def _validate_raw_mm_data(mmd: MultiModalData) -> MultiModalData: + return MultiModalData( + mm_hashes={k: list(v) for k, v in mmd.mm_hashes.items()}, + mm_placeholders={k: list(v) for k, v in mmd.mm_placeholders.items()}, + mm_items={ + modality: [_validate_raw_mm_item(item) for item in items] + for modality, items in mmd.mm_items.items() + }, + ) + + class MessageNode(StrictBaseModel): """One message in the graph: a message plus the tokens it adds to the cumulative sequence. Concatenating a root→leaf path's nodes reconstructs that branch's full token @@ -102,12 +144,19 @@ class MessageNode(StrictBaseModel): logprobs: list[float] = Field(default_factory=list) """Sampling logprobs for the sampled tokens — length equals the number of True entries in `mask`; empty for input messages.""" + finish_reason: FinishReason = None + """The response's finish reason (assistant nodes only) — kept for truncation detection.""" multi_modal_data: SkipJsonSchema[MultiModalData | None] = None - """The renderer items for the images this message's content introduces (pixel tensors, - grids, hashes, placeholders) — the only carrier of the pixels from the env server to the - trainer. `Branch.multi_modal_data` concatenates them along the path into the training - `mm_kwargs`. Rides the wire as raw bytes (msgpack `bin`) since pydantic can't JSON the numpy; - kept off disk by the dump-site `exclude` in prime-rl (the tensors bloat the rollout jsonl).""" + """The renderer items for images this message introduces. + + With the raw-image path, items are descriptors (hashes, grid metadata, and the inline + base64 image source), not image processor tensors. `Branch.multi_modal_data` concatenates + them along the path for the trainer. Old processed-payload sidecars are rejected. + """ + usage: Usage | None = None + """Provider-reported token usage for this message's response (assistant nodes). Preserved on + the wire and on disk so dashboards can show token counts and cost even when the endpoint + returns no token ids.""" routed_experts: SkipJsonSchema[np.ndarray | None] = None """This node's slice of the MoE expert-routing array — uint8 `[len(token_ids), layers, top_k]`, the expert ids inference selected for exactly this node's tokens. Attributed from @@ -125,10 +174,10 @@ class MessageNode(StrictBaseModel): @field_serializer("multi_modal_data") def serialize_multi_modal_data(self, mmd: MultiModalData | None) -> dict | None: - """`MultiModalData` -> msgpack-safe dict so the pixel tensors ride the wire; numpy - `mm_items` values become raw-bytes `__nd__` dicts (every renderer emits `return_tensors="np"`).""" + """`MultiModalData` -> msgpack-safe raw descriptor dict.""" if mmd is None: return None + mmd = _validate_raw_mm_data(mmd) return { "mm_hashes": {k: list(v) for k, v in mmd.mm_hashes.items()}, "mm_placeholders": { @@ -136,9 +185,7 @@ def serialize_multi_modal_data(self, mmd: MultiModalData | None) -> dict | None: for modality, ranges in mmd.mm_placeholders.items() }, "mm_items": { - modality: [ - {k: _encode_ndarray(v) for k, v in item.items()} for item in items - ] + modality: [dict(item) for item in items] for modality, items in mmd.mm_items.items() }, } @@ -146,25 +193,29 @@ def serialize_multi_modal_data(self, mmd: MultiModalData | None) -> dict | None: @field_validator("multi_modal_data", mode="before") @classmethod def deserialize_multi_modal_data(cls, value: Any) -> MultiModalData | None: - if value is None or isinstance(value, MultiModalData): + if value is None: return value + if isinstance(value, MultiModalData): + return _validate_raw_mm_data(value) if not isinstance(value, dict): raise TypeError(f"cannot build MultiModalData from {type(value).__name__}") - return MultiModalData( - mm_hashes={k: list(v) for k, v in (value.get("mm_hashes") or {}).items()}, - mm_placeholders={ - modality: [ - PlaceholderRange(offset=p["offset"], length=p["length"]) - for p in ranges - ] - for modality, ranges in (value.get("mm_placeholders") or {}).items() - }, - mm_items={ - modality: [ - {k: _decode_ndarray(v) for k, v in item.items()} for item in items - ] - for modality, items in (value.get("mm_items") or {}).items() - }, + return _validate_raw_mm_data( + MultiModalData( + mm_hashes={ + k: list(v) for k, v in (value.get("mm_hashes") or {}).items() + }, + mm_placeholders={ + modality: [ + PlaceholderRange(offset=p["offset"], length=p["length"]) + for p in ranges + ] + for modality, ranges in (value.get("mm_placeholders") or {}).items() + }, + mm_items={ + modality: list(items) + for modality, items in (value.get("mm_items") or {}).items() + }, + ) ) @field_serializer("routed_experts") @@ -332,6 +383,23 @@ def prompt_message_spans( for span in tail_spans ] + def previous_multi_modal_data(self) -> MultiModalData | None: + """Concatenate multimodal sidecars attached to the reusable prefix.""" + merged = MultiModalData() + found = False + for nid in self.prefix_node_ids: + mmd = self.trace.nodes[nid].multi_modal_data + if mmd is None or mmd.is_empty(): + continue + found = True + for modality, items in mmd.mm_items.items(): + merged.mm_items.setdefault(modality, []).extend(items) + for modality, hashes in mmd.mm_hashes.items(): + merged.mm_hashes.setdefault(modality, []).extend(hashes) + for modality, placeholders in mmd.mm_placeholders.items(): + merged.mm_placeholders.setdefault(modality, []).extend(placeholders) + return merged if found else None + def commit(self, response: Response, tools: list[Tool] | None = None) -> int: """Add this turn to the graph; returns the committed assistant node's id.""" assistant_id = _commit_turn(self, response) @@ -392,8 +460,9 @@ def _attribute_mm( renderer emits items per modality in prompt order (message order, then content-part order), so we walk the path advancing a per-modality cursor over every message's media but write only the nodes created this turn — `path[:num_reused]` is the reused prefix, already - attributed when first created. Item order is all training needs; placeholder offsets aren't - carried.""" + attributed when first created. Each node gets the hashes/items/placeholders for exactly the + media it introduced, preserving vLLM multimodal-list alignment when those node sidecars are + later merged for bridge or training.""" if mmd is None or mmd.is_empty(): return cursors: dict[str, int] = {} @@ -403,6 +472,7 @@ def _attribute_mm( continue node_items: dict[str, list] = {} node_hashes: dict[str, list] = {} + node_placeholders: dict[str, list[PlaceholderRange]] = {} for part in content: modality = _part_modality(part) if modality is None: @@ -414,13 +484,20 @@ def _attribute_mm( continue items = mmd.mm_items.get(modality) or [] hashes = mmd.mm_hashes.get(modality) or [] + placeholders = mmd.mm_placeholders.get(modality) or [] if k < len(items): node_items.setdefault(modality, []).append(items[k]) if k < len(hashes): node_hashes.setdefault(modality, []).append(hashes[k]) - if node_items: - trace.nodes[node_id].multi_modal_data = MultiModalData( - mm_items=node_items, mm_hashes=node_hashes + if k < len(placeholders): + node_placeholders.setdefault(modality, []).append(placeholders[k]) + if node_items or node_hashes or node_placeholders: + trace.nodes[node_id].multi_modal_data = _validate_raw_mm_data( + MultiModalData( + mm_items=node_items, + mm_hashes=node_hashes, + mm_placeholders=node_placeholders, + ) ) diff --git a/verifiers/v1/interception/server.py b/verifiers/v1/interception/server.py index 13a7e300f4..623519b519 100644 --- a/verifiers/v1/interception/server.py +++ b/verifiers/v1/interception/server.py @@ -33,9 +33,9 @@ from pydantic import TypeAdapter, ValidationError from pydantic_core import PydanticSerializationError, from_json, to_json +from verifiers.v1 import graph from verifiers.v1.dialects import DIALECTS, Dialect from verifiers.v1.dialects.base import is_sse_done_event -from verifiers.v1 import graph from verifiers.v1.errors import ( OverlongPromptError, ProviderError, diff --git a/verifiers/v1/legacy.py b/verifiers/v1/legacy.py index 9fe73f644a..0277648404 100644 --- a/verifiers/v1/legacy.py +++ b/verifiers/v1/legacy.py @@ -13,6 +13,7 @@ v1 stays importable without the v0 package present. """ +import asyncio import contextlib import logging from pathlib import Path @@ -190,6 +191,7 @@ def _to_v1_tokens(raw: Any) -> TurnTokens | None: prompt_ids=list(raw.get("prompt_ids") or []), completion_ids=list(raw.get("completion_ids") or []), completion_logprobs=list(raw.get("completion_logprobs") or []), + multi_modal_data=raw.get("multi_modal_data"), ) @@ -415,6 +417,20 @@ def _v0_client(self, client_config: ClientConfig, model: str): self._clients[key] = resolve_client(v0_config) return self._clients[key] + async def _state_output_with_live_trajectory(self, state: Any) -> dict: + """Build v0 rollout output metadata while preserving live trajectory sidecars. + + The JSON save path deltas ``tokens.multi_modal_data`` to avoid repeated + cumulative multimodal sidecars. Trace reconstruction needs the live, + cumulative sidecar for each turn so image descriptors align with the + full prompt the renderer saw. + """ + from verifiers.utils.save_utils import state_to_output + + out = await asyncio.to_thread(state_to_output, state, []) + out["trajectory"] = state.get("trajectory", []) + return out + async def _run_v0( self, task_idx: int, @@ -423,13 +439,13 @@ async def _run_v0( sampling: SamplingConfig, ) -> dict: client = self._v0_client(client_config, model) - return await self.env.run_rollout( + state = await self.env._run_rollout_state( input=dict(self.dataset[task_idx]), client=client, model=model, sampling_args=sampling.model_dump(exclude_none=True), - state_columns=["trajectory"], ) + return await self._state_output_with_live_trajectory(state) @staticmethod def _row(req: RunRequest) -> int: @@ -454,12 +470,14 @@ async def _run(self, req: RunRequest) -> RunResponse: async def _run_group(self, req: RunGroupRequest) -> RunGroupResponse: client = self._v0_client(req.client, req.model) # run_group scores the rollouts together so group/preference reward funcs apply. - outs = await self.env.run_group( + states = await self.env._run_group_states( group_inputs=[dict(self.dataset[req.task_idx]) for _ in range(req.n)], client=client, model=req.model, sampling_args=req.sampling.model_dump(exclude_none=True), - state_columns=["trajectory"], + ) + outs = await asyncio.gather( + *(self._state_output_with_live_trajectory(state) for state in states) ) traces = [ rollout_output_to_trace(out, req.task_idx).model_dump() for out in outs diff --git a/verifiers/v1/trace.py b/verifiers/v1/trace.py index 7e2e18720b..77b2d3ab04 100644 --- a/verifiers/v1/trace.py +++ b/verifiers/v1/trace.py @@ -170,7 +170,12 @@ def logprobs(self) -> list[float]: @property def multi_modal_data(self) -> MultiModalData | None: - """Node image data concatenated in token order for training; never persisted.""" + """The branch's multimodal sidecar — every node's images concatenated in path order. + + None when the branch has no images. The raw-image path carries lightweight descriptors + plus placeholder ranges, so downstream vLLM/training multimodal payloads can align hashes, + placeholders, and item refs without reprocessing images in the env worker. + """ merged = MultiModalData() found = False for node in self.nodes: @@ -182,6 +187,8 @@ def multi_modal_data(self) -> MultiModalData | None: merged.mm_items.setdefault(modality, []).extend(items) for modality, hashes in mmd.mm_hashes.items(): merged.mm_hashes.setdefault(modality, []).extend(hashes) + for modality, placeholders in mmd.mm_placeholders.items(): + merged.mm_placeholders.setdefault(modality, []).extend(placeholders) return merged if found else None @property @@ -262,7 +269,9 @@ def num_input_tokens(self) -> int: } } } -"""Raw tensor fields kept on the msgpack wire but excluded from JSON records.""" +"""Trainer-only sidecars kept on the msgpack wire but excluded from JSON records: raw +numpy tensors (``routed_experts``) can't round-trip JSON, while ``multi_modal_data`` +(renderer descriptors) and ``kept_tokens`` exist for training, not the rollout record.""" TRACE_VERSION = 2 diff --git a/verifiers/v1/types.py b/verifiers/v1/types.py index 7ee25f8502..483897957c 100644 --- a/verifiers/v1/types.py +++ b/verifiers/v1/types.py @@ -206,8 +206,8 @@ class TurnTokens(StrictBaseModel): default=None, exclude=True ) is_content: list[bool] | None = Field(default=None, exclude=True) - # Transient carrier (excluded): the renderer's multimodal sidecar (image tensors + offsets), - # attributed per node by the turn's `commit`, then dropped — never persisted. + # Transient carrier (excluded): the renderer's multimodal sidecar (raw-image descriptors, + # hashes, and placeholder offsets), attributed per node by the turn's `commit`, then dropped. multi_modal_data: MultiModalData | None = Field(default=None, exclude=True) # Transient carrier (excluded): the MoE expert-routing data from `generate` (expert ids # per token), attributed per node by the turn's `commit` into `MessageNode.routed_experts`,