From 88e5903f5e083ac791af498c692a223d4c6f5dcc Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Fri, 28 Aug 2026 09:24:58 -0700 Subject: [PATCH] feat(vllm-model): supply the previous call's exact training tokens Supply a verified parent's cumulative tokens to compatible generation backends and require generation-time proof before recording successful application. Store resolved prompts as bounded, reconstructible deltas so long external-harness trajectories remain exact without quadratic storage or materialization. Signed-off-by: Ananth Subramaniam --- nemo_gym/base_responses_api_model.py | 31 +- nemo_gym/token_id_capture/builder.py | 79 +++ nemo_gym/token_id_capture/config.py | 21 +- nemo_gym/token_id_capture/lineage.py | 132 +++- nemo_gym/token_id_capture/records.py | 44 +- nemo_gym/token_id_capture/sink.py | 98 ++- responses_api_models/vllm_model/app.py | 117 +++- .../configs/vllm_model_supply_prefix.yaml | 15 + .../vllm_model/tests/test_app.py | 645 ++++++++++++++++++ .../tests/test_token_capture_integration.py | 260 +++++++ tests/unit_tests/test_token_id_capture.py | 115 +++- tests/unit_tests/test_trajectory_builder.py | 171 ++++- 12 files changed, 1641 insertions(+), 87 deletions(-) create mode 100644 responses_api_models/vllm_model/configs/vllm_model_supply_prefix.yaml create mode 100644 responses_api_models/vllm_model/tests/test_token_capture_integration.py diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index 2f660132fe..30abed2d8c 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/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, ) diff --git a/nemo_gym/token_id_capture/builder.py b/nemo_gym/token_id_capture/builder.py index adbf14f5f9..7d66335e40 100644 --- a/nemo_gym/token_id_capture/builder.py +++ b/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,6 +304,7 @@ 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( @@ -243,6 +312,10 @@ def prefix_merging(entries: list[TokenEntry], terminal_call_id: str | None = Non 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)) diff --git a/nemo_gym/token_id_capture/config.py b/nemo_gym/token_id_capture/config.py index 539506a4a3..d5d999b293 100644 --- a/nemo_gym/token_id_capture/config.py +++ b/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//...``. @@ -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 diff --git a/nemo_gym/token_id_capture/lineage.py b/nemo_gym/token_id_capture/lineage.py index 29db6e44dd..a8fec207d3 100644 --- a/nemo_gym/token_id_capture/lineage.py +++ b/nemo_gym/token_id_capture/lineage.py @@ -293,6 +293,8 @@ class LineageNode: # A mid-rollout dialect switch can misalign it; verification then fails closed. context_len: int = 0 context_digest: str = "" + parent_call_id: str | None = None + prompt_is_delta: bool = False def stamp_continuation(entry: TokenEntry, request_items: list[dict]) -> TokenEntry: @@ -396,6 +398,8 @@ def add_entry(self, entry: TokenEntry, *, store_tokens: bool = True, entry_offse entry_offset=entry_offset, context_len=entry.continuation_context_len, context_digest=entry.continuation_context_digest, + parent_call_id=entry.parent_call_id, + prompt_is_delta=entry.prompt_is_delta, ) previous = self.by_call_id.get(entry.model_call_id) if previous is not None: @@ -530,6 +534,8 @@ class IncrementalLineageStore: ``_load_entry(rollout_id, ref)`` -> ``TokenEntry`` for one committed record. Optional hooks: + ``_load_entries(rollout_id, refs)`` — batch-load one parent chain + (default: call ``_load_entry`` for each reference). ``_read_locked(rollout_id)`` — context manager held around fetch+resolve for backends with a read-lock discipline (default: no lock). ``is_process_shared()`` — default ``True``; an external backend exists to @@ -539,14 +545,21 @@ class IncrementalLineageStore: class CursorReset(Exception): """The stored cursor no longer describes the backend; refetch from scratch.""" - def __init__(self, *, max_cached_rollouts: int = 65536) -> None: + def __init__(self, *, max_cached_rollouts: int = 65536, max_cached_tokens: int = 8_000_000) -> None: import threading if max_cached_rollouts < 1: raise ValueError("max_cached_rollouts must be positive") + if max_cached_tokens < 1: + raise ValueError("max_cached_tokens must be positive") # (cursor, refs, lineage): lineage stays at index 2 for diagnostics/tooling. self._cache: dict[str, tuple[Any, dict[str, Any], RolloutLineage]] = {} self._max_cached_rollouts = max_cached_rollouts + # Keep only the latest materialized parent for each rollout. + # The global token bound avoids recreating full-record memory growth. + self._materialized: dict[str, tuple[str, tuple[int, ...]]] = {} + self._materialized_tokens = 0 + self._max_cached_tokens = max_cached_tokens self._cache_guard = threading.Lock() # Fixed lock striping bounds synchronization metadata. # Hash collisions only serialize unrelated rollouts. @@ -559,6 +572,13 @@ def _fetch_new_entries(self, rollout_id: str, cursor: Any) -> tuple[list[tuple[T def _load_entry(self, rollout_id: str, ref: Any) -> TokenEntry: raise NotImplementedError + def _load_entries(self, rollout_id: str, refs: list[Any]) -> list[TokenEntry]: + """Load several committed entries. + + Backends can override this hook to fetch a parent chain in one operation. + """ + return [self._load_entry(rollout_id, ref) for ref in refs] + def _read_locked(self, rollout_id: str): from contextlib import nullcontext @@ -582,6 +602,27 @@ def _cache_put(self, rollout_id: str, value: tuple[Any, dict[str, Any], RolloutL if oldest == rollout_id: break self._cache.pop(oldest) + materialized = self._materialized.pop(oldest, None) + if materialized is not None: + self._materialized_tokens -= len(materialized[1]) + + def _cached_materialized(self, rollout_id: str) -> tuple[str, tuple[int, ...]] | None: + with self._cache_guard: + return self._materialized.get(rollout_id) + + def _remember_materialized(self, rollout_id: str, call_id: str, tokens: tuple[int, ...]) -> None: + with self._cache_guard: + previous = self._materialized.pop(rollout_id, None) + if previous is not None: + self._materialized_tokens -= len(previous[1]) + if len(tokens) > self._max_cached_tokens: + return + self._materialized[rollout_id] = (call_id, tokens) + self._materialized_tokens += len(tokens) + while self._materialized_tokens > self._max_cached_tokens: + oldest = next(iter(self._materialized)) + evicted = self._materialized.pop(oldest) + self._materialized_tokens -= len(evicted[1]) def _refresh(self, rollout_id: str) -> tuple[dict[str, Any], RolloutLineage]: with self._cache_guard: @@ -601,39 +642,60 @@ def _refresh(self, rollout_id: str) -> tuple[dict[str, Any], RolloutLineage]: def _materialize( self, rollout_id: str, node: LineageNode, refs: dict[str, Any], lineage: RolloutLineage - ) -> list[int]: + ) -> tuple[int, ...]: """Load one RESOLVED parent's cumulative tokens from the backend. + Read the chain in one batch and append each token segment once. Digest verification makes stale references fail closed. """ from nemo_gym.token_id_capture.records import compute_digest - def load(call_id: str) -> TokenEntry: - if call_id not in refs: - raise ValueError(f"lineage node for {call_id} has no backend ref") - entry = self._load_entry(rollout_id, refs[call_id]) - if entry.model_call_id != call_id: - raise ValueError(f"ref for {call_id} points at {entry.model_call_id}") - return entry - - # Walk delta suffixes back to a full-prompt anchor. - suffixes: list[tuple[list[int], list[int]]] = [] - current = load(node.call_id) - depth = 0 - while getattr(current, "prompt_is_delta", False): - depth += 1 - if depth > 10_000: + # Metadata carries enough lineage to collect every backend reference before loading tokens. + cached = self._cached_materialized(rollout_id) + chain: list[LineageNode] = [] + seen: set[str] = set() + current = node + cached_tokens: tuple[int, ...] | None = None + while True: + if cached is not None and current.call_id == cached[0]: + cached_tokens = cached[1] + break + if current.call_id in seen: + raise ValueError(f"delta chain for {node.call_id} contains a cycle") + if len(chain) >= 10_000: raise ValueError(f"delta chain for {node.call_id} exceeds sane depth") - suffixes.append((list(current.prompt_token_ids), list(current.generation_token_ids))) - if not current.parent_call_id or lineage.by_call_id.get(current.parent_call_id) is None: - raise ValueError(f"delta record {current.model_call_id} has no indexed parent") - current = load(current.parent_call_id) - tokens = cumulative_tokens(current) - for suffix, generation in reversed(suffixes): - tokens = tokens + suffix + generation + seen.add(current.call_id) + chain.append(current) + if not current.prompt_is_delta: + break + if not current.parent_call_id: + raise ValueError(f"delta record {current.call_id} has no parent call id") + parent = lineage.by_call_id.get(current.parent_call_id) + if parent is None: + raise ValueError(f"delta record {current.call_id} has no indexed parent") + current = parent + + ordered_nodes = list(reversed(chain)) + missing_ref = next((item.call_id for item in ordered_nodes if item.call_id not in refs), None) + if missing_ref is not None: + raise ValueError(f"lineage node for {missing_ref} has no backend ref") + entries = ( + self._load_entries(rollout_id, [refs[item.call_id] for item in ordered_nodes]) if ordered_nodes else [] + ) + + tokens = list(cached_tokens or ()) + for expected, entry in zip(ordered_nodes, entries, strict=True): + if entry.model_call_id != expected.call_id: + raise ValueError(f"ref for {expected.call_id} points at {entry.model_call_id}") + if entry.prompt_is_delta != expected.prompt_is_delta: + raise ValueError(f"metadata for {expected.call_id} disagrees with its stored entry") + tokens.extend(entry.prompt_token_ids) + tokens.extend(entry.generation_token_ids) if node.digest and compute_digest(tokens) != node.digest: raise ValueError(f"materialized tokens for {node.call_id} fail their digest") - return tokens + materialized = tuple(tokens) + self._remember_materialized(rollout_id, node.call_id, materialized) + return materialized async def resolve(self, rollout_id: str, request_items: list[dict]) -> LineageResolution: return await asyncio.to_thread(self._resolve, rollout_id, request_items) @@ -662,6 +724,8 @@ def is_process_shared(self) -> bool: async def close(self) -> None: with self._cache_guard: self._cache.clear() + self._materialized.clear() + self._materialized_tokens = 0 class FileLineageStore(IncrementalLineageStore): @@ -672,10 +736,16 @@ class FileLineageStore(IncrementalLineageStore): ``put`` is immediately visible. """ - def __init__(self, root: str | Path, *, max_cached_rollouts: int = 65536) -> None: + def __init__( + self, + root: str | Path, + *, + max_cached_rollouts: int = 65536, + max_cached_tokens: int = 8_000_000, + ) -> None: from nemo_gym.token_id_capture.store import TokenCaptureStore - super().__init__(max_cached_rollouts=max_cached_rollouts) + super().__init__(max_cached_rollouts=max_cached_rollouts, max_cached_tokens=max_cached_tokens) self._store = TokenCaptureStore(root) def _read_locked(self, rollout_id: str): @@ -711,3 +781,11 @@ def _load_entry(self, rollout_id: str, ref: Any) -> TokenEntry: with self._store.path_for(rollout_id).open("rb") as handle: handle.seek(ref) return TokenEntry.model_validate(orjson.loads(handle.readline())) + + def _load_entries(self, rollout_id: str, refs: list[Any]) -> list[TokenEntry]: + entries = [] + with self._store.path_for(rollout_id).open("rb") as handle: + for ref in refs: + handle.seek(ref) + entries.append(TokenEntry.model_validate(orjson.loads(handle.readline()))) + return entries diff --git a/nemo_gym/token_id_capture/records.py b/nemo_gym/token_id_capture/records.py index 238b4a3bb3..368d9da23f 100644 --- a/nemo_gym/token_id_capture/records.py +++ b/nemo_gym/token_id_capture/records.py @@ -145,11 +145,26 @@ class TokenEntry(BaseModel): # The context fields verify the request that produced this call. continuation_context_len: int = 0 continuation_context_digest: str = "" - # Persist the resolver's diagnostic reason. + + # The resolver's diagnostic reason for the parent decision above. + # Persisted so a resolution-rate regression is debuggable from records alone. parent_resolution_reason: str = "" - # Identify the continuation fingerprint algorithm. + # The fingerprint algorithm version that produced continuation_fingerprint. + # A resolver must not match records stamped by a different algorithm. fingerprint_version: int | None = None + # Delta storage is part of the initial record schema. + # A delta prompt stores only the suffix after the resolved parent. + # Root and unresolved records always store full prompts. + # This guarantees that every reconstructable chain has a full-prompt anchor. + prompt_is_delta: bool = False + + # Prefix supply fields are part of the initial record schema. + # Request intent is distinct from generation-time proof. + prefix_requested: bool = False + # This is true only when generation-time prompt_token_ids prove prefix application. + prefix_supplied: bool = False + @model_validator(mode="after") def _refuse_a_newer_record(self) -> "TokenEntry": """Accept older records and reject newer records. @@ -167,7 +182,8 @@ def _refuse_a_newer_record(self) -> "TokenEntry": if self.schema_version < TOKEN_ENTRY_MIN_SCHEMA_VERSION: raise ValueError( f"token record is schema_version {self.schema_version}, below the supported minimum " - f"{TOKEN_ENTRY_MIN_SCHEMA_VERSION}. Regenerate the rollout with a current writer." + f"{TOKEN_ENTRY_MIN_SCHEMA_VERSION}. Pre-lineage records were never written in " + "production; regenerate the rollout with a current writer." ) if len(self.generation_token_ids) != len(self.generation_log_probs): raise ValueError( @@ -183,11 +199,22 @@ def _refuse_a_newer_record(self) -> "TokenEntry": if self.parent_resolution in {ParentResolutionStatus.ROOT, ParentResolutionStatus.UNRESOLVED}: if self.parent_call_id is not None: raise ValueError(f"{self.parent_resolution.value} parent metadata cannot carry parent_call_id") + if self.prompt_is_delta and ( + self.parent_call_id is None or self.parent_resolution != ParentResolutionStatus.RESOLVED + ): + raise ValueError("a delta prompt requires a RESOLVED parent_call_id to reconstruct from") return self def cumulative_tokens(entry: TokenEntry) -> list[int]: - """The full sequence a child of this call must start with.""" + """The full sequence a child of this call must start with. + + Delta records require parent-chain reconstruction. + """ + if entry.prompt_is_delta: + raise ValueError( + f"model call {entry.model_call_id} stores a delta prompt; reconstruct through its parent chain" + ) return list(entry.prompt_token_ids) + list(entry.generation_token_ids) @@ -196,13 +223,18 @@ def stamp_lineage( parent_call_id: str | None, *, parent_resolution: ParentResolutionStatus | None = None, + cumulative: list[int] | None = None, ) -> TokenEntry: """Fill token lineage and the request-time parent decision. - ``cum_len`` and ``digest`` always describe this call. + ``cum_len`` and ``digest`` describe the full sequence. + A delta entry must pass ``cumulative`` explicitly. ``parent_resolution=None`` preserves records built by compatibility callers. """ - cumulative = cumulative_tokens(entry) + if cumulative is None: + cumulative = cumulative_tokens(entry) + elif entry.prompt_is_delta is False and cumulative != cumulative_tokens(entry): + raise ValueError("provided cumulative tokens disagree with the entry's own arrays") entry.cum_len = len(cumulative) entry.digest = compute_digest(cumulative) entry.parent_call_id = parent_call_id diff --git a/nemo_gym/token_id_capture/sink.py b/nemo_gym/token_id_capture/sink.py index 9f0f114d4f..d0a793f5d1 100644 --- a/nemo_gym/token_id_capture/sink.py +++ b/nemo_gym/token_id_capture/sink.py @@ -55,8 +55,8 @@ class CaptureContext: The context identifies the rollout and model call. ``token_sink`` receives the resulting record. A framework may provide any ``TokenSink`` implementation. - Every consumer shares the same per-call parent decision. - This keeps request-time resolution and capture metadata consistent. + Parent resolution runs once for each call. + Prefix supply and token capture read the same immutable decision. """ rollout_id: str @@ -68,6 +68,12 @@ class CaptureContext: model: str = "" # ``commit_entry`` sets this after another capture path records the call. committed: bool = False + # Store resolved continuations as parent-relative suffixes. + delta_records: bool = False + # This records the model server's intent to request prefix supply. + prefix_requested: bool = False + # This records proven application based on generation-time prompt_token_ids. + prefix_supplied: bool = False # Resolve the parent once before dispatch. # Downstream inference and capture share this immutable decision. parent_resolution: LineageResolution | None = None @@ -84,10 +90,12 @@ def parent_tokens(self) -> list[int]: _CAPTURE_CONTEXT: ContextVar[CaptureContext | None] = ContextVar("nemo_gym_capture_context", default=None) + +# Worker-level health counters are logged periodically. _STATS_LOCK = threading.Lock() _RESOLUTION_COUNTS = {"root": 0, "resolved": 0, "unresolved": 0} -_CAPTURE_FAILURES = 0 -_RESOLVER_UNAVAILABLE_NOTED = False +_CAPTURE_FAILURES = [0] +_RESOLVER_UNAVAILABLE_NOTED = [False] def _count_resolution(status_value: str) -> None: @@ -99,9 +107,9 @@ def _count_resolution(status_value: str) -> None: def capture_health_snapshot() -> dict: - """Return worker-level capture health counters.""" + """Return worker-level capture health for metrics endpoints.""" with _STATS_LOCK: - return {"resolutions": dict(_RESOLUTION_COUNTS), "capture_failures": _CAPTURE_FAILURES} + return {"resolutions": dict(_RESOLUTION_COUNTS), "capture_failures": _CAPTURE_FAILURES[0]} def set_token_sink(context: CaptureContext) -> Token: @@ -121,22 +129,6 @@ def reset_token_sink(token: Token) -> None: _CAPTURE_CONTEXT.reset(token) -async def register_call_intent() -> None: - """Record that the captured call is about to be dispatched. - - ``begin_call`` is an optional sink extension. - It lets a source detect a call whose entry was lost. - A failure happens before generation and must fail the model call. - The harness can retry without spending inference compute. - """ - context = _CAPTURE_CONTEXT.get() - if context is None or context.token_sink is None: - return - begin_call = getattr(context.token_sink, "begin_call", None) - if begin_call is not None: - await begin_call(context.rollout_id, context.model_call_id) - - async def resolve_parent(request_messages: list | None) -> None: """Resolve which recorded call this request continues. @@ -144,7 +136,8 @@ async def resolve_parent(request_messages: list | None) -> None: Resolve once before dialect conversion or dispatch. Prefix supply and capture then share one parent decision. Return without work for untagged traffic. - Every attempt records a root, resolved, or unresolved decision. + Every attempted resolution records a root, resolved, or unresolved decision. + An unresolved decision includes its reason. """ context = _CAPTURE_CONTEXT.get() if context is None or request_messages is None: @@ -157,12 +150,16 @@ async def resolve_parent(request_messages: list | None) -> None: ParentResolutionStatus.UNRESOLVED, reason="resolver_unavailable", ) - global _RESOLVER_UNAVAILABLE_NOTED + # Startup requires an explicit unresolved-continuation opt-in. + # Emit one warning and track later calls in the counters. with _STATS_LOCK: - first = not _RESOLVER_UNAVAILABLE_NOTED - _RESOLVER_UNAVAILABLE_NOTED = True + first = not _RESOLVER_UNAVAILABLE_NOTED[0] + _RESOLVER_UNAVAILABLE_NOTED[0] = True if first: - logger.warning("No lineage resolver is available. Every continuation will be unresolved and masked.") + logger.warning( + "No lineage resolver is available: every continuation resolves UNRESOLVED " + "and multi-call rollouts will be masked (allow_unresolved_continuations is set)." + ) else: context.parent_resolution = await context.lineage_store.resolve(context.rollout_id, request_messages) _count_resolution(context.parent_resolution.status.value) @@ -174,6 +171,24 @@ async def resolve_parent(request_messages: list | None) -> None: ) +async def register_call_intent() -> None: + """Record durable call intent before dispatch starts generation. + + ``begin_call`` is an optional sink extension. + A dangling intent identifies a lost entry. + Failure happens before generation and propagates to the caller. + The harness can retry without spending inference compute. + Sinks without ``begin_call`` cannot report a missing final entry this way. + """ + context = _CAPTURE_CONTEXT.get() + if context is None or context.token_sink is None: + return + begin = getattr(context.token_sink, "begin_call", None) + if begin is None: + return + await begin(context.rollout_id, context.model_call_id) + + async def capture_tokens( response: Any, request_messages: list | None = None, @@ -232,6 +247,8 @@ async def capture_tokens( # so the recorded id matches what the client received in every dialect. response_id=str(payload.get("id") or "") or None, created_at=time.time(), + prefix_requested=context.prefix_requested, + prefix_supplied=context.prefix_supplied, ) if request_messages is not None: stamp_continuation(entry, list(request_messages)) @@ -270,18 +287,38 @@ async def commit_entry( context.committed = True return try: + # Use the resolution decided before dispatch. + # Engine-side callers may pass their own resolution. resolution = parent_resolution or context.parent_resolution if resolution is None: resolution = LineageResolution( ParentResolutionStatus.UNRESOLVED, reason="not_attempted", ) - # The cumulative length and digest always describe this call. + # The digest always describes the full sequence. + # Delta storage changes representation, not lineage identity. # The parent decision is persisted with the same sink write. + cumulative = None + if ( + context.delta_records + and not entry.prompt_is_delta + and resolution.status == ParentResolutionStatus.RESOLVED + and resolution.match is not None + and resolution.match.cumulative_token_ids + ): + parent_cum = list(resolution.match.cumulative_token_ids) + prompt = list(entry.prompt_token_ids) + # Store a suffix only when the prompt extends the exact parent tokens. + # Otherwise retain the full prompt and preserve a safe reconstruction anchor. + if len(prompt) >= len(parent_cum) and prompt[: len(parent_cum)] == parent_cum: + cumulative = prompt + list(entry.generation_token_ids) + entry.prompt_token_ids = prompt[len(parent_cum) :] + entry.prompt_is_delta = True stamp_lineage( entry, resolution.match.model_call_id if resolution.match is not None else None, parent_resolution=resolution.status, + cumulative=cumulative, ) entry.parent_resolution_reason = resolution.reason or "" await context.token_sink.put(entry) @@ -297,10 +334,9 @@ async def _capture_failed(context: CaptureContext, stage: str) -> None: Mark the rollout so consumers can mask the sample. Call this only from an ``except`` block. """ - global _CAPTURE_FAILURES with _STATS_LOCK: - _CAPTURE_FAILURES += 1 - failures = _CAPTURE_FAILURES + _CAPTURE_FAILURES[0] += 1 + failures = _CAPTURE_FAILURES[0] if failures % 10 == 0: logger.error("Training-token capture has failed %d times in this worker.", failures) logger.warning( diff --git a/responses_api_models/vllm_model/app.py b/responses_api_models/vllm_model/app.py index f7943c9cde..5bc893a102 100644 --- a/responses_api_models/vllm_model/app.py +++ b/responses_api_models/vllm_model/app.py @@ -19,12 +19,13 @@ import logging import os from copy import deepcopy +from threading import Lock from time import monotonic, time, time_ns from typing import Any, ClassVar, Dict, List, Optional, Union from aiohttp.client_exceptions import ClientResponseError from fastapi import Request -from pydantic import Field +from pydantic import Field, PrivateAttr, model_validator from nemo_gym.base_responses_api_model import ( BaseResponsesAPIModelConfig, @@ -49,6 +50,7 @@ split_responses_input_output_items, # noqa: F401 ) from nemo_gym.server_utils import SESSION_ID_KEY, is_nemo_gym_fastapi_entrypoint +from nemo_gym.token_id_capture import current_capture_context LOG = logging.getLogger("nemo_gym.vllm_model") @@ -167,6 +169,12 @@ class VLLMModelConfig(BaseResponsesAPIModelConfig): # Whether or not the model can generate a reasoning output, and called again to produce additional reasoning output. sequential_reasoning_allowed: bool = True + # Opt in to supplying a verified parent's exact tokens to the engine. + # Prefix supply requires generation-time prompt_token_ids as proof. + # Stock vLLM does not support the required_prefix_token_ids extension. + # Prefix supply is incompatible with use_completions_api=true. + supply_prefix_token_ids: bool = False + # As of Feb 2026, we default this to False since majority of open source models aren't responses native with the exception of GPT-OSS is_responses_native: bool = False @@ -213,6 +221,16 @@ class VLLMModelConfig(BaseResponsesAPIModelConfig): # small without depending on vLLM's ``--allowed-local-media-path``. audio_root: Optional[str] = None + @model_validator(mode="after") + def _validate_prefix_supply(self) -> "VLLMModelConfig": + if self.supply_prefix_token_ids and not self.return_token_id_information: + raise ValueError("supply_prefix_token_ids requires return_token_id_information=true") + if self.supply_prefix_token_ids and self.use_completions_api: + raise ValueError("supply_prefix_token_ids is not supported with use_completions_api=true") + if self.supply_prefix_token_ids and self.is_responses_native: + raise ValueError("supply_prefix_token_ids is not supported with is_responses_native=true") + return self + # When True, outbound calls go to vLLM's /v1/completions endpoint instead # of /v1/chat/completions. The Gym /v1/responses and /v1/chat/completions # external endpoints continue to work; only the upstream call swaps. @@ -622,8 +640,104 @@ def _preprocess_chat_completion_create_params(self, request: Request, body_dict: self._apply_sampling_overrides(body_dict) self._validate_single_choice_token_request(body_dict) + body_dict = self._apply_prefix_supply(body_dict) + return body_dict + # Protect the ``[supplied, eligible, total]`` diagnostic counts. + # Eligible calls have a resolved parent. + _prefix_supply_counts: List[int] = PrivateAttr(default_factory=lambda: [0, 0, 0]) + _prefix_supply_lock: Any = PrivateAttr(default_factory=Lock) + + def _apply_prefix_supply(self, body_dict: Dict[str, Any]) -> Dict[str, Any]: + """Add a verified parent's exact tokens to a compatible engine request. + + Prefix supply is opt-in. + A unique parent match provides the cumulative token prefix. + The backend must implement the ``required_prefix_token_ids`` extension. + Stock vLLM does not implement this extension. + ``prefix_requested`` records that the request included the prefix. + Only generation-time ``prompt_token_ids`` can prove that the backend applied it. + That proof sets ``prefix_supplied``. + A missing or ambiguous parent leaves the request unchanged. + """ + if not self.config.supply_prefix_token_ids: + return body_dict + context = current_capture_context() + # The parent was resolved before dispatch from the request as received. + # Conversion and preprocessing may have reshaped this body. + # Re-resolving here could select against a representation never indexed. + parent_tokens = context.parent_tokens if context is not None else [] + with self._prefix_supply_lock: + self._prefix_supply_counts[2] += 1 + if parent_tokens: + self._prefix_supply_counts[1] += 1 + if context is None: + # An uncorrelated rollout call has no verified parent. + return body_dict + if not parent_tokens: + return body_dict + body_dict["required_prefix_token_ids"] = parent_tokens + # This records intent only. + # ``prefix_supplied`` remains false until generation-time prompt_token_ids prove application. + context.prefix_requested = True + return body_dict + + @staticmethod + def _generation_prompt_token_ids(response: dict) -> Any: + """Return the prompt token IDs reported by generation. + + Prefer the message-level token bundle over top-level transport fields. + Token capture uses the same source order. + """ + choices = response.get("choices") + choice = choices[0] if isinstance(choices, list) and choices and isinstance(choices[0], dict) else {} + message = choice.get("message") + if isinstance(message, dict) and message.get("prompt_token_ids") is not None: + return message["prompt_token_ids"] + return response.get("prompt_token_ids") + + def _verify_generation_prefix(self, body_dict: dict, response: dict) -> None: + """Require generation-time proof that the engine applied the requested prefix.""" + context = current_capture_context() + if context is None or not context.prefix_requested: + return + required = body_dict.get("required_prefix_token_ids") + if not required: + raise RuntimeError("A requested token prefix was removed before generation.") + tokens = self._generation_prompt_token_ids(response) + if not isinstance(tokens, list): + raise RuntimeError( + f"`{self.config.name}` (base_url={self.config.base_url}) requested " + "required_prefix_token_ids, but the generation response did not include prompt_token_ids " + "proving which prompt the engine used. The backend must implement the " + "required_prefix_token_ids extension and return generation-time prompt token ids. " + "Disabling supply_prefix_token_ids is the fallback." + ) + tokens = [int(token) for token in tokens] + if tokens[: len(required)] != list(required): + raise RuntimeError( + f"`{self.config.name}` (base_url={self.config.base_url}) returned generation " + "prompt_token_ids that do not start with required_prefix_token_ids. The backend must " + "implement the required_prefix_token_ids extension and return generation-time prompt " + "token ids that extend the supplied prefix. Disabling supply_prefix_token_ids is the fallback." + ) + # This proves that the served prompt extended the exact parent tokens. + # It does not prove how the backend produced that prompt. + # A prefix-stable re-render still satisfies the training invariant. + context.prefix_supplied = True + with self._prefix_supply_lock: + self._prefix_supply_counts[0] += 1 + supplied, eligible, total = self._prefix_supply_counts + if supplied % 10 == 0: + LOG.info( + "prefix supply: %d/%d eligible calls supplied (%.0f%%; %d enabled calls total)", + supplied, + eligible, + 100.0 * supplied / eligible, + total, + ) + async def chat_completions( self, request: Request, body: NeMoGymChatCompletionCreateParamsNonStreaming = Body() ) -> NeMoGymChatCompletion: @@ -741,6 +855,7 @@ async def chat_completions( ) choice_dict = chat_completion_dict["choices"][0] + self._verify_generation_prefix(body_dict, chat_completion_dict) if self.config.uses_reasoning_parser: # See the TODO wrt reasoning_content above reasoning_content = choice_dict["message"].get("reasoning_content") or choice_dict["message"].get( diff --git a/responses_api_models/vllm_model/configs/vllm_model_supply_prefix.yaml b/responses_api_models/vllm_model/configs/vllm_model_supply_prefix.yaml new file mode 100644 index 0000000000..a09ff8874d --- /dev/null +++ b/responses_api_models/vllm_model/configs/vllm_model_supply_prefix.yaml @@ -0,0 +1,15 @@ +# Opt in to prefix supply on the Gym model server. +# +# The server requests the exact tokens from a verified parent through required_prefix_token_ids. +# The backend must return generation-time prompt_token_ids that prove the prefix was applied. +# prefix_requested records intent only. +# prefix_supplied records proven application. +# +# Supply requires exactly one recorded parent that matches the incoming conversation. +# A missing or ambiguous parent leaves the request unchanged. +# Stock vLLM does not support the required_prefix_token_ids extension. +# Prefix supply is incompatible with use_completions_api=true. +policy_model: + responses_api_models: + vllm_model: + supply_prefix_token_ids: true diff --git a/responses_api_models/vllm_model/tests/test_app.py b/responses_api_models/vllm_model/tests/test_app.py index 9ee572e524..2a73ae236e 100644 --- a/responses_api_models/vllm_model/tests/test_app.py +++ b/responses_api_models/vllm_model/tests/test_app.py @@ -12,7 +12,9 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import asyncio import json +import logging from typing import Any, Union from unittest.mock import AsyncMock, MagicMock @@ -53,6 +55,14 @@ NeMoGymSummary, ) from nemo_gym.server_utils import SESSION_ID_KEY, ServerClient +from nemo_gym.token_id_capture import ( + CaptureContext, + InMemoryLineageStore, + TokenCaptureStore, + reset_token_sink, + resolve_parent, + set_token_sink, +) from responses_api_models.vllm_model.app import ( VLLMConverter, VLLMModel, @@ -67,6 +77,12 @@ FIXED_TIME = 1691418000 FIXED_UUID = "123" +_TEST_LINEAGE = InMemoryLineageStore() + + +def lineage_index(): + return _TEST_LINEAGE.index + def test_transport_io_writer_keeps_full_payload(monkeypatch: MonkeyPatch, tmp_path) -> None: log_path = tmp_path / "model-io-transport.jsonl" @@ -5190,3 +5206,632 @@ def test_unpublished_endpoint_grace_lifecycle(self, tmp_path, monkeypatch: Monke now = 1701.0 with raises(RuntimeError, match="no longer published"): server._maybe_rebind_endpoint() + + +class TestPrefixSupply: + """Supply a verified parent's exact tokens to the engine. + + Re-rendering can tokenize an assistant turn differently from generation. + A reasoning template can also omit earlier thinking. + Either case breaks the token chain. + """ + + @staticmethod + def _server(monkeypatch: MonkeyPatch, *, enabled: bool) -> VLLMModel: + config = VLLMModelConfig( + host="0.0.0.0", + port=8081, + base_url="http://api.openai.com/v1", + api_key="dummy_key", # pragma: allowlist secret + model="dummy_model", + entrypoint="", + name="", + return_token_id_information=enabled, + uses_reasoning_parser=False, + supply_prefix_token_ids=enabled, + ) + get_global_config_dict_mock = MagicMock(return_value={}) + monkeypatch.setattr(nemo_gym.server_utils, "get_global_config_dict", get_global_config_dict_mock) + return VLLMModel(config=config, server_client=MagicMock(spec=ServerClient, global_config_dict={})) + + @staticmethod + def _armed(rollout_id: str, tmp_path) -> CaptureContext: + return CaptureContext( + rollout_id=rollout_id, + model_call_id="call-x", + token_sink=TokenCaptureStore(tmp_path), + lineage_store=_TEST_LINEAGE, + ) + + def test_requires_token_information_for_generation_proof(self) -> None: + with raises(ValueError, match="return_token_id_information=true"): + VLLMModelConfig( + host="0.0.0.0", + port=8081, + base_url="http://api.openai.com/v1", + api_key="dummy_key", # pragma: allowlist secret + model="dummy_model", + entrypoint="", + name="", + return_token_id_information=False, + uses_reasoning_parser=False, + supply_prefix_token_ids=True, + ) + + def test_rejects_completions_api_prefix_supply(self) -> None: + with raises(ValueError, match="not supported with use_completions_api=true"): + VLLMModelConfig( + host="0.0.0.0", + port=8081, + base_url="http://api.openai.com/v1", + api_key="dummy_key", # pragma: allowlist secret + model="dummy_model", + entrypoint="", + name="", + return_token_id_information=True, + uses_reasoning_parser=False, + supply_prefix_token_ids=True, + use_completions_api=True, + ) + + def test_supplies_the_parents_cumulative_tokens(self, monkeypatch: MonkeyPatch, tmp_path) -> None: + server = self._server(monkeypatch, enabled=True) + first_turn = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}] + lineage_index().for_rollout("sup-0").record("parent", first_turn, [1, 2, 3, 4], "d") + + token = set_token_sink(self._armed("sup-0", tmp_path)) + try: + asyncio.run(resolve_parent(first_turn + [{"role": "user", "content": "next"}])) + out = server._apply_prefix_supply({"messages": first_turn + [{"role": "user", "content": "next"}]}) + finally: + reset_token_sink(token) + + assert out["required_prefix_token_ids"] == [1, 2, 3, 4] + + def test_off_by_default(self, monkeypatch: MonkeyPatch, tmp_path) -> None: + server = self._server(monkeypatch, enabled=False) + turn = [{"role": "assistant", "content": "hello"}] + lineage_index().for_rollout("sup-1").record("parent", turn, [1, 2], "d") + + token = set_token_sink(self._armed("sup-1", tmp_path)) + try: + asyncio.run(resolve_parent(turn)) + out = server._apply_prefix_supply({"messages": turn}) + finally: + reset_token_sink(token) + + assert "required_prefix_token_ids" not in out + + def test_uncorrelated_call_is_left_alone(self, monkeypatch: MonkeyPatch) -> None: + server = self._server(monkeypatch, enabled=True) + asyncio.run(resolve_parent([{"role": "assistant", "content": "hello"}])) + out = server._apply_prefix_supply({"messages": [{"role": "assistant", "content": "hello"}]}) + assert "required_prefix_token_ids" not in out + + def test_fingerprint_miss_falls_back_rather_than_supplying_something_wrong( + self, monkeypatch: MonkeyPatch, tmp_path + ) -> None: + """Leave a rewritten history untouched. + + The backend trusts the supplied prefix. + A wrong prefix would generate from a conversation that the harness never requested. + """ + server = self._server(monkeypatch, enabled=True) + lineage_index().for_rollout("sup-2").record("parent", [{"role": "assistant", "content": "hello"}], [1, 2], "d") + + token = set_token_sink(self._armed("sup-2", tmp_path)) + try: + asyncio.run(resolve_parent([{"role": "assistant", "content": "a summary"}])) + out = server._apply_prefix_supply({"messages": [{"role": "assistant", "content": "a summary"}]}) + finally: + reset_token_sink(token) + + assert "required_prefix_token_ids" not in out + + def test_ambiguous_parent_is_not_supplied(self, monkeypatch: MonkeyPatch, tmp_path) -> None: + server = self._server(monkeypatch, enabled=True) + turn = [{"role": "assistant", "content": "same"}] + lineage = lineage_index().for_rollout("sup-3") + lineage.record("a", turn, [1, 2], "da") + lineage.record("b", turn, [3, 4], "db") + + token = set_token_sink(self._armed("sup-3", tmp_path)) + try: + asyncio.run(resolve_parent(turn)) + out = server._apply_prefix_supply({"messages": turn}) + finally: + reset_token_sink(token) + + assert "required_prefix_token_ids" not in out + + def test_a_fork_gets_the_parents_prefix_not_the_previous_calls(self, monkeypatch: MonkeyPatch, tmp_path) -> None: + """Give each branch the verified parent's cumulative tokens. + + A running cursor would include the first branch's generation in the second branch. + The backend would then generate from a conversation that never happened. + """ + server = self._server(monkeypatch, enabled=True) + shared = [{"role": "user", "content": "q"}, {"role": "assistant", "content": "plan"}] + lineage = lineage_index().for_rollout("sup-4") + lineage.record("parent", shared, [1, 2, 3], "dp") + lineage.record( + "branch-a", + shared + [{"role": "user", "content": "a"}, {"role": "assistant", "content": "A"}], + [1, 2, 3, 9, 9], + "da", + ) + + token = set_token_sink(self._armed("sup-4", tmp_path)) + try: + asyncio.run(resolve_parent(shared + [{"role": "user", "content": "b"}])) + out = server._apply_prefix_supply({"messages": shared + [{"role": "user", "content": "b"}]}) + finally: + reset_token_sink(token) + + assert out["required_prefix_token_ids"] == [1, 2, 3] + + def test_reasoning_stripped_history_still_supplies_the_real_tokens( + self, monkeypatch: MonkeyPatch, tmp_path + ) -> None: + """Restore reasoning tokens omitted from the rendered history. + + A reasoning template can drop earlier thinking from a later prompt. + The rendered prompt then cannot extend the previous generation. + The verified parent's recorded tokens restore the chain. + """ + server = self._server(monkeypatch, enabled=True) + # The recorded turn included reasoning. + # The history returned by the harness does not. + recorded_turn = [{"role": "user", "content": "q"}, {"role": "assistant", "content": "answer"}] + real_tokens_including_reasoning = [1, 2, 3, 4, 5, 6, 7] + lineage_index().for_rollout("sup-5").record("parent", recorded_turn, real_tokens_including_reasoning, "d") + + token = set_token_sink(self._armed("sup-5", tmp_path)) + try: + asyncio.run(resolve_parent(recorded_turn + [{"role": "user", "content": "next"}])) + out = server._apply_prefix_supply({"messages": recorded_turn + [{"role": "user", "content": "next"}]}) + finally: + reset_token_sink(token) + + assert out["required_prefix_token_ids"] == real_tokens_including_reasoning + + +class TestPrefixSupplyAccounting: + """Distinguish requested prefixes from prefixes proven to be applied. + + Contiguous chains do not prove that supply fired. + ``prefix_requested`` records intent only. + ``prefix_supplied`` records generation-time proof. + A lock protects the diagnostic counters. + """ + + def test_supplied_call_is_only_requested_until_generation_proves_it( + self, monkeypatch: MonkeyPatch, tmp_path + ) -> None: + server = TestPrefixSupply._server(monkeypatch, enabled=True) + turn = [{"role": "user", "content": "q"}, {"role": "assistant", "content": "a"}] + lineage_index().for_rollout("acct-0").record("parent", turn, [1, 2, 3], "d") + + ctx = CaptureContext( + rollout_id="acct-0", + model_call_id="c", + token_sink=TokenCaptureStore(tmp_path), + lineage_store=_TEST_LINEAGE, + ) + token = set_token_sink(ctx) + try: + asyncio.run(resolve_parent(turn + [{"role": "user", "content": "next"}])) + out = server._apply_prefix_supply({"messages": turn + [{"role": "user", "content": "next"}]}) + finally: + reset_token_sink(token) + + assert out["required_prefix_token_ids"] == [1, 2, 3] + assert ctx.prefix_requested is True + assert ctx.prefix_supplied is False + # A resolved parent makes this call eligible. + # Generation proof has not landed yet. + assert server._prefix_supply_counts == [0, 1, 1] + + def test_fallback_counts_in_total_but_not_as_eligible(self, monkeypatch: MonkeyPatch, tmp_path) -> None: + server = TestPrefixSupply._server(monkeypatch, enabled=True) + lineage_index().for_rollout("acct-1").record("parent", [{"role": "assistant", "content": "a"}], [1, 2], "d") + + ctx = CaptureContext( + rollout_id="acct-1", + model_call_id="c", + token_sink=TokenCaptureStore(tmp_path), + lineage_store=_TEST_LINEAGE, + ) + token = set_token_sink(ctx) + try: + # A rewritten history has no unique verified parent. + # Prefix supply must decline. + asyncio.run(resolve_parent([{"role": "assistant", "content": "rewritten"}])) + out = server._apply_prefix_supply({"messages": [{"role": "assistant", "content": "rewritten"}]}) + finally: + reset_token_sink(token) + + assert "required_prefix_token_ids" not in out + assert ctx.prefix_supplied is False + # A call without a resolved parent is not eligible. + assert server._prefix_supply_counts == [0, 0, 1] + + +class TestPrefixSupplyReachesTokenize: + """Require generation-time proof for supplied prefixes. + + The generation request carries ``required_prefix_token_ids``. + A supplied request skips the separate tokenize call. + The generation response must return the actual ``prompt_token_ids``. + Those token IDs prove whether the backend applied the requested prefix. + """ + + @staticmethod + def _model() -> VLLMModel: + config = VLLMModelConfig( + host="0.0.0.0", + port=8080, + entrypoint="", + name="vllm_model", + base_url="http://localhost:9999/v1", + api_key="dummy_key", # pragma: allowlist secret + model="dummy_model", + return_token_id_information=True, + uses_reasoning_parser=False, + uses_interleaved_reasoning=False, + supply_prefix_token_ids=True, + ) + return VLLMModel(config=config, server_client=MagicMock(spec=ServerClient, global_config_dict={})) + + def test_supplied_prefix_uses_generation_prompt_proof(self, tmp_path) -> None: + model = self._model() + app = model.setup_webserver() + + turn = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}] + lineage_index().for_rollout("tok-0").record("parent", turn, [11, 12, 13], "d") + + chat_kwargs: dict[str, Any] = {} + tokenize_kwargs: dict[str, Any] = {} + + async def mock_create_chat_completion(**kwargs): + chat_kwargs.update(kwargs) + return { + "id": "c", + "object": "chat.completion", + "created": 0, + "model": "dummy_model", + "prompt_token_ids": [11, 12, 13, 77], + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "token_ids": [77], + "message": {"role": "assistant", "content": "ok"}, + "logprobs": { + "content": [{"token": "token_id:77", "logprob": -0.5, "bytes": None, "top_logprobs": []}] + }, + } + ], + } + + async def mock_create_tokenize(**kwargs): + tokenize_kwargs.update(kwargs) + return {"tokens": [11, 12, 13, 77]} + + mock_client = MagicMock(spec=NeMoGymAsyncOpenAI) + mock_client.create_chat_completion = AsyncMock(side_effect=mock_create_chat_completion) + mock_client.create_tokenize = AsyncMock(side_effect=mock_create_tokenize) + model._clients = [mock_client] + + sink = set_token_sink( + CaptureContext( + rollout_id="tok-0", + model_call_id="call-y", + token_sink=TokenCaptureStore(tmp_path), + lineage_store=_TEST_LINEAGE, + ) + ) + try: + client = TestClient(app) + response = client.post( + "/v1/chat/completions", + json={"messages": turn + [{"role": "user", "content": "next"}]}, + ) + finally: + reset_token_sink(sink) + + assert response.status_code == 200 + assert chat_kwargs["required_prefix_token_ids"] == [11, 12, 13] + assert tokenize_kwargs == {} + assert mock_client.create_tokenize.await_count == 0 + + def _supply_and_return_prompt( + self, tmp_path, prompt_tokens: list[int] | None, rollout: str + ) -> tuple[CaptureContext, RuntimeError | None]: + """Run one supplied call against an engine that returns ``prompt_tokens``.""" + model = self._model() + app = model.setup_webserver() + + turn = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}] + lineage_index().for_rollout(rollout).record("parent", turn, [11, 12, 13], "d") + + async def mock_create_chat_completion(**kwargs): + return { + "id": "c", + "object": "chat.completion", + "created": 0, + "model": "dummy_model", + "prompt_token_ids": prompt_tokens, + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "token_ids": [77], + "message": {"role": "assistant", "content": "ok"}, + "logprobs": { + "content": [{"token": "token_id:77", "logprob": -0.5, "bytes": None, "top_logprobs": []}] + }, + } + ], + } + + async def mock_create_tokenize(**kwargs): + return {"tokens": prompt_tokens} + + mock_client = MagicMock(spec=NeMoGymAsyncOpenAI) + mock_client.create_chat_completion = AsyncMock(side_effect=mock_create_chat_completion) + mock_client.create_tokenize = AsyncMock(side_effect=mock_create_tokenize) + model._clients = [mock_client] + + context = CaptureContext( + rollout_id=rollout, + model_call_id="call-v", + token_sink=TokenCaptureStore(tmp_path), + lineage_store=_TEST_LINEAGE, + ) + sink = set_token_sink(context) + error = None + try: + response = TestClient(app).post( + "/v1/chat/completions", + json={"messages": turn + [{"role": "user", "content": "next"}]}, + ) + except RuntimeError as caught: + error = caught + response = None + finally: + reset_token_sink(sink) + if response is not None: + assert response.status_code == 200 + return context, error + + def test_a_backend_that_ignores_the_prefix_is_recorded_as_not_supplied(self, tmp_path, caplog) -> None: + """Treat a requested prefix as intent until the generation response proves application. + + A backend can ignore an unsupported request field and answer normally. + Only generation-time prompt_token_ids can detect that behavior. + """ + with caplog.at_level(logging.ERROR): + context, error = self._supply_and_return_prompt(tmp_path, [900, 901, 77], "ver-ignored") + assert error is not None + assert context.prefix_supplied is False + assert "do not start with required_prefix_token_ids" in str(error) + + def test_a_backend_that_applies_the_prefix_is_recorded_as_supplied(self, tmp_path) -> None: + """Record supply when generation-time prompt_token_ids start with the requested prefix.""" + context, error = self._supply_and_return_prompt(tmp_path, [11, 12, 13, 77], "ver-applied") + assert error is None + assert context.prefix_supplied is True + + def test_a_backend_without_generation_prompt_proof_fails_closed(self, tmp_path) -> None: + context, error = self._supply_and_return_prompt(tmp_path, None, "ver-unproved") + assert error is not None + assert "did not include prompt_token_ids" in str(error) + assert context.prefix_supplied is False + + def test_tokenize_body_omits_the_prefix_when_supply_did_not_fire(self, tmp_path) -> None: + """Omit the prefix when no verified parent was resolved.""" + model = self._model() + app = model.setup_webserver() + + tokenize_kwargs: dict[str, Any] = {} + + async def mock_create_chat_completion(**kwargs): + return { + "id": "c", + "object": "chat.completion", + "created": 0, + "model": "dummy_model", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "ok"}, + "logprobs": { + "content": [{"token": "token_id:77", "logprob": -0.5, "bytes": None, "top_logprobs": []}] + }, + } + ], + } + + async def mock_create_tokenize(**kwargs): + tokenize_kwargs.update(kwargs) + return {"tokens": [5]} + + mock_client = MagicMock(spec=NeMoGymAsyncOpenAI) + mock_client.create_chat_completion = AsyncMock(side_effect=mock_create_chat_completion) + mock_client.create_tokenize = AsyncMock(side_effect=mock_create_tokenize) + model._clients = [mock_client] + + sink = set_token_sink( + CaptureContext( + rollout_id="tok-unknown", + model_call_id="call-z", + token_sink=TokenCaptureStore(tmp_path), + lineage_store=_TEST_LINEAGE, + ) + ) + try: + client = TestClient(app) + response = client.post("/v1/chat/completions", json={"messages": [{"role": "user", "content": "hi"}]}) + finally: + reset_token_sink(sink) + + assert response.status_code == 200 + assert "required_prefix_token_ids" not in tokenize_kwargs + + +class TestPrefixSupplyUsesTheRequestAsReceived: + """Resolve the parent before request conversion or preprocessing. + + Responses-to-Chat conversion and preprocessing can reshape the request body. + The lineage index records the request as received. + Resolving from the reshaped body could miss a verified parent. + """ + + def test_supply_does_not_fire_without_a_resolved_parent(self, monkeypatch, tmp_path): + server = TestPrefixSupply._server(monkeypatch, enabled=True) + turn = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}] + lineage_index().for_rollout("sup-x").record("parent", turn, [1, 2, 3], "d") + + token = set_token_sink(TestPrefixSupply._armed("sup-x", tmp_path)) + try: + # No parent resolution proved that this request continues another call. + out = server._apply_prefix_supply({"messages": turn}) + finally: + reset_token_sink(token) + + assert "required_prefix_token_ids" not in out + + def test_a_body_reshaped_after_resolution_still_supplies(self, monkeypatch, tmp_path): + """Preserve a verified parent decision across later request reshaping.""" + server = TestPrefixSupply._server(monkeypatch, enabled=True) + original = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}] + lineage_index().for_rollout("sup-y").record("parent", original, [7, 8, 9], "d") + + token = set_token_sink(TestPrefixSupply._armed("sup-y", tmp_path)) + try: + asyncio.run(resolve_parent(original)) + # Conversion reshaped the messages before prefix supply ran. + out = server._apply_prefix_supply( + {"messages": [{"role": "user", "content": "reshaped beyond recognition"}]} + ) + finally: + reset_token_sink(token) + + assert out["required_prefix_token_ids"] == [7, 8, 9] + + +class TestGenerationProofAcceptsBundleShape: + """Read generation-time prompt tokens from every supported response shape. + + A backend can return prompt token IDs in a message-level bundle or a top-level transport field. + Prefix verification uses the same priority order as token capture. + """ + + def test_a_backend_returning_a_message_bundle_passes_verification(self, tmp_path) -> None: + model = TestPrefixSupplyReachesTokenize._model() + app = model.setup_webserver() + + turn = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}] + lineage_index().for_rollout("bundle-0").record("parent", turn, [11, 12, 13], "d") + + async def mock_create_chat_completion(**kwargs): + return { + "id": "c", + "object": "chat.completion", + "created": 0, + "model": "dummy_model", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": { + "role": "assistant", + "content": "ok", + "prompt_token_ids": [11, 12, 13, 77], + "generation_token_ids": [77], + "generation_log_probs": [-0.5], + }, + } + ], + } + + mock_client = MagicMock(spec=NeMoGymAsyncOpenAI) + mock_client.create_chat_completion = AsyncMock(side_effect=mock_create_chat_completion) + mock_client.create_tokenize = AsyncMock() + model._clients = [mock_client] + + context = CaptureContext( + rollout_id="bundle-0", + model_call_id="call-b", + token_sink=TokenCaptureStore(tmp_path), + lineage_store=_TEST_LINEAGE, + ) + sink = set_token_sink(context) + try: + response = TestClient(app).post( + "/v1/chat/completions", + json={"messages": turn + [{"role": "user", "content": "next"}]}, + ) + finally: + reset_token_sink(sink) + + assert response.status_code == 200 + assert context.prefix_supplied is True + assert mock_client.create_tokenize.await_count == 0 + + def _verify(self, tmp_path, response: dict) -> CaptureContext: + model = TestPrefixSupplyReachesTokenize._model() + context = CaptureContext( + rollout_id="bundle-1", + model_call_id="call-b", + token_sink=TokenCaptureStore(tmp_path), + lineage_store=_TEST_LINEAGE, + ) + context.prefix_requested = True + token = set_token_sink(context) + try: + result = model._verify_generation_prefix({"required_prefix_token_ids": [11, 12, 13]}, response) + finally: + reset_token_sink(token) + assert result is None # The proof has no consumer; the function returns nothing. + return context + + def test_message_level_prompt_token_ids_outrank_top_level(self, tmp_path) -> None: + context = self._verify( + tmp_path, + { + "prompt_token_ids": [900, 901], + "choices": [{"index": 0, "message": {"role": "assistant", "prompt_token_ids": [11, 12, 13, 77]}}], + }, + ) + assert context.prefix_supplied is True + + def test_top_level_prompt_token_ids_remain_a_valid_source(self, tmp_path) -> None: + context = self._verify( + tmp_path, + { + "prompt_token_ids": [11, 12, 13, 77], + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}}], + }, + ) + assert context.prefix_supplied is True + + +class TestPrefixSupplyRejectsResponsesNative: + def test_rejected_at_construction(self) -> None: + with raises(ValueError, match="not supported with is_responses_native=true"): + VLLMModelConfig( + host="0.0.0.0", + port=8081, + base_url="http://api.openai.com/v1", + api_key="dummy_key", # pragma: allowlist secret + model="dummy_model", + entrypoint="", + name="", + return_token_id_information=True, + uses_reasoning_parser=False, + supply_prefix_token_ids=True, + is_responses_native=True, + ) diff --git a/responses_api_models/vllm_model/tests/test_token_capture_integration.py b/responses_api_models/vllm_model/tests/test_token_capture_integration.py new file mode 100644 index 0000000000..1f741fb61b --- /dev/null +++ b/responses_api_models/vllm_model/tests/test_token_capture_integration.py @@ -0,0 +1,260 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +from collections.abc import Callable +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +from fastapi.testclient import TestClient + +from nemo_gym.openai_utils import NeMoGymAsyncOpenAI +from nemo_gym.server_utils import ServerClient +from nemo_gym.token_id_capture import ( + CaptureContext, + FileLineageStore, + InMemoryLineageStore, + LineageStore, + ParentResolutionStatus, + TokenCaptureSnapshot, + TokenCaptureStore, + TokenEntry, + TokenSink, + TokenSource, + reset_token_sink, + set_token_sink, + trajectories_from_source, +) +from responses_api_models.vllm_model.app import VLLMModel, VLLMModelConfig + + +def _model(client: NeMoGymAsyncOpenAI) -> VLLMModel: + config = VLLMModelConfig( + host="0.0.0.0", + port=8080, + entrypoint="", + name="vllm_model", + base_url="http://localhost:9999/v1", + api_key="dummy_key", # pragma: allowlist secret + model="dummy_model", + return_token_id_information=True, + uses_reasoning_parser=False, + uses_interleaved_reasoning=False, + supply_prefix_token_ids=True, + ) + model = VLLMModel(config=config, server_client=MagicMock(spec=ServerClient, global_config_dict={})) + model._clients = [client] + return model + + +def _completion(prompt: list[int], generation: list[int], content: str) -> dict[str, Any]: + return { + "id": f"completion-{content}", + "object": "chat.completion", + "created": 0, + "model": "dummy_model", + "prompt_token_ids": prompt, + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "token_ids": generation, + "message": {"role": "assistant", "content": content}, + "logprobs": { + "content": [ + { + "token": f"token_id:{token_id}", + "logprob": -0.1, + "bytes": None, + "top_logprobs": [], + } + for token_id in generation + ] + }, + } + ], + } + + +class _ExternalBackend: + def __init__(self) -> None: + self.entries: dict[str, dict[str, TokenEntry]] = {} + self.incomplete: set[str] = set() + self.frozen: set[str] = set() + self.versions: dict[str, int] = {} + self.lineage = InMemoryLineageStore() + + async def commit(self, entry: TokenEntry) -> None: + if entry.rollout_id in self.frozen: + raise RuntimeError(f"Token capture for rollout {entry.rollout_id} is already frozen") + rollout = self.entries.setdefault(entry.rollout_id, {}) + previous = rollout.get(entry.model_call_id) + if previous is not None and previous != entry: + raise ValueError(f"Conflicting model call {entry.model_call_id}") + if previous is None: + rollout[entry.model_call_id] = entry + await self.lineage.put(entry) + self.versions[entry.rollout_id] = self.versions.get(entry.rollout_id, 0) + 1 + + +class _ExternalSink: + def __init__(self, backend: _ExternalBackend) -> None: + self.backend = backend + + async def put(self, entry: TokenEntry) -> None: + await self.backend.commit(entry) + + async def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: + self.backend.incomplete.add(rollout_id) + + async def close(self) -> None: + pass + + +class _ExternalLineageStore: + def __init__(self, backend: _ExternalBackend) -> None: + self.backend = backend + + async def resolve(self, rollout_id: str, request_items: list[dict]): + return await self.backend.lineage.resolve(rollout_id, request_items) + + def is_process_shared(self) -> bool: + return True + + async def close(self) -> None: + pass + + +class _ExternalSource: + def __init__(self, backend: _ExternalBackend) -> None: + self.backend = backend + + async def freeze(self, rollout_id: str) -> TokenCaptureSnapshot: + self.backend.frozen.add(rollout_id) + return TokenCaptureSnapshot( + rollout_id=rollout_id, + entries=tuple(self.backend.entries.get(rollout_id, {}).values()), + incomplete=rollout_id in self.backend.incomplete, + snapshot_id=f"snapshot-{rollout_id}", + version=self.backend.versions.get(rollout_id, 0), + ) + + async def drop(self, rollout_id: str, *, snapshot_id: str, version: int) -> bool: + if snapshot_id != f"snapshot-{rollout_id}" or version != self.backend.versions.get(rollout_id, 0): + return False + self.backend.entries.pop(rollout_id, None) + return True + + async def close(self) -> None: + pass + + +def _simulate_two_workers( + sink_factory: Callable[[], TokenSink], + lineage_factory: Callable[[], LineageStore], + source: TokenSource, + read_entries: Callable[[], list[TokenEntry]], +) -> None: + rollout_id = "simulated-rollout" + outbound_requests: list[dict[str, Any]] = [] + completions = [ + _completion([11, 12], [13, 14], "first answer"), + _completion([11, 12, 13, 14, 21], [22], "second answer"), + ] + + async def create_chat_completion(**kwargs): + outbound_requests.append(kwargs) + return completions[len(outbound_requests) - 1] + + def worker() -> TestClient: + client = MagicMock(spec=NeMoGymAsyncOpenAI) + client.create_chat_completion = AsyncMock(side_effect=create_chat_completion) + client.create_tokenize = AsyncMock() + return TestClient(_model(client).setup_webserver()) + + worker_a = worker() + worker_b = worker() + + def serve(client: TestClient, call_id: str, messages: list[dict]) -> tuple[dict, CaptureContext]: + context = CaptureContext( + rollout_id=rollout_id, + model_call_id=call_id, + token_sink=sink_factory(), + lineage_store=lineage_factory(), + ) + token = set_token_sink(context) + try: + response = client.post("/v1/chat/completions", json={"messages": messages}) + finally: + reset_token_sink(token) + assert response.status_code == 200, response.text + return response.json(), context + + first_request = [{"role": "user", "content": "first question"}] + first_response, first_context = serve(worker_a, "call-a", first_request) + first_message = first_response["choices"][0]["message"] + first_answer = {"role": first_message["role"], "content": first_message["content"]} + + second_request = first_request + [first_answer, {"role": "user", "content": "second question"}] + _, second_context = serve(worker_b, "call-b", second_request) + + entries = read_entries() + assert len(entries) == 2 + assert entries[0].parent_resolution == ParentResolutionStatus.ROOT + assert entries[0].prefix_requested is False + assert entries[0].prefix_supplied is False + assert entries[1].parent_resolution == ParentResolutionStatus.RESOLVED + assert entries[1].parent_call_id == entries[0].model_call_id + assert entries[1].prefix_requested is True + assert entries[1].prefix_supplied is True + + assert first_context.parent_resolution is not None + assert first_context.parent_resolution.status == ParentResolutionStatus.ROOT + assert second_context.parent_resolution is not None + assert second_context.parent_resolution.status == ParentResolutionStatus.RESOLVED + assert second_context.parent_tokens == [11, 12, 13, 14] + assert outbound_requests[0].get("required_prefix_token_ids") is None + assert outbound_requests[1]["required_prefix_token_ids"] == [11, 12, 13, 14] + + built = asyncio.run(trajectories_from_source(rollout_id, source)) + assert built is not None + assert built["mask_sample"] is False + assert built["metrics"]["roots"] == 1 + assert built["metrics"]["chains"] == 1 + assert built["metrics"]["delivered_fraction"] == 1.0 + assert built["metrics"]["unresolved_parent_calls"] == 0 + + output = built["rebuilt_response"]["output"] + assert len(output) == 2 + assert output[0]["prompt_token_ids"] == [11, 12] + assert output[0]["generation_token_ids"] == [13, 14] + assert output[1]["prompt_token_ids"] == [11, 12, 13, 14, 21] + assert output[1]["generation_token_ids"] == [22] + + snapshot = built["_capture_snapshot"] + assert asyncio.run( + source.drop( + rollout_id, + snapshot_id=snapshot["snapshot_id"], + version=snapshot["version"], + ) + ) + + +def test_local_store_capture_supply_and_rebuild_one_safe_trajectory(tmp_path) -> None: + _simulate_two_workers( + sink_factory=lambda: TokenCaptureStore(tmp_path), + lineage_factory=lambda: FileLineageStore(tmp_path), + source=TokenCaptureStore(tmp_path), + read_entries=lambda: TokenCaptureStore(tmp_path).read_entries("simulated-rollout"), + ) + + +def test_external_protocols_capture_supply_and_rebuild_one_safe_trajectory() -> None: + backend = _ExternalBackend() + _simulate_two_workers( + sink_factory=lambda: _ExternalSink(backend), + lineage_factory=lambda: _ExternalLineageStore(backend), + source=_ExternalSource(backend), + read_entries=lambda: list(backend.entries["simulated-rollout"].values()), + ) diff --git a/tests/unit_tests/test_token_id_capture.py b/tests/unit_tests/test_token_id_capture.py index 9800e42741..f8e0fb9e83 100644 --- a/tests/unit_tests/test_token_id_capture.py +++ b/tests/unit_tests/test_token_id_capture.py @@ -1277,6 +1277,43 @@ def _put_shared_file_entry( TokenCaptureStore(root).append(entry) +def _put_shared_file_delta_chain(root: str, depth: int) -> tuple[list[dict], list[int]]: + """Write a resolved delta chain and return its next request and cumulative tokens.""" + rollout_id = "delta-chain" + request = [{"role": "user", "content": "start"}] + cumulative: list[int] = [] + store = TokenCaptureStore(root) + for index in range(depth): + output = [{"role": "assistant", "content": f"answer-{index}"}] + parent_call_id = f"call-{index - 1}" if index else None + parent_resolution = ParentResolutionStatus.RESOLVED if index else ParentResolutionStatus.ROOT + prompt_tokens = [10 + index] + generation_tokens = [100 + index] + cumulative = cumulative + prompt_tokens + generation_tokens + entry = TokenEntry( + rollout_id=rollout_id, + model_call_id=f"call-{index}", + prompt_token_ids=prompt_tokens, + generation_token_ids=generation_tokens, + generation_log_probs=[-0.1], + output_items=output, + parent_call_id=parent_call_id, + parent_resolution=parent_resolution, + prompt_is_delta=index > 0, + ) + stamp_continuation(entry, request) + stamp_lineage( + entry, + parent_call_id, + parent_resolution=parent_resolution, + cumulative=cumulative, + ) + store.append(entry) + request.extend(output) + request.append({"role": "user", "content": f"continue-{index}"}) + return request, cumulative + + async def test_file_lineage_resolves_across_independent_worker_instances(tmp_path): reader = FileLineageStore(tmp_path) request = [{"role": "user", "content": "hello"}] @@ -1291,7 +1328,59 @@ async def test_file_lineage_resolves_across_independent_worker_instances(tmp_pat assert parent.match.cumulative_token_ids == (1, 2, 3) -async def test_file_lineage_cache_is_lru_and_metadata_only(tmp_path): +async def test_file_lineage_batch_loads_and_flattens_a_delta_chain(tmp_path): + request, expected_tokens = _put_shared_file_delta_chain(str(tmp_path), depth=100) + resolver = FileLineageStore(tmp_path) + + with patch.object(resolver, "_load_entries", wraps=resolver._load_entries) as load_entries: + resolution = await resolver.resolve("delta-chain", request) + + assert resolution.status == ParentResolutionStatus.RESOLVED + assert resolution.match is not None + assert resolution.match.model_call_id == "call-99" + assert resolution.match.cumulative_token_ids == tuple(expected_tokens) + load_entries.assert_called_once() + assert len(load_entries.call_args.args[1]) == 100 + + +async def test_file_lineage_reuses_the_latest_materialized_parent(tmp_path): + request, expected_tokens = _put_shared_file_delta_chain(str(tmp_path), depth=100) + resolver = FileLineageStore(tmp_path) + assert (await resolver.resolve("delta-chain", request)).status == ParentResolutionStatus.RESOLVED + + output = [{"role": "assistant", "content": "answer-100"}] + expected_tokens.extend([110, 200]) + entry = TokenEntry( + rollout_id="delta-chain", + model_call_id="call-100", + prompt_token_ids=[110], + generation_token_ids=[200], + generation_log_probs=[-0.1], + output_items=output, + parent_call_id="call-99", + parent_resolution=ParentResolutionStatus.RESOLVED, + prompt_is_delta=True, + ) + stamp_continuation(entry, request) + stamp_lineage( + entry, + "call-99", + parent_resolution=ParentResolutionStatus.RESOLVED, + cumulative=expected_tokens, + ) + TokenCaptureStore(tmp_path).append(entry) + request.extend(output) + request.append({"role": "user", "content": "continue-100"}) + + with patch.object(resolver, "_load_entries", wraps=resolver._load_entries) as load_entries: + resolution = await resolver.resolve("delta-chain", request) + + assert resolution.match is not None + assert resolution.match.cumulative_token_ids == tuple(expected_tokens) + assert len(load_entries.call_args.args[1]) == 1 + + +async def test_file_lineage_cache_is_lru_and_keeps_nodes_metadata_only(tmp_path): resolver = FileLineageStore(tmp_path, max_cached_rollouts=2) continuation = [ {"role": "user", "content": "hello"}, @@ -1318,6 +1407,22 @@ async def test_file_lineage_cache_is_lru_and_metadata_only(tmp_path): assert cold.match.cumulative_token_ids == (1, 2, 3) +async def test_file_lineage_materialized_cache_has_a_global_token_bound(tmp_path): + resolver = FileLineageStore(tmp_path, max_cached_tokens=4) + continuation = [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi"}, + {"role": "user", "content": "next"}, + ] + for rollout_id in ("r-a", "r-b"): + _put_shared_file_entry(str(tmp_path), rollout_id, f"{rollout_id}-c1") + assert (await resolver.resolve(rollout_id, continuation)).status == ParentResolutionStatus.RESOLVED + + assert "r-a" not in resolver._materialized + assert resolver._materialized["r-b"][1] == (1, 2, 3) + assert resolver._materialized_tokens == 3 + + def test_file_lineage_uses_bounded_striped_locks(tmp_path): resolver = FileLineageStore(tmp_path) @@ -2022,6 +2127,14 @@ def test_a_record_below_the_schema_floor_is_refused(): TokenEntry(**_entry_fields(schema_version=TOKEN_ENTRY_MIN_SCHEMA_VERSION - 1)) +def test_omitted_optional_schema_fields_use_safe_defaults(): + entry = TokenEntry(**_entry_fields()) + + assert entry.prompt_is_delta is False + assert entry.prefix_requested is False + assert entry.prefix_supplied is False + + def test_a_record_newer_than_this_reader_is_refused(): """Reject newer records hidden by ``extra="allow"``.""" with pytest.raises(ValidationError, match="this reader understands up to"): diff --git a/tests/unit_tests/test_trajectory_builder.py b/tests/unit_tests/test_trajectory_builder.py index cc6908d0c4..ed79dafa03 100644 --- a/tests/unit_tests/test_trajectory_builder.py +++ b/tests/unit_tests/test_trajectory_builder.py @@ -20,6 +20,9 @@ import pytest from nemo_gym.token_id_capture import ( + CaptureContext, + FileLineageStore, + InMemoryLineageStore, ParentResolutionStatus, TokenCaptureSnapshot, assert_prefix_contiguity, @@ -27,12 +30,16 @@ prefix_merging, project_chain_to_output_items, project_main_chain_response, + reset_token_sink, + set_token_sink, + stamp_continuation, stamp_lineage, token_id_capture_dirs_from_config, trajectories_for_rollout, trajectories_from_source, ) from nemo_gym.token_id_capture.records import TokenEntry +from nemo_gym.token_id_capture.sink import commit_entry, resolve_parent from nemo_gym.token_id_capture.store import TokenCaptureStore @@ -48,8 +55,8 @@ def _entry(mcid, prompt, gen, parent=None, lp=None, created_at=0.0): # Chain selection uses this value. created_at=created_at, ) - # Stamp the way a current writer does: every record carries a decision. - # (Pre-v3 unstamped records no longer exist and are refused by readers.) + # Stamp the way a current writer does. + # Every supported record carries a parent decision. status = ParentResolutionStatus.RESOLVED if parent is not None else ParentResolutionStatus.ROOT stamp_lineage(e, parent, parent_resolution=status) return e @@ -773,3 +780,163 @@ def test_a_chain_that_breaks_is_split_and_reported(): assert len(out.chains) > 1 assert out.notes.delivered_fraction < 1.0 assert_prefix_contiguity(project_main_chain_response("r0", out, model="m")) + + +def _commit_delta_entry( + store: TokenCaptureStore, + lineage: FileLineageStore, + rollout_id: str, + call_id: str, + prompt: list[int], + generation: list[int], + request: list[dict], +) -> CaptureContext: + entry = TokenEntry( + rollout_id=rollout_id, + model_call_id=call_id, + prompt_token_ids=prompt, + generation_token_ids=generation, + generation_log_probs=[-0.1] * len(generation), + output_items=[{"type": "message", "role": "assistant", "content": f"answer {call_id}"}], + ) + stamp_continuation(entry, request) + context = CaptureContext( + rollout_id=rollout_id, + model_call_id=call_id, + token_sink=store, + lineage_store=lineage, + delta_records=True, + ) + token = set_token_sink(context) + try: + asyncio.run(resolve_parent(request)) + asyncio.run(commit_entry(entry)) + finally: + reset_token_sink(token) + return context + + +def test_delta_records_store_suffixes_and_reconstruct_exact_prompts(tmp_path): + store = TokenCaptureStore(tmp_path) + lineage = FileLineageStore(tmp_path) + rollout_id = "delta" + request1 = [{"role": "user", "content": "q1"}] + _commit_delta_entry(store, lineage, rollout_id, "c1", [1, 2, 3], [4, 5], request1) + request2 = request1 + [ + {"role": "assistant", "content": "answer c1"}, + {"role": "user", "content": "q2"}, + ] + second = _commit_delta_entry( + store, + lineage, + rollout_id, + "c2", + [1, 2, 3, 4, 5, 6, 7], + [8], + request2, + ) + request3 = request2 + [ + {"role": "assistant", "content": "answer c2"}, + {"role": "user", "content": "q3"}, + ] + third = _commit_delta_entry( + store, + lineage, + rollout_id, + "c3", + [1, 2, 3, 4, 5, 6, 7, 8, 9], + [10, 11], + request3, + ) + + assert second.parent_resolution is not None + assert second.parent_resolution.status == ParentResolutionStatus.RESOLVED + assert third.parent_resolution is not None + assert third.parent_resolution.status == ParentResolutionStatus.RESOLVED + entries = {entry.model_call_id: entry for entry in store.read_entries(rollout_id)} + assert entries["c1"].prompt_is_delta is False + assert entries["c2"].prompt_is_delta is True + assert entries["c2"].prompt_token_ids == [6, 7] + assert entries["c3"].prompt_is_delta is True + assert entries["c3"].prompt_token_ids == [9] + + built = asyncio.run(trajectories_from_source(rollout_id, store)) + assert built["mask_sample"] is False + output = built["rebuilt_response"]["output"] + assert [item["generation_token_ids"] for item in output] == [[4, 5], [8], [10, 11]] + assert output[2]["prompt_token_ids"] == [1, 2, 3, 4, 5, 6, 7, 8, 9] + + +def test_delta_materialization_handles_a_thousand_turns_without_recursion(): + root = _entry("c0", [1], [2]) + entries = [root] + cumulative = [1, 2] + for turn in range(1, 1_101): + suffix = [10_000 + turn] + generation = [20_000 + turn] + full_prompt = cumulative + suffix + entry = _entry(f"c{turn}", full_prompt, generation, parent=f"c{turn - 1}") + entry.prompt_token_ids = suffix + entry.prompt_is_delta = True + entries.append(entry) + cumulative = full_prompt + generation + + out = prefix_merging(entries) + + assert len(out.chains) == 1 + assert len(out.chains[0].links) == len(entries) + assert out.chains[0].links[-1].entry.prompt_token_ids == cumulative[:-1] + + +def test_broken_delta_chain_masks_instead_of_guessing(tmp_path): + store = TokenCaptureStore(tmp_path) + orphan = TokenEntry( + rollout_id="orphan", + model_call_id="child", + prompt_token_ids=[9], + generation_token_ids=[10], + generation_log_probs=[-0.1], + output_items=[{"type": "message", "role": "assistant", "content": "a"}], + prompt_is_delta=True, + parent_call_id="missing-parent", + parent_resolution=ParentResolutionStatus.RESOLVED, + ) + stamp_continuation(orphan, [{"role": "user", "content": "q"}]) + stamp_lineage( + orphan, + "missing-parent", + parent_resolution=ParentResolutionStatus.RESOLVED, + cumulative=[1, 2, 9, 10], + ) + store.append(orphan) + + built = asyncio.run(trajectories_from_source("orphan", store)) + out = prefix_merging(store.read_entries("orphan")) + + assert built["mask_sample"] is True + assert out.notes.parent_link_failures["delta_chain_unreconstructable"] == 1 + assert out.notes.unresolved_parent_calls == ["child"] + + +def test_in_memory_lineage_refuses_delta_records(): + entry = TokenEntry( + rollout_id="r", + model_call_id="child", + prompt_token_ids=[9], + generation_token_ids=[10], + generation_log_probs=[-0.1], + output_items=[{"type": "message", "role": "assistant", "content": "answer"}], + prompt_is_delta=True, + parent_call_id="parent", + parent_resolution=ParentResolutionStatus.RESOLVED, + ) + stamp_continuation(entry, [{"role": "user", "content": "q"}]) + stamp_lineage( + entry, + "parent", + parent_resolution=ParentResolutionStatus.RESOLVED, + cumulative=[1, 9, 10], + ) + + with pytest.raises(ValueError, match="durable-log-backed"): + asyncio.run(InMemoryLineageStore().put(entry))