feat(vllm-model): supply the previous call's exact training tokens

nemo_gym/base_responses_api_model.py

@@ -1153,6 +1153,7 @@ def __init__(
         token_store: Any = None,
         configured_sink: Any = None,
         lineage_store: Any = None,
+        delta_records: bool = False,
         token_capture_enabled: bool = False,
     ) -> None:
         self._app = app
@@ -1163,6 +1164,7 @@ def __init__(
         # Built from token_id_capture.sink, once, in this process.
         self._configured_sink = configured_sink
         self._lineage_store = lineage_store
+        self._delta_records = delta_records
         # Capture may have no destination in this process.
         # A framework may stage records from its inference worker.
         # This process still resolves the capture identity.
@@ -1196,19 +1198,20 @@ async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None
         # Installed sinks are resolved for each request.
         token_sink = self._configured_sink or installed_token_sink() or self._token_store
         capture_wanted = token_capture_requested and (token_sink is not None or self._token_capture_enabled)
-        if token_capture_requested and dialect is None and token_sink is not None:
-            # This call cannot produce a capture record.
-            # Its output may still feed a later prompt.
-            # Mark the rollout incomplete before forwarding the request.
-            try:
-                await token_sink.mark_incomplete(rollout_from_path, "")
-            except Exception:
-                logger.warning(
-                    "Could not mark rollout %s incomplete for unobserved path %s.",
-                    rollout_from_path,
-                    path,
-                    exc_info=True,
-                )
+        if token_capture_requested and dialect is None:
+            # An uncapturable call must not look complete.
+            # Its output can still feed later prompts.
+            # Mark the rollout before forwarding so consumers retain that evidence.
+            if token_sink is not None:
+                try:
+                    await token_sink.mark_incomplete(rollout_from_path, "")
+                except Exception:
+                    logger.warning(
+                        "Could not mark rollout %s incomplete for unobserved path %s.",
+                        rollout_from_path,
+                        path,
+                        exc_info=True,
+                    )
         if (self._store is None and not capture_wanted) or rollout_from_path is None or dialect is None:
             await self._app(scope, receive, send)
             return
@@ -1229,6 +1232,7 @@ async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None
                     model_call_id=model_call_id,
                     token_sink=token_sink,
                     lineage_store=self._lineage_store,
+                    delta_records=self._delta_records,
                 )
             )
 
@@ -1459,6 +1463,7 @@ async def _capture_lifespan(application):
         token_store=token_store,
         configured_sink=configured_sink,
         lineage_store=lineage_store,
+        delta_records=(capture_settings.token_id_capture.delta_records if capture_settings is not None else False),
         token_capture_enabled=capture_settings.enabled if capture_settings is not None else False,
     )
 

nemo_gym/token_id_capture/builder.py

