verifiers 0.2.1.dev12__tar.gz → 0.2.1.dev20__tar.gz

{verifiers-0.2.1.dev12 → verifiers-0.2.1.dev20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: verifiers
-Version: 0.2.1.dev12
+Version: 0.2.1.dev20
 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.2.1.dev12 → verifiers-0.2.1.dev20}/verifiers/v1/__init__.py

@@ -114,6 +114,7 @@ from verifiers.v1.types import (
     ID,
     ImageUrlContentPart,
     ImageUrlSource,
+    KeptTokens,
     Message,
     MessageContent,
     Messages,
@@ -166,6 +167,7 @@ __all__ = [
     "MessageNode",
     "Branch",
     "TurnTokens",
+    "KeptTokens",
     "Timing",
     "TimeSpan",
     "Error",

{verifiers-0.2.1.dev12 → verifiers-0.2.1.dev20}/verifiers/v1/cli/dashboard/eval.py

@@ -40,6 +40,7 @@ _LABEL_WIDTH = len("timeouts")
 
 _STYLE = {
     "pending": "dim",
+    "boot": "orange3",
     "setup": "yellow",
     "running": "cyan",
     "finalize": "magenta",
@@ -49,6 +50,7 @@ _STYLE = {
 }
 _MARK_LABEL = {
     "pending": "pending",
+    "boot": "boot",
     "setup": "setup",
     "running": "rollout",
     "finalize": "finalize",
@@ -317,7 +319,7 @@ def _breakdown(done: list[Trace]) -> Table | None:
             if judge.cost is not None:
                 total_judge_cost += judge.cost
             have_judge = True
-        for phase in ("setup", "generation", "finalize", "scoring"):
+        for phase in ("boot", "setup", "generation", "finalize", "scoring"):
             span = getattr(trace.timing, phase)
             if span.end:  # phase was timed for this rollout
                 phase_secs[phase] = phase_secs.get(phase, 0.0) + span.duration
@@ -351,7 +353,7 @@ def _breakdown(done: list[Trace]) -> Table | None:
         grid.add_row("usage", "  ·  ".join(usage))
     time_segments = [
         f"{phase} {format_time(phase_secs[phase] / phase_count[phase])}"
-        for phase in ("setup", "generation", "finalize", "scoring")
+        for phase in ("boot", "setup", "generation", "finalize", "scoring")
         if phase_count.get(phase)
     ]
     if time_segments:
@@ -383,11 +385,12 @@ def _tokens(trace: Trace) -> tuple[int, int, int | None, int | None, int]:
 
 
 def _started(rollout: Rollout) -> float:
-    # Sort key: when a rollout began (its setup start). A still-pending rollout has no trace
-    # yet, so it sorts last (+inf) — behind everything already in flight, in task order.
-    return (
-        rollout.trace.timing.setup.start if rollout.trace is not None else float("inf")
-    )
+    # Sort key: when a rollout began (its boot start; setup for pre-boot-span traces on
+    # resume). A still-pending rollout has no trace yet, so it sorts last (+inf) — behind
+    # everything already in flight, in task order.
+    if rollout.trace is None:
+        return float("inf")
+    return rollout.trace.timing.boot.start or rollout.trace.timing.setup.start
 
 
 def _groups(rollouts: list[Rollout]) -> list[list[Rollout]]:
@@ -450,14 +453,14 @@ def Rows(groups: list[list[Rollout]], now: float, runtime_type: str) -> Table:
             )
             runtime = f"{runtime_type}({descriptor})" if descriptor else runtime_type
             turns = t.num_turns
-            start = t.timing.setup.start
+            start = t.timing.boot.start or t.timing.setup.start
             end = (
                 t.timing.scoring.end
                 or t.timing.finalize.end
                 or t.timing.generation.end
-                # a rollout that errored in setup has only setup.end — freeze there once done,
-                # else (still running) the timer would grow off `now` forever
-                or (t.timing.setup.end if t.is_completed else 0)
+                # a rollout that errored in boot/setup has only that span's end — freeze there
+                # once done, else (still running) the timer would grow off `now` forever
+                or ((t.timing.setup.end or t.timing.boot.end) if t.is_completed else 0)
                 or now
             )
             prompt, completion, cached, reasoning, nbranches = _tokens(t)

{verifiers-0.2.1.dev12 → verifiers-0.2.1.dev20}/verifiers/v1/cli/debug.py

@@ -138,15 +138,19 @@ def record_debug_error(
     setup_timeout: float | None,
     action_timeout: float | None,
 ) -> None:
-    if not trace.timing.setup.end:
-        trace.timing.setup.end = time.time()
-    if trace.timing.generation.start and not trace.timing.generation.end:
-        trace.timing.generation.end = time.time()
+    now = time.time()
+    for span in (trace.timing.boot, trace.timing.setup, trace.timing.generation):
+        if span.start and not span.end:
+            span.end = now
     in_action = bool(trace.timing.generation.start)
-    stage = "debug action" if in_action else "setup"
+    stage = (
+        "debug action" if in_action else "setup" if trace.timing.setup.start else "boot"
+    )
     timeout = action_timeout if in_action else setup_timeout
     error_start = (
-        trace.timing.generation.start if in_action else trace.timing.setup.start
+        trace.timing.generation.start
+        if in_action
+        else trace.timing.setup.start or trace.timing.boot.start
     )
     debug.update(error_info(error, error_start, timeout, stage))
     debug.setdefault("runtime", runtime_info(runtime))
@@ -226,8 +230,11 @@ async def debug_task(task: Task, config: DebugConfig) -> tuple[Trace, bool]:
         else task.data.timeout.setup
     )
     try:
-        trace.timing.setup.start = time.time()
+        trace.timing.boot.start = time.time()
         await runtime.start()
+        now = time.time()
+        trace.timing.boot.end = now
+        trace.timing.setup.start = now
         debug["runtime"] = runtime_info(runtime)
         await asyncio.wait_for(
             invoke(task.setup, {"trace": trace, "runtime": runtime}),

{verifiers-0.2.1.dev12 → verifiers-0.2.1.dev20}/verifiers/v1/clients/train.py

@@ -25,6 +25,7 @@ from verifiers.v1.graph import PendingTurn
 from verifiers.v1.types import (
     AssistantMessage,
     FinishReason,
+    KeptTokens,
     Response,
     SamplingConfig,
     Tool,
@@ -149,6 +150,9 @@ def response_from_generate(
             is_content=attribution.is_content if attribution is not None else None,
             multi_modal_data=result.get("multi_modal_data"),
             routed_experts=result.get("routed_experts"),
+            kept_tokens=KeptTokens(**kept)
+            if (kept := result.get("kept_tokens"))
+            else None,
         ),
     )
 

{verifiers-0.2.1.dev12 → verifiers-0.2.1.dev20}/verifiers/v1/graph.py

@@ -31,6 +31,7 @@ from renderers.base import MultiModalData, PlaceholderRange, RenderedTokens
 from verifiers.v1.types import (
     AssistantMessage,
     FinishReason,
+    KeptTokens,
     Message,
     Response,
     StrictBaseModel,
@@ -118,6 +119,12 @@ class MessageNode(StrictBaseModel):
     the turn's `generate` payload by `_attribute_routed_experts`; `Branch.routed_experts`
     concatenates these along the path into the trainer's router-replay input. Rides the wire as
     a raw-bytes `__nd__` dict; kept off disk by the dump-site `exclude` in prime-rl."""
+    kept_tokens: KeptTokens | None = None
+    """Kept-set sampling masks for this node's sampled tokens, decoded: `ids` flat int32
+    in position order, `counts` the per-token kept-set sizes (aligned with `logprobs`;
+    0 = no mask). Assistant nodes only; consumed via `Branch.kept_tokens` for
+    sampling-replay training. Rides the wire as raw-bytes `__nd__` dicts; kept off disk
+    by the dump-site `exclude` in prime-rl."""
 
     model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
 
@@ -166,18 +173,40 @@ class MessageNode(StrictBaseModel):
         )
 
     @field_serializer("routed_experts")
-    def serialize_routed_experts(self, re: np.ndarray | None) -> dict | None:
-        """uint8 routing array -> raw-bytes `__nd__` dict so it rides the wire (numpy can't JSON)."""
-        return None if re is None else _encode_ndarray(re)
+    def serialize_ndarray_field(self, arr: np.ndarray | None) -> dict | None:
+        """Integer array -> raw-bytes `__nd__` dict so it rides the wire (numpy can't JSON)."""
+        return None if arr is None else _encode_ndarray(arr)
 
     @field_validator("routed_experts", mode="before")
     @classmethod
-    def deserialize_routed_experts(cls, value: Any) -> np.ndarray | None:
+    def deserialize_ndarray_field(cls, value: Any) -> np.ndarray | None:
         if value is None or isinstance(value, np.ndarray):
             return value
         if isinstance(value, dict) and value.get("__nd__"):
             return _decode_ndarray(value)
-        raise TypeError(f"cannot build routed_experts from {type(value).__name__}")
+        raise TypeError(f"cannot build ndarray field from {type(value).__name__}")
+
+    @field_serializer("kept_tokens")
+    def serialize_kept_tokens(self, kept: KeptTokens | None) -> dict | None:
+        """`KeptTokens` -> dict of raw-bytes `__nd__` entries so the arrays ride the wire."""
+        if kept is None:
+            return None
+        return {
+            "ids": _encode_ndarray(kept.ids),
+            "counts": _encode_ndarray(kept.counts),
+        }
+
+    @field_validator("kept_tokens", mode="before")
+    @classmethod
+    def deserialize_kept_tokens(cls, value: Any) -> KeptTokens | None:
+        if value is None or isinstance(value, KeptTokens):
+            return value
+        if isinstance(value, dict):
+            return KeptTokens(
+                ids=_decode_ndarray(value["ids"]),
+                counts=_decode_ndarray(value["counts"]),
+            )
+        raise TypeError(f"cannot build KeptTokens from {type(value).__name__}")
 
 
 def _canonical_tool_arguments(arguments: str) -> str:
@@ -434,6 +463,23 @@ def _attribute_routed_experts(
         off = end
 
 
+def _attribute_kept_tokens(
+    trace: Trace, assistant_id: int, payload: KeptTokens | None
+) -> None:
+    """Attach this turn's kept-set sampling masks to the assistant node (the payload
+    covers exactly the turn's completion tokens, so no path arithmetic). A payload
+    that doesn't line up with the node's sampled tokens is dropped, not misaligned."""
+    if payload is None:
+        return
+    counts = np.frombuffer(binascii.a2b_base64(payload.counts), dtype=np.int32)
+    ids = np.frombuffer(binascii.a2b_base64(payload.ids), dtype=np.int32)
+    node = trace.nodes[assistant_id]
+    if len(counts) != sum(node.mask) or int(counts.sum()) != len(ids):
+        return
+    # Own the buffers — the payload views reference the turn's response bytes.
+    node.kept_tokens = KeptTokens(ids=ids.copy(), counts=counts.copy())
+
+
 def _commit_turn(turn: PendingTurn, response: Response) -> None:
     trace = turn.trace
     prompt = turn.prompt
@@ -540,6 +586,10 @@ def _commit_turn(turn: PendingTurn, response: Response) -> None:
         trace, new_node_ids, path_len, tokens.routed_experts if tokens else None
     )
 
+    # Attribute this turn's kept-set sampling masks onto the assistant node (they are
+    # completion-aligned, so only the sampled node carries them).
+    _attribute_kept_tokens(trace, assistant_id, tokens.kept_tokens if tokens else None)
+
 
 # --- walking the graph (views) ---------------------------------------------------------
 

{verifiers-0.2.1.dev12 → verifiers-0.2.1.dev20}/verifiers/v1/rollout.py

@@ -37,6 +37,7 @@ logger = logging.getLogger(__name__)
 
 class Phase(StrEnum):
     PENDING = "pending"
+    BOOT = "boot"
     SETUP = "setup"
     RUNNING = "running"
     FINALIZE = "finalize"
@@ -108,8 +109,8 @@ class Rollout:
             state=state_cls(type(self.task))(),
         )
         self.trace = trace  # expose for the --rich dashboard
-        self.phase = Phase.SETUP  # leaving the queue: provisioning starts now
-        trace.timing.setup.start = time.time()
+        self.phase = Phase.BOOT  # leaving the queue: the runtime boots now
+        trace.timing.boot.start = time.time()
         self.runtime = make_runtime(self.runtime_config, name=trace.id)
         runtime = self.runtime
         trace.runtime = runtime.info
@@ -125,6 +126,10 @@ class Rollout:
         try:
             session = RolloutSession(ctx, trace, stops, self.limits)
             await runtime.start()
+            now = time.time()
+            trace.timing.boot.end = now
+            trace.timing.setup.start = now
+            self.phase = Phase.SETUP
             # Task setup and harness provisioning share one setup-stage deadline.
             setup_deadline = (
                 None
@@ -228,6 +233,7 @@ class Rollout:
             trace.is_completed = True
             now = time.time()
             for span in (
+                trace.timing.boot,
                 trace.timing.setup,
                 trace.timing.generation,
                 trace.timing.finalize,

{verifiers-0.2.1.dev12 → verifiers-0.2.1.dev20}/verifiers/v1/trace.py

@@ -22,6 +22,7 @@ from verifiers.v1.state import State, StateT
 from verifiers.v1.task import DataT, WireTaskData
 from verifiers.v1.types import (
     AssistantMessage,
+    KeptTokens,
     Messages,
     StrictBaseModel,
     Tool,
@@ -46,6 +47,7 @@ class TimeSpan(StrictBaseModel):
 
 class Timing(StrictBaseModel):
     start: float = Field(default_factory=time.time)
+    boot: TimeSpan = Field(default_factory=TimeSpan)
     setup: TimeSpan = Field(default_factory=TimeSpan)
     generation: TimeSpan = Field(default_factory=TimeSpan)
     finalize: TimeSpan = Field(default_factory=TimeSpan)
@@ -139,6 +141,32 @@ class Branch(StrictBaseModel):
         total = sum(len(n.token_ids) for n in nodes)
         return merged if merged.shape[0] == total else None
 
+    @property
+    def kept_tokens(self) -> KeptTokens | None:
+        """The branch's kept-set sampling masks: `counts` is int32 aligned 1:1 with
+        `token_ids` (0 = no mask, safe under partial coverage — unlike `routed_experts`
+        this is not all-or-nothing), `ids` the flat int32 concatenation of the kept
+        sets in position order. None when no node carries kept-set data."""
+        if all(n.kept_tokens is None for n in self.nodes):
+            return None
+        # `_attribute_kept_tokens` validates counts/ids against the node's sampled
+        # tokens before setting the field, so this is a straight scatter+concat
+        # (a corrupted node would fail loudly on the scatter shape mismatch).
+        ids_parts: list[np.ndarray] = []
+        counts_parts: list[np.ndarray] = []
+        for node in self.nodes:
+            counts = np.zeros(len(node.mask), dtype=np.int32)
+            if node.kept_tokens is not None and len(node.kept_tokens.counts):
+                counts[np.nonzero(node.mask)[0]] = node.kept_tokens.counts
+                ids_parts.append(node.kept_tokens.ids)
+            counts_parts.append(counts)
+        ids = (
+            np.concatenate(ids_parts).astype(np.int32, copy=False)
+            if ids_parts
+            else np.zeros(0, dtype=np.int32)
+        )
+        return KeptTokens(ids=ids, counts=np.concatenate(counts_parts))
+
     @property
     def num_total_tokens(self) -> int:
         return sum(len(n.token_ids) for n in self.nodes)
@@ -172,7 +200,13 @@ class Branch(StrictBaseModel):
 
 
 _NODE_DUMP_EXCLUDE: dict = {
-    "nodes": {"__all__": {"multi_modal_data", "routed_experts"}}
+    "nodes": {
+        "__all__": {
+            "multi_modal_data",
+            "routed_experts",
+            "kept_tokens",
+        }
+    }
 }
 """Raw tensor fields kept on the msgpack wire but excluded from JSON records."""
 

{verifiers-0.2.1.dev12 → verifiers-0.2.1.dev20}/verifiers/v1/types.py

@@ -1,4 +1,5 @@
 from collections.abc import Iterable
+from dataclasses import dataclass
 from typing import Annotated, Any, Literal
 
 from pydantic import AfterValidator, AliasChoices, BaseModel, ConfigDict, Field
@@ -178,6 +179,18 @@ class RoutedExperts(TypedDict):
     start: int
 
 
+@dataclass
+class KeptTokens:
+    """Kept-set sampling masks for sampling replay: `ids` (every kept set concatenated
+    in position order) and `counts` (kept-set size per completion token; 0 = no usable
+    mask). Base64 blobs straight off the `generate` response on the `TurnTokens`
+    carrier; decoded to flat int32 arrays on `MessageNode` (`len(ids) == sum(counts)`,
+    row boundaries recovered from `counts`)."""
+
+    ids: Any
+    counts: Any
+
+
 class TurnTokens(StrictBaseModel):
     """Training tokens from renderer tokenization or provider-returned token IDs."""
 
@@ -200,6 +213,11 @@ class TurnTokens(StrictBaseModel):
     # per token), attributed per node by the turn's `commit` into `MessageNode.routed_experts`,
     # then dropped. None unless the engine ran with `enable_return_routed_experts`.
     routed_experts: RoutedExperts | None = Field(default=None, exclude=True)
+    # Transient carrier (excluded): the kept-set sampling masks from `generate` (token ids
+    # surviving top-p/top-k truncation, per completion token), attributed to the assistant
+    # node by the turn's `commit`, then dropped. None unless the engine ran with
+    # `enable_return_kept_tokens`.
+    kept_tokens: KeptTokens | None = Field(default=None, exclude=True)
 
 
 class Response(StrictBaseModel):