diff --git a/verifiers/types.py b/verifiers/types.py index 87aa10a7b2..a0bbbc68d4 100644 --- a/verifiers/types.py +++ b/verifiers/types.py @@ -216,7 +216,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. @@ -263,7 +263,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 2c8d347cb7..4b7ca05469 100644 --- a/verifiers/v1/clients/train.py +++ b/verifiers/v1/clients/train.py @@ -12,6 +12,7 @@ from openai import OpenAIError from renderers import OverlongPromptError as RendererOverlongPromptError from renderers import RenderedTokens, Renderer, RendererConfig +from renderers.base import is_multimodal from verifiers.v1.clients.base import build_async_openai from verifiers.v1.clients.client import SESSION_ID_HEADER, Client @@ -172,16 +173,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 - - @dataclass class RendererSlot: """One renderer and the rollouts currently holding it. Encoding mutates a fast @@ -366,22 +357,23 @@ async def get_response( async with pool.acquire() as slot: renderer = slot.renderer # 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) - ) + # 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 slot.run(bridge) diff --git a/verifiers/v1/graph.py b/verifiers/v1/graph.py index f039e80cf6..b3bb96431f 100644 --- a/verifiers/v1/graph.py +++ b/verifiers/v1/graph.py @@ -31,12 +31,14 @@ from verifiers.v1.types import ( AssistantMessage, + FinishReason, KeptTokens, Message, Response, TextContentPart, Tool, ToolMessage, + Usage, ) if TYPE_CHECKING: @@ -61,6 +63,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(BaseModel): """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 @@ -101,17 +143,24 @@ class MessageNode(BaseModel): 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.""" advantages: list[float] | None = None """Per-token credit over the sampled tokens, same layout as `logprobs`. `None` until a consumer's RL algorithm assigns it, which is not the same as a credit of zero: a group whose rewards were all equal is assigned zeros and carries no gradient, while an unassigned node was never scored at all.""" 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 @@ -129,10 +178,10 @@ class MessageNode(BaseModel): @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": { @@ -140,9 +189,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() }, } @@ -150,25 +197,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") @@ -368,6 +419,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) @@ -428,8 +496,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] = {} @@ -439,6 +508,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: @@ -450,13 +520,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/legacy.py b/verifiers/v1/legacy.py index 0dd81f02cb..49ec1329b1 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 @@ -193,6 +194,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"), ) @@ -428,6 +430,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, @@ -436,13 +452,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: @@ -467,12 +483,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 52e479bcbd..c9005b5eff 100644 --- a/verifiers/v1/trace.py +++ b/verifiers/v1/trace.py @@ -241,7 +241,12 @@ def advantages(self) -> list[float] | None: @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: @@ -253,6 +258,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 diff --git a/verifiers/v1/types.py b/verifiers/v1/types.py index dcb99f7067..00fa5e7856 100644 --- a/verifiers/v1/types.py +++ b/verifiers/v1/types.py @@ -202,8 +202,8 @@ class TurnTokens(BaseModel): 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`,