@@ -198,6 +198,69 @@ def _resolve_parent(
     return None, False, "missing_resolution"
 
 
+def _materialize_delta_prompts(entries: list[TokenEntry]) -> tuple[list[TokenEntry], list[str]]:
+    """Rebuild full prompts for delta records by walking their parent chains.
+
+    Return full-prompt entries and call ids with broken chains.
+    """
+    by_id = {entry.model_call_id: entry for entry in entries}
+    cumulative: dict[str, list[int] | None] = {}
+
+    def cum_of(call_id: str) -> list[int] | None:
+        if call_id in cumulative:
+            return cumulative[call_id]
+
+        path: list[TokenEntry] = []
+        seen: set[str] = set()
+        current_id = call_id
+        while current_id not in cumulative:
+            if current_id in seen or len(path) > 10_000:
+                cumulative[current_id] = None
+                break
+            seen.add(current_id)
+            entry = by_id.get(current_id)
+            if entry is None:
+                cumulative[current_id] = None
+                break
+            if not entry.prompt_is_delta:
+                cumulative[current_id] = list(entry.prompt_token_ids) + list(entry.generation_token_ids)
+                break
+            path.append(entry)
+            if entry.parent_call_id is None:
+                cumulative[current_id] = None
+                break
+            current_id = entry.parent_call_id
+
+        value = cumulative[current_id]
+        if value is None:
+            for entry in path:
+                cumulative[entry.model_call_id] = None
+            return None
+        for entry in reversed(path):
+            value = value + list(entry.prompt_token_ids) + list(entry.generation_token_ids)
+            cumulative[entry.model_call_id] = value
+        return cumulative[call_id]
+
+    materialized: list[TokenEntry] = []
+    broken: list[str] = []
+    for entry in entries:
+        if not entry.prompt_is_delta:
+            materialized.append(entry)
+            continue
+        cum = cum_of(entry.model_call_id)
+        if cum is None:
+            broken.append(entry.model_call_id)
+            continue
+        full_prompt = cum[: len(cum) - len(entry.generation_token_ids)]
+        rebuilt = entry.model_copy(update={"prompt_token_ids": full_prompt, "prompt_is_delta": False})
+        # Verify the reconstructed full sequence.
+        if rebuilt.digest and compute_digest(cum) != rebuilt.digest:
+            broken.append(entry.model_call_id)
+            continue
+        materialized.append(rebuilt)
+    return materialized, broken
+
+
 class _NullPrefixIndex:
     """Avoid building a token-prefix index when every parent is present.
 
@@ -228,6 +291,11 @@ def prefix_merging(entries: list[TokenEntry], terminal_call_id: str | None = Non
             duplicate_conflicts.append(candidate.model_call_id)
     entries = list(deduped.values())
 
+    # Chain construction assumes full prompts.
+    # Materialize delta records before sorting or checking parent digests.
+    # Exclude a record when its parent chain cannot be reconstructed exactly.
+    entries, unreconstructable = _materialize_delta_prompts(entries)
+
     # A call without generated tokens has no training signal.
     # Its cumulative sequence equals its prompt.
     # Keeping it would make it the parent of another call with the same prompt.
@@ -236,13 +304,18 @@ def prefix_merging(entries: list[TokenEntry], terminal_call_id: str | None = Non
     empty_generation = [e.model_call_id for e in entries if not e.generation_token_ids]
     entries = [e for e in entries if e.generation_token_ids]
     if not entries:
+        # Report delta reconstruction failures even when no entries remain.
         return BuildOutput(
             chains=[],
             notes=BuildNotes(
                 builder="prefix_merging",
                 empty_generation_calls=empty_generation,
                 terminal_call_id=terminal_call_id,
                 terminal_chain="not_captured" if terminal_call_id else "",
+                parent_link_failures=(
+                    {"delta_chain_unreconstructable": len(unreconstructable)} if unreconstructable else {}
+                ),
+                unresolved_parent_calls=list(unreconstructable),
             ),
         )
 
@@ -267,6 +340,12 @@ def prefix_merging(entries: list[TokenEntry], terminal_call_id: str | None = Non
             parent_link_failures.get("duplicate_call_id_conflict", 0) + 1
         )
         unresolved_parent_calls.append(call_id)
+    for call_id in unreconstructable:
+        parent_link_failures["delta_chain_unreconstructable"] = (
+            parent_link_failures.get("delta_chain_unreconstructable", 0) + 1
+        )
+        unresolved_parent_calls.append(call_id)
+
     for entry in ordered:
         prompt = list(entry.prompt_token_ids)
         node = _Node(entry=entry, cumulative=prompt + list(entry.generation_token_ids))

nemo_gym/token_id_capture/config.py

@@ -22,6 +22,9 @@
       enabled: true
       dir: /tmp/nemo_gym_token_id_captures  # The writer and consumer share this node-local directory.
       sink: my_pkg.sinks:MyDataPlaneSink   # This optional sink replaces the file store.
+      lineage_store: my_pkg.sinks:MyResolver  # Required with a custom sink (same backend namespace).
+      delta_records: true                  # Store RESOLVED continuations as parent-relative suffixes.
+      max_mask_fraction: 0.5               # Abort a run that is mostly producing masked rollouts.
 ```
 
 Evaluation capture uses ``/ng-rollout/<id>/...``.
@@ -114,8 +117,13 @@ class TokenIdCaptureSettings(BaseModel):
     # Without one, every multi-call continuation is unresolved and masked.
     # This flag permits that degraded behavior explicitly.
     allow_unresolved_continuations: bool = False
-    # Abort once enough finalized rollouts exceed this masked fraction.
-    # ``None`` disables the limit.
+    # Store resolved prompts as suffixes of their verified parent tokens.
+    # This avoids repeatedly storing the growing full prompt.
+    # Root and unresolved records remain full-prompt reconstruction anchors.
+    delta_records: bool = False
+    # Abort when the finalized-rollout masked fraction exceeds this limit.
+    # Enforcement begins after ``mask_fraction_min_samples`` observations.
+    # ``None`` disables the kill switch.
     max_mask_fraction: float | None = None
     mask_fraction_min_samples: int = 50
 
@@ -172,14 +180,15 @@ def _require_resolver(block: TokenIdCaptureSettings) -> None:
             return
         if block.allow_unresolved_continuations:
             logger.warning(
-                "token_id_capture has a custom sink and no lineage_store. "
-                "Every continuation will be unresolved and masked."
+                "token_id_capture has a custom sink and no lineage_store: every continuation "
+                "will resolve UNRESOLVED and multi-call rollouts will be masked."
             )
             return
         raise ValueError(
-            "token_id_capture has a custom sink but no lineage_store. Configure "
+            "token_id_capture has a custom sink but no lineage_store, so no continuation can "
+            "resolve its parent and every multi-call rollout will be masked. Configure "
             "token_id_capture.lineage_store on the same backend as the sink, or set "
-            "token_id_capture.allow_unresolved_continuations: true to accept unresolved continuations."
+            "token_id_capture.allow_unresolved_continuations: true to accept the loss."
         )
 
     @property