diff --git a/tests/unit/test_graph.py b/tests/unit/test_graph.py index 4df44699ed..3c70574358 100644 --- a/tests/unit/test_graph.py +++ b/tests/unit/test_graph.py @@ -740,3 +740,60 @@ def test_prompt_supplied_assistant_messages_are_not_sampled_turns(): assert [n.sampled for n in trace.nodes] == [False, False, False, True] assert trace.num_turns == 1 assert trace.assistant_messages == [response] + + +def test_payload_segments_attributed_per_node_and_rebased_per_branch(): + """By-handle rows tile the nodes a turn creates. A bridged turn's first row belongs to the + prior sampled turn's final token; a reused non-sampled prefix row is dropped. Each branch + rebases its nodes' segments onto its own token positions.""" + trace = vf.Trace( + agent=vf.AgentInfo(config=vf.AgentConfig()), + task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="x")), + ) + user = vf.UserMessage(content="u1") + + def commit(messages, content, prompt_ids, spans, pos, rows): + segment = vf.PayloadSegment( + field="routed_experts", + file="f", + offset=pos * 2, + pos=pos, + rows=rows, + dtype="uint8", + shape=[2], + ) + graph.prepare_turn(trace, messages).commit( + vf.Response( + id=content, + created=0, + model="t", + message=vf.AssistantMessage(content=content), + finish_reason="stop", + tokens=TurnTokens( + prompt_ids=prompt_ids, + completion_ids=[40, 41], + message_spans=spans, + payload=[segment], + ), + ) + ) + + def rows(branch): + return [(s.pos, s.rows, s.offset) for s in branch.payload] + + commit([user], "a1", [10, 11, 12], [(0, 3)], 0, 4) + # Bridged: request row 4 is a1's unforwarded final token. + commit( + [user, vf.AssistantMessage(content="a1"), vf.UserMessage(content="u2")], + "a2", + [10, 11, 12, 40, 41, 30, 31], + [(0, 3), None, (5, 7)], + 4, + 4, + ) + # Fork off the shared user prompt: re-forwarded rows 0-2 are dropped. + commit([user], "b1", [10, 11, 12], [(0, 3)], 0, 4) + + a2, b1 = (b for b in trace.branches if len(b.token_ids) in (9, 5)) + assert rows(a2) == [(0, 3, 0), (3, 1, 6), (4, 1, 8), (5, 2, 10), (7, 1, 14)] + assert rows(b1) == [(0, 3, 0), (3, 1, 6)] diff --git a/verifiers/v1/__init__.py b/verifiers/v1/__init__.py index 265f69286b..9fd2f371c0 100644 --- a/verifiers/v1/__init__.py +++ b/verifiers/v1/__init__.py @@ -153,6 +153,8 @@ Message, MessageContent, Messages, + NodeSegment, + PayloadSegment, Request, Response, SamplingMask, @@ -260,8 +262,10 @@ "ModelCall": "verifiers.v1.trace", "ModelContext": "verifiers.v1.clients", "NetworkPolicyConfig": "verifiers.v1.configs.runtime", + "NodeSegment": "verifiers.v1.types", "PATCH_CAP_BYTES": "verifiers.v1.utils.git", "ParentLink": "verifiers.v1.semantic", + "PayloadSegment": "verifiers.v1.types", "PolicyEvent": "verifiers.v1.trace", "PolicySpan": "verifiers.v1.episode", "PrimeConfig": "verifiers.v1.configs.runtime", @@ -440,6 +444,8 @@ "Branch", "TurnTokens", "SamplingMask", + "PayloadSegment", + "NodeSegment", "TopLogprobs", "Timing", "TimeSpan", diff --git a/verifiers/v1/clients/train.py b/verifiers/v1/clients/train.py index 67ddb99b29..d3224fe30e 100644 --- a/verifiers/v1/clients/train.py +++ b/verifiers/v1/clients/train.py @@ -25,6 +25,7 @@ from verifiers.v1.types import ( AssistantMessage, FinishReason, + PayloadSegment, Response, SamplingMask, Tool, @@ -193,6 +194,9 @@ def response_from_generate( sampling_mask=SamplingMask.from_sampling_mask(mask) if (mask := result.get("sampling_mask")) else None, + payload=[PayloadSegment(**segment) for segment in segments] + if (segments := result.get("payload")) + else None, weight_versions=[WeightVersionSpan(**span) for span in spans] if (spans := result.get("weight_versions")) else None, diff --git a/verifiers/v1/graph.py b/verifiers/v1/graph.py index c87be9302c..402fc6be45 100644 --- a/verifiers/v1/graph.py +++ b/verifiers/v1/graph.py @@ -48,6 +48,8 @@ FloatStream, IntStream, Message, + NodeSegment, + PayloadSegment, Response, SamplingMask, TextContentPart, @@ -778,6 +780,43 @@ def _attribute_routed_experts( off = end +def _attribute_payload( + trace: Trace, + prefix_node_ids: list[int], + new_node_ids: list[int], + path_len: int, + segments: list[PayloadSegment] | None, +) -> None: + """Append this turn's by-handle rows to `Trace.payload`, node by node. The request's + positions tile `[path_len:]` over the nodes created this turn. The one position before, + the prefix's final token, belongs to the prefix's last node when that node was sampled: + the previous turn never forwarded its final token, so this prefill is the first to + produce its row. Rows of other reused prefix positions (a turn that re-forwarded its + history) are dropped. Rows are never rewritten, only appended.""" + if not segments: + return + # (node, request lo, request hi, node start) + spans: list[tuple[int, int, int, int]] = [] + last = next( + (nid for nid in reversed(prefix_node_ids) if len(trace.nodes[nid].token_ids)), + None, + ) + if last is not None and trace.nodes[last].sampled: + spans.append( + (last, path_len - 1, path_len, path_len - len(trace.nodes[last].token_ids)) + ) + start = path_len + for nid in new_node_ids: + end = start + len(trace.nodes[nid].token_ids) + spans.append((nid, start, end, start)) + start = end + for segment in segments: + for nid, lo, hi, node_start in spans: + part = segment.clip(lo, hi, shift=-node_start) + if part is not None: + trace.payload.append(NodeSegment(node=nid, segment=part)) + + def _attribute_sampling_mask( trace: Trace, assistant_id: int, payload: SamplingMask | None ) -> None: @@ -1092,6 +1131,10 @@ def _commit_turn(turn: PendingTurn, response: Response) -> int: trace, assistant_id, tokens.top_logprobs if tokens else None ) + _attribute_payload( + trace, prefix, new_node_ids, path_len, tokens.payload if tokens else None + ) + return assistant_id diff --git a/verifiers/v1/serve/delta.py b/verifiers/v1/serve/delta.py index 24b67c543c..f713c0fd95 100644 --- a/verifiers/v1/serve/delta.py +++ b/verifiers/v1/serve/delta.py @@ -46,6 +46,7 @@ LIST_FIELDS = ( "nodes", "calls", + "payload", "errors", "extra_usage", "request_rewrites", diff --git a/verifiers/v1/trace.py b/verifiers/v1/trace.py index fca578ac62..b7ed382827 100644 --- a/verifiers/v1/trace.py +++ b/verifiers/v1/trace.py @@ -29,6 +29,8 @@ FinishReason, Message, Messages, + NodeSegment, + PayloadSegment, SamplingMask, Tool, ToolMessage, @@ -48,9 +50,11 @@ "sampling_mask", "top_logprobs", } - } + }, + "payload": True, } -"""Raw tensor fields kept on the msgpack wire but excluded from disk serialization.""" +"""Raw tensor fields and short-lived payload handles, kept on the msgpack wire but excluded +from disk serialization.""" class TimeSpan(BaseModel): @@ -209,6 +213,9 @@ class Branch(BaseModel): """Whether this physical path contributes a training sample.""" mm_token_type_id_map: dict[int, int] = Field(default_factory=dict) """The trace's `mm_token_type_id_map`, carried so `mm_token_type_ids` is self-contained.""" + payload: list[PayloadSegment] = Field(default_factory=list) + """The trace's payload segments on this path, positions relative to `token_ids`, in + the order they were appended (a later segment wins where two overlap).""" @property def messages(self) -> Messages: @@ -413,6 +420,10 @@ class Trace(BaseModel, Generic[DataT, StateT, AgentConfigT]): """The message graph, including physical and semantic parent links.""" calls: list[ModelCall] = Field(default_factory=list) """Every model call; automatically recorded at intercept time + linked into `nodes`.""" + payload: list[NodeSegment] = Field(default_factory=list) + """By-handle per-token side arrays (router-replay ids, sampling masks), append-only: + each turn's `commit` attributes the inference server's segments to the nodes they cover. + Handles point into a short-lived shared directory; kept off disk records.""" mm_token_type_id_map: dict[int, int] = Field(default_factory=dict) """Special-token id -> modality marker (1 = image placeholder, 2 = video placeholder) from the renderer that tokenized this trace, stamped at turn commit. Applied to a @@ -562,6 +573,11 @@ def branches(self) -> list[Branch]: path.append(nid) nid = self.nodes[nid].parent path.reverse() + node_starts: dict[int, int] = {} + offset = 0 + for n in path: + node_starts[n] = offset + offset += len(self.nodes[n].token_ids) is_compaction_attempt = any( link.type == "compaction_attempt" for link in self.nodes[leaf].semantic_parents @@ -576,6 +592,13 @@ def branches(self) -> list[Branch]: or leaf in accepted_compaction_attempts ), mm_token_type_id_map=self.mm_token_type_id_map, + payload=[ + entry.segment.model_copy( + update={"pos": entry.segment.pos + node_starts[entry.node]} + ) + for entry in self.payload + if entry.node in node_starts + ], ) ) return branches diff --git a/verifiers/v1/types.py b/verifiers/v1/types.py index 9d79d837ab..c4f19882d6 100644 --- a/verifiers/v1/types.py +++ b/verifiers/v1/types.py @@ -299,6 +299,48 @@ def _decode_packed(packed: dict[str, Any], dtype: type) -> np.ndarray: return array.astype(dtype, copy=False) +class PayloadSegment(BaseModel): + """`rows` consecutive rows of one per-token side array (`field`, e.g. `routed_experts`, + `sampling_mask`) that the inference server wrote to `file` at byte `offset`, as a + C-contiguous `[rows, *shape]` `dtype` array. `pos` is the token position of the first + row in whatever sequence holds the segment: the request on `TurnTokens`, the node on + `Trace.payload`, the branch on `Branch.payload`. Only the handle travels; the trainer + reads the rows.""" + + field: str + file: str + offset: int + pos: int + rows: int + dtype: str + shape: list[int] + + @property + def end(self) -> int: + return self.pos + self.rows + + def clip(self, lo: int, hi: int, shift: int = 0) -> "PayloadSegment | None": + """The rows inside positions `[lo, hi)`, moved by `shift`; None when empty.""" + start, end = max(self.pos, lo), min(self.end, hi) + if start >= end: + return None + row_bytes = np.dtype(self.dtype).itemsize * int(np.prod(self.shape)) + return self.model_copy( + update={ + "offset": self.offset + (start - self.pos) * row_bytes, + "pos": start + shift, + "rows": end - start, + } + ) + + +class NodeSegment(BaseModel): + """A payload segment attributed to `Trace.nodes[node]`, positions node-local.""" + + node: int + segment: PayloadSegment + + @dataclass class TopLogprobs: """Top-k sampling head stored as flat `ids`, `logprobs` and `counts` arrays. @@ -369,6 +411,9 @@ class TurnTokens(BaseModel): # Transient carrier (excluded): per-completion-token sampling masks, # attributed to the assistant node by the turn's `commit`, then dropped. sampling_mask: SamplingMask | None = Field(default=None, exclude=True) + # Transient carrier (excluded): by-handle side arrays from `generate`, request-absolute + # positions, attributed per node onto `Trace.payload` by the turn's `commit`. + payload: list[PayloadSegment] | None = Field(default=None, exclude=True) # Transient carrier (excluded): the engine's weight-version spans over the completion, # attributed to the assistant node by the turn's `commit`, then dropped. weight_versions: list[WeightVersionSpan] | None = Field(default=None, exclude=True)