verifiers 0.3.2.dev101__tar.gz → 0.3.2.dev103__tar.gz

{verifiers-0.3.2.dev101 → verifiers-0.3.2.dev103}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: verifiers
-Version: 0.3.2.dev101
+Version: 0.3.2.dev103
 Summary: Verifiers: Environments for LLM Reinforcement Learning
 Project-URL: Homepage, https://github.com/primeintellect-ai/verifiers
 Project-URL: Documentation, https://github.com/primeintellect-ai/verifiers

{verifiers-0.3.2.dev101 → verifiers-0.3.2.dev103}/tests/v1/test_graph.py

@@ -91,6 +91,82 @@ def test_routed_experts_attributed_and_aligned_across_turns():
     )
 
 
+def test_bridged_turn_recovers_the_prior_turn_unforwarded_row():
+    """A turn never forwards its own final sampled token, so the engine returns one routing row
+    fewer than the turn has prompt plus completion positions, and that final position is filled
+    with a copy of its predecessor: a placeholder, effectively. The next turn's prefill does
+    forward it and reports it as row 0, so the placeholder must be replaced by that real row
+    rather than discarded."""
+    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")
+
+    # Start first turn
+    graph.prepare_turn(trace, [user]).commit(
+        vf.Response(
+            id="a",
+            created=0,
+            model="t",
+            message=vf.AssistantMessage(content="a1"),
+            finish_reason="stop",
+            tokens=TurnTokens(
+                prompt_ids=[10, 11, 12],
+                completion_ids=[20, 21],
+                message_spans=[(0, 2)],
+                routed_experts=_routed_payload(4, 0, 0),
+            ),
+        )
+    )
+    sampled_node = trace.nodes[-1]
+    assert sampled_node.routed_experts is not None
+    # First turn has the placeholder routing assignments:
+    assert bool(
+        (sampled_node.routed_experts[-1] == sampled_node.routed_experts[-2]).all()
+    )
+    forwarded_rows = sampled_node.routed_experts[:-1].copy()
+
+    # Start second turn
+    bridged = _routed_payload(4, 4, 100)
+    true_row = np.frombuffer(base64.b64decode(bridged["data"]), dtype=np.uint8).reshape(
+        bridged["shape"]
+    )[0]
+    graph.prepare_turn(
+        trace,
+        [user, vf.AssistantMessage(content="a1"), vf.UserMessage(content="u2")],
+    ).commit(
+        vf.Response(
+            id="b",
+            created=0,
+            model="t",
+            message=vf.AssistantMessage(content="a2"),
+            finish_reason="stop",
+            tokens=TurnTokens(
+                prompt_ids=[10, 11, 12, 20, 21, 30, 31],
+                completion_ids=[40, 41],
+                message_spans=[(0, 2), None, (5, 7)],
+                routed_experts=bridged,
+            ),
+        )
+    )
+
+    recovered = sampled_node.routed_experts
+    assert recovered is not None
+    # Did the first turn's final placeholder get corrected?
+    assert bool((recovered[-1] == true_row).all()), (
+        f"expected turn 1's final position to hold {true_row.tolist()}, "
+        f"got {recovered[-1].tolist()}"
+    )
+    assert bool((recovered[:-1] == forwarded_rows).all()), (
+        "only the fabricated final row may change"
+    )
+    assert recovered.flags.owndata
+    branch = trace.branches[-1]
+    assert branch.routed_experts is not None
+    assert branch.routed_experts.shape[0] == len(branch.token_ids)
+
+
 def test_routed_experts_none_when_absent():
     """No routing captured (engine ran without `enable_return_routed_experts`) -> the branch
     reports None and the trainer simply skips replay."""

{verifiers-0.3.2.dev101 → verifiers-0.3.2.dev103}/tests/v1/test_serve_delta.py

@@ -4,6 +4,7 @@ the episode the worker finished with."""
 import asyncio
 import copy
 
+import numpy as np
 import pytest
 
 import verifiers.v1 as vf
@@ -176,7 +177,8 @@ async def test_pending_preview_streams_and_clears_on_commit():
 
 
 @pytest.mark.asyncio
-async def test_deltas_stream_once_and_reassemble_the_episode():
+@pytest.mark.parametrize("routing_dtype", [np.uint8, np.uint16])
+async def test_deltas_stream_once_and_reassemble_the_episode(routing_dtype):
     slot = Slot()
     frames: list[bytes] = []
 
@@ -192,8 +194,15 @@ async def test_deltas_stream_once_and_reassemble_the_episode():
         await settle()
         assert len(frames) == 1, "the mint and the boot span coalesce into one delta"
         add_turn(trace, "a1")
+        trace.nodes[1].routed_experts = np.arange(6, dtype=np.uint8).reshape(3, 2, 1)
         await settle()
         add_turn(trace, "a2")
+        # The next prefill corrects the previous turn's last routing row, potentially
+        # widening its dtype. The previous node has already crossed the wire.
+        trace.nodes[1].routed_experts = np.array(
+            [[[0], [1]], [[2], [3]], [[100], [np.iinfo(routing_dtype).max]]],
+            dtype=routing_dtype,
+        )
         trace.nodes[0].semantic_parents.append(ParentLink(node=1, type="reply"))
         await settle()
         trace.record_reward("match", 1.0)
@@ -224,6 +233,10 @@ async def test_deltas_stream_once_and_reassemble_the_episode():
     rebuilt = WireEpisode.model_validate(unpack(pack(assembly.finish(head, summaries))))
     assert rebuilt.ok and rebuilt.traces[0].id == trace.id
     assert rebuilt.traces[0].messages == trace.messages
+    np.testing.assert_array_equal(
+        rebuilt.traces[0].nodes[1].routed_experts, trace.nodes[1].routed_experts
+    )
+    assert rebuilt.traces[0].nodes[1].routed_experts.dtype == routing_dtype
     assert (
         rebuilt.traces[0].nodes[0].semantic_parents == trace.nodes[0].semantic_parents
     )

{verifiers-0.3.2.dev101 → verifiers-0.3.2.dev103}/verifiers/v1/graph.py

@@ -673,8 +673,41 @@ def _attribute_mm(
             )
 
 
+def _replace_placeholder_routing_row(
+    trace: Trace, prefix_node_ids: list[int], arr: np.ndarray, off: int
+) -> None:
+    """Replace the prefix's placeholder routing row with the one this turn's prefill forwarded."""
+    if not 1 <= off <= arr.shape[0]:
+        return
+    # Only assistant (`sampled`) nodes are affected: the model forward isn't run on  the final
+    # generated token for such turns, meaning its routing decisions are fundamentally unavailable.
+    # But, because routing needs one row per token, that row instead receives an inaccurate
+    # placeholder, attempt to fix up below.
+    node_with_placeholder = None
+    for nid in reversed(prefix_node_ids):
+        if trace.nodes[nid].token_ids:
+            node_with_placeholder = trace.nodes[nid]
+            break
+    if node_with_placeholder is None or not node_with_placeholder.sampled:
+        return
+    node_rows = node_with_placeholder.routed_experts
+    if (
+        node_rows is None
+        or node_rows.shape[0] == 0
+        or node_rows.shape[1:] != arr.shape[1:]
+    ):
+        return
+    # Row `i` of this turn's array is sequence position `start + i`, so the prefix's final position
+    # is `arr[off - 1]`. Concatenating widens the node when this turn serialized `uint16`, where an
+    # in-place write would truncate.
+    node_with_placeholder.routed_experts = np.concatenate(
+        [node_rows[:-1], arr[off - 1 : off]], axis=0
+    )
+
+
 def _attribute_routed_experts(
     trace: Trace,
+    prefix_node_ids: list[int],
     new_node_ids: list[int],
     path_len: int,
     payload: Any,
@@ -683,7 +716,8 @@ def _attribute_routed_experts(
     payload's array covers the turn's prompt+completion from `payload["start"]` (0 = from token
     0); the nodes created this turn tile sequence positions `[path_len:]` in creation order, so
     we hand each node `arr[off : off+len(node.token_ids)]` and advance. Reused-prefix nodes keep
-    the routing attributed when they were first created. A node whose slice falls outside the
+    the routing attributed when they were first created, except for the one position this turn
+    corrects (see `_replace_placeholder_routing_row`). A node whose slice falls outside the
     array (a `start` past `path_len`, e.g. an unexpected prefix-cache delta) is left unset — the
     branch then reports no routing rather than misaligning."""
     if payload is None:
@@ -693,6 +727,7 @@ def _attribute_routed_experts(
         payload["shape"]
     )
     off = path_len - int(payload.get("start", 0) or 0)
+    _replace_placeholder_routing_row(trace, prefix_node_ids, arr, off)
     needed = off + sum(len(trace.nodes[nid].token_ids) for nid in new_node_ids)
     for nid in new_node_ids:
         n = len(trace.nodes[nid].token_ids)
@@ -701,8 +736,8 @@ def _attribute_routed_experts(
             # Own only this node's rows; a view would retain the turn's full-context array.
             trace.nodes[nid].routed_experts = arr[off:end].copy()
         elif n and arr.shape[0] and 0 <= off and end == needed == arr.shape[0] + 1:
-            # The engine omits the turn's final position because no forward pass follows it.
-            # Pad only the final node's suffix instead of copying the full-context array.
+            # No forward pass follows the turn's final position, so it gets a placeholder: a
+            # copy of the previous row, appended to this node's slice of the array.
             trace.nodes[nid].routed_experts = np.concatenate(
                 [arr[off:], arr[-1:]], axis=0
             )
@@ -813,7 +848,7 @@ def _commit_turn(turn: PendingTurn, response: Response) -> int:
     parent = prefix[-1] if prefix else None
     # cursor: in prompt_ids, the end of the previous *new* message's tokens
     cursor: int | None = None
-    # Track new nodes separately so routed-expert attribution does not need this full path.
+    # Track new nodes separately so routed-expert attribution needs only node ids, not this path.
     new_node_ids: list[int] = []
     # Materialize the reused message path only for multimodal cursor attribution.
     mm_path: list[tuple[int, Message]] | None = None
@@ -874,8 +909,9 @@ def _commit_turn(turn: PendingTurn, response: Response) -> int:
 
     # Attribute this turn's expert-routing array onto the nodes created this turn (new input
     # nodes in creation order, then the assistant node), each getting the routing for its tokens.
+    # The prefix goes in too, so the position the previous turn could only pad can be corrected.
     _attribute_routed_experts(
-        trace, new_node_ids, path_len, tokens.routed_experts if tokens else None
+        trace, prefix, new_node_ids, path_len, tokens.routed_experts if tokens else None
     )
 
     # Sampling masks are completion-aligned, so only the sampled node carries them.

{verifiers-0.3.2.dev101 → verifiers-0.3.2.dev103}/verifiers/v1/serve/delta.py

@@ -4,12 +4,12 @@ The worker streams a served episode as it grows. Each trace announces its own ch
 (`Trace.notify`, fired by the rollout at every phase change and by the interception proxy
 after every recorded turn); the `DeltaStreamer` then diffs the run's live traces against
 what it has already sent and ships only the new part — the trace header once, then
-appended nodes / calls / errors, semantic links landing on earlier nodes, the scalar
+appended nodes / calls / errors, semantic links and routing-row repairs on earlier nodes, the scalar
 fields whose value changed (timing spans, stop condition, rewards, ...), and the
 `pending` preview — the messages of the request in flight that no node holds yet, so a
 watcher sees a tool result before the model has answered it. The `Trace` is
-append-only at turn granularity (a turn's nodes are committed complete, with their
-tokens), so apart from the preview, which a committed turn repeats, every byte of the
+append-only at turn granularity except for links and the last routing row of a node,
+which the next prefill can repair. Apart from these and the preview, every byte of the
 episode crosses the wire once and the stream costs about what a single reply would; the
 reply that ends the run carries only the episode head and per-trace counts the client
 checks its assembly against. A cursor advances only once its delta is on the wire, so a
@@ -29,8 +29,10 @@ from collections.abc import Awaitable, Callable
 from typing import TYPE_CHECKING, Any, Self
 
 import msgpack
+import numpy as np
 from pydantic import BaseModel
 
+from verifiers.v1.graph import _decode_ndarray, _encode_ndarray
 from verifiers.v1.serve.encoding import msgpack_encoder
 
 if TYPE_CHECKING:
@@ -73,7 +75,7 @@ def pack(payload: Any) -> bytes:
 
 
 def unpack(data: bytes) -> Any:
-    # `links` is keyed by node index (int).
+    # Node updates are keyed by node index (int).
     return msgpack.unpackb(data, raw=False, strict_map_key=False)
 
 
@@ -90,6 +92,8 @@ class TraceCursor:
         self.sent = dict.fromkeys(LIST_FIELDS, 0)
         self.links: list[int] = []
         """Per sent node, how many of its semantic links went out with or after it."""
+        self.final_rows: dict[int, bytes] = {}
+        """Per node that carries routing, its final row as packed when last sent."""
         self.scalars: dict[str, bytes] = {}
         self.pending: bytes = pack([])
 
@@ -157,6 +161,31 @@ class DeltaStreamer:
                 else:
                     self.cursors[trace_id] = cursor
 
+    def _maybe_add_routing_repairs(
+        self,
+        delta: dict[str, Any],
+        trace: Trace,
+        cursor: TraceCursor,
+        sent_nodes: int,
+    ) -> None:
+        """Add final rows repaired since they were sent, keyed by node index.
+
+        Record every node's current row on the cursor. `sent_nodes` is the pre-flush
+        count, so a node first sent in this delta is never reported as a repair.
+        """
+        repairs: dict[int, dict] = {}
+        for index, node in enumerate(trace.nodes):
+            if node.routed_experts is None:
+                continue
+            row = _encode_ndarray(node.routed_experts[-1:])
+            packed = pack(row)
+            if cursor.final_rows.get(index) != packed:
+                cursor.final_rows[index] = packed
+                if index < sent_nodes:
+                    repairs[index] = row
+        if repairs:
+            delta["routing_repairs"] = repairs
+
     def diff(self) -> list[tuple[str, dict, TraceCursor | None]]:
         """Each trace's delta against its sent cursor, with the cursor as it stands once
         that delta is sent (None for a discard). Nothing here is committed: `flush`
@@ -183,6 +212,7 @@ class DeltaStreamer:
                     cursor.links[index] = len(node_links)
             if links:
                 delta["links"] = links
+            self._maybe_add_routing_repairs(delta, trace, cursor, cursor.sent["nodes"])
             for field in LIST_FIELDS:
                 items = getattr(trace, field)
                 sent = cursor.sent[field]
@@ -220,6 +250,20 @@ class EpisodeAssembly:
     def __init__(self) -> None:
         self.traces: dict[str, dict] = {}
 
+    def _maybe_apply_routing_repairs(self, delta: dict[str, Any], trace: dict) -> None:
+        """Replace repaired final rows in nodes the client already holds.
+
+        Rebuild each array: a repair can widen its dtype, and decoded rows alias
+        the original delta's bytes, which must remain unchanged for consumers.
+        """
+        for index, row in (delta.get("routing_repairs") or {}).items():
+            node = trace["nodes"][int(index)]
+            node["routed_experts"] = _encode_ndarray(
+                np.concatenate(
+                    [_decode_ndarray(node["routed_experts"])[:-1], _decode_ndarray(row)]
+                )
+            )
+
     def apply(self, delta: dict) -> None:
         trace_id = delta["trace"]
         if delta.get("discard"):
@@ -234,6 +278,7 @@ class EpisodeAssembly:
             }
         for index, links in (delta.get("links") or {}).items():
             trace["nodes"][int(index)]["semantic_parents"].extend(links)
+        self._maybe_apply_routing_repairs(delta, trace)
         # a later `links` delta grows a node's semantic_parents in place, so the node
         # is copied: the delta stays as it was when the caller received it
         if "nodes" in delta: